commonwarexyz · patrick-ogrady · May 30, 2025 · May 30, 2025 · May 30, 2025 · May 30, 2025
diff --git a/runtime/src/iouring/mod.rs b/runtime/src/iouring/mod.rs
@@ -9,7 +9,8 @@ use io_uring::{
     types::Timespec,
     IoUring,
 };
-use std::{collections::HashMap, time::Duration};
+use prometheus_client::{metrics::gauge::Gauge, registry::Registry};
+use std::{collections::HashMap, sync::Arc, time::Duration};
 
 /// Reserved ID for a CQE that indicates an operation timed out.
 const TIMEOUT_WORK_ID: u64 = u64::MAX;
@@ -21,6 +22,31 @@ const SHUTDOWN_TIMEOUT_WORK_ID: u64 = u64::MAX - 1;
 /// wake up to check for new work.
 const POLL_WORK_ID: u64 = u64::MAX - 2;
 
+#[derive(Debug)]
+/// Tracks io_uring metrics.
+pub struct Metrics {
+    /// Number of operations submitted to the io_uring whose CQEs haven't
+    /// yet been processed. Note this metric doesn't include timeouts,
+    /// which are generated internally by the io_uring event loop.
+    /// It's only updated before `submit_and_wait` is called, so it may
+    /// temporarily vary from the actual number of pending operations.
+    pending_operations: Gauge,
+}
+
+impl Metrics {
+    pub fn new(registry: &mut Registry) -> Self {
+        let metrics = Self {
+            pending_operations: Gauge::default(),
+        };
+        registry.register(
+            "pending_operations",
+            "Number of operations submitted to the io_uring whose CQEs haven't yet been processed",
+            metrics.pending_operations.clone(),
+        );
+        metrics
+    }
+}
+
 #[derive(Clone, Debug)]
 /// Configuration for an io_uring instance.
 /// See `man io_uring`.
