Rust中支持工作线程自主调度任务的并行算法实现问询

针对你需要工作线程自主调度任务、主线程/工作线程均可向队列添加任务的需求，以下是几种可行方案，基于Rayon或其他Rust线程库实现：

方案1：Rayon + 无锁任务队列 #

Rayon本身基于工作窃取调度，但如果需要自定义任务提交逻辑，可以结合crossbeam::queue提供的无锁队列实现动态任务添加：

1. 依赖准备 #

在Cargo.toml中添加依赖：

[dependencies]
rayon = "1.8"
crossbeam = "0.8"

2. 核心实现代码 #

use rayon::prelude::*;
use std::sync::Arc;
use crossbeam::queue::SegQueue;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;

// 定义可跨线程执行的任务类型
type Task = Box<dyn FnOnce() + Send + 'static>;

fn main() {
    // 线程安全的任务队列，支持多生产者多消费者
    let task_queue = Arc::new(SegQueue::new());
    // 控制工作线程运行状态的原子标志
    let is_running = Arc::new(AtomicBool::new(true));

    // 启动Rayon线程池的工作循环
    rayon::scope(|s| {
        // 为每个Rayon工作线程启动任务循环
        for _ in 0..rayon::current_num_threads() {
            let queue = task_queue.clone();
            let running = is_running.clone();
            s.spawn(move |_| {
                while running.load(Ordering::Relaxed) {
                    // 尝试从队列取任务执行
                    if let Ok(task) = queue.pop() {
                        task();
                    } else {
                        // 队列为空时短暂休眠，避免空转消耗CPU
                        std::thread::sleep(Duration::from_millis(10));
                    }
                }
            });
        }

        // 主线程添加初始任务
        task_queue.push(Box::new(|| {
            println!("执行初始任务");
            // 任务执行过程中动态添加新任务
            let queue = task_queue.clone();
            queue.push(Box::new(|| println!("执行动态添加的子任务")));
        }));

        // 主线程轮询队列，可根据业务逻辑添加更多任务或监控状态
        std::thread::sleep(Duration::from_secs(1));
        task_queue.push(Box::new(|| println!("主线程添加的额外任务")));

        // 所有任务提交完成后，终止工作线程
        is_running.store(false, Ordering::Relaxed);
    });
}

关键细节 #

• SegQueue是无锁队列，支持多线程同时读写，性能优于标准库的mpsc::channel（后者是单生产者模型）。
• 用原子布尔值AtomicBool控制工作线程的退出逻辑，避免僵尸线程。
• Rayon的scope确保所有工作线程在主线程退出前完成任务。

方案2：自定义线程池（完全自主调度） #

如果Rayon的线程池模型不符合需求，可以直接用标准库std::thread创建自定义线程池，配合线程安全队列：

use std::sync::Arc;
use crossbeam::queue::SegQueue;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use std::thread;

type Task = Box<dyn FnOnce() + Send + 'static>;

fn main() {
    let task_queue = Arc::new(SegQueue::new());
    let is_running = Arc::new(AtomicBool::new(true));
    let thread_count = 4;

    // 创建自定义工作线程
    for _ in 0..thread_count {
        let queue = task_queue.clone();
        let running = is_running.clone();
        thread::spawn(move || {
            while running.load(Ordering::Relaxed) {
                if let Ok(task) = queue.pop() {
                    task();
                } else {
                    std::thread::sleep(Duration::from_millis(10));
                }
            }
        });
    }

    // 添加任务
    task_queue.push(Box::new(|| println!("自定义线程池执行任务")));
    task_queue.push(Box::new(|| {
        let queue = task_queue.clone();
        queue.push(Box::new(|| println!("自定义线程池执行动态子任务")));
    }));

    // 等待任务完成后终止线程
    std::thread::sleep(Duration::from_secs(1));
    is_running.store(false, Ordering::Relaxed);
}

方案3：使用tokio（异步场景） #

如果你的任务包含IO密集型操作，异步runtimetokio的任务调度更高效，天然支持动态任务提交：

use tokio::task;

#[tokio::main]
async fn main() {
    // 提交初始任务
    let handle = task::spawn(async {
        println!("执行初始异步任务");
        // 动态提交子任务
        task::spawn(async {
            println!("执行动态异步子任务");
        }).await.unwrap();
    });

    handle.await.unwrap();
}

选择建议 #

• CPU密集型任务：优先选Rayon+无锁队列，Rayon的工作窃取调度能最大化CPU利用率。
• 需要完全自定义调度逻辑：选自定义线程池+crossbeam::queue。
• IO密集型任务：选tokio或async-std异步runtime。

内容的提问来源于stack exchange，提问作者pnadeau