从零写一个 Instrumented

[dependencies]
tokio = { version = "=1.54.1", features = ["full"] }
# 只用来对照，mini 本身不依赖它
tokio-metrics = "=0.5.2"

use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll, Wake, Waker};
use std::time::{Duration, Instant};

/// 写的一方：只累加。
#[derive(Default)]
struct Counters {
  polls: AtomicU64,
  poll_ns: AtomicU64,
  slow_polls: AtomicU64,
  scheduled: AtomicU64,
  scheduled_ns: AtomicU64,
}

struct Shared {
  counters: Counters,
  epoch: Instant,
  slow: Duration,
}

impl Shared {
  /// 相对 epoch 的纳秒数，加 1 让 0 专门表示「没被唤醒」。
  fn now_ns(&self) -> u64 {
    self.epoch.elapsed().as_nanos() as u64 + 1
  }
}

/// 包装 waker：记下第一次唤醒的时刻，再转给真正的 waker。
struct WakeState {
  woke_at: AtomicU64,
  real: Mutex<Option<Waker>>,
  shared: Arc<Shared>,
}

impl Wake for WakeState {
  fn wake(self: Arc<Self>) {
    self.wake_by_ref();
  }

  fn wake_by_ref(self: &Arc<Self>) {
    let now = self.shared.now_ns();
    // 两次 poll 之间只记第一次唤醒
    let _ = self.woke_at.compare_exchange(0, now, Relaxed, Relaxed);
    if let Some(w) = self.real.lock().unwrap_or_else(|e| e.into_inner()).as_ref() {
      w.wake_by_ref();
    }
  }
}

#[derive(Clone)]
pub struct Monitor {
  shared: Arc<Shared>,
}

impl Monitor {
  pub fn new(slow: Duration) -> Self {
    let shared = Shared { counters: Counters::default(), epoch: Instant::now(), slow };
    Monitor { shared: Arc::new(shared) }
  }

  pub fn instrument<F: Future>(&self, fut: F) -> Instrumented<F> {
    let state = Arc::new(WakeState {
      woke_at: AtomicU64::new(0),
      real: Mutex::new(None),
      shared: self.shared.clone(),
    });
    let waker = Waker::from(state.clone());
    Instrumented { fut: Box::pin(fut), state, waker }
  }
  // snapshot、intervals 见下一节
}

pub struct Instrumented<F> {
  fut: Pin<Box<F>>,
  state: Arc<WakeState>,
  waker: Waker,
}

impl<F: Future> Future for Instrumented<F> {
  type Output = F::Output;

  fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<F::Output> {
    let this = self.get_mut();
    let shared = &this.state.shared;
    let c = &shared.counters;
    let start = Instant::now();
    let start_ns = shared.now_ns();

    // 1. 排队：被唤醒到这次 poll 开始
    let woke_at = this.state.woke_at.swap(0, Relaxed);
    if woke_at > 0 {
      c.scheduled.fetch_add(1, Relaxed);
      c.scheduled_ns.fetch_add(start_ns.saturating_sub(woke_at), Relaxed);
    }

    // 2. 存下真正的 waker，用自己的 waker 去 poll 内层
    {
      let mut real = this.state.real.lock().unwrap_or_else(|e| e.into_inner());
      if !real.as_ref().is_some_and(|w| w.will_wake(cx.waker())) {
        *real = Some(cx.waker().clone());
      }
    }
    let mut inner_cx = Context::from_waker(&this.waker);
    let res = this.fut.as_mut().poll(&mut inner_cx);

    // 3. poll 时长，按阈值分快慢
    let d = start.elapsed();
    c.polls.fetch_add(1, Relaxed);
    c.poll_ns.fetch_add(d.as_nanos() as u64, Relaxed);
    if d >= shared.slow {
      c.slow_polls.fetch_add(1, Relaxed);
    }
    res
  }
}

impl Monitor {
  pub fn snapshot(&self) -> Snapshot {
    let c = &self.shared.counters;
    Snapshot {
      polls: c.polls.load(Relaxed),
      poll_ns: c.poll_ns.load(Relaxed),
      slow_polls: c.slow_polls.load(Relaxed),
      scheduled: c.scheduled.load(Relaxed),
      scheduled_ns: c.scheduled_ns.load(Relaxed),
    }
  }

  pub fn intervals(&self) -> Intervals {
    Intervals { monitor: self.clone(), prev: Snapshot::default() }
  }
}