@@ -135,6 +161,7 @@ fn handle_cqe(waiters: &mut HashMap<u64, oneshot::Sender<i32>>, cqe: CqueueEntry
 /// It should be run in a separate task.
 pub(crate) async fn run(
     cfg: Config,
+    metrics: Arc<Metrics>,
     mut receiver: mpsc::Receiver<(SqueueEntry, oneshot::Sender<i32>)>,
 ) {
     let mut ring = new_ring(&cfg).expect("unable to create io_uring instance");
@@ -237,6 +264,7 @@ pub(crate) async fn run(
         // even if it's there before this call. That is, a completion
         // that arrived before this call will be counted and cause this
         // call to return. Note that waiters.len() > 0 here.
+        metrics.pending_operations.set(waiters.len() as _);
         ring.submit_and_wait(1).expect("unable to submit to ring");
     }
 }
@@ -296,13 +324,18 @@ mod tests {
         opcode,
         types::{Fd, Timespec},
     };
-    use std::os::{fd::AsRawFd, unix::net::UnixStream};
+    use prometheus_client::registry::Registry;
     use std::time::Duration;
+    use std::{
+        os::{fd::AsRawFd, unix::net::UnixStream},
+        sync::Arc,
+    };
 
     async fn recv_then_send(cfg: Config, should_succeed: bool) {
         // Create a new io_uring instance
         let (mut submitter, receiver) = channel(0);
-        let handle = tokio::spawn(super::run(cfg, receiver));
+        let metrics = Arc::new(super::Metrics::new(&mut Registry::default()));
+        let handle = tokio::spawn(super::run(cfg, metrics.clone(), receiver));
 
         let (left_pipe, right_pipe) = UnixStream::pair().unwrap();
 
@@ -317,10 +350,12 @@ mod tests {
             .await
             .expect("failed to send work");
 
+        while metrics.pending_operations.get() == 0 {
+            // Wait for the read to be submitted
+            tokio::time::sleep(Duration::from_millis(100)).await;
+        }
+
         // Submit a write that satisfies the read.
-        // Note that since the channel capacity is 0, we can only successfully send
-        // the write after the event loop has reached receiver.await(), which implies
-        // the event loop is parked in submit_and_wait when the send below is called.
         let write =
             opcode::Write::new(Fd(right_pipe.as_raw_fd()), msg.as_ptr(), msg.len() as _).build();
         let (write_tx, write_rx) = oneshot::channel();
@@ -381,7 +416,8 @@ mod tests {
             ..Default::default()
         };
         let (mut submitter, receiver) = channel(1);
-        let handle = tokio::spawn(super::run(cfg, receiver));
+        let metrics = Arc::new(super::Metrics::new(&mut Registry::default()));
+        let handle = tokio::spawn(super::run(cfg, metrics, receiver));
 
         // Submit a work item that will time out (because we don't write to the pipe)
         let (pipe_left, _pipe_right) = UnixStream::pair().unwrap();
@@ -408,7 +444,8 @@ mod tests {
             ..Default::default()
         };
         let (mut submitter, receiver) = channel(1);
-        let handle = tokio::spawn(super::run(cfg, receiver));
+        let metrics = Arc::new(super::Metrics::new(&mut Registry::default()));
+        let handle = tokio::spawn(super::run(cfg, metrics, receiver));
 
         // Submit an operation that will complete after shutdown
         let timeout = Timespec::new().sec(3);
@@ -433,7 +470,8 @@ mod tests {
             ..Default::default()
         };
         let (mut submitter, receiver) = channel(1);
-        let handle = tokio::spawn(super::run(cfg, receiver));
+        let metrics = Arc::new(super::Metrics::new(&mut Registry::default()));
+        let handle = tokio::spawn(super::run(cfg, metrics, receiver));
 
         // Submit an operation that will complete long after shutdown starts
         let timeout = Timespec::new().sec(5_000);

diff --git a/runtime/src/network/iouring.rs b/runtime/src/network/iouring.rs
@@ -6,6 +6,7 @@ use futures::{
     SinkExt as _,
 };
 use io_uring::{squeue::Entry as SqueueEntry, types::Fd};
+use prometheus_client::registry::Registry;
 use std::{
     net::SocketAddr,
     os::fd::{AsRawFd, OwnedFd},
@@ -31,17 +32,24 @@ impl Network {
     /// large enough, given the number of connections that will be maintained.
     /// Each ongoing send/recv to/from each connection will consume a slot in the io_uring.
     /// The io_uring `size` should be a multiple of the number of expected connections.
-    pub(crate) fn start(cfg: iouring::Config) -> Result<Self, crate::Error> {
+    pub(crate) fn start(
+        cfg: iouring::Config,
+        registry: &mut Registry,
+    ) -> Result<Self, crate::Error> {
         // Create an io_uring instance to handle send operations.
         let (send_submitter, rx) = mpsc::channel(cfg.size as usize);
         std::thread::spawn({
             let cfg = cfg.clone();
-            move || block_on(iouring::run(cfg, rx))
+            let registry = registry.sub_registry_with_prefix("iouring_sender");
+            let metrics = Arc::new(iouring::Metrics::new(registry));
+            move || block_on(iouring::run(cfg, metrics, rx))
         });
 
         // Create an io_uring instance to handle receive operations.
         let (recv_submitter, rx) = mpsc::channel(cfg.size as usize);
-        std::thread::spawn(|| block_on(iouring::run(cfg, rx)));
+        let registry = registry.sub_registry_with_prefix("iouring_receiver");
+        let metrics = Arc::new(iouring::Metrics::new(registry));
+        std::thread::spawn(|| block_on(iouring::run(cfg, metrics, rx)));
 
         Ok(Self {
             send_submitter: send_submitter.clone(),

diff --git a/runtime/src/storage/iouring.rs b/runtime/src/storage/iouring.rs
@@ -9,6 +9,7 @@ use futures::{
     SinkExt as _,
 };
 use io_uring::{opcode, squeue::Entry as SqueueEntry, types};
+use prometheus_client::registry::Registry;
 use std::fs::{self, File};
 use std::io::Error as IoError;
 use std::os::fd::AsRawFd;
@@ -32,15 +33,16 @@ pub struct Storage {
 
 impl Storage {
     /// Returns a new `Storage` instance.
-    pub fn start(cfg: Config) -> Self {
+    pub fn start(cfg: Config, registry: &mut Registry) -> Self {
         let (io_sender, receiver) =
             mpsc::channel::<(SqueueEntry, oneshot::Sender<i32>)>(cfg.ring_config.size as usize);
 
         let storage = Storage {
             storage_directory: cfg.storage_directory.clone(),
             io_sender,
         };
-        std::thread::spawn(|| block_on(iouring::run(cfg.ring_config, receiver)));
+        let metrics = Arc::new(iouring::Metrics::new(registry));
+        std::thread::spawn(|| block_on(iouring::run(cfg.ring_config, metrics, receiver)));
         storage
     }
 }
@@ -301,10 +303,13 @@ mod tests {
         let storage_directory =
             env::temp_dir().join(format!("commonware_iouring_storage_{}", rng.gen::<u64>()));
 
-        let storage = Storage::start(Config {
-            storage_directory: storage_directory.clone(),
-            ring_config: Default::default(),
-        });
+        let storage = Storage::start(
+            Config {
+                storage_directory: storage_directory.clone(),
+                ring_config: Default::default(),
+            },
+            &mut Registry::default(),
+        );
         (storage, storage_directory)
     }
 

diff --git a/runtime/src/tokio/runtime.rs b/runtime/src/tokio/runtime.rs
@@ -242,11 +242,12 @@ impl crate::Runner for Runner {
 
         cfg_if::cfg_if! {
             if #[cfg(feature = "iouring-storage")] {
+                let iouring_registry = runtime_registry.sub_registry_with_prefix("iouring_storage");
                 let storage = MeteredStorage::new(
                     IoUringStorage::start(IoUringConfig {
                         storage_directory: self.cfg.storage_directory.clone(),
                         ring_config: Default::default(),
-                    }),
+                    }, iouring_registry),
                     runtime_registry,
                 );
             } else {
@@ -262,8 +263,9 @@ impl crate::Runner for Runner {
 
         cfg_if::cfg_if! {
             if #[cfg(feature = "iouring-network")] {
+                let iouring_registry = runtime_registry.sub_registry_with_prefix("iouring_network");
                 let network = MeteredNetwork::new(
-                    IoUringNetwork::start(crate::iouring::Config::default()).unwrap(),
+                    IoUringNetwork::start(crate::iouring::Config::default(),iouring_registry).unwrap(),
                     runtime_registry,
                 );
             } else {