TaskMonitorCore 与 FutureMonitor

#[derive(Clone, Debug)]
pub struct TaskMonitor {
    base: Arc<TaskMonitorCore>,
}

impl Deref for TaskMonitor {
    type Target = TaskMonitorCore;
    // ...
}
// ...
#[derive(Debug)]
pub struct TaskMonitorCore {
    metrics: RawMetrics,
    record_scheduling_log: bool,
}

static BACKGROUND: TaskMonitorCore = TaskMonitorCore::new();

    pub fn instrument<F>(&'static self, task: F) -> Instrumented<F, &'static Self> {
        Self::instrument_with(task, self)
    }

// Deliberately not `Clone`: cloning would share the underlying span and let two
// instrumented futures be measured as one, which is exactly the reuse `instrument`
// (by consuming `self`) is meant to prevent.
#[derive(Debug, Default)]
pub struct FutureMonitor {
    monitor: TaskMonitor,
    span: Arc<SpanState>,
}

impl FutureMonitor {
    // ...
    pub fn instrument<F>(self, task: F) -> MonitoredFuture<F> {
        MonitoredFuture {
            inner: self.monitor.instrument(task),
            monitor: self.monitor,
            span: self.span,
        }
    }
}

impl<F: Future> Future for MonitoredFuture<F> {
    type Output = (F::Output, FutureMetrics);

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.project();
        // ...
        this.span
            .record_poll_start(TaskScheduling::try_current().unwrap_or_default());
        let ret = this.inner.poll(cx);
        // ...
        this.span.record_poll_end();
        match ret {
            Poll::Ready(output) => {
                let metrics = build_future_metrics(this.monitor, this.span);
                Poll::Ready((output, metrics))
            }
            Poll::Pending => Poll::Pending,
        }
    }
}

thread_local! {
    /// The scheduling log of the innermost [`Instrumented`] task currently being
    /// polled on this thread, if any.
    ///
    /// This is a hand-rolled task-local built on a `thread_local!`: it is only
    /// set for the duration of a single synchronous poll (see
    /// [`SchedulingLogGuard`]) and never held across an `.await`, so the usual
    /// hazard of thread-locals in async code does not apply.
    static CURRENT_SCHEDULING_LOG: RefCell<Option<Arc<SchedulingLog>>> =
        const { RefCell::new(None) };
}

use std::time::{Duration, Instant};
use tokio_metrics::{FutureMonitor, TaskMonitor, TaskMonitorCore, TaskScheduling};

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

// const fn 构造，可以直接放进 static，不需要 OnceLock / lazy_static
static BACKGROUND: TaskMonitorCore = TaskMonitorCore::new();

// 一个「请求」：先在阻塞线程池上查一次「数据库」（2ms），再算 1ms
async fn handle(id: u32) -> u32 {
  tokio::task::spawn_blocking(|| std::thread::sleep(Duration::from_millis(2))).await.unwrap();
  spin(Duration::from_millis(1));
  id
}

async fn serve(root: TaskMonitor, label: &'static str) {
  let task = tokio::spawn(root.instrument(async move {
    println!("{label}: try_current inside = {:?}", TaskScheduling::try_current().is_some());
    for id in 0..2 {
      let (out, m) = FutureMonitor::new().instrument(handle(id)).await;
      println!(
        "{label}: req {out} polls={} slow={} idle={:.1?} total={:.1?} scheduled={} sched={:.1?}",
        m.poll_count, m.slow_poll_count, m.total_idle_duration, m.total_duration,
        m.scheduled_count, m.total_scheduled_duration
      );
    }
  }));
  // 同一个线程上再放一个忙 5ms 的任务，制造排队
  tokio::spawn(async { spin(Duration::from_millis(5)) });
  task.await.unwrap();
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
  tokio::spawn(BACKGROUND.instrument(async { spin(Duration::from_micros(100)) })).await.unwrap();
  println!("static core: polls={} slow={}", BACKGROUND.cumulative().total_poll_count, BACKGROUND.cumulative().total_slow_poll_count);
  println!("outside any task: try_current = {:?}", TaskScheduling::try_current());

  let mut b = TaskMonitor::builder();
  b.publish_scheduling_delay();
  serve(b.build(), "opt-in ").await;
  serve(TaskMonitor::new(), "default").await;
}

static core: polls=1 slow=1
outside any task: try_current = None
opt-in : try_current inside = true
opt-in : req 0 polls=2 slow=1 idle=2.6ms total=6.0ms scheduled=1 sched=2.5ms
opt-in : req 1 polls=2 slow=1 idle=2.5ms total=3.5ms scheduled=1 sched=4.8µs
default: try_current inside = false
default: req 0 polls=2 slow=1 idle=2.5ms total=6.0ms scheduled=0 sched=0.0ns
default: req 1 polls=2 slow=1 idle=2.5ms total=3.5ms scheduled=0 sched=0.0ns