Rust in Everything
基于 tokio 1.54.1

select! 的取消安全：分支输了，数据丢了

主线位置：主线之外，用户层 API。select! 在当前任务里同时 poll 几个 future，一个完成，其余的被 drop。

这个问题没有报错。表现是：select! 循环里偶尔少一条消息、少半行数据；或者某个分支永远等不到结果。

flowchart TB
  lp[循环开始] --> mk[创建各分支的 future]
  mk --> sel[select! poll]
  sel -->|分支 A 先完成| run[执行 A 的处理代码]
  run --> drop[分支 B 的 future 被 drop]
  drop --> lost{B 做到一半了吗}
  lost -->|是| bad[已读的数据丢失]
  lost -->|否| ok[无影响]
  drop --> lp

宏怎么展开、分支什么时候被 drop、tokio 文档对 cancel safety 的定义和哪些方法是 cancel safe 的清单，见 select! 与 join!。那一页用 channel 演示过「取到消息后被取消」。这里看另外两种常见形态。

触发它的代码

丢半行：read_line

对端把一行分两次发出，中间隔 30ms。读端在 select! 里 read_line，另一个分支是每 10ms 一次的心跳：

use std::time::Duration;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, DuplexStream};
use tokio::time::{interval, sleep};

// The peer sends one line in two pieces, 30ms apart, then a second line.
fn peer() -> DuplexStream {
  let (client, mut server) = tokio::io::duplex(64);
  tokio::spawn(async move {
    server.write_all(b"hello ").await.unwrap();
    sleep(Duration::from_millis(30)).await;
    server.write_all(b"world\nsecond line\n").await.unwrap();
  });
  client
}

async fn buggy() {
  let mut reader = BufReader::new(peer());
  let mut tick = interval(Duration::from_millis(10));
  let mut line = String::new();
  loop {
    tokio::select! {
      n = reader.read_line(&mut line) => {
        if n.unwrap() == 0 { break }
        println!("  got {line:?}");
        line.clear();
      }
      _ = tick.tick() => {} // heartbeat, metrics, shutdown check...
    }
  }
}

read_line inside select!:
  got "world\n"
  got "second line\n"

"hello " 没了。read_line 已经把它从 BufReader 里取走，还没等到换行符，心跳分支先完成，read_line 的 future 被 drop，读到的字节跟着一起销毁。下一轮新建的 read_line 从 "world\n" 开始读。

永远完成不了：每轮重建的长操作

一个要 50ms 才能完成的操作，和 10ms 的心跳一起放进 select!：

// Something that takes 50ms and must not be restarted (a handshake, a big download...)
async fn fetch_snapshot() -> &'static str {
  sleep(Duration::from_millis(50)).await;
  "snapshot"
}

let start = Instant::now();
let mut tick = interval(Duration::from_millis(10));
let mut ticks = 0;
loop {
  tokio::select! {
    s = fetch_snapshot() => { println!("in-loop future: got {s} at {:?}", start.elapsed()); break }
    _ = tick.tick() => {
      ticks += 1;
      if ticks == 30 { println!("in-loop future: gave up after {ticks} ticks ({:?})", start.elapsed()); break }
    }
  }
}

in-loop future: gave up after 30 ticks (293.364083ms)

每 10ms 心跳赢一次，fetch_snapshot() 被 drop，下一轮从头开始，永远走不到 50ms。这里没有数据丢失，但如果这个操作有副作用（发了请求、扣了配额），每一轮都会重复一次。

源码里在哪

问题不在 select!，在被 drop 的 future 把状态存在了哪里。

read_line 创建 future 时，把调用方的 String 整个拿走，读到的字节存在 future 自己的 ReadLine 结构里：

pub(crate) fn read_line<'a, R>(reader: &'a mut R, string: &'a mut String) -> ReadLine<'a, R>
where
    R: AsyncBufRead + ?Sized + Unpin,
{
    ReadLine {
        reader,
        buf: mem::take(string).into_bytes(),
        output: string,
        read: 0,
        _pin: PhantomPinned,
    }
}

