Synchronous/blocking JoinSet

join_set.rs

// This is not considered production quality, just provided to accompany a blog post.
use std::sync::{Arc, atomic::AtomicUsize, mpsc};

pub struct JoinSet<T> {
    tx: mpsc::Sender<std::thread::Result<T>>,
    rx: mpsc::Receiver<std::thread::Result<T>>,
    threads: Arc<AtomicUsize>,
}

impl<T> JoinSet<T>
where
    T: Send + 'static,
{
    pub fn new() -> Self {
        let (tx, rx) = mpsc::channel();
        JoinSet {
            tx,
            rx,
            threads: Default::default(),
        }
    }

    pub fn spawn<F>(&mut self, f: F)
    where
        F: FnOnce() -> T + Send + 'static,
    {
        let threads = self.threads.clone();
        threads.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        let tx = self.tx.clone();
        std::thread::spawn(move || {
            let handle = std::thread::spawn(f);
            tx.send(handle.join()).ok();
            threads.fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
        });
    }

    pub fn join_next(&mut self) -> Option<std::thread::Result<T>> {
        if self.threads.load(std::sync::atomic::Ordering::Relaxed) == 0 {
            None
        } else {
            Some(self.rx.recv().unwrap())
        }
    }
}