Rust 异步编程中的取消安全性：CancellationToken 与 select! 的正确使用方法-拓冰建站

Rust 异步编程中的取消安全性：CancellationToken 与 select! 的正确使用方法

一、隐式取消的破坏性：当 Future 在 await 点被静默丢弃

Rust 的 async 生态中，取消（Cancellation）是隐式且普遍的。当一个select!分支完成时，其他分支的 Future 被直接 drop。这看似无害——Future 的 drop 会递归释放所有资源——但问题在于：被取消的 Future 可能已经完成了部分副作用。

具体场景：一个 Raft 节点在select!中同时等待心跳超时和AppendEntries消息。心跳超时先触发，触发选举，但AppendEntriesFuture 正在反序列化日志条目并写入磁盘。Future 被 drop 后，部分写入的日志条目留在磁盘上，破坏了日志的完整性。

这个问题被称为"取消安全问题"（Cancel Safety）。一个 Future 是取消安全的，当且仅当它在任意.await点被 drop 后，系统状态仍保持一致。大部分异步代码不是取消安全的，因为副作用操作（写磁盘、发网络包、修改共享状态）不可逆。

Tokio 的文档明确指出：TcpStream::read是取消安全的（读不到数据只是回到调用前状态），但mpsc::Sender::send不是（消息可能已被消费但未送出）。

二、CancellationToken 的传播模型

sequenceDiagram participant Root as Root Task participant CT as CancellationToken participant Child1 as Child Future A participant Child2 as Child Future B participant DB as Database Root->>CT: new() CT->>CT: create child_token() CT->>CT: create child_token() Root->>Child1: spawn(child_token_a) Root->>Child2: spawn(child_token_b) Child1->>DB: begin_transaction() DB-->>Child1: tx_handle Note over Root: 超时触发 Root->>CT: cancel() CT->>Child1: is_cancelled() == true Child1->>DB: rollback() DB-->>Child1: rolled back CT->>Child2: is_cancelled() == true Child2->>Child2: cleanup() Child1-->>Root: JoinHandle::await (cancelled) Child2-->>Root: JoinHandle::await (cancelled) Note over Root,DB: 系统恢复到一致状态

CancellationToken的设计遵循树状传播模型：根 token 被取消时，所有子 token 都会收到通知。这与 Rust 的所有权树一致——父任务拥有的 token 可以在 spawn 子任务时派生子 token。

关键语义：CancellationToken::cancelled()返回一个 Future，它在 token 被取消前永远 Pending。这允许在select!中优雅地处理取消：

tokio::select! { result = do_work() => { /* 正常完成 */ } _ = token.cancelled() => { /* 优雅取消 */ } }

这与直接 drop Future 的根本区别在于：被取消的 Future 有机会执行清理逻辑（回滚事务、关闭连接、释放锁）。

三、生产级取消安全模式实现

