Rust in Everything
基于 async-stream 0.3.6

为什么需要它

主线位置：主线之前。本页不看 async-stream 的源码，只看它要解决的问题，以及同一个问题的另外几种解法各自在什么时候执行。
flowchart TB
  need["想写一个 Stream"] --> manual["手写 poll_next：自己的状态机"]
  need --> unfold["futures::stream::unfold：状态进出闭包"]
  need --> chan["通道 + spawn + ReceiverStream"]
  need --> macro["stream!：在 async 块里写 yield"]
  manual --> lazy["消费者 poll 时才执行"]
  unfold --> lazy
  macro --> lazy
  chan --> eager["另一个任务提前执行，跑到通道满"]

问题：async fn 只能返回一个值

async fn 把「一段会等待的代码」变成一个 Future，编译器替你生成状态机：每个 .await 是一个状态，局部变量存在状态机里。但 Future 只产出一个值。要产出一串值，得实现 Stream：

pub trait Stream {
    type Item;
    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>>;
}

poll_next 是同步函数，不能 .await。每次被调用，它要从上次停下的地方接着干，所以「上次停在哪里」「局部变量是什么」都得由你自己存在 Self 里，等待中的 future 也要自己存着、自己 poll。这正是 async fn 替你做掉的那部分工作，到了 Stream 这里又得手写一遍。

Rust 的语言层面有生成器（coroutine），可以写 yield，但在 stable 上还不能用：0.3.6 的 tests/ui 里，编译器对裸 yield 的报错就是「yield can only be used in #[coroutine] closures, or gen blocks」。async-stream 在 stable 上用宏把 yield 改写掉，让你像写 async fn 一样写 Stream。原生 gen 块的进展放到 03 设计取舍再说。

四种写法

下面把同一个流写四遍：依次产出 5 个值，每个值要先 await 一个异步函数 fetch 才能拿到。fetch 产出时打印一行，好看清它在什么时候执行。

async-stream = "=0.3.6"
futures = "0.3"
tokio = { version = "=1.54.1", features = ["full"] }
tokio-stream = "0.1"

async fn fetch(tag: &'static str, i: u32) -> u32 {
  tokio::time::sleep(Duration::from_millis(1)).await;
  println!("  {tag}: produced {i}");
  i * 10
}

手写。状态是「下一个编号」加「正在等的那个 future」。fetch 返回的 future 类型写不出名字，只能装箱：

struct Manual {
  next: u32,
  in_flight: Option<Pin<Box<dyn Future<Output = u32> + Send>>>,
}

impl Stream for Manual {
  type Item = u32;

  fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<u32>> {
    if self.in_flight.is_none() {
      if self.next == 5 {
        return Poll::Ready(None);
      }
      let i = self.next;
      self.next += 1;
      self.in_flight = Some(Box::pin(fetch("manual", i)));
    }
    let fut = self.in_flight.as_mut().expect("set above");
    let v = std::task::ready!(fut.as_mut().poll(cx));
    self.in_flight = None;
    Poll::Ready(Some(v))
  }
}

stream!。状态就是局部变量 i 和 for 循环本身，fetch 直接 .await：

fn with_macro() -> impl Stream<Item = u32> {
  stream! {
    for i in 0..5 {
      yield fetch("stream!", i).await;
    }
  }
}

futures::stream::unfold。不用宏。状态显式地传进闭包，闭包返回一个 future，产出 (值, 新状态)：

fn with_unfold() -> impl Stream<Item = u32> {
  futures::stream::unfold(0u32, |i| async move {
    if i == 5 {
      return None;
    }
    Some((fetch("unfold", i).await, i + 1))
  })
}

通道 + spawn。生产者是另一个任务，把值塞进容量为 1 的 mpsc，消费者拿到的是 tokio-stream 的 ReceiverStream：

fn with_channel() -> impl Stream<Item = u32> {
  let (tx, rx) = mpsc::channel(1);
  tokio::spawn(async move {
    for i in 0..5 {
      if tx.send(fetch("channel", i).await).await.is_err() {
        break;
      }
    }
  });
  ReceiverStream::new(rx)
}

实验：谁在什么时候干活

