反应器

/// I/O driver, backed by Mio.
pub(crate) struct Driver {
    /// True when an event with the signal token is received
    signal_ready: bool,

    /// Reuse the `mio::Events` value across calls to poll.
    events: mio::Events,

    /// Buffer for polls that do not wait, if
    /// `Builder::max_io_events_per_busy_tick` is set.
    events_busy: Option<mio::Events>,

    /// The system event queue.
    poll: mio::Poll,
}

/// A reference to an I/O driver.
pub(crate) struct Handle {
    /// Registers I/O resources.
    registry: mio::Registry,

    /// Tracks all registrations
    registrations: RegistrationSet,

    /// State that should be synchronized
    synced: Mutex<registration_set::Synced>,

    /// Used to wake up the reactor from a call to `turn`.
    /// Not supported on `Wasi` due to lack of threading support.
    #[cfg(not(target_os = "wasi"))]
    waker: mio::Waker,

    pub(crate) metrics: IoDriverMetrics,
    // ...
}

const TOKEN_WAKEUP: mio::Token = mio::Token(0);
const TOKEN_SIGNAL: mio::Token = mio::Token(1);

fn turn(&mut self, handle: &Handle, max_wait: Option<Duration>) {
    debug_assert!(!handle.registrations.is_shutdown(&handle.synced.lock()));

    handle.release_pending_registrations();

    // A poll that does not wait takes the busy batch. Events it leaves
    // behind stay queued in the kernel, so the next poll returns them.
    let events = match (&mut self.events_busy, max_wait) {
        (Some(busy), Some(wait)) if wait.is_zero() => busy,
        _ => &mut self.events,
    };

    // Block waiting for an event to happen, peeling out how many events
    // happened.
    match self.poll.poll(events, max_wait) {
        Ok(()) => {}
        Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
        // ...
        Err(e) => panic!("unexpected error when polling the I/O driver: {e:?}"),
    }

    // Process all the events that came in, dispatching appropriately
    let mut ready_count = 0;
    for event in events.iter() {
        let token = event.token();

        if token == TOKEN_WAKEUP {
            // Nothing to do, the event is used to unblock the I/O driver
        } else if token == TOKEN_SIGNAL {
            self.signal_ready = true;
        } else {
            let ready = Ready::from_mio(event);

            // Safety: we ensure that the pointers used as tokens are not freed
            // until they are both deregistered from mio **and** we know the I/O
            // driver is not concurrently polling. The I/O driver holds ownership of
            // an `Arc<ScheduledIo>` so we can safely cast this to a ref.
            unsafe { dispatch_event(token, ready) };

            ready_count += 1;
        }
    }

    // ...

    handle.metrics.incr_ready_count_by(ready_count);
}

pub(crate) fn token(&self) -> mio::Token {
    mio::Token(super::EXPOSE_IO.expose_provenance(self))
}

pub(super) unsafe fn dispatch_event(token: mio::Token, ready: Ready) {
    let ptr = super::EXPOSE_IO.from_exposed_addr(token.0);

    // SAFETY: the caller guarantees that the token refers to a live `ScheduledIo`.
    let io: &ScheduledIo = unsafe { &*ptr };

    io.set_readiness(Tick::Set, |curr| curr | ready);
    io.wake(ready);
}

pub(super) struct RegistrationSet {
    num_pending_release: AtomicUsize,
}

pub(super) struct Synced {
    // True when the I/O driver shutdown. At this point, no more registrations
    // should be added to the set.
    is_shutdown: bool,

    // List of all registrations tracked by the set
    registrations: LinkedList<Arc<ScheduledIo>>,

    // Registrations that are pending drop. When a `Registration` is dropped, it
    // stores its `ScheduledIo` in this list. The I/O driver is responsible for
    // dropping it. This ensures the `ScheduledIo` is not freed while it can
    // still be included in an I/O event.
    pending_release: Vec<Arc<ScheduledIo>>,
}

let scheduled_io = self.registrations.allocate(&mut self.synced.lock())?;
let token = scheduled_io.token();

// we should remove the `scheduled_io` from the `registrations` set if registering
// the `source` with the OS fails. Otherwise it will leak the `scheduled_io`.
if let Err(e) = self.registry.register(source, token, interest.to_mio()) {
    // safety: `scheduled_io` is part of the `registrations` set.
    unsafe {
        self.registrations
            .remove(&mut self.synced.lock(), &scheduled_io)
    };

    return Err(e);
}

// Deregister the source with the OS poller **first**.
//
// If `deregister` fails (for example, because the caller closed the raw
// file descriptor before dropping `AsyncFd`), we must NOT release the
// `ScheduledIo` registration. On Linux, if the file descriptor was
// duplicated before the original fd number was closed, the open file
// description remains registered in epoll with the `ScheduledIo` token
// pointer even though `EPOLL_CTL_DEL` fails with `EBADF`.
self.registry.deregister(source)?;