只有读到换行、future 完成时，才把结果写回 output。中途被 drop，buf 随 future 销毁，调用方的 String 是空的，BufReader 里那部分数据已经被消费。它的文档写明了这一点：「This method is not cancel safe. ... some data may have been partially read, and this data is lost.」

对照 Lines：读了一半的字节存在 Lines 结构体自己的 buf 里，不在 future 里。next_line 的 future 被 drop，Lines 还在，下次调用接着读。所以它的文档写的是「This method is cancel safe.」

📝NOTE
冷知识：Lines 的 buf 到 1.54.0 才是字节

1.53 及更早，Lines 有三个字段 buf: String、bytes: Vec<u8>、read，通过 read_line 的内部函数往 String 里攒。进度一样存在 Lines 里，取消安全不受影响；出问题的是 IO 错误。#8400（1.54.0）的说明写道：一行读到一半遇到 IO 错误，如果这半行恰好断在一个多字节 UTF-8 字符中间，String 里放不下它；下一次 poll 会触发 debug_assert!(output.is_empty())，或者在 put_back_original_data 里 vector.len() - num_bytes_read 下溢，release 构建也会在后面的 .expect 上 panic。改法是 Lines 只留一个 buf: Vec<u8>，像 Split 那样直接调 read_until_internal，整行读完才转成 String。

规律是：future 被 drop 时，已经产生的进度存在哪里。存在 future 外面（Lines、channel、BufReader 的缓冲区）就安全；存在 future 里面（局部变量、ReadLine 的 buf、read_exact 已经填了多少字节的计数）就会丢。

❗IMPORTANT
动脑筋：换成 read_until，第一行打印什么？

把上面 buggy 里的 read_line 换成 read_until，其他不变（同一个 peer()，同样 10ms 的心跳，#[tokio::main]）：

let mut reader = BufReader::new(peer());
let mut tick = interval(Duration::from_millis(10));
let mut buf = Vec::new();
loop {
  tokio::select! {
    n = reader.read_until(b'\n', &mut buf) => {
      let n = n.unwrap();
      if n == 0 { break }
      println!("n = {n}, buf = {:?}", String::from_utf8_lossy(&buf));
      buf.clear();
    }
    _ = tick.tick() => {}
  }
}

"hello " 还会丢吗？第一行的 n 是几？答案在本页末尾。

怎么发现

在代码里找满足下面两条的 select!：

  1. 在循环里（loop、while）。只执行一次的 select!，输掉的分支本来就不要了。
  2. 分支里的 future 是在 select! 里现场创建的，而且是下面之一：
    • 自己写的 async fn 或 async 块，里面有不止一个 .await。第一个 .await 完成后、最后一个完成前被 drop，前面的效果就丢了。
    • tokio 文档里 # Cancel safety 一节写着 not cancel safe 的方法：read_exact、read_to_end、read_to_string、read_line、write_all，以及各种锁和信号量的等待（丢的是排队位置）。完整清单在 select! 与 join! 的表里。

查不到文档的第三方方法，按「中途被 drop 会不会丢进度」自己判断，判断不了就按不安全处理。

怎么修

换成 cancel safe 的方法

read_line 换成 lines().next_line()，进度存在 Lines 里：

async fn with_lines() {
  let mut lines = BufReader::new(peer()).lines();
  let mut tick = interval(Duration::from_millis(10));
  loop {
    tokio::select! {
      l = lines.next_line() => match l.unwrap() {
        Some(l) => println!("  got {l:?}"),
        None => break,
      },
      _ = tick.tick() => {}
    }
  }
}

lines().next_line() inside select!:
  got "hello world"
  got "second line"

类似的替换：read_exact 换成 read/read_buf 加自己的缓冲区和计数；write_all 换成 write 加自己的偏移量；按帧读写用 tokio_util::codec 的 FramedRead，进度存在 Framed 结构里。

把 future 固定在循环外面

操作不能重来，就在循环外创建一次，用 tokio::pin! 固定，分支里写 &mut fut。select! drop 的只是这个引用，future 本身留着，下一轮接着 poll：

