从 TcpStream::read 到挂起

pub(crate) unsafe fn poll_read<'a>(
    &'a self,
    cx: &mut Context<'_>,
    buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>>
where
    &'a E: io::Read + 'a,
{
    use std::io::Read;

    loop {
        let evt = ready!(self.registration.poll_read_ready(cx))?;

        let b = unsafe { &mut *(buf.unfilled_mut() as *mut [std::mem::MaybeUninit<u8>] as *mut [u8]) };

        // used only when the cfgs below apply
        #[allow(unused_variables)]
        let len = b.len();

        match self.io.as_ref().unwrap().read(b) {
            Ok(n) => {
                // When mio is using the epoll or kqueue selector, reading a partially full
                // buffer is sufficient to show that the socket buffer has been drained.
                //
                // This optimization does not work for level-triggered selectors such as
                // windows or when poll is used.
                //
                // Read more:
                // https://github.com/tokio-rs/tokio/issues/5866
                #[cfg(all(
                    // keep in sync with poll_write
                    not(mio_unsupported_force_poll_poll),
                    any(
                        // epoll
                        target_os = "android",
                        // ...
                        target_os = "macos",
                        // ...
                    )
                ))]
                if 0 < n && n < len {
                    self.registration.clear_readiness(evt);
                }

                // Safety: We trust `TcpStream::read` to have filled up `n` bytes in the
                // buffer.
                unsafe { buf.assume_init(n) };
                buf.advance(n);
                return Poll::Ready(Ok(()));
            },
            Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
                self.registration.clear_readiness(evt);
            }
            Err(e) => return Poll::Ready(Err(e)),
        }
    }
}

fn poll_io<R>(
    &self,
    cx: &mut Context<'_>,
    direction: Direction,
    mut f: impl FnMut() -> io::Result<R>,
) -> Poll<io::Result<R>> {
    loop {
        let ev = ready!(self.poll_ready(cx, direction))?;

        match f() {
            Ok(ret) => {
                return Poll::Ready(Ok(ret));
            }
            Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                self.clear_readiness(ev);
            }
            Err(e) => return Poll::Ready(Err(e)),
        }
    }
}

fn poll_ready(
    &self,
    cx: &mut Context<'_>,
    direction: Direction,
) -> Poll<io::Result<ReadyEvent>> {
    ready!(crate::trace::trace_leaf(cx));
    // Keep track of task budget
    let coop = ready!(crate::task::coop::poll_proceed(cx));
    let ev = ready!(self.shared.poll_readiness(cx, direction));

    if ev.is_shutdown {
        return Poll::Ready(Err(gone()));
    }

    coop.made_progress();
    Poll::Ready(Ok(ev))
}

pub(super) fn poll_readiness(
    &self,
    cx: &mut Context<'_>,
    direction: Direction,
) -> Poll<ReadyEvent> {
    let curr = self.readiness.load(Acquire);

    let ready = direction.mask() & Ready::from_usize(READINESS.unpack(curr));
    let is_shutdown = SHUTDOWN.unpack(curr) != 0;

    if ready.is_empty() && !is_shutdown {
        // Update the task info
        let mut waiters = self.waiters.lock();
        let waker = match direction {
            Direction::Read => &mut waiters.reader,
            Direction::Write => &mut waiters.writer,
        };

        // Avoid cloning the waker if one is already stored that matches the
        // current task.
        match waker {
            Some(waker) => waker.clone_from(cx.waker()),
            None => *waker = Some(cx.waker().clone()),
        }

        // Try again, in case the readiness was changed while we were
        // taking the waiters lock
        let curr = self.readiness.load(Acquire);
        let ready = direction.mask() & Ready::from_usize(READINESS.unpack(curr));
        let is_shutdown = SHUTDOWN.unpack(curr) != 0;
        if is_shutdown {
            Poll::Ready(ReadyEvent {
                tick: TICK.unpack(curr) as u16,
                ready: direction.mask(),
                is_shutdown,
            })
        } else if ready.is_empty() {
            Poll::Pending
        } else {
            Poll::Ready(ReadyEvent {
                tick: TICK.unpack(curr) as u16,
                ready,
                is_shutdown,
            })
        }
    } else {
        Poll::Ready(ReadyEvent {
            tick: TICK.unpack(curr) as u16,
            ready,
            is_shutdown,
        })
    }
}

use std::future::Future;
use std::io::Write;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll, Wake, Waker};
use std::time::{Duration, Instant};
use tokio::io::{AsyncRead, ReadBuf};
use tokio::net::{TcpListener, TcpStream};