消费者对每种写法只取第一项，然后去忙别的 30ms，再把流丢掉：

async fn take_one(name: &str, s: impl Stream<Item = u32>) {
  let mut s = Box::pin(s);
  println!("{name}: first = {:?}", s.next().await);
  // 消费者拿到一项之后去忙别的
  tokio::time::sleep(Duration::from_millis(30)).await;
  println!("{name}: done waiting");
}

#[tokio::main]
async fn main() {
  take_one("manual", Manual { next: 0, in_flight: None }).await;
  take_one("stream!", with_macro()).await;
  take_one("unfold", with_unfold()).await;
  take_one("channel", with_channel()).await;
}

输出（跑 3 次相同）：

  manual: produced 0
manual: first = Some(0)
manual: done waiting
  stream!: produced 0
stream!: first = Some(0)
stream!: done waiting
  unfold: produced 0
unfold: first = Some(0)
unfold: done waiting
  channel: produced 0
channel: first = Some(0)
  channel: produced 1
  channel: produced 2
channel: done waiting

前三种写法在消费者等待的 30ms 里什么也没做：没人 poll，流就不往前走。通道那种不一样，生产者是一个独立的任务，消费者拿走 0 之后它接着跑：1 放进通道的唯一空位，2 生产出来后卡在 send 上等空位。消费者不来取，它就停在那里，这是通道容量带来的背压。

sequenceDiagram
  participant C as 消费者
  participant S as stream! 的 async 块
  participant P as 生产者任务
  participant Q as mpsc（容量 1）
  Note over C,S: stream!：同一个任务里
  C->>S: next()：poll
  S-->>C: Ready(Some(0))
  Note over S: 没人 poll，停在 yield 之后
  Note over C,Q: 通道：两个任务
  P->>Q: send(0)
  Q-->>C: next() 拿到 0
  P->>Q: send(1)：占满空位
  P->>Q: send(2)：等空位

怎么选

手写stream!unfold通道 + spawn
本例的代码行数（不含空行）217811
状态放在哪Self 的字段async 块的局部变量闭包的参数和返回值生产者任务的局部变量
什么时候执行消费者 poll 时消费者 poll 时消费者 poll 时生产者任务被调度时，和消费者并发
额外开销本例里每项一次 Box没有分配没有分配一个任务、一个通道，值要跨任务搬运
依赖无proc-macrofutures-utiltokio 运行时、tokio-stream
能否借用调用方的局部变量能能能不能，spawn 要求 'static

几点说明：

  • 「没有分配」指流本身：stream! 把 async 块整个存在 AsyncStream 里，yield 也不经过任何堆上的缓冲区，具体见 02 运行时。上面的实验为了能统一写 take_one，用 Box::pin 把流装箱了一次，那是消费者自己的选择。
  • 惰性执行和并发执行没有好坏之分。希望生产和消费重叠（比如边下载边解析）时，通道的提前执行正是你要的；只想把一段逻辑包装成流交给别人消费时，惰性写法不多占一个任务，消费者丢掉流，生产也就停了。
  • unfold 适合状态简单的流。状态一多（几个计数器、一个缓冲区、一个连接），每次都要整个传进去再整个传出来，读起来就不如 stream! 里的普通局部变量。

tokio-stream 自己的文档在「从函数返回一个 Stream」一节里给的也是这两条路：要么把通道的发送端传进去、把接收端当流，要么用 async-stream（tokio-stream lib.rs 第 47–53 行）。它的 StreamMap 和 StreamExt::merge 的文档示例也直接用 async_stream::stream! 造输入流。

小结

  • Stream::poll_next 是同步函数，状态机要自己写；async fn 替你做的那些事，写 Stream 时又得手写一遍。
  • stream! 让你在 async 块里写 yield，状态就是局部变量；stable 上没有原生生成器，这是宏改写出来的。
  • 手写、stream!、unfold 都在消费者 poll 时才执行；通道 + spawn 的生产者在另一个任务里提前跑，跑到通道满为止。实测消费者只取一项时，前三种只生产了 0，通道版生产到了 2。
  • tokio-stream 的文档推荐的也是「通道」或「async-stream」两条路。

💬 关联讨论

发起讨论
加载中...