Rust中如何复用已生成的线程对象？降低线程创建开销

问题分析 #

原reduce函数每次调用都会通过std::thread::scope创建新线程，数千次调用的线程创建/销毁开销极高。我们需要仅修改reduce函数，实现线程复用，且不改动reducer等其他代码。

解决方案 #

利用全局静态变量初始化一次线程池，线程持续等待任务；每次调用reduce时将任务参数发送给线程，执行完成后同步结果。以下是修改后的代码：

use std::sync::{Arc, AtomicBool, Barrier, OnceLock};
use std::sync::mpsc::{Sender, Receiver};
use std::thread;

// 任务结构体，封装reducer所需的所有参数
#[derive(Clone)]
struct Task {
    heap: Arc<Heap>,
    prog: Arc<Program>,
    tids: Arc<[usize]>,
    stop: Arc<AtomicBool>,
    barr: Arc<Barrier>,
    locs: Arc<Vec<AtomicU64>>,
    root: u64,
    tid: usize,
    debug: bool,
}

// 全局工作线程的任务发送器
static WORKERS: OnceLock<Vec<Sender<Task>>> = OnceLock::new();
// 全局线程完成信号接收器
static COMPLETION_RXS: OnceLock<Vec<Receiver<()>>> = OnceLock::new();

pub fn reduce(heap: &Heap, prog: &Program, tids: &[usize], root: u64, debug: bool) -> Ptr {
    // 首次调用时初始化线程池
    let workers = WORKERS.get_or_init(|| {
        let mut senders = Vec::with_capacity(tids.len());
        let mut comp_recvs = Vec::with_capacity(tids.len());
        
        for &tid in tids {
            let (task_tx, task_rx) = mpsc::channel();
            let (comp_tx, comp_rx) = mpsc::channel();
            
            senders.push(task_tx);
            comp_recvs.push(comp_rx);
            
            // 启动常驻工作线程
            thread::spawn(move || {
                while let Ok(task) = task_rx.recv() {
                    // 调用原reducer函数，参数直接从任务中解引用
                    reducer(
                        &task.heap,
                        &task.prog,
                        &task.tids,
                        &task.stop,
                        &task.barr,
                        &task.locs,
                        task.root,
                        task.tid,
                        task.debug,
                    );
                    // 发送任务完成信号
                    let _ = comp_tx.send(());
                }
            });
        }
        
        // 保存完成信号接收器
        COMPLETION_RXS.set(comp_recvs).unwrap();
        senders
    });
    
    let comp_rxs = COMPLETION_RXS.get().unwrap();

    // 创建本次调用专属的同步变量（和原逻辑一致）
    let stop = Arc::new(AtomicBool::new(false));
    let barr = Arc::new(Barrier::new(tids.len()));
    let locs = Arc::new(
        tids.iter()
            .map(|_| AtomicU64::new(u64::MAX))
            .collect::<Vec<AtomicU64>>()
    );

    // 将输入引用转为Arc以安全传递给线程（若原参数已用Arc包裹，直接clone即可）
    // 注：若Heap/Program未实现Clone，可使用unsafe Arc::from_raw（需确保原引用生命周期覆盖任务执行）
    let heap_arc = Arc::new(heap.clone());
    let prog_arc = Arc::new(prog.clone());
    let tids_arc = Arc::from(tids);

    // 向所有工作线程发送任务
    for (i, worker) in workers.iter().enumerate() {
        let task = Task {
            heap: heap_arc.clone(),
            prog: prog_arc.clone(),
            tids: tids_arc.clone(),
            stop: stop.clone(),
            barr: barr.clone(),
            locs: locs.clone(),
            root,
            tid: tids[i],
            debug,
        };
        worker.send(task).unwrap();
    }

    // 等待所有线程完成任务
    for rx in comp_rxs.iter() {
        rx.recv().unwrap();
    }

    // 返回最终结果（和原逻辑一致）
    load_ptr(heap, root)
}

关键细节说明 #

1. 线程池初始化：通过OnceLock确保线程池仅在首次调用reduce时创建，后续调用直接复用已启动的线程。
2. 任务传递：用mpsc通道将任务参数发送给线程，每个线程执行完任务后发送完成信号，主线程等待所有信号后继续执行。
3. 生命周期处理：将输入引用转为Arc，确保参数在线程执行期间保持有效；若原参数已用Arc管理，直接克隆即可，无需额外处理。
4. 兼容性：完全保留原reducer函数的调用方式，无需改动其他代码逻辑。

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