let start = Instant::now();
let mut tick = interval(Duration::from_millis(10));
let mut ticks = 0;
let fut = fetch_snapshot();
tokio::pin!(fut);
loop {
  tokio::select! {
    s = &mut fut => { println!("pinned future:  got {s} at {:?} after {ticks} ticks", start.elapsed()); break }
    _ = tick.tick() => ticks += 1,
  }
}

pinned future:  got snapshot at 53.590084ms after 6 ticks

它完成后不能再 poll。循环还要继续时，要么 break，要么换一个新的 future（fut.set(..)），要么给分支加 if !done 前置条件把它禁用。

这种写法要求 future 不借用循环里还要用的变量。read_line(&mut line) 借用了 reader 和 line，固定在外面以后，循环体里就不能再碰它们，所以 read_line 更适合上一种改法。

把工作挪进独立任务

读写放进一个单独的任务，它自己按顺序 .await，不参与 select!；结果通过 channel 送出来。主循环 select! 的是 rx.recv()，它是 cancel safe 的：

async fn with_task() {
  let (tx, mut rx) = tokio::sync::mpsc::channel(16);
  tokio::spawn(async move {
    let mut reader = BufReader::new(peer());
    let mut line = String::new();
    while reader.read_line(&mut line).await.unwrap() > 0 {
      if tx.send(std::mem::take(&mut line)).await.is_err() { break }
    }
  });
  let mut tick = interval(Duration::from_millis(10));
  loop {
    tokio::select! {
      l = rx.recv() => match l {
        Some(l) => println!("  got {l:?}"),
        None => break,
      },
      _ = tick.tick() => {}
    }
  }
}

reader task + mpsc:
  got "hello world\n"
  got "second line\n"

三段 read_line 相关的输出来自同一次运行（buggy、with_lines、with_task 依次调用），两次运行结果一致。

选哪个

情况改法
有现成的 cancel safe 版本（next_line、read、recv、FramedRead）直接换
一个长操作，必须完成且只做一次tokio::pin! 固定在循环外，&mut fut
一串有顺序的异步步骤（读一帧、解析、回写）放进独立任务，用 channel 和主循环通信
程序正在关闭，丢掉半截数据无所谓不用改，取消就是想要的效果

动脑筋答案

程序在 exp-hf-pitfalls/src/bin/read_until_select.rs，M2 Max、rustc 1.99.0，连跑 3 次，输出一致：

n = 6, buf = "hello world\n"
n = 12, buf = "second line\n"
  • "hello " 没丢。 read_until_internal 每读到一段就 buf.extend_from_slice(..)，直接追加进调用方传进来的 Vec<u8>。心跳赢了、future 被 drop，"hello " 已经在 buf 里了。下一轮新建的 read_until 接着往同一个 buf 后面追加 "world\n"。read_line 做不到这一点，是因为它要保证 String 里永远是合法的 UTF-8，读了一半的字节只能先放在 future 自己的 buf 里。
  • n = 6，不是 12。 返回值是 future 里的计数器 read，它随第一个 future 一起被 drop 了，新 future 从 0 开始，只数到了 "world\n" 的 6 个字节。read_until 的文档在 # Cancel safety 一节写了这一点：取消后再调用，计数器会重置。所以在 select! 里用它，要看 buf.len()，不能看返回值。
  • n = 12：第二行没有被打断，一次读完 "second line\n"。

小结

  • select! 输掉的分支的 future 会被 drop。future 里存着的进度随之消失，下一轮从头开始。
  • 典型后果有两种：读了一半的数据丢失（read_line、read_exact、自己写的多步 async fn），以及长操作每轮重来、永远完成不了。
  • 判断一个操作是否 cancel safe，看它的进度存在 future 里面还是外面；tokio 的方法在文档 # Cancel safety 一节写明。
  • 改法：换 cancel safe 的方法；把 future 用 tokio::pin! 固定在循环外；或者把有顺序的工作挪进独立任务，select! 只等 channel。

💬 关联讨论

发起讨论
加载中...