monitor tests, more docs

Author: Dav1dde

CHANGELOG.md

@@ -1,5 +1,11 @@
 # Changelog
 
+## Unreleased
+
+**Internal**:
+
+- Track an utilization metric for internal services. ([#4501](https://github.com/getsentry/relay/pull/4501))
+
 ## 25.2.0
 
 - Allow log ingestion behind a flag, only for internal use currently. ([#4471](https://github.com/getsentry/relay/pull/4471))

relay-system/src/runtime/handle.rs

@@ -14,10 +14,10 @@ struct HandleInner {
     tokio_cb_metrics: Arc<TokioCallbackMetrics>,
 }
 
-/// Handle to the [`Runtime`].
+/// Handle to the [`Runtime`](crate::Runtime).
 ///
 /// The handle is internally reference-counted and can be freely cloned.
-/// A handle can be obtained using the [`Runtime::handle`] method.
+/// A handle can be obtained using the [`Runtime::handle`](crate::Runtime::handle) method.
 #[derive(Debug, Clone)]
 pub struct Handle {
     inner: Arc<HandleInner>,

relay-system/src/runtime/spawn.rs

@@ -38,7 +38,7 @@ where
 
 /// Spawns a new asynchronous task in a specific runtime, returning a [`JoinHandle`] for it.
 ///
-/// This is in instrumented spawn variant of Tokio's [`Handle::spawn`].
+/// This is in instrumented spawn variant of Tokio's [`Handle::spawn`](crate::Handle::spawn).
 #[allow(clippy::disallowed_methods)]
 pub fn spawn_in<F>(
     handle: &tokio::runtime::Handle,
@@ -166,8 +166,8 @@ mod tests {
         #[cfg(windows)]
         assert_debug_snapshot!(captures, @r###"
         [
-            "runtime.task.spawn.created:1|c|#id:relay-system\\src\\runtime\\spawn.rs:124,file:relay-system\\src\\runtime\\spawn.rs,line:124",
-            "runtime.task.spawn.terminated:1|c|#id:relay-system\\src\\runtime\\spawn.rs:124,file:relay-system\\src\\runtime\\spawn.rs,line:124",
+            "runtime.task.spawn.created:1|c|#id:relay-system\\src\\runtime\\spawn.rs:124,file:relay-system\\src\\runtime\\spawn.rs,line:155",
+            "runtime.task.spawn.terminated:1|c|#id:relay-system\\src\\runtime\\spawn.rs:124,file:relay-system\\src\\runtime\\spawn.rs,line:155",
         ]
         "###);
     }

relay-system/src/service/monitor.rs

@@ -3,16 +3,14 @@ use std::pin::Pin;
 use std::sync::atomic::{AtomicU64, AtomicU8, Ordering};
 use std::sync::Arc;
 use std::task::{Context, Poll};
-use std::time::{Duration, Instant};
-
-use crate::service::registry::ServiceId;
+use tokio::time::{Duration, Instant};
 
 /// Minimum interval when utilization is recalculated.
 const UTILIZATION_UPDATE_THRESHOLD: Duration = Duration::from_secs(5);
 
 pin_project_lite::pin_project! {
+    /// A service monitor tracks service metrics.
     pub struct ServiceMonitor<F> {
-        id: ServiceId,
         #[pin]
         inner: F,
 
@@ -24,9 +22,9 @@ pin_project_lite::pin_project! {
 }
 
 impl<F> ServiceMonitor<F> {
-    pub fn wrap(id: ServiceId, inner: F) -> Self {
+    /// Wraps a service future with a monitor.
+    pub fn wrap(inner: F) -> Self {
         Self {
-            id,
             inner,
             metrics: Arc::new(RawMetrics {
                 poll_count: AtomicU64::new(0),
@@ -38,6 +36,7 @@ impl<F> ServiceMonitor<F> {
         }
     }
 
+    /// Provides access to the raw metrics tracked in this monitor.
     pub fn metrics(&self) -> &Arc<RawMetrics> {
         &self.metrics
     }
@@ -87,9 +86,55 @@ where
     }
 }
 
+/// The raw metrics extracted from a [`ServiceMonitor`].
+///
+/// All access outside the [`ServiceMonitor`] must be *read* only.
 #[derive(Debug)]
 pub struct RawMetrics {
+    /// Amount of times the service was polled.
     pub poll_count: AtomicU64,
+    /// The total time the service spent in its poll function.
     pub total_duration_ns: AtomicU64,
+    /// Estimated utilization percentage `[0-100]`
     pub utilization: AtomicU8,
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test(start_paused = true)]
+    async fn test_monitor() {
+        let mut monitor = ServiceMonitor::wrap(Box::pin(async {
+            loop {
+                tokio::time::advance(Duration::from_millis(500)).await;
+            }
+        }));
+        let metrics = Arc::clone(monitor.metrics());
+
+        assert_eq!(metrics.poll_count.load(Ordering::Relaxed), 0);
+        let _ = futures::poll!(&mut monitor);
+        assert_eq!(metrics.poll_count.load(Ordering::Relaxed), 1);
+        let _ = futures::poll!(&mut monitor);
+        assert_eq!(metrics.poll_count.load(Ordering::Relaxed), 2);
+        let _ = futures::poll!(&mut monitor);
+        assert_eq!(metrics.poll_count.load(Ordering::Relaxed), 3);
+
+        assert_eq!(metrics.utilization.load(Ordering::Relaxed), 0);
+        assert_eq!(
+            metrics.total_duration_ns.load(Ordering::Relaxed),
+            1500000000
+        );
+
+        // Advance time just enough to perfectly hit the update threshold.
+        tokio::time::advance(UTILIZATION_UPDATE_THRESHOLD - Duration::from_secs(2)).await;
+
+        let _ = futures::poll!(&mut monitor);
+        assert_eq!(metrics.poll_count.load(Ordering::Relaxed), 4);
+        assert_eq!(metrics.utilization.load(Ordering::Relaxed), 40);
+        assert_eq!(
+            metrics.total_duration_ns.load(Ordering::Relaxed),
+            2000000000
+        );
+    }
+}

relay-system/src/service/registry.rs

@@ -12,7 +12,7 @@ use crate::{ServiceObj, TaskId};
 pub struct ServicesMetrics(BTreeMap<ServiceId, ServiceMetrics>);
 
 impl ServicesMetrics {
-    /// Returns an iterator of all service identifiers and their [`Metrics`].
+    /// Returns an iterator of all service identifiers and their [`ServiceMetrics`].
     pub fn iter(&self) -> impl Iterator<Item = (ServiceId, ServiceMetrics)> + '_ {
         self.0.iter().map(|(id, metrics)| (*id, *metrics))
     }
@@ -104,13 +104,16 @@ impl Inner {
         let task_id = TaskId::from(&service);
         let group = self.services.entry(task_id).or_default();
 
+        // Cleanup group, evicting all terminated services, while we're at it.
+        group.instances.retain(|s| !s.handle.is_finished());
+
         let id = ServiceId {
             task: task_id,
             instance_id: group.next_instance_id,
         };
         group.next_instance_id += 1;
 
-        let future = ServiceMonitor::wrap(id, service.future);
+        let future = ServiceMonitor::wrap(service.future);
         let metrics = Arc::clone(future.metrics());
 
         let jh = crate::runtime::spawn_in(handle, task_id, future);
@@ -158,6 +161,5 @@ struct ServiceGroup {
 struct Service {
     instance_id: u32,
     metrics: Arc<RawMetrics>,
-    #[expect(unused, reason = "not yet implemented")]
     handle: ServiceStatusJoinHandle,
 }