jsonrpsee integration (#8783)

* Add tokio * No need to map CallError to CallError * jsonrpsee proc macros (#9673) * port error types to `JsonRpseeError` * migrate chain module to proc macro api * make it compile with proc macros * update branch * update branch * update to jsonrpsee master * port system rpc * port state rpc * port childstate & offchain * frame system rpc * frame transaction payment * bring back CORS hack to work with polkadot UI * port babe rpc * port manual seal rpc * port frame mmr rpc * port frame contracts rpc * port finality grandpa rpc * port sync state rpc * resolve a few TODO + no jsonrpc deps * Update bin/node/rpc-client/src/main.rs * Update bin/node/rpc-client/src/main.rs * Update bin/node/rpc-client/src/main.rs * Update bin/node/rpc-client/src/main.rs * Port over system_ rpc tests * Make it compile * Use prost 0.8 * Use prost 0.8 * Make it compile * Ignore more failing tests * Comment out WIP tests * fix nit in frame system api * Update lockfile * No more juggling tokio versions * No more wait_for_stop ? * Remove browser-testing * Arguments must be arrays * Use same argument names * Resolve todo: no wait_for_stop for WS server Add todo: is parse_rpc_result used? Cleanup imports * fmt * log * One test passes * update jsonrpsee * update jsonrpsee * cleanup rpc-servers crate * jsonrpsee: add host and origin filtering (#9787) * add access control in the jsonrpsee servers * use master * fix nits * rpc runtime_version safe * fix nits * fix grumbles * remove unused files * resolve some todos * jsonrpsee more cleanup (#9803) * more cleanup * resolve TODOs * fix some unwraps * remove type hints * update jsonrpsee * downgrade zeroize * pin jsonrpsee rev * remove unwrap nit * Comment out more tests that aren't ported * Comment out more tests * Fix tests after merge * Subscription test * Invalid nonce test * Pending exts * WIP removeExtrinsic test * Test remove_extrinsic * Make state test: should_return_storage work * Uncomment/fix the other non-subscription related state tests * test: author_insertKey * test: author_rotateKeys * Get rest of state tests passing * asyncify a little more * Add todo to note #msg change * Crashing test for has_session_keys * Fix error conversion to avoid stack overflows Port author_hasSessionKeys test fmt * test author_hasKey * Add two missing tests Add a check on the return type Add todos for James's concerns * RPC tests for state, author and system (#9859) * Fix test runner * Impl Default for SubscriptionTaskExecutor * Keep the minimul amount of code needed to compile tests * Re-instate `RpcSession` (for now) * cleanup * Port over RPC tests * Add tokio * No need to map CallError to CallError * Port over system_ rpc tests * Make it compile * Use prost 0.8 * Use prost 0.8 * Make it compile * Ignore more failing tests * Comment out WIP tests * Update lockfile * No more juggling tokio versions * No more wait_for_stop ? * Remove browser-testing * Arguments must be arrays * Use same argument names * Resolve todo: no wait_for_stop for WS server Add todo: is parse_rpc_result used? Cleanup imports * fmt * log * One test passes * Comment out more tests that aren't ported * Comment out more tests * Fix tests after merge * Subscription test * Invalid nonce test * Pending exts * WIP removeExtrinsic test * Test remove_extrinsic * Make state test: should_return_storage work * Uncomment/fix the other non-subscription related state tests * test: author_insertKey * test: author_rotateKeys * Get rest of state tests passing * asyncify a little more * Add todo to note #msg change * Crashing test for has_session_keys * Fix error conversion to avoid stack overflows Port author_hasSessionKeys test fmt * test author_hasKey * Add two missing tests Add a check on the return type Add todos for James's concerns * offchain rpc tests * Address todos * fmt Co-authored-by: James Wilson <james@jsdw.me> * fix drop in state test * update jsonrpsee * fix ignored system test * fix chain tests * remove some boiler plate * Port BEEFY RPC (#9883) * Merge master * Port beefy RPC (ty @niklas!) * trivial changes left over from merge * Remove unused code * Update jsonrpsee * fix build * make tests compile again * beefy update jsonrpsee * fix: respect rpc methods policy * update cargo.lock * update jsonrpsee * update jsonrpsee * downgrade error logs * update jsonrpsee * Fix typo * remove unused file * Better name * Port Babe RPC tests * Put docs back * Resolve todo * Port tests for System RPCs * Resolve todo * fix build * Updated jsonrpsee to current master * fix: port finality grandpa rpc tests * Move .into() outside of the match * more review grumbles * jsonrpsee: add `rpc handlers` back (#10245) * add back RpcHandlers * cargo fmt * fix docs * fix grumble: remove needless alloc * resolve TODO * fmt * Fix typo * grumble: Use constants based on BASE_ERROR * grumble: DRY whitelisted listening addresses grumble: s/JSONRPC/JSON-RPC/ * cleanup * grumbles: Making readers aware of the possibility of gaps * review grumbles * grumbles * remove notes from niklasad1 * Update `jsonrpsee` * fix: jsonrpsee features * jsonrpsee: fallback to random port in case the specified port failed (#10304) * jsonrpsee: fallback to random port * better comment * Update client/rpc-servers/src/lib.rs Co-authored-by: Maciej Hirsz <1096222+maciejhirsz@users.noreply.github.com> * Update client/rpc-servers/src/lib.rs Co-authored-by: Maciej Hirsz <1096222+maciejhirsz@users.noreply.github.com> * address grumbles * cargo fmt * addrs already slice Co-authored-by: Maciej Hirsz <1096222+maciejhirsz@users.noreply.github.com> * Update jsonrpsee to 092081a0a2b8904c6ebd2cd99e16c7bc13ffc3ae * lockfile * update jsonrpsee * fix warning * Don't fetch jsonrpsee from crates * make tests compile again * fix rpc tests * remove unused deps * update tokio * fix rpc tests again * fix: test runner `HttpServerBuilder::builder` fails unless it's called within tokio runtime * cargo fmt * grumbles: fix subscription aliases * make clippy happy * update remaining subscriptions alias * cleanup * cleanup * fix chain subscription: less boiler plate (#10285) * fix chain subscription: less boiler plate * fix bad merge * cargo fmt * Switch to jsonrpsee 0.5 * fix build * add missing features * fix nit: remove needless Box::pin * Integrate jsonrpsee metrics (#10395) * draft metrics impl * Use latest api * Add missing file * Http server metrics * cleanup * bump jsonrpsee * Remove `ServerMetrics` and use a single middleware for both connection counting (aka sessions) and call metrics. * fix build * remove needless Arc::clone * Update to jsonrpsee 0.6 * lolz * fix metrics * Revert "lolz" This reverts commit eed6c6a56e78d8e307b4950f4c52a1c3a2322ba1. * fix: in-memory rpc support subscriptions * commit Cargo.lock * Update tests to 0.7 * fix TODOs * ws server: generate subscriptionIDs as Strings Some libraries seems to expect the subscription IDs to be Strings, let's not break this in this PR. * Increase timeout * Port over tests * cleanup * Using error codes from the spec * fix clippy * cargo fmt * update jsonrpsee * fix nits * fix: rpc_query * enable custom subid gen through spawn_tasks * remove unsed deps * unify tokio deps * Revert "enable custom subid gen through spawn_tasks" This reverts commit 5c5eb70328fe39d154fdb55c56e637b4548cf470. * fix bad merge of `test-utils` * fix more nits * downgrade wasm-instrument to 0.1.0 * [jsonrpsee]: enable custom RPC subscription ID generatation (#10731) * enable custom subid gen through spawn_tasks * fix nits * Update client/service/src/builder.rs Co-authored-by: David <dvdplm@gmail.com> * add Poc; needs jsonrpsee pr * update jsonrpsee * add re-exports * add docs Co-authored-by: David <dvdplm@gmail.com> * cargo fmt * fmt * port RPC-API dev * Remove unused file * fix nit: remove async trait * fix doc links * fix merge nit: remove jsonrpc deps * kill namespace on rpc apis * companion for jsonrpsee v0.10 (#11158) * companion for jsonrpsee v0.10 * update versions v0.10.0 * add some fixes * spelling * fix spaces Co-authored-by: Niklas Adolfsson <niklasadolfsson1@gmail.com> * send error before subs are closed * fix unsubscribe method names: chain * fix tests * jsonrpc server: print binded local address * grumbles: kill SubscriptionTaskExecutor * Update client/sync-state-rpc/src/lib.rs Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com> * Update client/rpc/src/chain/chain_full.rs Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com> * Update client/rpc/src/chain/chain_full.rs Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com> * sync-state-rpc: kill anyhow * no more anyhow * remove todo * jsonrpsee: fix bad params in subscriptions. (#11251) * update jsonrpsee * fix error responses * revert error codes * dont do weird stuff in drop impl * rpc servers: remove needless clone * Remove silly constants * chore: update jsonrpsee v0.12 * commit Cargo.lock * deps: downgrade git2 * feat: CLI flag max subscriptions per connection * metrics: use old logging format * fix: read WS address from substrate output (#11379) Co-authored-by: Niklas Adolfsson <niklasadolfsson1@gmail.com> Co-authored-by: James Wilson <james@jsdw.me> Co-authored-by: Maciej Hirsz <hello@maciej.codes> Co-authored-by: Maciej Hirsz <1096222+maciejhirsz@users.noreply.github.com> Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>
2026-06-12 20:31:13 +00:00 · 2022-05-10 10:52:19 +02:00
parent e45d53552d
commit 29c0c6a4a8
93 changed files with 3813 additions and 5094 deletions
@@ -20,213 +20,193 @@

 #![warn(missing_docs)]

-mod middleware;
+use jsonrpsee::{
+	http_server::{AccessControlBuilder, HttpServerBuilder, HttpServerHandle},
+	ws_server::{WsServerBuilder, WsServerHandle},
+	RpcModule,
+};
+use std::{error::Error as StdError, net::SocketAddr};

-use jsonrpc_core::{IoHandlerExtension, MetaIoHandler};
-use log::error;
-use prometheus_endpoint::{register, Counter, PrometheusError, Registry, U64};
-use pubsub::PubSubMetadata;
-use std::io;
+pub use crate::middleware::{RpcMetrics, RpcMiddleware};
+pub use jsonrpsee::core::{
+	id_providers::{RandomIntegerIdProvider, RandomStringIdProvider},
+	traits::IdProvider,
+};

 const MEGABYTE: usize = 1024 * 1024;

 /// Maximal payload accepted by RPC servers.
 pub const RPC_MAX_PAYLOAD_DEFAULT: usize = 15 * MEGABYTE;

-/// Maximal buffer size in WS server.
-pub const WS_MAX_BUFFER_CAPACITY_DEFAULT: usize = 16 * MEGABYTE;
-
 /// Default maximum number of connections for WS RPC servers.
 const WS_MAX_CONNECTIONS: usize = 100;

-/// The RPC IoHandler containing all requested APIs.
-pub type RpcHandler<T> = pubsub::PubSubHandler<T, RpcMiddleware>;
+/// Default maximum number subscriptions per connection for WS RPC servers.
+const WS_MAX_SUBS_PER_CONN: usize = 1024;

-pub use middleware::{method_names, RpcMetrics, RpcMiddleware};
+pub mod middleware;

-/// Construct rpc `IoHandler`
-pub fn rpc_handler<M: PubSubMetadata>(
-	extension: impl IoHandlerExtension<M>,
-	rpc_middleware: RpcMiddleware,
-) -> RpcHandler<M> {
-	let io_handler = MetaIoHandler::with_middleware(rpc_middleware);
-	let mut io = pubsub::PubSubHandler::new(io_handler);
-	extension.augment(&mut io);
-
-	// add an endpoint to list all available methods.
-	let mut methods = io.iter().map(|x| x.0.clone()).collect::<Vec<String>>();
-	io.add_method("rpc_methods", {
-		methods.sort();
-		let methods = serde_json::to_value(&methods)
-			.expect("Serialization of Vec<String> is infallible; qed");
-
-		move |_| {
-			let methods = methods.clone();
-			async move {
-				Ok(serde_json::json!({
-					"version": 1,
-					"methods": methods,
-				}))
-			}
-		}
-	});
-	io
-}
-
-/// 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(
-							"substrate_rpc_sessions_opened",
-							"Number of persistent RPC sessions opened",
-						)?,
-						r,
-					)?
-					.into(),
-					session_closed: register(
-						Counter::new(
-							"substrate_rpc_sessions_closed",
-							"Number of persistent RPC sessions closed",
-						)?,
-						r,
-					)?
-					.into(),
-				})
-			})
-			.unwrap_or_else(|| Ok(Default::default()))
-	}
-}
-
-/// Type alias for ipc server
-pub type IpcServer = ipc::Server;
 /// Type alias for http server
