Sleep 的一生

pub fn sleep(duration: Duration) -> Sleep {
    let deadline = Instant::now() + safe_delay(duration);
    Sleep::new_timeout(deadline, trace::caller_location())
}

pub(crate) fn new_timeout(
    deadline: Instant,
    location: Option<&'static Location<'static>>,
) -> Sleep {
    let handle = scheduler::Handle::current();
    // Panic if the time driver is not enabled (backwards compat)
    _ = handle.driver().time();
    // ...
    Sleep {
        deadline,
        driver: handle,
        inner,
        timer: None,
    }
}

pub struct Sleep {
    deadline: Instant,
    driver: scheduler::Handle,
    inner: Inner,
    #[pin]
    timer: Option<Timer>,
}

fn poll_elapsed(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Result<(), Error>> {
    ready!(crate::trace::trace_leaf(cx));
    let mut this = self.project();

    // ...
    let coop = ready!(crate::task::coop::poll_proceed(cx));

    let timer = match this.timer.as_mut().as_pin_mut() {
        Some(timer) => timer,
        None => {
            let handle = this.driver;
            let time_source = handle.driver().time().time_source();
            let deadline = time_source.deadline_to_tick(*this.deadline);

            // ...

            let timer = Timer::new(handle.clone(), deadline);
            this.timer.set(Some(timer));
            let mut timer = this.timer.as_pin_mut().unwrap();
            timer.as_mut().init(deadline);
            timer
        }
    };

    let result = timer.poll_elapsed(cx).map(move |r| {
        coop.made_progress();
        r
    });

    // ...
    return result;
}

pub(crate) struct TimerShared {
    pointers: linked_list::Pointers<TimerShared>,
    registered_when: AtomicU64,
    state: StateCell,
    _p: PhantomPinned,
}

pub(super) struct StateCell {
    state: AtomicU64,
    result: UnsafeCell<TimerResult>,
    waker: AtomicWaker,
}

pub(in crate::runtime::time) const STATE_DEREGISTERED: u64 = u64::MAX;
const STATE_PENDING_FIRE: u64 = STATE_DEREGISTERED - 1;
const STATE_MIN_VALUE: u64 = STATE_PENDING_FIRE;
/// The largest safe integer to use for ticks.
pub(super) const MAX_SAFE_MILLIS_DURATION: u64 = STATE_MIN_VALUE - 1;

fn poll(&self, waker: &Waker) -> Poll<TimerResult> {
    // We must register first. This ensures that either `fire` will
    // observe the new waker, or we will observe a racing fire to have set
    // the state, or both.
    self.waker.register_by_ref(waker);

    self.read_state()
}

fn read_state(&self) -> Poll<TimerResult> {
    let cur_state = self.state.load(Ordering::Acquire);

    if cur_state == STATE_DEREGISTERED {
        // SAFETY: The driver has fired this timer; this involves writing
        // the result, and then writing (with release ordering) the state
        // field.
        Poll::Ready(unsafe { self.result.with(|p| *p) })
    } else {
        Poll::Pending
    }
}

pub(crate) fn reset(self: Pin<&mut Self>, deadline: u64) {
    if self.inner.extend_expiration(deadline).is_ok() {
        return;
    }

    unsafe {
        self.driver()
            .reregister(&self.driver.driver().io, deadline, (&self.inner).into());
    }
}

unsafe fn fire(&self, result: TimerResult) -> Option<Waker> {
    // Quick initial check to see if the timer is already fired. ...
    let cur_state = self.state.load(Ordering::Relaxed);
    if cur_state == STATE_DEREGISTERED {
        return None;
    }

    // ...
    unsafe { self.result.with_mut(|p| *p = result) };

    self.state.store(STATE_DEREGISTERED, Ordering::Release);

    self.waker.take_waker()
}

pub(self) unsafe fn clear_entry(&self, entry: NonNull<TimerShared>) {
    unsafe {
        let mut lock = self.inner.lock();

        if entry.as_ref().might_be_registered() {
            lock.wheel.remove(entry);
        }

        let waker = entry.as_ref().handle().fire(Ok(()));
        drop(lock);
        drop(waker);
    }
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
  for d in [Duration::from_micros(1), Duration::from_micros(100), Duration::from_micros(1500), Duration::from_millis(2)] {
    let mut v = Vec::new();
    for _ in 0..5 {
      let t = Instant::now();
      sleep(d).await;
      v.push(format!("{:.3}", t.elapsed().as_secs_f64() * 1000.0));
    }
    println!("sleep({d:?}) 实际(ms): {}", v.join(" "));
  }

  // 连续 100 次 sleep(100µs)，总共睡了多久
  let t = Instant::now();
  for _ in 0..100 {
    sleep(Duration::from_micros(100)).await;
  }
  println!("100 x sleep(100µs) 总计 {:.1}ms", t.elapsed().as_secs_f64() * 1000.0);
}

sleep(1µs) 实际(ms): 1.169 1.144 1.087 1.134 1.143
sleep(100µs) 实际(ms): 1.140 1.140 1.143 1.137 1.131
sleep(1.5ms) 实际(ms): 3.386 3.387 2.260 3.387 2.273
sleep(2ms) 实际(ms): 3.225 3.403 3.393 3.396 3.389
100 x sleep(100µs) 总计 123.0ms

pub fn timeout<F>(duration: Duration, future: F) -> Timeout<F::IntoFuture>
where
    F: IntoFuture,
{
    // Closures don't preserve `#[track_caller]`.
    #[allow(clippy::manual_map)]
    Timeout {
        value: future.into_future(),
        delay: match Instant::now().checked_add(duration) {
            Some(deadline) => Some(Sleep::new_timeout(deadline, trace::caller_location())),
            None => None,
        },
    }
}

fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
    let me = self.project();

    let had_budget_before = coop::has_budget_remaining();

    // First, try polling the future
    if let Poll::Ready(v) = me.value.poll(cx) {
        return Poll::Ready(Ok(v));
    }

    match me.delay.as_pin_mut() {
        Some(delay) => poll_delay(had_budget_before, delay, cx).map(Err),
        None => Poll::Pending,
    }
}

