[Rust 小練習] m3u8分片下載場景任務調度

場景:

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);
        });
    }
}
發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章