metrics: Expose litep2p metrics in an agnostic manner

src/config.rs

@@ -23,6 +23,7 @@
 use crate::{
     crypto::ed25519::Keypair,
     executor::{DefaultExecutor, Executor},
+    metrics::MetricsRegistry,
     protocol::{
         libp2p::{bitswap, identify, kademlia, ping},
         mdns::Config as MdnsConfig,
@@ -83,6 +84,9 @@ pub struct ConfigBuilder {
     #[cfg(feature = "websocket")]
     websocket: Option<WebSocketConfig>,
 
+    /// Metrics registry
+    metrics_registry: Option<MetricsRegistry>,
+
     /// Keypair.
     keypair: Option<Keypair>,
 
@@ -143,6 +147,7 @@ impl ConfigBuilder {
             webrtc: None,
             #[cfg(feature = "websocket")]
             websocket: None,
+            metrics_registry: None,
             keypair: None,
             ping: None,
             identify: None,
@@ -187,6 +192,12 @@ impl ConfigBuilder {
         self
     }
 
+    /// Add metrics registry.
+    pub fn with_metrics_registry(mut self, registry: MetricsRegistry) -> Self {
+        self.metrics_registry = Some(registry);
+        self
+    }
+
     /// Add keypair.
     ///
     /// If no keypair is specified, litep2p creates a new keypair.
@@ -295,6 +306,7 @@ impl ConfigBuilder {
             webrtc: self.webrtc.take(),
             #[cfg(feature = "websocket")]
             websocket: self.websocket.take(),
+            metrics_registry: self.metrics_registry.take(),
             ping: self.ping.take(),
             identify: self.identify.take(),
             kademlia: self.kademlia.take(),
@@ -328,6 +340,9 @@ pub struct Litep2pConfig {
     #[cfg(feature = "websocket")]
     pub(crate) websocket: Option<WebSocketConfig>,
 
+    /// Prometheus metrics registry.
+    pub(crate) metrics_registry: Option<MetricsRegistry>,
+
     /// Keypair.
     pub(crate) keypair: Keypair,
 

src/error.rs

@@ -127,6 +127,8 @@ pub enum Error {
     ConnectionLimit(ConnectionLimitsError),
     #[error("Failed to dial peer immediately")]
     ImmediateDialError(#[from] ImmediateDialError),
+    #[error("Invalid metric: `{0}`")]
+    MetricError(String),
 }
 
 /// Error type for address parsing.

src/lib.rs

@@ -72,6 +72,7 @@ pub mod config;
 pub mod crypto;
 pub mod error;
 pub mod executor;
+pub mod metrics;
 pub mod protocol;
 pub mod substream;
 pub mod transport;
@@ -164,7 +165,8 @@ impl Litep2p {
             bandwidth_sink.clone(),
             litep2p_config.max_parallel_dials,
             litep2p_config.connection_limits,
-        );
+            litep2p_config.metrics_registry.clone(),
+        )?;
 
         // add known addresses to `TransportManager`, if any exist
         if !litep2p_config.known_addresses.is_empty() {
@@ -188,9 +190,14 @@ impl Litep2p {
                 litep2p_config.keep_alive_timeout,
             );
             let executor = Arc::clone(&litep2p_config.executor);
-            litep2p_config.executor.run(Box::pin(async move {
-                NotificationProtocol::new(service, config, executor).run().await
-            }));
+            let notification = NotificationProtocol::new(
+                service,
+                config,
+                executor,
+                litep2p_config.metrics_registry.clone(),
+            )?;
+
+            litep2p_config.executor.run(Box::pin(async move { notification.run().await }));
         }
 
         // start request-response protocol event loops
@@ -207,9 +214,15 @@ impl Litep2p {
                 config.codec,
                 litep2p_config.keep_alive_timeout,
             );
-            litep2p_config.executor.run(Box::pin(async move {
-                RequestResponseProtocol::new(service, config).run().await
-            }));
+            let request_response = RequestResponseProtocol::new(
+                service,
+                config,
+                litep2p_config.metrics_registry.clone(),
+            )?;
+
+            litep2p_config
+                .executor
+                .run(Box::pin(async move { request_response.run().await }));
         }
 
         // start user protocol event loops
@@ -241,9 +254,13 @@ impl Litep2p {
                 ping_config.codec,
                 litep2p_config.keep_alive_timeout,
             );
-            litep2p_config.executor.run(Box::pin(async move {
-                Ping::new(service, ping_config).run().await
-            }));
+            let ping = Ping::new(
+                service,
+                ping_config,
+                litep2p_config.metrics_registry.clone(),
+            )?;
+
+            litep2p_config.executor.run(Box::pin(async move { ping.run().await }));
         }
 
         // start kademlia protocol event loop if enabled
@@ -264,8 +281,14 @@ impl Litep2p {
                 kademlia_config.codec,
                 litep2p_config.keep_alive_timeout,
             );
+            let kad = Kademlia::new(
+                service,
+                kademlia_config,
+                litep2p_config.metrics_registry.clone(),
+            )?;
+
             litep2p_config.executor.run(Box::pin(async move {
-                let _ = Kademlia::new(service, kademlia_config).run().await;
+                let _ = kad.run().await;
             }));
         }
 
@@ -313,8 +336,11 @@ impl Litep2p {
         // enable tcp transport if the config exists
         if let Some(config) = litep2p_config.tcp.take() {
             let handle = transport_manager.transport_handle(Arc::clone(&litep2p_config.executor));
-            let (transport, transport_listen_addresses) =
-                <TcpTransport as TransportBuilder>::new(handle, config)?;
+            let (transport, transport_listen_addresses) = <TcpTransport as TransportBuilder>::new(
+                handle,
+                config,
+                litep2p_config.metrics_registry.clone(),
+            )?;
 
             for address in transport_listen_addresses {
                 transport_manager.register_listen_address(address.clone());
@@ -330,8 +356,11 @@ impl Litep2p {
         #[cfg(feature = "quic")]
         if let Some(config) = litep2p_config.quic.take() {
             let handle = transport_manager.transport_handle(Arc::clone(&litep2p_config.executor));
-            let (transport, transport_listen_addresses) =
-                <QuicTransport as TransportBuilder>::new(handle, config)?;
+            let (transport, transport_listen_addresses) = <QuicTransport as TransportBuilder>::new(
+                handle,
+                config,
+                litep2p_config.metrics_registry.clone(),
+            )?;
 
             for address in transport_listen_addresses {
                 transport_manager.register_listen_address(address.clone());
@@ -348,7 +377,11 @@ impl Litep2p {
         if let Some(config) = litep2p_config.webrtc.take() {
             let handle = transport_manager.transport_handle(Arc::clone(&litep2p_config.executor));
             let (transport, transport_listen_addresses) =
-                <WebRtcTransport as TransportBuilder>::new(handle, config)?;
+                <WebRtcTransport as TransportBuilder>::new(
+                    handle,
+                    config,
+                    litep2p_config.metrics_registry.clone(),
+                )?;
 
             for address in transport_listen_addresses {
                 transport_manager.register_listen_address(address.clone());
@@ -365,7 +398,11 @@ impl Litep2p {
         if let Some(config) = litep2p_config.websocket.take() {
             let handle = transport_manager.transport_handle(Arc::clone(&litep2p_config.executor));
             let (transport, transport_listen_addresses) =
-                <WebSocketTransport as TransportBuilder>::new(handle, config)?;
+                <WebSocketTransport as TransportBuilder>::new(
+                    handle,
+                    config,
+                    litep2p_config.metrics_registry.clone(),
+                )?;
 
             for address in transport_listen_addresses {
                 transport_manager.register_listen_address(address.clone());
@@ -390,7 +427,11 @@ impl Litep2p {
         // if identify was enabled, give it the enabled protocols and listen addresses and start it
         if let Some((service, mut identify_config)) = identify_info.take() {
             identify_config.protocols = transport_manager.protocols().cloned().collect();
-            let identify = Identify::new(service, identify_config);
+            let identify = Identify::new(
+                service,
+                identify_config,
+                litep2p_config.metrics_registry.clone(),
+            )?;
 
             litep2p_config.executor.run(Box::pin(async move {
                 let _ = identify.run().await;

src/metrics.rs

@@ -0,0 +1,150 @@
+// Copyright 2020 Parity Technologies (UK) Ltd.
+// Copyright 2024 litep2p developers
+//
+// Permission is hereby granted, free of charge, to any person obtaining a
+// copy of this software and associated documentation files (the "Software"),
+// to deal in the Software without restriction, including without limitation
+// the rights to use, copy, modify, merge, publish, distribute, sublicense,
+// and/or sell copies of the Software, and to permit persons to whom the
+// Software is furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+// DEALINGS IN THE SOFTWARE.
+
+//! A generic module for handling the metrics exposed by litep2p.
+//!
+//! Contains the traits and types that are used to define and interact with metrics.
+
+use crate::{utils::futures_stream::FuturesStream, Error};
+use futures::{Stream, StreamExt};
+use std::{
+    future::Future,
+    pin::Pin,
+    sync::Arc,
+    task::{Context, Poll},
+};
+
+pub type MetricCounter = Arc<dyn MetricCounterT>;
+
+pub type MetricGauge = Arc<dyn MetricGaugeT>;
+
+pub type MetricsRegistry = Arc<dyn MetricsRegistryT>;
+
+/// Represents a metric that can only go up.
+pub trait MetricCounterT: Send + Sync {
+    /// Increment the counter by `value`.
+    fn inc(&self, value: u64);
+}
+
+/// Represents a metric that can arbitrarily go up and down.
+pub trait MetricGaugeT: Send + Sync {
+    /// Set the gauge to `value`.
+    fn set(&self, value: u64);
+
+    /// Increment the gauge.
+    fn inc(&self);
+
+    /// Decrement the gauge.
+    fn dec(&self);
+
+    /// Add `value` to the gauge.
+    fn add(&self, value: u64);
+
+    /// Subtract `value` from the gauge.
+    fn sub(&self, value: u64);
+}
+
+/// A registry for metrics.
+pub trait MetricsRegistryT: Send + Sync {
+    /// Register a new counter.
+    fn register_counter(&self, name: String, help: String) -> Result<MetricCounter, Error>;
+
+    /// Register a new gauge.
+    fn register_gauge(&self, name: String, help: String) -> Result<MetricGauge, Error>;
+}
+
+/// A scope for metrics that modifies a provided gauge in an RAII fashion.
+///
+/// The gauge is incremented when constructed and decremented when the object is dropped.
+#[derive(Clone)]
+pub struct ScopeGaugeMetric {
+    inner: MetricGauge,
+}
+
+impl ScopeGaugeMetric {
+    /// Create a new [`ScopeGaugeMetric`].
+    pub fn new(inner: MetricGauge) -> Self {
+        inner.inc();
+        ScopeGaugeMetric { inner }
+    }
+}
+
+impl Drop for ScopeGaugeMetric {
+    fn drop(&mut self) {
+        self.inner.dec();
+    }
+}
+
+/// Wrapper around `FuturesStream` that provides information to the given metric.
+#[derive(Default)]
+pub struct MeteredFuturesStream<F> {
+    stream: FuturesStream<F>,
+    metric: Option<MetricGauge>,
+}
+
+impl<F> MeteredFuturesStream<F> {
+    pub fn new(metric: Option<MetricGauge>) -> Self {
+        MeteredFuturesStream {
+            stream: FuturesStream::new(),
+            metric,
+        }
+    }
+
+    pub fn push(&mut self, future: F) {
+        if let Some(ref metric) = self.metric {
+            metric.inc();
+        }
+
+        self.stream.push(future);
+    }
+
+    /// Number of futures in the stream.
+    pub fn len(&self) -> usize {
+        self.stream.len()
+    }
+
+    /// Returns `true` if the stream is empty.
+    pub fn is_empty(&self) -> bool {
+        self.stream.len() == 0
+    }
+}
+
+impl<F: Future> Stream for MeteredFuturesStream<F> {
+    type Item = <F as Future>::Output;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let result = self.stream.poll_next_unpin(cx);
+        if result.is_ready() {
+            if let Some(ref metric) = self.metric {
+                metric.dec();
+            }
+        }
+        result
+    }
+}
+
+impl<F> Drop for MeteredFuturesStream<F> {
+    fn drop(&mut self) {
+        if let Some(ref metric) = self.metric {
+            metric.sub(self.len() as u64);
+        }
+    }
+}