Remove ServerMetrics and use a single middleware for both connectio…

…n counting (aka sessions) and call metrics.

client/rpc-servers/src/lib.rs

@@ -25,7 +25,7 @@ use jsonrpsee::{
 	ws_server::{WsServerBuilder, WsServerHandle},
 	RpcModule,
 };
-use prometheus_endpoint::{register, Counter, PrometheusError, Registry, U64};
+use prometheus_endpoint::Registry;
 use std::net::SocketAddr;
 
 use crate::middleware::{RpcMetrics, RpcMiddleware};
@@ -43,44 +43,6 @@ const WS_MAX_CONNECTIONS: usize = 100;
 
 pub mod middleware;
 
-// TODO: (dp) Wire up the ServerMetrics as well
-/// RPC server-specific prometheus metrics.
-#[derive(Debug, Clone, Default)]
-pub struct ServerMetrics {
-	/// Number of sessions opened.
-	session_opened: Option<Counter<U64>>,
-	/// Number of sessions closed.
-	session_closed: Option<Counter<U64>>,
-}
-
-impl ServerMetrics {
-	/// Create new WebSocket RPC server metrics.
-	pub fn new(registry: Option<&Registry>) -> Result<Self, PrometheusError> {
-		registry
-			.map(|r| {
-				Ok(Self {
-					session_opened: register(
-						Counter::new(
-							"rpc_sessions_opened",
-							"Number of persistent RPC sessions opened",
-						)?,
-						r,
-					)?
-					.into(),
-					session_closed: register(
-						Counter::new(
-							"rpc_sessions_closed",
-							"Number of persistent RPC sessions closed",
-						)?,
-						r,
-					)?
-					.into(),
-				})
-			})
-			.unwrap_or_else(|| Ok(Default::default()))
-	}
-}
-
 /// Type alias for http server
 pub type HttpServer = HttpServerHandle;
 /// Type alias for ws server
@@ -121,7 +83,7 @@ pub fn start_http<M: Send + Sync + 'static>(
 		let server = tokio::task::block_in_place(|| rt.block_on(async { builder.build(addrs) }))?;
 		server.start(rpc_api)?
 	} else {
-		let server= tokio::task::block_in_place(|| rt.block_on(async { builder.build(addrs) }))?;
+		let server = tokio::task::block_in_place(|| rt.block_on(async { builder.build(addrs) }))?;
 		server.start(rpc_api)?
 	};
 
@@ -164,7 +126,7 @@ pub fn start_ws<M: Send + Sync + 'static>(
 		let server = tokio::task::block_in_place(|| rt.block_on(builder.build(addrs)))?;
 		server.start(rpc_api)?
 	} else {
-		let server= tokio::task::block_in_place(|| rt.block_on(builder.build(addrs)))?;
+		let server = tokio::task::block_in_place(|| rt.block_on(builder.build(addrs)))?;
 		server.start(rpc_api)?
 	};
 

client/rpc-servers/src/middleware.rs

@@ -20,18 +20,28 @@
 
 use jsonrpsee::types::middleware::Middleware;
 use prometheus_endpoint::{
-	register, CounterVec, HistogramOpts, HistogramVec, Opts, PrometheusError, Registry, U64,
+	register, Counter, CounterVec, HistogramOpts, HistogramVec, Opts, PrometheusError, Registry,
+	U64,
 };
 
 /// Metrics for RPC middleware storing information about the number of requests started/completed,
 /// calls started/completed and their timings.
 #[derive(Debug, Clone)]
 pub struct RpcMetrics {
+	/// Number of RPC requests received since the server started.
 	requests_started: CounterVec<U64>,
+	/// Number of RPC requests completed since the server started.
 	requests_finished: CounterVec<U64>,
+	/// Histogram over RPC execution times.
 	calls_time: HistogramVec,
+	/// Number of calls started.
 	calls_started: CounterVec<U64>,
+	/// Number of calls completed.
 	calls_finished: CounterVec<U64>,
+	/// Number of Websocket sessions opened (Websocket only).
+	ws_sessions_opened: Option<Counter<U64>>,
+	/// Number of Websocket sessions closed (Websocket only).
+	ws_sessions_closed: Option<Counter<U64>>,
 }
 
 impl RpcMetrics {
@@ -85,6 +95,16 @@ impl RpcMetrics {
 				)?,
 				metrics_registry,
 			)?,
+			ws_sessions_opened: register(
+				Counter::new("rpc_sessions_opened", "Number of persistent RPC sessions opened")?,
+				metrics_registry,
+			)?
+			.into(),
+			ws_sessions_closed: register(
+				Counter::new("rpc_sessions_closed", "Number of persistent RPC sessions closed")?,
+				metrics_registry,
+			)?
+			.into(),
 		})
 	}
 }
@@ -106,6 +126,10 @@ impl RpcMiddleware {
 impl Middleware for RpcMiddleware {
 	type Instant = std::time::Instant;
 
+	fn on_connect(&self) {
+		self.metrics.ws_sessions_opened.as_ref().map(|counter| counter.inc());
+	}
+
 	fn on_request(&self) -> Self::Instant {
 		let now = std::time::Instant::now();
 		self.metrics.requests_started.with_label_values(&[self.transport_label]).inc();
@@ -140,4 +164,8 @@ impl Middleware for RpcMiddleware {
 		log::trace!(target: "rpc_metrics", "[{}] on_response started_at={:?}", self.transport_label, started_at);
 		self.metrics.requests_finished.with_label_values(&[self.transport_label]).inc();
 	}
+
+	fn on_disconnect(&self) {
+		self.metrics.ws_sessions_closed.as_ref().map(|counter| counter.inc());
+	}
 }