metrics: add a new metric for budget exhaustion yields

tokio/src/runtime/coop.rs

@@ -177,12 +177,35 @@ cfg_coop! {
                 cell.set(budget);
                 Poll::Ready(restore)
             } else {
+                inc_budget_forced_yield_count();
+
                 cx.waker().wake_by_ref();
                 Poll::Pending
             }
         }).unwrap_or(Poll::Ready(RestoreOnPending(Cell::new(Budget::unconstrained()))))
     }
 
+    cfg_rt! {
+        cfg_metrics! {
+            #[inline(always)]
+            fn inc_budget_forced_yield_count() {
+                if let Ok(handle) = context::try_current() {
+                    handle.scheduler_metrics().inc_budget_forced_yield_count();
+                }
+            }
+        }
+
+        cfg_not_metrics! {
+            #[inline(always)]
+            fn inc_budget_forced_yield_count() {}
+        }
+    }
+
+    cfg_not_rt! {
+        #[inline(always)]
+        fn inc_budget_forced_yield_count() {}
+    }
+
     impl Budget {
         /// Decrements the budget. Returns `true` if successful. Decrementing fails
         /// when there is not enough remaining budget.

tokio/src/runtime/metrics/runtime.rs

@@ -124,6 +124,21 @@ impl RuntimeMetrics {
             .load(Relaxed)
     }
 
+    /// Returns the number of times that tasks have been forced to yield back to the scheduler
+    /// after exhausting their task budgets.
+    ///
+    /// This count starts at zero when the runtime is created and increases by one each time a task yields due to exhausting its budget.
+    ///
+    /// The counter is monotonically increasing. It is never decremented or
+    /// reset to zero.
+    pub fn budget_forced_yield_count(&self) -> u64 {
+        self.handle
+            .inner
+            .scheduler_metrics()
+            .budget_forced_yield_count
+            .load(Relaxed)
+    }
+
     /// Returns the total number of times the given worker thread has parked.
     ///
     /// The worker park count starts at zero when the runtime is created and

tokio/src/runtime/metrics/scheduler.rs

@@ -11,17 +11,24 @@ use crate::loom::sync::atomic::{AtomicU64, Ordering::Relaxed};
 pub(crate) struct SchedulerMetrics {
     /// Number of tasks that are scheduled from outside the runtime.
     pub(super) remote_schedule_count: AtomicU64,
+    pub(super) budget_forced_yield_count: AtomicU64,
 }
 
 impl SchedulerMetrics {
     pub(crate) fn new() -> SchedulerMetrics {
         SchedulerMetrics {
             remote_schedule_count: AtomicU64::new(0),
+            budget_forced_yield_count: AtomicU64::new(0),
         }
     }
 
     /// Increment the number of tasks scheduled externally
     pub(crate) fn inc_remote_schedule_count(&self) {
         self.remote_schedule_count.fetch_add(1, Relaxed);
     }
+
+    /// Increment the number of tasks forced to yield due to budget exhaustion
+    pub(crate) fn inc_budget_forced_yield_count(&self) {
+        self.budget_forced_yield_count.fetch_add(1, Relaxed);
+    }
 }

tokio/tests/rt_metrics.rs

@@ -1,9 +1,13 @@
 #![warn(rust_2018_idioms)]
 #![cfg(all(feature = "full", tokio_unstable, not(tokio_wasi)))]
 
+use std::future::Future;
 use std::sync::{Arc, Mutex};
+use std::task::Poll;
+use tokio::macros::support::poll_fn;
 
 use tokio::runtime::Runtime;
+use tokio::task::consume_budget;
 use tokio::time::{self, Duration};
 
 #[test]
@@ -433,6 +437,33 @@ fn worker_local_queue_depth() {
     });
 }
 
+#[test]
+fn budget_exhaustion_yield() {
+    let rt = current_thread();
+    let metrics = rt.metrics();
+
+    assert_eq!(0, metrics.budget_forced_yield_count());
+
+    let mut did_yield = false;
+
+    // block on a task which consumes budget until it yields
+    rt.block_on(poll_fn(|cx| loop {
+        if did_yield {
+            return Poll::Ready(());
+        }
+
+        let fut = consume_budget();
+        tokio::pin!(fut);
+
+        if fut.poll(cx).is_pending() {
+            did_yield = true;
+            return Poll::Pending;
+        }
+    }));
+
+    assert_eq!(1, rt.metrics().budget_forced_yield_count());
+}
+
 #[cfg(any(target_os = "linux", target_os = "macos"))]
 #[test]
 fn io_driver_fd_count() {