#[derive(Debug, Default, Clone, Copy)]
pub struct Snapshot {
  pub polls: u64,
  pub poll_ns: u64,
  pub slow_polls: u64,
  pub scheduled: u64,
  pub scheduled_ns: u64,
}

impl Snapshot {
  pub fn mean_poll(&self) -> Duration {
    Duration::from_nanos(self.poll_ns.checked_div(self.polls).unwrap_or(0))
  }
  pub fn mean_scheduled(&self) -> Duration {
    Duration::from_nanos(self.scheduled_ns.checked_div(self.scheduled).unwrap_or(0))
  }
  fn delta(&self, prev: &Snapshot) -> Snapshot {
    Snapshot {
      polls: self.polls.wrapping_sub(prev.polls),
      poll_ns: self.poll_ns.wrapping_sub(prev.poll_ns),
      slow_polls: self.slow_polls.wrapping_sub(prev.slow_polls),
      scheduled: self.scheduled.wrapping_sub(prev.scheduled),
      scheduled_ns: self.scheduled_ns.wrapping_sub(prev.scheduled_ns),
    }
  }
}

/// 读的一方：记住上一次，返回差。
pub struct Intervals {
  monitor: Monitor,
  prev: Snapshot,
}

impl Iterator for Intervals {
  type Item = Snapshot;
  fn next(&mut self) -> Option<Snapshot> {
    let now = self.monitor.snapshot();
    let d = now.delta(&self.prev);
    self.prev = now;
    Some(d)
  }
}

use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::{Duration, Instant};

use mini::Monitor;

fn spin(d: Duration) {
  let t = Instant::now();
  while t.elapsed() < d {}
}

/// 手写的让出：唤醒自己、返回一次 Pending
struct WakeSelf(bool);
impl Future for WakeSelf {
  type Output = ();
  fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
    if self.0 {
      return Poll::Ready(());
    }
    self.0 = true;
    cx.waker().wake_by_ref();
    Poll::Pending
  }
}

async fn victim(yield_now: bool) {
  for _ in 0..10 {
    if yield_now {
      tokio::task::yield_now().await;
    } else {
      WakeSelf(false).await;
    }
  }
}

async fn hog() {
  for _ in 0..10 {
    spin(Duration::from_millis(5));
    WakeSelf(false).await;
  }
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
  let slow = Duration::from_micros(50);

  // A：三次「忙 5ms，再睡 1ms」
  let m = Monitor::new(slow);
  let mut iv = m.intervals();
  m.instrument(async {
    for _ in 0..3 {
      spin(Duration::from_millis(5));
      tokio::time::sleep(Duration::from_millis(1)).await;
    }
  })
  .await;
  let a = iv.next().unwrap();
  println!("A mini: polls={} slow={} mean_poll={:?} scheduled={} mean_sched={:?}", a.polls, a.slow_polls, a.mean_poll(), a.scheduled, a.mean_scheduled());

  // B：同一个线程上，hog 每次 poll 忙 5ms，victim 让出 10 次
  for yield_now in [false, true] {
    let label = if yield_now { "yield_now" } else { "wake_self" };
    let m = Monitor::new(slow);
    let h = tokio::spawn(hog());
    let v = tokio::spawn(m.instrument(victim(yield_now)));
    let _ = tokio::join!(h, v);
    let s = m.snapshot();
    println!("B mini   {label}: polls={} scheduled={} mean_sched={:?}", s.polls, s.scheduled, s.mean_scheduled());

    let tm = tokio_metrics::TaskMonitor::new();
    let h = tokio::spawn(hog());
    let v = tokio::spawn(tm.instrument(victim(yield_now)));
    let _ = tokio::join!(h, v);
    let t = tm.cumulative();
    println!(
      "B tokio-metrics {label}: polls={} scheduled={} mean_sched={:?} idled={} mean_idle={:?}",
      t.total_poll_count, t.total_scheduled_count, t.mean_scheduled_duration(), t.total_idled_count, t.mean_idle_duration()
    );
  }
}

A mini: polls=4 slow=3 mean_poll=3.772739ms scheduled=3 mean_sched=333ns
B mini   wake_self: polls=11 scheduled=10 mean_sched=4.501349ms
B tokio-metrics wake_self: polls=11 scheduled=10 mean_sched=4.501766ms idled=0 mean_idle=0ns
B mini   yield_now: polls=11 scheduled=10 mean_sched=433ns
B tokio-metrics yield_now: polls=11 scheduled=10 mean_sched=624ns idled=10 mean_idle=4.517054ms