go_spawn/

lib.rs

//! `go-spawn` is a library that provides macros to spawn and join threads with minimal boilerplate.

pub mod error;
#[cfg(test)]
mod test;

use state::LocalStorage;
use std::{cell::RefCell, collections::VecDeque, thread::JoinHandle};

#[doc(hidden)]
pub static __THREAD_HANDLE_INITIALIZED: LocalStorage<RefCell<VecDeque<JoinHandle<()>>>> =
    LocalStorage::new();

#[cfg(test)]
fn drop_all_handles() {
    if let Some(handles) = __THREAD_HANDLE_INITIALIZED.try_get() {
        handles.borrow_mut().clear();
    }
}

/// Spawns a thread to execute the given code. Any captured variables will be moved to the spawned
/// thread.
///
/// The thread's handle will be stored in thread-local storage until joined by `join!`
/// or `join_all!`.
///
/// # Examples
/// ```
/// # fn expensive_background_work() {}
/// # fn do_work() {
/// use go_spawn::go;
/// use std::sync::{
///     atomic::{AtomicI64, Ordering},
///     Arc,
/// };
///
/// go!(expensive_background_work());
///
/// let counter = Arc::new(AtomicI64::new(0));
/// let copy_of_counter = counter.clone();
///
/// go! {
///     for _ in 0..1_000_000 {
///         copy_of_counter.fetch_add(1, Ordering::SeqCst);
///     }
/// }
/// # }
/// ```
#[macro_export]
macro_rules! go {
    ( $( $tokens:tt )+ ) => {
        go_spawn::__go_helper!(@move $( $tokens )+ )
    };
}

/// Spawns a thread to execute the given code. Any captured variables will be borrowed by the
/// spawned thread.
///
/// The thread's handle will be stored in thread-local storage until joined by
/// `join!` or `join_all!`.
///
/// # Examples
/// ```
/// # fn expensive_background_work() {}
/// # fn do_work() {
/// use go_spawn::go_ref;
/// use std::sync::{
///     atomic::{AtomicI64, Ordering},
/// };
///
/// go_ref!(expensive_background_work());
///
/// static COUNTER: AtomicI64 = AtomicI64::new(0);
///
/// go_ref! {
///     for _ in 0..1_000_000 {
///         COUNTER.fetch_add(1, Ordering::SeqCst);
///     }
/// }
/// # }
/// ```
#[macro_export]
macro_rules! go_ref {
    ( $( $tokens:tt )+ ) => {
        go_spawn::__go_helper!(@ref $( $tokens )+ )
    };
}

/// Joins the most recent not-yet-joined thread spawned from the current thread by either `go!` or
/// `go_ref!`.
///
/// This will always yield a `go_spawn::error::Result`, which indicates whether a thread
/// was successfully joined.
///
/// # Examples
/// ```
/// # fn expensive_background_work() {}
/// # fn do_work() {
/// use go_spawn::{
///     error::{JoinError, Result},
///     go, go_ref, join,
/// };
/// use std::sync::{
///     atomic::{AtomicI64, Ordering},
/// };
///
/// go_ref!(expensive_background_work());
///
/// static COUNTER: AtomicI64 = AtomicI64::new(0);
///
/// go_ref! {
///     for _ in 0..1_000_000 {
///         COUNTER.fetch_add(1, Ordering::SeqCst);
///     }
/// }
///
/// assert!(join!().is_ok());
/// assert_eq!(COUNTER.load(Ordering::SeqCst), 1_000_0000);
///
/// go!(panic!("{}", 4));
///
/// let result = join!().unwrap_err();
/// let expected = Some("4".to_string());
/// assert!(
///     matches!(
///         result,
///         JoinError::ThreadPanic(value) if value.downcast_ref::<String>() == expected.as_ref(),
///     )
/// );
///
/// go!(std::panic::panic_any(4_i32));
///
/// let result = join!().unwrap_err();
/// let expected = Some(4_i32);
/// assert!(
///     matches!(
///         result,
///         JoinError::ThreadPanic(value) if value.downcast_ref::<i32>() == expected.as_ref(),
///     )
/// );
/// # }
/// ```
#[macro_export]
macro_rules! join {
    () => {
        go_spawn::__THREAD_HANDLE_INITIALIZED
            .try_get()
            .and_then(|handles| handles.borrow_mut().pop_back())
            .map(|handle| {
                handle
                    .join()
                    .map_err(go_spawn::error::JoinError::ThreadPanic)
            })
            .unwrap_or_else(|| Err(go_spawn::error::JoinError::NoHandleFound))
    };
}

/// Joins all pending threads spawned from the current thread by either `go!` or `go_ref!`.
///
/// Threads will be joined starting with the most recent and working backwards until the earliest
/// thread spawned.
///
/// # Handling panics
///
/// `join_all` optionally takes a callback to invoke on the error values returned when joining
/// threads that panicked. The callback will be invoked on each error value returned when joining
/// threads before proceeding to join the next thread, so it's recommended to avoid doing heavy
/// work in these error handlers (e.g. by using a channel to pass the error somewhere else to be
/// processed).
///
/// See the documentation for [`JoinHandle::join`](https://doc.rust-lang.org/std/thread/struct.JoinHandle.html#method.join) for more details.
///
/// # Examples
/// ```
/// # fn expensive_background_work() {}
/// # fn do_work() {
/// use go_spawn::{go, go_ref, join_all};
/// use std::sync::{
///     atomic::{AtomicI64, Ordering},
/// };
///
/// static COUNTER: AtomicI64 = AtomicI64::new(0);
///
/// for _ in 0..4 {
///     go_ref! {
///         COUNTER.fetch_add(1, Ordering::SeqCst);
///     }
/// }
///
/// join_all!();
/// assert_eq!(COUNTER.load(Ordering::SeqCst), 4);
///
/// for i in 0_i64..4 {
///     go! {
///         std::panic::panic_any(i);
///     }
/// }
///
/// static ERR_COUNTER: AtomicI64 = AtomicI64::new(0);
///
/// join_all!(|error| {
///     let value: &i64 = error.downcast_ref().unwrap();
///     assert_eq!(ERR_COUNTER.fetch_add(1, Ordering::SeqCst), *value);
/// });
/// # }
/// ```
#[macro_export]
macro_rules! join_all {
    () => {
        go_spawn::join_all!(std::mem::drop)
    };

    ($error_handler:expr) => {
        if let Some(handles) = go_spawn::__THREAD_HANDLE_INITIALIZED.try_get() {
            let mut handles = handles.borrow_mut();

            handles
                .drain(..)
                .rev()
                .filter_map(|handle| handle.join().err())
                .for_each($error_handler);
        }
    };
}

#[doc(hidden)]
#[macro_export]
macro_rules! __go_helper {
    (@move $( $tokens:tt )+ ) => {
        go_spawn::__go_helper!(@spawn move || { let _ = { $( $tokens )+ }; } )
    };

    (@ref $( $tokens:tt )+ ) => {
        go_spawn::__go_helper!(@spawn || { let _ = { $( $tokens )+ }; } )
    };

    (@spawn $expression:expr) => {{
        go_spawn::__THREAD_HANDLE_INITIALIZED.set(|| {
            std::cell::RefCell::new(std::collections::VecDeque::new())
        });
        go_spawn::__THREAD_HANDLE_INITIALIZED
            .get()
            .borrow_mut()
            .push_back(std::thread::spawn($expression));
    }};
}