fn poll_delay(
    had_budget_before: bool,
    delay: Pin<&mut Sleep>,
    cx: &mut task::Context<'_>,
) -> Poll<Elapsed> {
    let delay_poll = || match delay.poll(cx) {
        Poll::Ready(()) => Poll::Ready(Elapsed::new()),
        Poll::Pending => Poll::Pending,
    };

    let has_budget_now = coop::has_budget_remaining();

    if let (true, false) = (had_budget_before, has_budget_now) {
        // if it is the underlying future that exhausted the budget, we poll
        // the `delay` with an unconstrained one. ...
        coop::with_unconstrained(delay_poll)
    } else {
        delay_poll()
    }
}

struct NaiveTimeout<F> {
  inner: Pin<Box<F>>,
  delay: Pin<Box<Sleep>>,
}

impl<F: Future> Future for NaiveTimeout<F> {
  type Output = Result<F::Output, &'static str>;
  fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
    if let Poll::Ready(v) = self.inner.as_mut().poll(cx) {
      return Poll::Ready(Ok(v));
    }
    self.delay.as_mut().poll(cx).map(|_| Err("Elapsed"))
  }
}

async fn busy_reader() -> u64 {
  let (tx, mut rx) = mpsc::unbounded_channel::<u64>();
  let start = Instant::now();
  let mut n = 0u64;
  loop {
    tx.send(n).unwrap();
    rx.recv().await;
    n += 1;
    if start.elapsed() > Duration::from_millis(500) {
      return n;
    }
  }
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
  let start = Instant::now();
  let res = timeout(Duration::from_millis(50), busy_reader()).await;
  println!("timeout:      {:?}，用时 {}ms", res.map_err(|_| "Elapsed"), start.elapsed().as_millis());

  let start = Instant::now();
  let res = NaiveTimeout { inner: Box::pin(busy_reader()), delay: Box::pin(sleep(Duration::from_millis(50))) }.await;
  println!("NaiveTimeout: {:?}，用时 {}ms", res.map(|_| "内层自己结束"), start.elapsed().as_millis());
}

timeout:      Err("Elapsed")，用时 50ms
NaiveTimeout: Ok("内层自己结束")，用时 500ms

pub fn poll_tick(&mut self, cx: &mut Context<'_>) -> Poll<Instant> {
    // Wait for the delay to be done
    ready!(Pin::new(&mut self.delay).poll(cx));

    // Get the time when we were scheduled to tick
    let timeout = self.delay.deadline();

    let now = Instant::now();

    // ...
    let next = if now > timeout + Duration::from_millis(5) {
        self.missed_tick_behavior
            .next_timeout(timeout, now, self.period)
    } else {
        timeout + self.period
    };

    // When we arrive here, the internal delay returned `Poll::Ready`.
    // Reset the delay but do not register it. It should be registered with
    // the next call to [`poll_tick`].
    self.delay.as_mut().reset_without_timer(next);

    // Return the time when we were scheduled to tick
    Poll::Ready(timeout)
}

fn next_timeout(&self, timeout: Instant, now: Instant, period: Duration) -> Instant {
    match self {
        Self::Burst => timeout + period,
        Self::Delay => now + period,
        Self::Skip => {
            let offset = Duration::from_nanos(
                ((now - timeout).as_nanos() % period.as_nanos())
                    .try_into()
                    // ...
                    .expect(
                        "too much time has elapsed since the interval was supposed to tick",
                    ),
            );
            now + (period - offset)
        }
    }
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
  for behavior in [MissedTickBehavior::Burst, MissedTickBehavior::Delay, MissedTickBehavior::Skip] {
    let start = Instant::now();
    let mut iv = interval(Duration::from_millis(100));
    iv.set_missed_tick_behavior(behavior);
    let mut line = format!("{behavior:?}:");
    for i in 0..6 {
      let scheduled = iv.tick().await;
      let actual = start.elapsed().as_millis();
      line += &format!(" #{i} 计划{}/实际{}", (scheduled - start).as_millis(), actual);
      if i == 1 {
        // 第 2 次 tick 之后卡住 350ms，期间错过 3 个 tick
        std::thread::sleep(Duration::from_millis(350));
      }
    }
    println!("{line}");
  }
}

Burst: #0 计划0/实际1 #1 计划100/实际102 #2 计划200/实际453 #3 计划300/实际453 #4 计划400/实际453 #5 计划500/实际502
Delay: #0 计划0/实际1 #1 计划100/实际101 #2 计划200/实际455 #3 计划555/实际557 #4 计划655/实际657 #5 计划755/实际757
Skip: #0 计划0/实际1 #1 计划100/实际101 #2 计划200/实际456 #3 计划500/实际502 #4 计划600/实际602 #5 计划700/实际702

#[tokio::main(flavor = "current_thread", start_paused = true)]
async fn main() {
  let wall = std::time::Instant::now();
  let t = Instant::now();
  sleep(Duration::from_secs(3600)).await;
  println!("tokio 时钟走了 {:?}，墙上时间走了 {:?}", t.elapsed(), wall.elapsed());
}

tokio 时钟走了 3600s，墙上时间走了 65.375µs

10 x sleep(100µs): tokio 时钟走了 10ms
10 x sleep(1.5ms): tokio 时钟走了 20ms
10 x sleep(2ms): tokio 时钟走了 20ms