// 包住任务的 waker：被调用时打印是哪个线程在叫醒
struct LoggingWaker {
  inner: Waker,
  start: Instant,
}

impl Wake for LoggingWaker {
  fn wake(self: Arc<Self>) {
    self.wake_by_ref();
  }
  fn wake_by_ref(self: &Arc<Self>) {
    let t = std::thread::current();
    println!(
      "{:>6.1?} wake() 在线程 {}",
      self.start.elapsed(),
      t.name().unwrap_or("?")
    );
    self.inner.wake_by_ref();
  }
}

// 手写的 read future：每次 poll 打印 poll_read 的结果
struct LoggedRead<'a> {
  stream: &'a mut TcpStream,
  buf: [u8; 64],
  polls: u32,
  start: Instant,
}

impl Future for LoggedRead<'_> {
  type Output = usize;
  fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<usize> {
    let me = &mut *self;
    me.polls += 1;
    let waker = Waker::from(Arc::new(LoggingWaker {
      inner: cx.waker().clone(),
      start: me.start,
    }));
    let mut cx2 = Context::from_waker(&waker);
    let mut rb = ReadBuf::new(&mut me.buf);
    let res = Pin::new(&mut *me.stream).poll_read(&mut cx2, &mut rb);
    let name = std::thread::current().name().unwrap_or("?").to_string();
    match res {
      Poll::Pending => {
        println!("{:>6.1?} 第 {} 次 poll：Pending（线程 {name}）", me.start.elapsed(), me.polls);
        Poll::Pending
      }
      Poll::Ready(r) => {
        r.unwrap();
        let n = rb.filled().len();
        println!(
          "{:>6.1?} 第 {} 次 poll：Ready，读到 {n} 字节 {:?}（线程 {name}）",
          me.start.elapsed(),
          me.polls,
          String::from_utf8_lossy(rb.filled())
        );
        Poll::Ready(n)
      }
    }
  }
}

fn main() {
  let rt = tokio::runtime::Builder::new_multi_thread()
    .worker_threads(1)
    .enable_all()
    .build()
    .unwrap();
  rt.block_on(async {
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    // 对端用 std 的阻塞 socket，放在普通线程里写，和运行时无关
    let peer = std::thread::spawn(move || std::net::TcpStream::connect(addr).unwrap());
    let (mut stream, _) = listener.accept().await.unwrap();
    let mut peer = peer.join().unwrap();

    let start = Instant::now();
    let reader = tokio::spawn(async move {
      let mut fut = LoggedRead { stream: &mut stream, buf: [0; 64], polls: 0, start };
      (&mut fut).await
    });

    std::thread::spawn(move || {
      std::thread::sleep(Duration::from_millis(50));
      println!("{:>6.1?} 对端写入 5 字节（普通线程）", start.elapsed());
      peer.write_all(b"hello").unwrap();
    });

    reader.await.unwrap();
  });
}

86.9µs 第 1 次 poll：Pending（线程 tokio-rt-worker）
52.2ms 对端写入 5 字节（普通线程）
52.4ms wake() 在线程 tokio-rt-worker
52.5ms 第 2 次 poll：Ready，读到 5 字节 "hello"（线程 tokio-rt-worker）

use std::io::Write;
use std::pin::Pin;
use std::task::Poll;
use std::time::Duration;
use tokio::io::{AsyncRead, ReadBuf};
use tokio::net::TcpListener;

#[tokio::main(flavor = "current_thread")]
async fn main() {
  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 (mut s, _) = listener.accept().await.unwrap();
  let mut peer = peer.join().unwrap();
  peer.write_all(&[7u8; 4096]).unwrap();
  tokio::time::sleep(Duration::from_millis(50)).await; // 等 4096 字节都进了接收缓冲区

  let reads_per_poll = tokio::spawn(async move {
    let mut total = 0;
    let mut log = Vec::new();
    std::future::poll_fn(|cx| {
      let mut ok = 0;
      loop {
        let mut b = [0u8; 8];
        let mut rb = ReadBuf::new(&mut b);
        match Pin::new(&mut s).poll_read(cx, &mut rb) {
          Poll::Ready(r) => {
            r.unwrap();
            ok += 1;
            total += rb.filled().len();
            if total == 4096 {
              log.push(ok);
              return Poll::Ready(());
            }
          }
          Poll::Pending => {
            log.push(ok);
            return Poll::Pending;
          }
        }
      }
    })
    .await;
    log
  })
  .await
  .unwrap();
  println!("每次 poll 读成功的次数：{reads_per_poll:?}");
}

每次 poll 读成功的次数：[128, 128, 128, 128]