場景:
Producer解析m3u8文件,獲取ts列表, Consumer中併發處理TS文件的下載,最大的併發任務數量爲X 。 一開始想的是用Vec<Future> 保存任務,然後tokio::spawn處理,但是有個缺點是一個任務集合完全結束後纔能有新的任務加入。 我想實現的是,最大X個任務併發,當併發任務數<X時,加入新的任務。 實現如下:
use std::{
error::Error,
sync::{Arc, Mutex},
time::Duration,
};
use rand::{thread_rng, Rng};
use tokio::{sync::mpsc, time::sleep};
#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
// 最多併發數 同時也是channel背壓的buffer大小
const BUFF_SIZE: usize = 5;
// 任務的數量
const TASK_SIZE: u32 = 33;
let (tx, mut rx) = mpsc::channel(BUFF_SIZE);
let counter = Arc::new(Mutex::new(0_u32));
let _counter = counter.clone();
tokio::spawn(async move {
println!("start.............");
for i in 0..TASK_SIZE {
// 如果計數器的值大於等於最多併發數 , 進入Pending狀態
while { *_counter.lock().unwrap() } > BUFF_SIZE as u32 {}
println!("send {i}");
tx.send(i).await.unwrap();
// 任務發送到隊列後,計數器需要+1
let mut num_guard = _counter.lock().unwrap();
*num_guard = *num_guard + 1;
}
});
loop {
let _counter = counter.clone();
let Some(i) = rx.recv().await else { return Ok(()) };
println!("got = {} ", i);
tokio::spawn(async move {
// 隨機的延時
let rnd_delay = { Duration::from_secs((thread_rng()).gen_range(0..10_u64)) };
sleep(rnd_delay).await;
// 任務處理完畢,計數器-1
let mut num_guard = _counter.lock().unwrap();
*num_guard = *num_guard - 1;
println!("{} done! num={}", i, *num_guard);
});
}
}