-pub type HttpServer = http::Server;
+pub type HttpServer = HttpServerHandle;
 /// Type alias for ws server
-pub type WsServer = ws::Server;
+pub type WsServer = WsServerHandle;

-impl ws::SessionStats for ServerMetrics {
-	fn open_session(&self, _id: ws::SessionId) {
-		self.session_opened.as_ref().map(|m| m.inc());
-	}
+/// WebSocket specific settings on the server.
+pub struct WsConfig {
+	/// Maximum connections.
+	pub max_connections: Option<usize>,
+	/// Maximum subscriptions per connection.
+	pub max_subs_per_conn: Option<usize>,
+	/// Maximum rpc request payload size.
+	pub max_payload_in_mb: Option<usize>,
+	/// Maximum rpc response payload size.
+	pub max_payload_out_mb: Option<usize>,
+}

-	fn close_session(&self, _id: ws::SessionId) {
-		self.session_closed.as_ref().map(|m| m.inc());
+impl WsConfig {
+	// Deconstructs the config to get the finalized inner values.
+	//
+	// `Payload size` or `max subs per connection` bigger than u32::MAX will be truncated.
+	fn deconstruct(self) -> (u32, u32, u64, u32) {
+		let max_conns = self.max_connections.unwrap_or(WS_MAX_CONNECTIONS) as u64;
+		let max_payload_in_mb = payload_size_or_default(self.max_payload_in_mb) as u32;
+		let max_payload_out_mb = payload_size_or_default(self.max_payload_out_mb) as u32;
+		let max_subs_per_conn = self.max_subs_per_conn.unwrap_or(WS_MAX_SUBS_PER_CONN) as u32;
+
+		(max_payload_in_mb, max_payload_out_mb, max_conns, max_subs_per_conn)
 	}
 }

 /// Start HTTP server listening on given address.
-pub fn start_http<M: pubsub::PubSubMetadata + Default + Unpin>(
-	addr: &std::net::SocketAddr,
+pub async fn start_http<M: Send + Sync + 'static>(
+	addrs: [SocketAddr; 2],
 	cors: Option<&Vec<String>>,
-	io: RpcHandler<M>,
-	maybe_max_payload_mb: Option<usize>,
-	tokio_handle: tokio::runtime::Handle,
-) -> io::Result<http::Server> {
-	let max_request_body_size = maybe_max_payload_mb
-		.map(|mb| mb.saturating_mul(MEGABYTE))
-		.unwrap_or(RPC_MAX_PAYLOAD_DEFAULT);
+	max_payload_in_mb: Option<usize>,
+	max_payload_out_mb: Option<usize>,
+	metrics: Option<RpcMetrics>,
+	rpc_api: RpcModule<M>,
+	rt: tokio::runtime::Handle,
+) -> Result<HttpServerHandle, Box<dyn StdError + Send + Sync>> {
+	let max_payload_in = payload_size_or_default(max_payload_in_mb);
+	let max_payload_out = payload_size_or_default(max_payload_out_mb);

-	http::ServerBuilder::new(io)
-		.threads(1)
-		.event_loop_executor(tokio_handle)
-		.health_api(("/health", "system_health"))
-		.allowed_hosts(hosts_filtering(cors.is_some()))
-		.rest_api(if cors.is_some() { http::RestApi::Secure } else { http::RestApi::Unsecure })
-		.cors(map_cors::<http::AccessControlAllowOrigin>(cors))
-		.max_request_body_size(max_request_body_size)
-		.start_http(addr)
-}
+	let mut acl = AccessControlBuilder::new();

-/// Start IPC server listening on given path.
-pub fn start_ipc<M: pubsub::PubSubMetadata + Default>(
-	addr: &str,
-	io: RpcHandler<M>,
-	server_metrics: ServerMetrics,
-) -> io::Result<ipc::Server> {
-	let builder = ipc::ServerBuilder::new(io);
-	#[cfg(target_os = "unix")]
-	builder.set_security_attributes({
-		let security_attributes = ipc::SecurityAttributes::empty();
-		security_attributes.set_mode(0o600)?;
-		security_attributes
-	});
-	builder.session_stats(server_metrics).start(addr)
+	if let Some(cors) = cors {
+		// Whitelist listening address.
+		// NOTE: set_allowed_hosts will whitelist both ports but only one will used.
+		acl = acl.set_allowed_hosts(format_allowed_hosts(&addrs[..]))?;
+		acl = acl.set_allowed_origins(cors)?;
+	};
+
+	let builder = HttpServerBuilder::new()
+		.max_request_body_size(max_payload_in as u32)
+		.max_response_body_size(max_payload_out as u32)
+		.set_access_control(acl.build())
+		.custom_tokio_runtime(rt);
+
+	let rpc_api = build_rpc_api(rpc_api);
+	let (handle, addr) = if let Some(metrics) = metrics {
+		let middleware = RpcMiddleware::new(metrics, "http".into());
+		let builder = builder.set_middleware(middleware);
+		let server = builder.build(&addrs[..]).await?;
+		let addr = server.local_addr();
+		(server.start(rpc_api)?, addr)
+	} else {
+		let server = builder.build(&addrs[..]).await?;
+		let addr = server.local_addr();
+		(server.start(rpc_api)?, addr)
+	};
+
+	log::info!(
+		"Running JSON-RPC HTTP server: addr={}, allowed origins={:?}",
+		addr.map_or_else(|_| "unknown".to_string(), |a| a.to_string()),
+		cors
+	);
+
+	Ok(handle)
 }

 /// Start WS server listening on given address.
-pub fn start_ws<
-	M: pubsub::PubSubMetadata + From<futures::channel::mpsc::UnboundedSender<String>>,
->(
-	addr: &std::net::SocketAddr,
-	max_connections: Option<usize>,
+pub async fn start_ws<M: Send + Sync + 'static>(
+	addrs: [SocketAddr; 2],
 	cors: Option<&Vec<String>>,
-	io: RpcHandler<M>,
-	maybe_max_payload_mb: Option<usize>,
-	maybe_max_out_buffer_capacity_mb: Option<usize>,
-	server_metrics: ServerMetrics,
-	tokio_handle: tokio::runtime::Handle,
-) -> io::Result<ws::Server> {
-	let max_payload = maybe_max_payload_mb
-		.map(|mb| mb.saturating_mul(MEGABYTE))
-		.unwrap_or(RPC_MAX_PAYLOAD_DEFAULT);
-	let max_out_buffer_capacity = maybe_max_out_buffer_capacity_mb
-		.map(|mb| mb.saturating_mul(MEGABYTE))
-		.unwrap_or(WS_MAX_BUFFER_CAPACITY_DEFAULT);
+	ws_config: WsConfig,
+	metrics: Option<RpcMetrics>,
+	rpc_api: RpcModule<M>,
+	rt: tokio::runtime::Handle,
+	id_provider: Option<Box<dyn IdProvider>>,
+) -> Result<WsServerHandle, Box<dyn StdError + Send + Sync>> {
+	let (max_payload_in, max_payload_out, max_connections, max_subs_per_conn) =
+		ws_config.deconstruct();

-	if max_payload > max_out_buffer_capacity {
-		log::warn!(
-			"maximum payload ({}) is more than maximum output buffer ({}) size in ws server, the payload will actually be limited by the buffer size",
-			max_payload,
-			max_out_buffer_capacity,
-		)
-	}
+	let mut builder = WsServerBuilder::new()
+		.max_request_body_size(max_payload_in)
+		.max_response_body_size(max_payload_out)
+		.max_connections(max_connections)
+		.max_subscriptions_per_connection(max_subs_per_conn)
+		.custom_tokio_runtime(rt);

-	ws::ServerBuilder::with_meta_extractor(io, |context: &ws::RequestContext| {
-		context.sender().into()
-	})
-	.event_loop_executor(tokio_handle)
-	.max_payload(max_payload)
-	.max_connections(max_connections.unwrap_or(WS_MAX_CONNECTIONS))
-	.max_out_buffer_capacity(max_out_buffer_capacity)
-	.allowed_origins(map_cors(cors))
-	.allowed_hosts(hosts_filtering(cors.is_some()))
-	.session_stats(server_metrics)
-	.start(addr)
-	.map_err(|err| match err {
-		ws::Error::Io(io) => io,
-		ws::Error::ConnectionClosed => io::ErrorKind::BrokenPipe.into(),
-		e => {
-			error!("{}", e);
-			io::ErrorKind::Other.into()
-		},
-	})
-}
-
-fn map_cors<T: for<'a> From<&'a str>>(cors: Option<&Vec<String>>) -> http::DomainsValidation<T> {
-	cors.map(|x| x.iter().map(AsRef::as_ref).map(Into::into).collect::<Vec<_>>())
-		.into()
-}
-
-fn hosts_filtering(enable: bool) -> http::DomainsValidation<http::Host> {
-	if enable {
-		// NOTE The listening address is whitelisted by default.
-		// Setting an empty vector here enables the validation
-		// and allows only the listening address.
-		http::DomainsValidation::AllowOnly(vec![])
+	if let Some(provider) = id_provider {
+		builder = builder.set_id_provider(provider);
 	} else {
-		http::DomainsValidation::Disabled
+		builder = builder.set_id_provider(RandomStringIdProvider::new(16));
+	};
+
+	if let Some(cors) = cors {
+		// Whitelist listening address.
+		// NOTE: set_allowed_hosts will whitelist both ports but only one will used.
+		builder = builder.set_allowed_hosts(format_allowed_hosts(&addrs[..]))?;
+		builder = builder.set_allowed_origins(cors)?;
 	}
+
+	let rpc_api = build_rpc_api(rpc_api);
+	let (handle, addr) = if let Some(metrics) = metrics {
+		let middleware = RpcMiddleware::new(metrics, "ws".into());
+		let builder = builder.set_middleware(middleware);
+		let server = builder.build(&addrs[..]).await?;
+		let addr = server.local_addr();
+		(server.start(rpc_api)?, addr)
+	} else {
+		let server = builder.build(&addrs[..]).await?;
+		let addr = server.local_addr();
+		(server.start(rpc_api)?, addr)
+	};
+
+	log::info!(
+		"Running JSON-RPC WS server: addr={}, allowed origins={:?}",
+		addr.map_or_else(|_| "unknown".to_string(), |a| a.to_string()),
+		cors
+	);
+
+	Ok(handle)
+}
+
+fn format_allowed_hosts(addrs: &[SocketAddr]) -> Vec<String> {
+	let mut hosts = Vec::with_capacity(addrs.len() * 2);
+	for addr in addrs {
+		hosts.push(format!("localhost:{}", addr.port()));
+		hosts.push(format!("127.0.0.1:{}", addr.port()));
+	}
+	hosts
+}
+
+fn build_rpc_api<M: Send + Sync + 'static>(mut rpc_api: RpcModule<M>) -> RpcModule<M> {
+	let mut available_methods = rpc_api.method_names().collect::<Vec<_>>();
+	available_methods.sort_unstable();
+
+	rpc_api
+		.register_method("rpc_methods", move |_, _| {
+			Ok(serde_json::json!({
+				"version": 1,
+				"methods": available_methods,
+			}))
+		})
+		.expect("infallible all other methods have their own address space; qed");
+
+	rpc_api
+}
+
+fn payload_size_or_default(size_mb: Option<usize>) -> usize {
+	size_mb.map_or(RPC_MAX_PAYLOAD_DEFAULT, |mb| mb.saturating_mul(MEGABYTE))
 }
@@ -16,34 +16,38 @@
 // You should have received a copy of the GNU General Public License
 // along with this program. If not, see <https://www.gnu.org/licenses/>.

-//! Middleware for RPC requests.
+//! RPC middlware to collect prometheus metrics on RPC calls.

-use std::collections::HashSet;
-
-use jsonrpc_core::{FutureOutput, FutureResponse, Metadata, Middleware as RequestMiddleware};
+use jsonrpsee::core::middleware::Middleware;
 use prometheus_endpoint::{
-	register, CounterVec, HistogramOpts, HistogramVec, Opts, PrometheusError, Registry, U64,
+	register, Counter, CounterVec, HistogramOpts, HistogramVec, Opts, PrometheusError, Registry,
+	U64,
 };

-use futures::{future::Either, Future, FutureExt};
-use pubsub::PubSubMetadata;
-
-use crate::RpcHandler;
-
-/// Metrics for RPC middleware
+/// 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 {
 	/// Create an instance of metrics
 	pub fn new(metrics_registry: Option<&Registry>) -> Result<Option<Self>, PrometheusError> {
-		if let Some(r) = metrics_registry {
+		if let Some(metrics_registry) = metrics_registry {
 			Ok(Some(Self {
 				requests_started: register(
 					CounterVec::new(
@@ -53,7 +57,7 @@ impl RpcMetrics {
 						),
 						&["protocol"],
 					)?,
-					r,
+					metrics_registry,
 				)?,
 				requests_finished: register(
 					CounterVec::new(
@@ -63,7 +67,7 @@ impl RpcMetrics {
 						),
 						&["protocol"],
 					)?,
-					r,
+					metrics_registry,
 				)?,
 				calls_time: register(
 					HistogramVec::new(
@@ -73,7 +77,7 @@ impl RpcMetrics {
 						),
 						&["protocol", "method"],
 					)?,
-					r,
+					metrics_registry,
 				)?,
 				calls_started: register(
 					CounterVec::new(
@@ -83,7 +87,7 @@ impl RpcMetrics {
 						),
 						&["protocol", "method"],
 					)?,
-					r,
+					metrics_registry,
 				)?,
 				calls_finished: register(
 					CounterVec::new(
@@ -93,8 +97,24 @@ impl RpcMetrics {
 						),
 						&["protocol", "method", "is_error"],
 					)?,
-					r,
+					metrics_registry,
 				)?,
+				ws_sessions_opened: register(
+					Counter::new(
+						"substrate_rpc_sessions_opened",
+						"Number of persistent RPC sessions opened",
+					)?,
+					metrics_registry,
+				)?
+				.into(),
+				ws_sessions_closed: register(
+					Counter::new(
+						"substrate_rpc_sessions_closed",
+						"Number of persistent RPC sessions closed",
+					)?,
+					metrics_registry,
+				)?
+				.into(),
 			}))
 		} else {
 			Ok(None)
@@ -102,140 +122,71 @@ impl RpcMetrics {
 	}
 }

-/// Instantiates a dummy `IoHandler` given a builder function to extract supported method names.
-pub fn method_names<F, M, E>(gen_handler: F) -> Result<HashSet<String>, E>
-where
-	F: FnOnce(RpcMiddleware) -> Result<RpcHandler<M>, E>,
-	M: PubSubMetadata,
-{
-	let io = gen_handler(RpcMiddleware::new(None, HashSet::new(), "dummy"))?;
-	Ok(io.iter().map(|x| x.0.clone()).collect())
-}
-
+#[derive(Clone)]
 /// Middleware for RPC calls
 pub struct RpcMiddleware {
-	metrics: Option<RpcMetrics>,
-	known_rpc_method_names: HashSet<String>,
-	transport_label: String,
+	metrics: RpcMetrics,
+	transport_label: &'static str,
 }

 impl RpcMiddleware {
-	/// Create an instance of middleware.
-	///
-	/// - `metrics`: Will be used to report statistics.
-	/// - `transport_label`: The label that is used when reporting the statistics.
-	pub fn new(
-		metrics: Option<RpcMetrics>,
-		known_rpc_method_names: HashSet<String>,
-		transport_label: &str,
-	) -> Self {
-		RpcMiddleware { metrics, known_rpc_method_names, transport_label: transport_label.into() }
+	/// Create a new [`RpcMiddleware`] with the provided [`RpcMetrics`].
+	pub fn new(metrics: RpcMetrics, transport_label: &'static str) -> Self {
+		Self { metrics, transport_label }
 	}
 }

-impl<M: Metadata> RequestMiddleware<M> for RpcMiddleware {
-	type Future = FutureResponse;
-	type CallFuture = FutureOutput;
+impl Middleware for RpcMiddleware {
+	type Instant = std::time::Instant;

-	fn on_request<F, X>(
-		&self,
-		request: jsonrpc_core::Request,
-		meta: M,
-		next: F,
-	) -> Either<Self::Future, X>
-	where
-		F: Fn(jsonrpc_core::Request, M) -> X + Send + Sync,
-		X: Future<Output = Option<jsonrpc_core::Response>> + Send + 'static,
-	{
-		let metrics = self.metrics.clone();
-		let transport_label = self.transport_label.clone();
-		if let Some(ref metrics) = metrics {
-			metrics.requests_started.with_label_values(&[transport_label.as_str()]).inc();
-		}
-		let r = next(request, meta);
-		Either::Left(
-			async move {
-				let r = r.await;
-				if let Some(ref metrics) = metrics {
-					metrics.requests_finished.with_label_values(&[transport_label.as_str()]).inc();
-				}
-				r
-			}
-			.boxed(),
-		)
+	fn on_connect(&self) {
+		self.metrics.ws_sessions_opened.as_ref().map(|counter| counter.inc());
 	}

-	fn on_call<F, X>(
-		&self,
-		call: jsonrpc_core::Call,
-		meta: M,
-		next: F,
-	) -> Either<Self::CallFuture, X>
-	where
-		F: Fn(jsonrpc_core::Call, M) -> X + Send + Sync,
-		X: Future<Output = Option<jsonrpc_core::Output>> + Send + 'static,
-	{
-		let start = std::time::Instant::now();
-		let name = call_name(&call, &self.known_rpc_method_names).to_owned();
-		let metrics = self.metrics.clone();
-		let transport_label = self.transport_label.clone();
-		log::trace!(target: "rpc_metrics", "[{}] {} call: {:?}", transport_label, name, &call);
-		if let Some(ref metrics) = metrics {
-			metrics
-				.calls_started
-				.with_label_values(&[transport_label.as_str(), name.as_str()])
-				.inc();
-		}
-		let r = next(call, meta);
-		Either::Left(
-			async move {
-				let r = r.await;
-				let micros = start.elapsed().as_micros();
-				if let Some(ref metrics) = metrics {
-					metrics
-						.calls_time
-						.with_label_values(&[transport_label.as_str(), name.as_str()])
-						.observe(micros as _);
-					metrics
-						.calls_finished
-						.with_label_values(&[
-							transport_label.as_str(),
-							name.as_str(),
-							if is_success(&r) { "true" } else { "false" },
-						])
-						.inc();
-				}
-				log::debug!(
-					target: "rpc_metrics",
-					"[{}] {} call took {} μs",
-					transport_label,
-					name,
-					micros,
-				);
-				r
-			}
-			.boxed(),
-		)
+	fn on_request(&self) -> Self::Instant {
+		let now = std::time::Instant::now();
+		self.metrics.requests_started.with_label_values(&[self.transport_label]).inc();
+		now
+	}
+
+	fn on_call(&self, name: &str) {
+		log::trace!(target: "rpc_metrics", "[{}] on_call name={}", self.transport_label, name);
+		self.metrics
+			.calls_started
+			.with_label_values(&[self.transport_label, name])
+			.inc();
+	}
+
+	fn on_result(&self, name: &str, success: bool, started_at: Self::Instant) {
+		let micros = started_at.elapsed().as_micros();
+		log::debug!(
+			target: "rpc_metrics",
+			"[{}] {} call took {} μs",
+			self.transport_label,
+			name,
+			micros,
+		);
+		self.metrics
+			.calls_time
+			.with_label_values(&[self.transport_label, name])
+			.observe(micros as _);
+
+		self.metrics
+			.calls_finished
+			.with_label_values(&[
+				self.transport_label,
+				name,
+				if success { "true" } else { "false" },
+			])
+			.inc();
+	}
+
+	fn on_response(&self, started_at: Self::Instant) {
+		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());
 	}
 }
-
-fn call_name<'a>(call: &'a jsonrpc_core::Call, known_methods: &HashSet<String>) -> &'a str {
-	// To prevent bloating metric with all invalid method names we filter them out here.
-	let only_known = |method: &'a String| {
-		if known_methods.contains(method) {
-			method.as_str()
-		} else {
-			"invalid method"
-		}
-	};
-
-	match call {
-		jsonrpc_core::Call::Invalid { .. } => "invalid call",
-		jsonrpc_core::Call::MethodCall(ref call) => only_known(&call.method),
-		jsonrpc_core::Call::Notification(ref notification) => only_known(&notification.method),
-	}
-}
-
-fn is_success(output: &Option<jsonrpc_core::Output>) -> bool {
-	matches!(output, Some(jsonrpc_core::Output::Success(..)))
-}