Register metrics for the notification handles (#2562)

Add metrics for notification handles so substream events are correctly reported to Prometheus
paritytech · Nov 30, 2023 · 180a6ad · 180a6ad
1 parent e41a0e7
commit 180a6ad
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 9 deletions.
diff --git a/substrate/client/network/src/protocol.rs b/substrate/client/network/src/protocol.rs
@@ -43,7 +43,7 @@ use sp_runtime::traits::Block as BlockT;
 
 use std::{collections::HashSet, iter, task::Poll};
 
-use notifications::{Notifications, NotificationsOut};
+use notifications::{metrics, Notifications, NotificationsOut};
 
 pub(crate) use notifications::ProtocolHandle;
 
@@ -95,7 +95,7 @@ impl<B: BlockT> Protocol<B> {
 			// NOTE: Block announcement protocol is still very much hardcoded into
 			// `Protocol`. 	This protocol must be the first notification protocol given to
 			// `Notifications`
-			let (protocol_configs, handles): (Vec<_>, Vec<_>) = iter::once({
+			let (protocol_configs, mut handles): (Vec<_>, Vec<_>) = iter::once({
 				let config = notifications::ProtocolConfig {
 					name: block_announces_protocol.protocol_name().clone(),
 					fallback_names: block_announces_protocol.fallback_names().cloned().collect(),
@@ -122,11 +122,16 @@ impl<B: BlockT> Protocol<B> {
 			}))
 			.unzip();
 
+			let metrics = registry.as_ref().and_then(|registry| metrics::register(&registry).ok());
+			handles.iter_mut().for_each(|handle| {
+				handle.set_metrics(metrics.clone());
+			});
+
 			(
 				Notifications::new(
 					protocol_controller_handles,
 					from_protocol_controllers,
-					registry,
+					metrics,
 					protocol_configs.into_iter(),
 				),
 				installed_protocols,

diff --git a/substrate/client/network/src/protocol/notifications.rs b/substrate/client/network/src/protocol/notifications.rs
@@ -25,7 +25,7 @@ pub use self::{
 	service::{notification_service, ProtocolHandlePair},
 };
 
-pub(crate) use self::service::ProtocolHandle;
+pub(crate) use self::service::{metrics, ProtocolHandle};
 
 mod behaviour;
 mod handler;

diff --git a/substrate/client/network/src/protocol/notifications/behaviour.rs b/substrate/client/network/src/protocol/notifications/behaviour.rs
@@ -40,7 +40,6 @@ use libp2p::{
 };
 use log::{debug, error, trace, warn};
 use parking_lot::RwLock;
-use prometheus_endpoint::Registry;
 use rand::distributions::{Distribution as _, Uniform};
 use sc_utils::mpsc::TracingUnboundedReceiver;
 use smallvec::SmallVec;
@@ -405,7 +404,7 @@ impl Notifications {
 	pub(crate) fn new(
 		protocol_controller_handles: Vec<protocol_controller::ProtocolHandle>,
 		from_protocol_controllers: TracingUnboundedReceiver<Message>,
-		registry: &Option<Registry>,
+		metrics: Option<metrics::Metrics>,
 		notif_protocols: impl Iterator<
 			Item = (
 				ProtocolConfig,
@@ -429,7 +428,6 @@ impl Notifications {
 			.unzip();
 		assert!(!notif_protocols.is_empty());
 
-		let metrics = registry.as_ref().and_then(|registry| metrics::register(&registry).ok());
 		let (mut protocol_handles, command_streams): (Vec<_>, Vec<_>) = protocol_handles
 			.into_iter()
 			.enumerate()
@@ -2452,7 +2450,7 @@ mod tests {
 			Notifications::new(
 				vec![handle],
 				from_controller,
-				&None,
+				None,
 				iter::once((
 					ProtocolConfig {
 						name: "/foo".into(),

diff --git a/substrate/client/network/src/protocol/notifications/tests.rs b/substrate/client/network/src/protocol/notifications/tests.rs
@@ -99,7 +99,7 @@ fn build_nodes() -> (Swarm<CustomProtoWithAddr>, Swarm<CustomProtoWithAddr>) {
 			inner: Notifications::new(
 				vec![controller_handle],
 				from_controller,
-				&None,
+				None,
 				iter::once((
 					ProtocolConfig {
 						name: "/foo".into(),