#1046·crossbeam

使用 `WaitGroup` 实现线程安全的共享状态并发

作者: siennathesane创建于 2023年12月6日更新于 2026年2月21日
标签featurecrossbeam-utils

我有一个复杂的用例,用于 crossbeam_utils::sync::WaitGroup。我有一个线程安全的结构,用于为数据库构建 SSTables,并且它在内存中非常大,因此无法轻易克隆。作为一部分,表构建器需要能够在完成后对块进行压缩和加密,但这项工作是异步执行的。我想使用 WaitGroup 在将 SSTable 写入磁盘之前提供最终同步点。以下是我当前的 MVCE([playground 链接](https://play.Rust-lang.org/?version=stable&mode=debug&edition=2021&gist=85d6c0f2a47e86faa0fc2c95d8b0029f)):

rust
use bytes::Bytes;
use crossbeam_deque::Worker;
use crossbeam_utils::sync::WaitGroup;
use parking_lot::Mutex;
use std::{sync::Arc, thread, thread::available_parallelism};

/// Block on disk
struct Block {
    data: Bytes, // several megabytes
}

/// Thread-safe SSTable Builder
struct Builder {
    wg: Arc<WaitGroup>,
    done: Arc<Mutex<bool>>,
    blocks: Arc<Mutex<Vec<Block>>>, // could be 100s of MiBs until it's flushed
    work_queue: Arc<Mutex<Worker<usize>>>,
}

impl Builder {
    /// Returns an Arc<Builder> to ensure that cloning doesn't clone hundreds
    /// of megabytes
    pub fn new() -> Arc<Self> {
        let f = Arc::new(Builder {
            wg: Arc::new(WaitGroup::new()),
            done: Arc::new(Mutex::new(false)),
            blocks: Arc::new(Mutex::new(vec![])),
            work_queue: Arc::new(Mutex::new(Worker::<usize>::new_lifo())),
        });

        // spin up the internal data workers
        let p_count = available_parallelism().unwrap().get();
        for _ in 0..=p_count {
            let f_alias = f.clone();
            thread::spawn(move || {
                f_alias.worker();
            });
        }

        f
    }

    pub fn add(&self, _key: Bytes, _value: Bytes) {
        // add to block
    }

    pub fn complete(&self) {
        // inform workers there's no more work
        {
            let mut done = self.done.lock();
            *done = true;
        }

        // ensure the work is complete
        self.wg.wait();

        // flush blocks to disk
    }

    // a worker thread
    fn worker(&self) {
        let wg = self.wg.clone();

        while !*self.done.lock() {
            let stealer = self.work_queue.lock().stealer();
            let idx = stealer.steal().success().unwrap();

            let mut block_list_ref = self.blocks.lock();
            let _block = &mut block_list_ref[idx];

            // compress and encrypt the block stolen from the queue
            // this modifies the block vec in-place
        }

        drop(wg);
    }
}

fn main() {
    // we need a new sstable
    let f = Builder::new();

    // some other threads will call this
    // f.add(key, value);

    // sstable is complete
    // finalize all blocks
    // write to disk
    f.complete();
}

内容来源: crossbeam-rs/crossbeam