chore: add metrics (#40)

* chore: add metrics

* nit

Author: ctian1

Cargo.toml

@@ -17,6 +17,7 @@ members = [
     "crates/artifacts",
     "crates/calibrator",
     "crates/rpc",
+    "crates/metrics",
     "programs/fibonacci",
     "programs/is-prime",
 ]
@@ -74,6 +75,7 @@ tonic-build = "0.12.0"
 # misc
 anyhow = "1.0.71"
 async-trait = "0.1.68"
+axum = "0.7.4"
 backoff = { version = "0.4", features = ["futures", "tokio"] }
 bytes = "1.6.0"
 config = "0.14.0"
@@ -95,4 +97,13 @@ sqlx = { version = "0.8", features = [
 thiserror = "1.0.63"
 tokio = { version = "1.0", features = ["full"] }
 url = "2.5.0"
-reqwest = "0.12.0"
+reqwest = "0.12.0"
+
+# metrics
+metrics = "0.24.0"
+metrics-derive = "0.1"
+metrics-exporter-prometheus = { version = "0.16.0", default-features = false }
+metrics-process = "2.1.0"
+metrics-util = { default-features = false, version = "0.18.0" }
+once_cell = "1.18.0"
+socket2 = { version = "0.5", default-features = false }

crates/metrics/Cargo.toml

@@ -0,0 +1,28 @@
+[package]
+name = "spn-metrics"
+description = ""
+readme = "README.md"
+version = { workspace = true }
+edition = { workspace = true }
+license = { workspace = true }
+repository = { workspace = true }
+keywords = { workspace = true }
+categories = { workspace = true }
+
+[dependencies]
+metrics = { workspace = true }
+eyre = { workspace = true }
+metrics-exporter-prometheus = { workspace = true }
+metrics-process = { workspace = true }
+metrics-util = { workspace = true }
+metrics-derive = { workspace = true }
+tokio = { workspace = true }
+tracing = { workspace = true }
+once_cell = { workspace = true }
+
+[target.'cfg(target_os = "linux")'.dependencies]
+procfs = "0.16.0"
+
+[dev-dependencies]
+reqwest = { workspace = true }
+socket2 = { workspace = true }

crates/metrics/src/hooks.rs

@@ -0,0 +1,87 @@
+use metrics_process::Collector;
+use once_cell::sync::Lazy;
+use std::{fmt, sync::Arc, time::SystemTime};
+
+pub(crate) trait Hook: Fn() + Send + Sync {}
+impl<T: Fn() + Send + Sync> Hook for T {}
+
+impl fmt::Debug for Hooks {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let hooks_len = self.inner.len();
+        f.debug_struct("Hooks")
+            .field("inner", &format!("Arc<Vec<Box<dyn Hook>>>, len: {}", hooks_len))
+            .finish()
+    }
+}
+
+/// Helper type for managing hooks.
+#[derive(Clone)]
+pub struct Hooks {
+    inner: Arc<Vec<Box<dyn Hook<Output = ()>>>>,
+}
+
+impl Hooks {
+    /// Create a new set of hooks.
+    pub fn new() -> Self {
+        let collector = Collector::default();
+        let hooks: Vec<Box<dyn Hook<Output = ()>>> = vec![
+            Box::new(move || collector.collect()),
+            Box::new(collect_memory_stats),
+            Box::new(collect_io_stats),
+            Box::new(collect_uptime_seconds),
+        ];
+        Self { inner: Arc::new(hooks) }
+    }
+
+    pub(crate) fn iter(&self) -> impl Iterator<Item = &Box<dyn Hook<Output = ()>>> {
+        self.inner.iter()
+    }
+}
+
+impl Default for Hooks {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+fn collect_memory_stats() {}
+
+#[cfg(target_os = "linux")]
+fn collect_io_stats() {
+    use metrics::counter;
+    use tracing::error;
+
+    let Ok(process) = procfs::process::Process::myself()
+        .map_err(|error| error!(%error, "failed to get currently running process"))
+    else {
+        return;
+    };
+
+    let Ok(io) = process.io().map_err(
+        |error| error!(%error, "failed to get io stats for the currently running process"),
+    ) else {
+        return;
+    };
+
+    counter!("io.rchar").absolute(io.rchar);
+    counter!("io.wchar").absolute(io.wchar);
+    counter!("io.syscr").absolute(io.syscr);
+    counter!("io.syscw").absolute(io.syscw);
+    counter!("io.read_bytes").absolute(io.read_bytes);
+    counter!("io.write_bytes").absolute(io.write_bytes);
+    counter!("io.cancelled_write_bytes").absolute(io.cancelled_write_bytes);
+}
+
+#[cfg(not(target_os = "linux"))]
+const fn collect_io_stats() {}
+
+/// Global start time of the process.
+static START_TIME: Lazy<SystemTime> = Lazy::new(SystemTime::now);
+
+/// Collects and records the process uptime in seconds.
+fn collect_uptime_seconds() {
+    use metrics::gauge;
+    let uptime = START_TIME.elapsed().unwrap_or_default().as_secs();
+    let gauge = gauge!("process_uptime_seconds");
+    gauge.set(uptime as f64);
+}

crates/metrics/src/lib.rs

@@ -0,0 +1,13 @@
+/// The metrics hooks for prometheus.
+pub mod hooks;
+pub mod recorder;
+/// The metric server serving the metrics.
+pub mod server;
+pub mod version;
+
+pub use metrics_exporter_prometheus::*;
+pub use metrics_process::*;
+
+pub use metrics_derive::Metrics;
+
+pub use metrics;

crates/metrics/src/recorder.rs

@@ -0,0 +1,58 @@
+//! Prometheus recorder
+
+use metrics_exporter_prometheus::{PrometheusBuilder, PrometheusHandle};
+use metrics_util::layers::{PrefixLayer, Stack};
+use std::sync::OnceLock;
+
+/// Installs the Prometheus recorder as the global recorder.
+pub fn get_or_init_prometheus(service_name: &str) -> &'static PrometheusHandle {
+    PROMETHEUS_RECORDER_HANDLE.get_or_init(|| PrometheusRecorder::install(service_name))
+}
+
+/// The default Prometheus recorder handle. We use a global static to ensure that it is only
+/// installed once.
+static PROMETHEUS_RECORDER_HANDLE: OnceLock<PrometheusHandle> = OnceLock::new();
+
+/// Prometheus recorder installer.
+#[derive(Debug)]
+pub struct PrometheusRecorder;
+
+impl PrometheusRecorder {
+    /// Installs Prometheus as the metrics recorder.
+    pub fn install(service_name: &str) -> PrometheusHandle {
+        let builder = PrometheusBuilder::new().add_global_label("service", service_name);
+        let recorder = builder.build_recorder();
+        let handle = recorder.handle();
+
+        // Build metrics stack
+        Stack::new(recorder)
+            .push(PrefixLayer::new("spn"))
+            .install()
+            .expect("Couldn't set metrics recorder.");
+
+        handle
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    // Dependencies using different version of the `metrics` crate (to be exact, 0.21 vs 0.22)
+    // may not be able to communicate with each other through the global recorder.
+    //
+    // This test ensures that `metrics-process` dependency plays well with the current
+    // `metrics-exporter-prometheus` dependency version.
+    #[test]
+    fn process_metrics() {
+        // initialize the lazy handle
+        let _ = PROMETHEUS_RECORDER_HANDLE.get_or_init(|| PrometheusRecorder::install("spn"));
+
+        let process = metrics_process::Collector::default();
+        process.describe();
+        process.collect();
+
+        let metrics =
+            PROMETHEUS_RECORDER_HANDLE.get_or_init(|| PrometheusRecorder::install("spn")).render();
+        assert!(metrics.contains("process_cpu_seconds_total"), "{metrics:?}");
+    }
+}