use std::sync::Arc; use tokio::sync::{Mutex, oneshot}; use tokio_util::sync::CancellationToken; /// 取消安全的 Raft 日志复制器 /// /// 设计原则： /// 1. 每个副作用操作前检查取消状态 /// 2. 使用事务语义确保操作原子性 /// 3. 取消时执行完整的回滚路径 struct CancelSafeLogReplicator { /// 取消令牌（外部注入，通常是 Raft 节点级别的） cancel_token: CancellationToken, /// 日志存储句柄 log_store: Arc<Mutex<Vec<LogEntry>>>, } impl CancelSafeLogReplicator { /// 复制一批日志条目到 Follower /// /// 取消安全性保证： /// - 检查点 A：序列化之前检查取消 /// - 检查点 B：发送网络请求之前检查取消 /// - 检查点 C：写入本地确认之前检查取消 /// 任一点被取消都不会留下不一致的中间状态 async fn replicate_entries( &self, entries: &[LogEntry], follower_addr: &str, ) -> Result<u64, ReplicationError> { // 检查点 A：序列化阶段 if self.cancel_token.is_cancelled() { return Err(ReplicationError::Cancelled); } // 序列化（纯 CPU，不会长时间阻塞） let payload = serde_json::to_vec(entries) .map_err(|e| ReplicationError::Serialization(e.to_string()))?; // 检查点 B：网络发送阶段 if self.cancel_token.is_cancelled() { return Err(ReplicationError::Cancelled); } // 使用 select! 在发送和取消之间竞争 let send_result = tokio::select! { result = Self::send_rpc(follower_addr, &payload) => result, _ = self.cancel_token.cancelled() => { // 取消时立即返回，不等待网络超时 return Err(ReplicationError::Cancelled); } }; match send_result { Ok(response) => { // 检查点 C：写入本地状态 if self.cancel_token.is_cancelled() { return Err(ReplicationError::Cancelled); } let mut store = self.log_store.lock().await; // 扩展日志，记录匹配索引 store.extend_from_slice(entries); Ok(response.match_index) } Err(e) => { // 发送失败：无需回滚（未修改本地状态） Err(ReplicationError::Network(e.to_string())) } } } async fn send_rpc(addr: &str, payload: &[u8]) -> Result<RpcResponse, Box<dyn std::error::Error>> { // 简化的 RPC 发送实现 tokio::time::sleep(std::time::Duration::from_millis(10)).await; Ok(RpcResponse { match_index: 42 }) } } struct LogEntry { index: u64, term: u64, data: Vec<u8>, } struct RpcResponse { match_index: u64, } #[derive(Debug)] enum ReplicationError { Cancelled, Serialization(String), Network(String), } /// 取消安全的分阶段操作模式 /// 使用状态机确保所有中间状态都是可回滚的 enum TwoPhaseCommitState { Prepare, Prepared { tx_id: u64 }, Committed, } struct CancelSafeTwoPc { state: TwoPhaseCommitState, cancel_token: CancellationToken, } impl CancelSafeTwoPc { async fn prepare(&mut self, tx_id: u64) -> Result<(), &'static str> { // 阶段 1：准备（可回滚） self.cancel_token.is_cancelled().then(|| Err("cancelled")).unwrap_or(Ok(()))?; self.state = TwoPhaseCommitState::Prepared { tx_id }; Ok(()) } async fn commit(&mut self) -> Result<(), &'static str> { // 阶段 2：提交（不可逆，仅在确认取消状态后执行） if self.cancel_token.is_cancelled() { self.state = TwoPhaseCommitState::Prepare; return Err("cancelled before commit"); } self.state = TwoPhaseCommitState::Committed; Ok(()) } } /// select! 中取消安全的并发操作 /// 使用 fuse() 避免已完成的 Future 被重复 poll async fn safe_select_example(token: CancellationToken) { let mut work = Box::pin(do_important_work()); loop { tokio::select! { result = &mut work => { println!("Work completed"); break; } _ = token.cancelled() => { println!("Cancelled, cleaning up..."); // work Future 被 drop，但我们在 drop 前执行了清理 // 注意：work 的 drop 必须也是取消安全的 break; } } } } async fn do_important_work() -> u64 { tokio::time::sleep(std::time::Duration::from_millis(100)).await; 42 } /// 生产环境中的取消传播模式 async fn run_with_timeout(token: CancellationToken, timeout_ms: u64) -> Result<u64, &'static str> { // 创建子 token，超时时自动取消 let child_token = token.child_token(); let result = tokio::time::timeout( std::time::Duration::from_millis(timeout_ms), async { // 定时检查取消状态 for i in 0..10 { if child_token.is_cancelled() { return Err("cancelled mid-operation"); } tokio::time::sleep(std::time::Duration::from_millis(10)).await; } Ok(42u64) } ).await; match result { Ok(inner) => inner, Err(_elapsed) => { child_token.cancel(); Err("timeout") } } } #[tokio::main] async fn main() { let token = CancellationToken::new(); let handle = tokio::spawn(run_with_timeout(token.clone(), 50)); // 模拟外部超时 tokio::time::sleep(std::time::Duration::from_millis(30)).await; token.cancel(); match handle.await.unwrap() { Ok(val) => println!("Result: {}", val), Err(e) => println!("Error: {}", e), } }

TwoPhaseCommitState状态机是解决取消安全问题的最优模式：将不可逆操作隔离在commit阶段，prepare阶段的所有操作都是可回滚的。在select!中被取消时，只需 drop 状态机，未到达Committed状态的操作自动视为未发生。

CancellationToken::child_token()的关键性质：子 token 取消不会影响父 token。这允许局部超时控制——某个子任务超时时只取消它自己，不影响其他并发任务。

四、取消安全的判定清单与常见反模式

判断一个 Future 是否取消安全：

它在.await后是否修改了外部共享状态？→ 不安全
它是否发送了网络请求且响应可能丢失？→ 通常不安全
它是否持有锁跨越了.await？→ 极其危险

常见不安全模式：

mpsc::Sender::send().await：消息可能已被消费但未入队，被取消后丢失
TcpStream::write_all().await：部分数据已写入，被取消后连接处于不确定状态
Mutex::lock().await在临界区内.await：被取消时锁已持有但后续操作未执行

修复策略（按优先级）：

使用CancellationToken的cancelled()替代直接 drop
在取消时执行回滚（需要事务状态机支持）
将副作用操作移到spawn_blocking中（不会被取消）
接受偶尔的不一致，通过更高层协议（如 Raft 的日志校验）修复

五、总结

Rust 异步代码中 Future 被 drop 即取消，这种隐式取消会破坏已产生副作用的操作，导致系统进入不一致状态。
CancellationToken提供显式取消信号，cancelled()返回一个可被select!等待的 Future，允许被取消的 Future 执行完整的清理路径。
取消安全的最佳工程实践是使用两阶段状态机：prepare阶段可回滚，commit阶段不可逆，在取消时只需 drop 状态机。
CancellationToken::child_token()支持树状取消传播，子 token 取消不影响父 token，实现局部超时控制。
排查取消安全问题的顺序：检查.await点前后的外部状态修改 → 检查网络 I/O 的部分完成 → 检查锁的持有范围是否跨越.await。