feat: metrics collection

Cargo.toml

@@ -45,7 +45,7 @@ tempfile = "3.4"
 thiserror = "1"
 tokio = { version = "1", features = ["full"] }
 tokio-stream = "0.1"
-tokio-util = { version = "0.7" }
+tokio-util = { version = "0.7", features = ["io-util", "io"] }
 tracing = "0.1"
 tracing-futures = "0.2.5"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
@@ -55,7 +55,10 @@ x509-parser = "0.14"
 zeroize = "1.5"
 bao-tree = { version = "0.1.5", features = ["tokio_io"], default-features = false }
 range-collections = "0.4.0"
-
+once_cell = { version = "1.17.1", optional = true }
+prometheus-client = { version = "0.18.0", optional = true }
+paste = { version = "1.0.12", optional = true }
+hyper = { version = "0.14.16", features = ["server", "http1", "tcp"], optional = true }
 [dev-dependencies]
 proptest = "1.0.0"
 rand = "0.7"
@@ -64,8 +67,9 @@ regex = { version = "1.7.1", features = ["std"] }
 nix = "0.26.2"
 
 [features]
-default = ["cli"]
+default = ["cli", "metrics"]
 cli = ["clap", "console", "indicatif", "data-encoding", "multibase"]
+metrics = ["paste", "hyper", "prometheus-client", "once_cell"]
 test = []
 
 [[bin]]

src/lib.rs

@@ -3,6 +3,8 @@
 #![deny(rustdoc::broken_intra_doc_links)]
 pub mod blobs;
 pub mod get;
+#[cfg(feature = "metrics")]
+pub mod metrics;
 pub mod net;
 pub mod progress;
 pub mod protocol;

src/main.rs

@@ -27,6 +27,9 @@ use main_util::Blake3Cid;
 
 use crate::main_util::{iroh_data_root, pathbuf_from_name};
 
+#[cfg(feature = "metrics")]
+use iroh::metrics::init_metrics;
+
 const DEFAULT_RPC_PORT: u16 = 0x1337;
 const RPC_ALPN: [u8; 17] = *b"n0/provider-rpc/1";
 const MAX_RPC_CONNECTIONS: u32 = 16;
@@ -42,6 +45,10 @@ struct Cli {
     /// Log SSL pre-master key to file in SSLKEYLOGFILE environment variable.
     #[clap(long)]
     keylog: bool,
+    /// Bind address on which to serve Prometheus metrics
+    #[cfg(feature = "metrics")]
+    #[clap(long)]
+    metrics_addr: Option<SocketAddr>,
 }
 
 #[derive(Debug, Clone)]
@@ -440,6 +447,25 @@ fn print_add_response(hash: Hash, entries: Vec<ProvideResponseEntry>) {
 const PROGRESS_STYLE: &str =
     "{msg}\n{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})";
 
+#[cfg(feature = "metrics")]
+fn init_metrics_collection(
+    metrics_addr: Option<SocketAddr>,
+) -> Option<tokio::task::JoinHandle<()>> {
+    init_metrics();
+    // doesn't start the server if the address is None
+    if let Some(metrics_addr) = metrics_addr {
+        return Some(tokio::spawn(async move {
+            iroh::metrics::start_metrics_server(metrics_addr)
+                .await
+                .unwrap_or_else(|e| {
+                    eprintln!("Failed to start metrics server: {}", e);
+                });
+        }));
+    }
+    tracing::info!("Metrics server not started, no address provided");
+    None
+}
+
 fn main() -> Result<()> {
     let rt = tokio::runtime::Builder::new_multi_thread()
         .enable_all()
@@ -460,7 +486,10 @@ async fn main_impl() -> Result<()> {
 
     let cli = Cli::parse();
 
-    match cli.command {
+    #[cfg(feature = "metrics")]
+    let metrics_fut = init_metrics_collection(cli.metrics_addr);
+
+    let r = match cli.command {
         Commands::Get {
             hash,
             peer,
@@ -665,7 +694,14 @@ async fn main_impl() -> Result<()> {
             println!("Listening addresses: {:?}", response.addrs);
             Ok(())
         }
+    };
+
+    #[cfg(feature = "metrics")]
+    if let Some(metrics_fut) = metrics_fut {
+        metrics_fut.abort();
+        drop(metrics_fut);
     }
+    r
 }
 
 async fn provide(

src/metrics.rs

@@ -0,0 +1,47 @@
+//! Metrics collection
+//!
+//! Enables and manages a global registry of metrics.
+//! Divided up into modules, each module has its own metrics.
+//! Starting the metrics service will expose the metrics on a OpenMetrics http endpoint.
+//!
+//! To enable metrics collection, call `init_metrics()` before starting the service.
+//!
+//! To record a metric, use the `record!` macro with the metric and the value to record.
+//! To increment a metric by 1, use the `inc!` macro with the metric.
+//! To observe a metric, use the `observe!` macro with the metric and the value to observe.
+//! To expose the metrics, start the metrics service with `start_metrics_server()`.
+//!
+//! # Example:
+//! ```
+//! use iroh::metrics::init_metrics;
+//! use iroh::metrics::iroh::IrohMetrics;
+//! use crate::iroh::metrics::core::MRecorder;
+//!
+//! init_metrics();
+//! iroh::record!(IrohMetrics::RequestsTotal, 2);
+//! iroh::inc!(IrohMetrics::RequestsTotal);
+//! ```
+use std::net::SocketAddr;
+
+use hyper::Error;
+
+use self::core::CORE;
+
+#[macro_use]
+mod macros;
+
+/// Expose core types and traits
+pub mod core;
+/// Expose iroh metrics
+pub mod iroh;
+mod service;
+
+/// Enables metrics collection, otherwise all inc!, record! & observe! calls are no-ops
+pub fn init_metrics() {
+    CORE.set_enabled(true);
+}
+
+/// Start a server to serve the OpenMetrics endpoint.
+pub async fn start_metrics_server(addr: SocketAddr) -> Result<(), Error> {
+    self::service::run(addr).await
+}