switch to CounterPair struct, improve reliability of flaky test

Author: conradludgate

tokio/src/runtime/metrics/mod.rs

@@ -19,7 +19,7 @@ cfg_metrics! {
 
     mod runtime;
     #[allow(unreachable_pub)] // rust-lang/rust#57411
-    pub use runtime::RuntimeMetrics;
+    pub use runtime::{RuntimeMetrics, CounterPair};
 
     mod scheduler;
     pub(crate) use scheduler::SchedulerMetrics;

tokio/src/runtime/metrics/runtime.rs

@@ -15,6 +15,31 @@ pub struct RuntimeMetrics {
     handle: Handle,
 }
 
+/// A gauge represented as two counters.
+///
+/// Instead of decrementing a gauge, we increment a decrements counter.
+/// This is beneficial as it allows you to observe activity spikes that occur
+/// inbetween a scrape interval
+#[derive(Copy, Clone, Debug, PartialEq)]
+pub struct CounterPair {
+    /// Tracks how many times this gauge was incremented
+    pub inc: u64,
+    /// Tracks how many times this gauge was decremeneted
+    pub dec: u64,
+}
+
+impl CounterPair {
+    /// Determines the current length of the pair
+    pub fn len(&self) -> usize {
+        (self.inc - self.dec) as usize
+    }
+
+    /// Determines if the counter pair represents an empty collection
+    pub fn is_empty(&self) -> bool {
+        self.inc == self.dec
+    }
+}
+
 impl RuntimeMetrics {
     pub(crate) fn new(handle: Handle) -> RuntimeMetrics {
         RuntimeMetrics { handle }
@@ -85,10 +110,10 @@ impl RuntimeMetrics {
     /// }
     /// ```
     pub fn active_tasks_count(&self) -> usize {
-        self.handle.inner.active_tasks_count()
+        self.handle.inner.task_counts().len()
     }
 
-    /// Returns the number of started tasks in the runtime.
+    /// Returns a counter pair representing the number of started and completed tasks in the runtime.
     ///
     /// # Examples
     ///
@@ -100,30 +125,13 @@ impl RuntimeMetrics {
     ///    let metrics = Handle::current().metrics();
     ///
     ///     let n = metrics.start_tasks_count();
-    ///     println!("Runtime has {} active tasks", n);
-    /// }
-    /// ```
-    pub fn start_tasks_count(&self) -> u64 {
-        self.handle.inner.start_tasks_count()
-    }
-
-    /// Returns the number of finished tasks in the runtime.
-    ///
-    /// # Examples
-    ///
-    /// ```
-    /// use tokio::runtime::Handle;
-    ///
-    /// #[tokio::main]
-    /// async fn main() {
-    ///    let metrics = Handle::current().metrics();
-    ///
-    ///     let n = metrics.stop_tasks_count();
-    ///     println!("Runtime has {} active tasks", n);
+    ///     println!("Runtime has {} started tasks", n.inc);
+    ///     println!("Runtime has {} completed tasks", n.dec);
+    ///     println!("Runtime has {} active tasks", n.len());
     /// }
     /// ```
-    pub fn stop_tasks_count(&self) -> u64 {
-        self.handle.inner.stop_tasks_count()
+    pub fn task_counts(&self) -> CounterPair {
+        self.handle.inner.task_counts()
     }
 
     /// Returns the number of idle threads, which have spawned by the runtime

tokio/src/runtime/mod.rs

@@ -250,7 +250,7 @@ cfg_rt! {
 
     cfg_metrics! {
         mod metrics;
-        pub use metrics::{RuntimeMetrics, HistogramScale};
+        pub use metrics::{RuntimeMetrics, HistogramScale, CounterPair};
 
         pub(crate) use metrics::{MetricsBatch, SchedulerMetrics, WorkerMetrics, HistogramBuilder};
 

tokio/src/runtime/scheduler/current_thread/mod.rs

@@ -2,6 +2,7 @@ use crate::future::poll_fn;
 use crate::loom::sync::atomic::AtomicBool;
 use crate::loom::sync::Arc;
 use crate::runtime::driver::{self, Driver};
+use crate::runtime::metrics::CounterPair;
 use crate::runtime::scheduler::{self, Defer, Inject};
 use crate::runtime::task::{self, JoinHandle, OwnedTasks, Schedule, Task};
 use crate::runtime::{blocking, context, Config, MetricsBatch, SchedulerMetrics, WorkerMetrics};
@@ -539,16 +540,8 @@ cfg_metrics! {
             self.blocking_spawner.queue_depth()
         }
 
-        pub(crate) fn active_tasks_count(&self) -> usize {
-            self.shared.owned.active_tasks_count()
-        }
-
-        pub(crate) fn start_tasks_count(&self) -> u64 {
-            self.shared.owned.start_tasks_count()
-        }
-
-        pub(crate) fn stop_tasks_count(&self) -> u64 {
-            self.shared.owned.stop_tasks_count()
+        pub(crate) fn tasks_count(&self) -> CounterPair {
+            self.shared.owned.tasks_count()
         }
     }
 }

tokio/src/runtime/scheduler/mod.rs

@@ -164,7 +164,7 @@ cfg_rt! {
     }
 
     cfg_metrics! {
-        use crate::runtime::{SchedulerMetrics, WorkerMetrics};
+        use crate::runtime::{SchedulerMetrics, WorkerMetrics, metrics::CounterPair};
 
         impl Handle {
             pub(crate) fn num_workers(&self) -> usize {
@@ -185,16 +185,8 @@ cfg_rt! {
                 match_flavor!(self, Handle(handle) => handle.num_idle_blocking_threads())
             }
 
-            pub(crate) fn active_tasks_count(&self) -> usize {
-                match_flavor!(self, Handle(handle) => handle.active_tasks_count())
-            }
-
-            pub(crate) fn start_tasks_count(&self) -> u64 {
-                match_flavor!(self, Handle(handle) => handle.start_tasks_count())
-            }
-
-            pub(crate) fn stop_tasks_count(&self) -> u64 {
-                match_flavor!(self, Handle(handle) => handle.stop_tasks_count())
+            pub(crate) fn task_counts(&self) -> CounterPair {
+                match_flavor!(self, Handle(handle) => handle.tasks_count())
             }
 
             pub(crate) fn scheduler_metrics(&self) -> &SchedulerMetrics {

tokio/src/runtime/scheduler/multi_thread/handle/metrics.rs

@@ -1,6 +1,6 @@
 use super::Handle;
 
-use crate::runtime::{SchedulerMetrics, WorkerMetrics};
+use crate::runtime::{SchedulerMetrics, WorkerMetrics, metrics::CounterPair};
 
 impl Handle {
     pub(crate) fn num_workers(&self) -> usize {
@@ -15,16 +15,8 @@ impl Handle {
         self.blocking_spawner.num_idle_threads()
     }
 
-    pub(crate) fn active_tasks_count(&self) -> usize {
-        self.shared.owned.active_tasks_count()
-    }
-
-    pub(crate) fn start_tasks_count(&self) -> u64 {
-        self.shared.owned.start_tasks_count()
-    }
-
-    pub(crate) fn stop_tasks_count(&self) -> u64 {
-        self.shared.owned.stop_tasks_count()
+    pub(crate) fn tasks_count(&self) -> CounterPair {
+        self.shared.owned.tasks_count()
     }
 
     pub(crate) fn scheduler_metrics(&self) -> &SchedulerMetrics {

tokio/src/runtime/scheduler/multi_thread_alt/handle/metrics.rs

@@ -1,6 +1,6 @@
 use super::Handle;
 
-use crate::runtime::{SchedulerMetrics, WorkerMetrics};
+use crate::runtime::{SchedulerMetrics, WorkerMetrics, metrics::CounterPair};
 
 impl Handle {
     pub(crate) fn num_workers(&self) -> usize {
@@ -15,16 +15,8 @@ impl Handle {
         self.blocking_spawner.num_idle_threads()
     }
 
-    pub(crate) fn active_tasks_count(&self) -> usize {
-        self.shared.owned.active_tasks_count()
-    }
-
-    pub(crate) fn start_tasks_count(&self) -> u64 {
-        self.shared.owned.start_tasks_count()
-    }
-
-    pub(crate) fn stop_tasks_count(&self) -> u64 {
-        self.shared.owned.stop_tasks_count()
+    pub(crate) fn tasks_count(&self) -> CounterPair {
+        self.shared.owned.tasks_count()
     }
 
     pub(crate) fn scheduler_metrics(&self) -> &SchedulerMetrics {

tokio/src/runtime/task/list.rs

@@ -9,6 +9,7 @@
 use crate::future::Future;
 use crate::loom::cell::UnsafeCell;
 use crate::loom::sync::Mutex;
+use crate::runtime::metrics::CounterPair;
 use crate::runtime::task::{JoinHandle, LocalNotified, Notified, Schedule, Task};
 use crate::util::linked_list::{CountedLinkedList, Link, LinkedList};
 
@@ -166,16 +167,12 @@ impl<S: 'static> OwnedTasks<S> {
         }
     }
 
-    pub(crate) fn active_tasks_count(&self) -> usize {
-        self.inner.lock().list.count()
-    }
-
-    pub(crate) fn start_tasks_count(&self) -> u64 {
-        self.inner.lock().list.added()
-    }
-
-    pub(crate) fn stop_tasks_count(&self) -> u64 {
-        self.inner.lock().list.removed()
+    pub(crate) fn tasks_count(&self) -> CounterPair {
+        let lock = self.inner.lock();
+        CounterPair {
+            inc: lock.list.added(),
+            dec: lock.list.removed(),
+        }
     }
 
     pub(crate) fn remove(&self, task: &Task<S>) -> Option<Task<S>> {

tokio/src/util/linked_list.rs

@@ -272,12 +272,6 @@ impl<L: Link> CountedLinkedList<L, L::Target> {
         val
     }
 
-    pub(crate) fn count(&self) -> usize {
-        // this subtraction can't underflow.
-        // this cast can't overflow because the length of the linked list can't exceed usize.
-        (self.added - self.removed) as usize
-    }
-
     pub(crate) fn added(&self) -> u64 {
         self.added
     }

tokio/tests/rt_metrics.rs

@@ -6,6 +6,7 @@ use std::sync::{Arc, Mutex};
 use std::task::Poll;
 use tokio::macros::support::poll_fn;
 
+use tokio::runtime::CounterPair;
 use tokio::runtime::Runtime;
 use tokio::task::consume_budget;
 use tokio::time::{self, Duration};
@@ -104,36 +105,32 @@ fn active_tasks_count() {
 fn active_tasks_count_pairs() {
     let rt = current_thread();
     let metrics = rt.metrics();
-    assert_eq!(0, metrics.start_tasks_count());
-    assert_eq!(0, metrics.stop_tasks_count());
+    assert_eq!(CounterPair { inc: 0, dec: 0 }, metrics.task_counts());
 
     rt.block_on(rt.spawn(async move {
-        assert_eq!(1, metrics.start_tasks_count());
-        assert_eq!(0, metrics.stop_tasks_count());
+        assert_eq!(CounterPair { inc: 1, dec: 0 }, metrics.task_counts());
     }))
     .unwrap();
 
-    assert_eq!(1, rt.metrics().start_tasks_count());
-    assert_eq!(1, rt.metrics().stop_tasks_count());
+    assert_eq!(CounterPair { inc: 1, dec: 1 }, rt.metrics().task_counts());
 
     let rt = threaded();
     let metrics = rt.metrics();
-    assert_eq!(0, metrics.start_tasks_count());
-    assert_eq!(0, metrics.stop_tasks_count());
+    assert_eq!(CounterPair { inc: 0, dec: 0 }, metrics.task_counts());
 
     rt.block_on(rt.spawn(async move {
-        assert_eq!(1, metrics.start_tasks_count());
-        assert_eq!(0, metrics.stop_tasks_count());
+        assert_eq!(CounterPair { inc: 1, dec: 0 }, metrics.task_counts());
     }))
     .unwrap();
 
-    // for some reason, sometimes the stop count doesn't get a chance to incremenet before we get here.
-    // Only observed on single-cpu systems. Most likely the worker thread doesn't a chance to clean up
-    // the spawned task yet. We yield to give it an opportunity.
-    std::thread::yield_now();
-
-    assert_eq!(1, rt.metrics().start_tasks_count());
-    assert_eq!(1, rt.metrics().stop_tasks_count());
+    for _ in 0..100 {
+        if rt.metrics().task_counts() == (CounterPair { inc: 1, dec: 1 }) {
+            return;
+        }
+        // on single threaded machines (like in CI), we need to force the OS to run the runtime threads
+        std::thread::yield_now();
+    }
+    panic!("runtime didn't decrement active task gauge")
 }
 
 #[test]