Rust Tokio取消任務的幾種模式,你知道嗎?
Rust提供了對異步編程的支持,它可以生成異步任務,然后通過運行時執行器在操作系統線程之間調度執行。
與Rust中的所有東西一樣,異步編程必須是內存安全的,因此需要確保借用檢查器可以編譯通過。
這篇文章是關于任務取消模式的,下面我們來介紹Tokio任務的取消模式。
Select 和 Channels
所有這些模式的核心是兩個tokio特性:
- channel:用于任務間通信
- select:用于等待多個異步計算(不一定是任務!)
Tokio channel看起來有點復雜,但同時就程序的內存安全和彈性而言,它很強大。Tokio channel創建了兩個不同的對象,用于任務之間的通信,不能同時使用一個通道對象來接收和發送。
Tokio提供的頻道實際上有四種:
- mpsc:多個生產者,單一消費者
- oneshot:用于發送和接收單個值,發送后,通道關閉。
- broadcast:多個發送者,多個消費者
- watch:單一生產者,多個消費者
Drop JoinHandle不會取消任務
JoinHandle在刪除關聯的任務時將其分離,這意味著不再有任何任務句柄,也沒有辦法對其進行連接。
每次在tokio中生成任務時,都會返回JoinHandle??梢允褂胘oin句柄來等待任務完成,但是認為可以使用它來簡單地通過刪除任務來強制終止任務是錯誤的。這里有一個愚蠢的例子:
use tokio::time::{self, Duration};
#[tokio::main]
async fn main() {
let handle = tokio::spawn(async {
// do some work
tokio::time::sleep(Duration::from_secs(10)).await;
println!("Task completed");
});
// 100毫秒后取消任務
time::sleep(Duration::from_millis(100)).await;
drop(handle);
println!("Task was cancelled");
}
丟棄句柄并不會取消正在運行的任務!
Abort任務
這是取消任務的最極端的方式,沒有清理的空間:
use tokio::time::{self, Duration};
#[tokio::main]
async fn main() {
let handle = tokio::spawn(async {
// do some work
tokio::time::sleep(Duration::from_secs(1)).await;
println!("Task completed");
});
// 100毫秒后取消任務
time::sleep(Duration::from_millis(100)).await;
handle.abort();
time::sleep(Duration::from_secs(2)).await;
println!("Task was cancelled");
}
使用oneshot channel
oneshot channel允許通道上的發送單個值,可以由多個接收器偵聽。與drop模式不同,此模式允許通道執行一些清理工作。這里有一個例子:
use tokio::sync::oneshot;
use tokio::time::Duration;
#[tokio::main]
async fn main() {
let (tx, rx) = oneshot::channel();
let task = tokio::spawn(async move {
tokio::select! {
_ = rx => {
println!("Task is cancelling...");
}
_ = tokio::time::sleep(Duration::from_secs(10)) => {
println!("Task completed normally");
}
}
println!("Task is cleaning up");
});
tokio::time::sleep(Duration::from_millis(100)).await;
// 發送取消信號
let _ = tx.send(());
// 等待任務完成
let _ = task.await;
}
運行結果如下:
Task is cancelling...
Task is cleaning up
oneshot channel的限制是你不能用它來取消多個任務。
使用broadcast channel取消多個任務
如果要取消多個任務,可以使用broad channel。可以有多個生產者向通道發送信息,也可以有多個消費者從通道接收信息。每個接收方都可以看到在通道上發送的每個值。
這里有一個簡單的例子,來演示如何使用它來取消多個任務:
use tokio::sync::broadcast;
use tokio::time::Duration;
#[tokio::main]
async fn main() {
let (tx, mut rx1) = broadcast::channel(1);
let mut rx2 = tx.subscribe();
let task1 = tokio::spawn(async move {
tokio::select! {
_ = rx1.recv() => {
println!("Task 1 is cancelling...");
}
_ = tokio::time::sleep(Duration::from_secs(10)) => {
println!("Task 1 completed normally");
}
}
println!("Task 1 is cleaning up");
});
let task2 = tokio::spawn(async move {
tokio::select! {
_ = rx2.recv() => {
println!("Task 2 is cancelling...");
}
_ = tokio::time::sleep(Duration::from_secs(10)) => {
println!("Task 2 completed normally");
}
}
println!("Task 2 is cleaning up");
});
tokio::time::sleep(Duration::from_millis(100)).await;
// 發送取消信號
let _ = tx.send(());
// 等待任務完成
let _ = tokio::join!(task1, task2);
}
運行結果如下:
Task 2 is cancelling...
Task 2 is cleaning up
Task 1 is cancelling...
Task 1 is cleaning up
取消的順序可能會有所不同,因為任務可能會以不同的順序取消!
如果只想從單個任務向多個任務發送取消信號,那么broad channel可能有點過度,因為它提供了在多個任務之間傳遞消息的所有機制。
如果既需要消息傳遞又需要消息取消,這很方便。但如果只需要消息取消,還有更好的方法,開銷更少:watch channel。
使用watch channel取消多個任務
watch channel是多個消費者頻道的單一生產者。watch channel給了任務清理自己的機會。缺點是,消費者只能看到通道上發送的最近的值——這意味著如果任務在通道上發送了一個值之后啟動,它可能會錯過它,因此不會被取消,所以要小心這一點。這里有一個簡單的例子:
use tokio::sync::watch;
use tokio::time::Duration;
#[tokio::main]
async fn main() {
let (tx, mut rx1) = watch::channel(false);
let mut rx2 = tx.subscribe();
let task1 = tokio::spawn(async move {
loop {
tokio::select! {
_ = rx1.changed() => {
if *rx1.borrow() {
println!("Task 1 is cancelling...");
break;
}
}
_ = tokio::time::sleep(Duration::from_secs(10)) => {
println!("Task 1 completed normally");
break;
}
}
}
println!("Task 1 is cleaning up");
});
let task2 = tokio::spawn(async move {
loop {
tokio::select! {
_ = rx2.changed() => {
if *rx2.borrow() {
println!("Task 2 is cancelling...");
break;
}
}
_ = tokio::time::sleep(Duration::from_secs(10)) => {
println!("Task 2 completed normally");
break;
}
}
}
println!("Task 2 is cleaning up");
});
tokio::time::sleep(Duration::from_millis(100)).await;
// 發送取消信號
let _ = tx.send(true);
// 等待任務完成
let _ = tokio::join!(task1, task2);
}
取消令牌
官方的Tokio文檔中列出了一種名為CancellationToken的東西,用于優雅關機。這在tokio crate本身中不可用,但在相關的toko_util crate中可用。
use tokio::time::{sleep, Duration};
use tokio_util::sync::CancellationToken;
#[tokio::main]
async fn main() {
// Create a CancellationToken
let token = CancellationToken::new();
let token1 = token.clone();
let token2 = token.clone();
let task1 = tokio::spawn(async move {
loop {
tokio::select! {
_ = token1.cancelled() => {
println!("Task 1 is cancelling...");
break;
}
_ = tokio::time::sleep(Duration::from_secs(10)) => {
println!("Task 1 completed normally");
break;
}
}
}
println!("Task 1 is cleaning up");
});
let task2 = tokio::spawn(async move {
loop {
tokio::select! {
_ = token2.cancelled() => {
println!("Task 2 is cancelling...");
break;
}
_ = tokio::time::sleep(Duration::from_secs(10)) => {
println!("Task 2 completed normally");
break;
}
}
}
println!("Task 2 is cleaning up");
});
sleep(Duration::from_millis(100)).await;
// 發送取消信號
token.cancel();
// 等待任務完成
let _ = tokio::join!(task1, task2);
}
請注意我們是如何克隆令牌的,以便將其移動到各個異步任務中。