if self
    .registrations
    .deregister(&mut self.synced.lock(), registration)
{
    self.unpark();
}

pub(super) fn deregister(&self, synced: &mut Synced, registration: &Arc<ScheduledIo>) -> bool {
    synced.pending_release.push(registration.clone());

    let len = synced.pending_release.len();
    self.num_pending_release.store(len, Release);

    len == NOTIFY_AFTER
}

use std::io::Write;
use std::sync::atomic::{AtomicU64, Ordering::SeqCst};
use std::sync::Arc;
use std::task::{Wake, Waker};
use std::time::Duration;
use tokio::net::TcpListener;

// 已经 poll 过的「忙任务」个数
static POLLS: AtomicU64 = AtomicU64::new(0);

// 一个忙任务：自增计数，空转一小会儿，再 spawn 一个自己，然后结束。
// 运行队列里因此一直有活，worker 永远不会因为没事做而 park。
fn busy() {
  tokio::spawn(async {
    POLLS.fetch_add(1, SeqCst);
    let t = std::time::Instant::now();
    while t.elapsed() < Duration::from_micros(20) {}
    busy();
  });
}

struct Mark(Waker, Arc<AtomicU64>);
impl Wake for Mark {
  fn wake(self: Arc<Self>) {
    self.wake_by_ref();
  }
  fn wake_by_ref(self: &Arc<Self>) {
    self.1.store(POLLS.load(SeqCst), SeqCst);
    self.0.wake_by_ref();
  }
}

fn run(interval: u32) {
  let rt = tokio::runtime::Builder::new_current_thread()
    .event_interval(interval)
    .enable_all()
    .build()
    .unwrap();
  let mut gaps = Vec::new();
  rt.block_on(async {
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    let peer = std::thread::spawn(move || std::net::TcpStream::connect(addr).unwrap());
    let (s, _) = listener.accept().await.unwrap();
    let mut peer = peer.join().unwrap();
    peer.set_nodelay(true).unwrap();
    let s = Arc::new(s);
    let mut b = [0u8; 8];
    let _ = s.try_read(&mut b); // 清掉 accept 时假定的可读
    busy();
    // 写端是一个普通线程：收到信号后等 1ms，记下当时已经跑了多少个忙任务，再写 1 字节
    let (go_tx, go_rx) = std::sync::mpsc::channel::<()>();
    let (at_tx, at_rx) = std::sync::mpsc::channel::<u64>();
    std::thread::spawn(move || {
      for () in go_rx {
        std::thread::sleep(Duration::from_millis(1));
        let at = POLLS.load(SeqCst);
        peer.write_all(b"x").unwrap();
        at_tx.send(at).unwrap();
      }
    });
    for _ in 0..200 {
      let woke_at = Arc::new(AtomicU64::new(0));
      let s2 = s.clone();
      let w2 = woke_at.clone();
      let reader = tokio::spawn(async move {
        std::future::poll_fn(|cx| {
          let waker = Waker::from(Arc::new(Mark(cx.waker().clone(), w2.clone())));
          let mut cx2 = std::task::Context::from_waker(&waker);
          s2.poll_read_ready(&mut cx2)
        })
        .await
        .unwrap();
        POLLS.load(SeqCst)
      });
      go_tx.send(()).unwrap();
      let ran_at = reader.await.unwrap();
      let wrote_at = at_rx.recv().unwrap();
      let woke = woke_at.load(SeqCst);
      assert!(woke >= wrote_at, "读者没有经过 wake 就返回了");
      gaps.push((woke - wrote_at, ran_at - woke));
      // 读走这 1 字节，再读一次拿到 WouldBlock，把就绪位清掉，下一轮才会真的挂起
      let mut b = [0u8; 8];
      assert_eq!(s.try_read(&mut b).unwrap(), 1);
      assert!(s.try_read(&mut b).is_err());
    }
  });
  let max_wake = gaps.iter().map(|g| g.0).max().unwrap();
  let min_wake = gaps.iter().map(|g| g.0).min().unwrap();
  let avg_wake = gaps.iter().map(|g| g.0).sum::<u64>() as f64 / gaps.len() as f64;
  let max_run = gaps.iter().map(|g| g.1).max().unwrap();
  println!(
    "event_interval = {interval:>2}: 写入到 wake 之间跑了 {min_wake}..={max_wake} 个任务（平均 {avg_wake:.1}），wake 到读者被 poll 之间最多 {max_run} 个"
  );
}

fn main() {
  run(61);
  run(8);
}

event_interval = 61: 写入到 wake 之间跑了 1..=62 个任务（平均 46.5），wake 到读者被 poll 之间最多 1 个
event_interval =  8: 写入到 wake 之间跑了 0..=9 个任务（平均 4.8），wake 到读者被 poll 之间最多 1 个