[runtime] Track RSS (#1457)

Author: patrick-ogrady

Cargo.lock

@@ -86,6 +86,7 @@ io-uring = "0.7.4"
 rayon = "1.10.0"
 async-lock = "3.4.0"
 libc = "0.2.172"
+sysinfo = "0.33.0"
 zeroize = "1.5.7"
 blst = "0.3.13"
 p256 = "0.13.2"

Cargo.toml

@@ -29,6 +29,7 @@ rayon = { workspace = true }
 async-lock = { workspace = true }
 io-uring = { workspace = true, optional = true }
 libc = { workspace = true }
+sysinfo = { workspace = true }
 
 [features]
 default = []

runtime/Cargo.toml

@@ -39,6 +39,7 @@ cfg_if::cfg_if! {
     }
 }
 mod network;
+mod process;
 mod storage;
 pub mod telemetry;
 mod utils;
@@ -1672,6 +1673,37 @@ mod tests {
         test_metrics_label(executor);
     }
 
+    #[test]
+    fn test_tokio_process_rss_metric() {
+        let executor = tokio::Runner::default();
+        executor.start(|context| async move {
+            loop {
+                // Wait for RSS metric to be available
+                let metrics = context.encode();
+                if !metrics.contains("runtime_process_rss") {
+                    context.sleep(Duration::from_millis(100)).await;
+                    continue;
+                }
+
+                // Verify the RSS value is eventually populated (greater than 0)
+                for line in metrics.lines() {
+                    if line.starts_with("runtime_process_rss")
+                        && !line.starts_with("runtime_process_rss{")
+                    {
+                        let parts: Vec<&str> = line.split_whitespace().collect();
+                        if parts.len() >= 2 {
+                            let rss_value: i64 =
+                                parts[1].parse().expect("Failed to parse RSS value");
+                            if rss_value > 0 {
+                                return;
+                            }
+                        }
+                    }
+                }
+            }
+        });
+    }
+
     #[test]
     fn test_tokio_telemetry() {
         let executor = tokio::Runner::default();

runtime/src/lib.rs

@@ -0,0 +1,116 @@
+//! Process metrics collection.
+
+use prometheus_client::{metrics::gauge::Gauge, registry::Registry};
+use std::{future::Future, time::Duration};
+use sysinfo::{ProcessRefreshKind, ProcessesToUpdate, System};
+
+/// The interval at which to update process metrics.
+const TICK_INTERVAL: Duration = Duration::from_secs(10);
+
+/// Process metrics collector.
+pub struct Metrics {
+    /// Resident set size in bytes.
+    pub rss: Gauge,
+    /// Virtual memory size in bytes.
+    pub virtual_memory: Gauge,
+
+    /// Process ID.
+    pid: sysinfo::Pid,
+    /// System information handle.
+    system: System,
+}
+
+impl Metrics {
+    /// Initialize process metrics and register them with the given registry.
+    pub fn init(registry: &mut Registry) -> Self {
+        let metrics = Self {
+            pid: sysinfo::Pid::from_u32(std::process::id()),
+            rss: Gauge::default(),
+            virtual_memory: Gauge::default(),
+
+            system: System::new(),
+        };
+
+        // Register all metrics
+        registry.register(
+            "process_rss",
+            "Resident set size of the current process",
+            metrics.rss.clone(),
+        );
+        registry.register(
+            "process_virtual_memory",
+            "Virtual memory size of the current process",
+            metrics.virtual_memory.clone(),
+        );
+
+        metrics
+    }
+
+    /// Update all process metrics.
+    fn update(&mut self) {
+        // Refresh process information
+        self.system.refresh_processes_specifics(
+            ProcessesToUpdate::Some(&[self.pid]),
+            false,
+            ProcessRefreshKind::nothing().with_memory(),
+        );
+
+        // If the process exists, update the metrics
+        if let Some(process) = self.system.process(self.pid) {
+            self.rss.set(process.memory() as i64);
+            self.virtual_memory.set(process.virtual_memory() as i64);
+        }
+    }
+
+    /// Update process metrics periodically.
+    ///
+    /// This function takes a sleep function as a parameter to allow different runtimes
+    /// to provide their own implementation.
+    pub async fn collect<F, Fut>(mut self, sleep_fn: F)
+    where
+        F: Fn(Duration) -> Fut,
+        Fut: Future<Output = ()>,
+    {
+        loop {
+            self.update();
+            sleep_fn(TICK_INTERVAL).await;
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_process_metrics_init() {
+        let mut registry = Registry::default();
+        let mut metrics = Metrics::init(&mut registry);
+
+        // Update metrics
+        metrics.update();
+
+        // Check that RSS is reasonable (> 1MB for a running process)
+        let rss = metrics.rss.get();
+        assert!(rss > 0);
+
+        // Check that virtual memory is >= RSS
+        let virt = metrics.virtual_memory.get();
+        assert!(virt >= rss);
+
+        // Allocate some memory
+        let mut vec = vec![0; 10 * 1024 * 1024]; // 10MB
+        vec.push(1);
+
+        // Update metrics
+        metrics.update();
+
+        // Check that the metrics are updated
+        let new_rss = metrics.rss.get();
+        assert!(new_rss > rss, "RSS should be > {rss}");
+
+        // Check that virtual memory is updated
+        let new_virt = metrics.virtual_memory.get();
+        assert!(new_virt > virt, "Virtual memory should be > {virt}");
+    }
+}

runtime/src/process/metered.rs

@@ -0,0 +1,3 @@
+//! Process implementations.
+
+pub mod metered;

runtime/src/process/mod.rs

@@ -10,8 +10,8 @@ use crate::{
     network::iouring::{Config as IoUringNetworkConfig, Network as IoUringNetwork},
 };
 use crate::{
-    network::metered::Network as MeteredNetwork, signal::Signal,
-    storage::metered::Storage as MeteredStorage, telemetry::metrics::task::Label,
+    network::metered::Network as MeteredNetwork, process::metered::Metrics as MeteredProcess,
+    signal::Signal, storage::metered::Storage as MeteredStorage, telemetry::metrics::task::Label,
     utils::signal::Stopper, Clock, Error, Handle, SinkOf, StreamOf, METRICS_PREFIX,
 };
 use commonware_macros::select;
@@ -253,6 +253,14 @@ impl crate::Runner for Runner {
             .build()
             .expect("failed to create Tokio runtime");
 
+        // Collect process metrics.
+        //
+        // We prefer to collect process metrics outside of `Context` because
+        // we are using `runtime_registry` rather than the one provided by `Context`.
+        let process = MeteredProcess::init(runtime_registry);
+        runtime.spawn(process.collect(tokio::time::sleep));
+
+        // Initialize storage
         cfg_if::cfg_if! {
             if #[cfg(feature = "iouring-storage")] {
                 let iouring_registry = runtime_registry.sub_registry_with_prefix("iouring_storage");
@@ -274,6 +282,7 @@ impl crate::Runner for Runner {
             }
         }
 
+        // Initialize network
         cfg_if::cfg_if! {
             if #[cfg(feature = "iouring-network")] {
                 let iouring_registry = runtime_registry.sub_registry_with_prefix("iouring_network");