Add a new executor module to run functions in thread pools

Author: and-cb

Cargo.toml

@@ -18,6 +18,7 @@ zerocopy = { version = "0.8.24", features = ["derive"] }
 parking_lot = { version = "0.12.4", features = ["send_guard"] }
 fxhash = "0.2.1"
 static_assertions = "1.1.0"
+rayon = "1.10.0"
 
 [dev-dependencies]
 criterion = "0.6.0"

src/executor/futures.rs

@@ -0,0 +1,213 @@
+use crate::executor::Wait;
+use std::{
+    fmt,
+    sync::{Arc, OnceLock},
+};
+
+#[derive(Debug)]
+enum FutureStatus<T> {
+    Completed(T),
+    Poisoned,
+}
+
+use FutureStatus::*;
+
+/// A placeholder for a value that will be computed at a later time.
+///
+/// `Future`s are the result of running functions in separate threads using
+/// [`Executor::spawn`](crate::executor::Executor::spawn): calling `spawn()` in fact returns
+/// immediately, even though the function will complete at a later time. The `Future` returned by
+/// `spawn()` allows retrieving the result of the function once it completes.
+///
+/// # Poisoning
+///
+/// A `Future` may be in a "poisoned" status if the execution of the function that produced it
+/// failed with a panic.
+pub struct Future<T> {
+    cell: Arc<OnceLock<FutureStatus<T>>>,
+}
+
+impl<T> Future<T> {
+    #[inline]
+    #[must_use]
+    pub(super) fn pending() -> Self {
+        Self { cell: Arc::new(OnceLock::new()) }
+    }
+
+    /// Creates a new `Future` that is already completed with the given value.
+    #[inline]
+    #[must_use]
+    pub fn ready(value: T) -> Self {
+        let this = Self::pending();
+        this.complete(value);
+        this
+    }
+
+    #[inline]
+    pub(super) fn complete(&self, value: T) {
+        self.try_complete(value).unwrap_or_else(|err| panic!("{err}"))
+    }
+
+    #[inline]
+    pub(super) fn try_complete(&self, value: T) -> Result<(), ReadyError> {
+        self.cell.set(Completed(value)).map_err(|_| ReadyError)
+    }
+
+    // There's no `poison()` method simply because it's not used internally.
+
+    #[inline]
+    pub(super) fn try_poison(&self) -> Result<(), ReadyError> {
+        self.cell.set(Poisoned).map_err(|_| ReadyError)
+    }
+
+    /// Returns the value of the `Future`, or `None` if this `Future` was not completed yet.
+    ///
+    /// # Panics
+    ///
+    /// If the `Future` is poisoned. See [`Future::try_get()`] for a non-panicking version of this
+    /// method.
+    #[inline]
+    #[must_use]
+    pub fn get(&self) -> Option<&T> {
+        self.try_get().map(|result| result.expect("Future is poisoned"))
+    }
+
+    /// Returns the value of the `Future`, `None` if this `Future` was not completed yet, or an
+    /// error if this `Future` is poisoned.
+    #[inline]
+    #[must_use]
+    pub fn try_get(&self) -> Option<Result<&T, PoisonError>> {
+        match self.cell.get() {
+            None => None,
+            Some(Completed(ref value)) => Some(Ok(value)),
+            Some(Poisoned) => Some(Err(PoisonError)),
+        }
+    }
+}
+
+impl<T> Wait for Future<T> {
+    type Output = T;
+
+    #[inline]
+    fn wait(&self) -> &Self::Output {
+        match self.cell.wait() {
+            Completed(value) => value,
+            Poisoned => panic!("{PoisonError}"),
+        }
+    }
+}
+
+impl<T> Clone for Future<T> {
+    fn clone(&self) -> Self {
+        Self { cell: Arc::clone(&self.cell) }
+    }
+}
+
+impl<T: fmt::Debug> fmt::Debug for Future<T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        struct Pending;
+
+        impl fmt::Debug for Pending {
+            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+                f.write_str("<pending>")
+            }
+        }
+
+        struct Poisoned;
+
+        impl fmt::Debug for Poisoned {
+            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+                f.write_str("<poisoned>")
+            }
+        }
+
+        f.debug_tuple("Future")
+            .field(match self.try_get() {
+                None => &Pending,
+                Some(Ok(value)) => value,
+                Some(Err(PoisonError)) => &Poisoned,
+            })
+            .finish()
+    }
+}
+
+#[derive(Clone, PartialEq, Eq, Debug)]
+pub(super) struct ReadyError;
+
+impl fmt::Display for ReadyError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_str("attempted to complete or poison the Future twice")
+    }
+}
+
+impl std::error::Error for ReadyError {}
+
+#[derive(Clone, PartialEq, Eq, Debug)]
+pub struct PoisonError;
+
+impl fmt::Display for PoisonError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_str("execution of the closure for this Future resulted in a panic")
+    }
+}
+
+impl std::error::Error for PoisonError {}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::{panic, sync::Barrier, thread, time::Duration};
+
+    #[test]
+    fn pending_to_completed() {
+        let f = Future::<u32>::pending();
+
+        assert_eq!(f.get(), None);
+        assert_eq!(f.try_get(), None);
+
+        f.complete(123);
+
+        assert_eq!(f.get(), Some(&123));
+        assert_eq!(f.try_get(), Some(Ok(&123)));
+    }
+
+    #[test]
+    fn pending_to_poisoned() {
+        let f = Future::<u32>::pending();
+
+        assert_eq!(f.get(), None);
+        assert_eq!(f.try_get(), None);
+
+        f.try_poison().expect("poison failed");
+
+        panic::catch_unwind(|| f.get()).expect_err("get() should have panicked");
+        assert_eq!(f.try_get(), Some(Err(PoisonError)));
+    }
+
+    #[test]
+    fn wait() {
+        let f = Future::<u32>::pending();
+        let g = f.clone();
+        let barrier = Barrier::new(2);
+
+        thread::scope(|s| {
+            s.spawn(|| {
+                barrier.wait();
+                thread::sleep(Duration::from_secs(1));
+                g.complete(123);
+            });
+
+            assert_eq!(f.get(), None);
+            assert_eq!(f.try_get(), None);
+
+            barrier.wait();
+
+            assert_eq!(f.wait(), &123);
+            assert_eq!(f.get(), Some(&123));
+            assert_eq!(f.try_get(), Some(Ok(&123)));
+
+            // Waiting twice or more should return the same value
+            assert_eq!(f.wait(), &123);
+        });
+    }
+}

src/executor/inline.rs

@@ -0,0 +1,29 @@
+use crate::executor::{Executor, Future};
+
+/// A dummy executor that executes all functions in the same thread.
+#[derive(Copy, Clone, Debug)]
+pub struct Inline;
+
+impl Executor for Inline {
+    #[inline]
+    fn defer<F, T>(&self, f: F) -> Future<T>
+    where
+        F: FnOnce() -> T + Send + 'static,
+        T: Send + Sync + 'static,
+    {
+        Future::ready(f())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::executor::Wait;
+
+    #[test]
+    fn defer() {
+        let inline = Inline;
+        let future = inline.defer(|| 123);
+        assert_eq!(future.wait(), &123);
+    }
+}

src/executor/mod.rs

@@ -0,0 +1,33 @@
+//! Concurrent execution.
+//!
+//! This module provides structures and traits to run functions in separate threads and obtain
+//! their results later on. The implementation is currently based on the popular [`rayon`] crate.
+//! This module also provides a dummy implementation called [`Inline`] that does not spawn any
+//! actual thread but executes everything serially.
+//!
+//! # Examples
+//!
+//! ```
+//! use triedb::executor::{threadpool, Executor, Wait};
+//!
+//! // Create a thread pool
+//! let pool = threadpool::builder().build().unwrap();
+//!
+//! // Run some closures in the background.
+//! let future1 = pool.defer(|| 1 + 1);
+//! let future2 = pool.defer(|| 2 + 2);
+//!
+//! // Wait for the closures to return a result.
+//! assert_eq!(future1.wait(), &2);
+//! assert_eq!(future2.wait(), &4);
+//! ```
+
+mod futures;
+mod inline;
+mod traits;
+
+pub mod threadpool;
+
+pub use futures::{Future, PoisonError};
+pub use inline::Inline;
+pub use traits::{Executor, Wait};

src/executor/threadpool.rs

@@ -0,0 +1,77 @@
+use crate::executor::{Executor, Future};
+use rayon::{ThreadPool, ThreadPoolBuilder};
+use std::thread;
+
+/// A wrapper around [`ThreadPoolBuilder::new()`] which sets some default values to make the
+/// resulting [`ThreadPool`] play nicely with the [`Executor`] trait.
+pub fn builder() -> ThreadPoolBuilder {
+    ThreadPoolBuilder::new()
+        .thread_name(|num| format!("executor-thread-{num:03}"))
+        // The default behavior for `rayon` is to abort in case of panic, causing the whole program
+        // to crash. We instead want to catch individual panics and poison the relevant `Future`
+        // when those occur.
+        //
+        // This panic hanlder does nothing, so that the abort behavior is suppressed. We don't need
+        // to explicitly print an error message or the backtrace because this will be already taken
+        // care of.
+        .panic_handler(|_| {})
+}
+
+impl Executor for ThreadPool {
+    fn defer<F, T>(&self, f: F) -> Future<T>
+    where
+        F: FnOnce() -> T + Send + 'static,
+        T: Send + Sync + 'static,
+    {
+        let sender_future = Future::pending();
+        let receiver_future = sender_future.clone();
+
+        self.spawn(move || {
+            // Create the guard panic first, then run the closure. The guard will be dropped at the
+            // end of this scope. If the function succeeds, the guard won't do anything; if it
+            // panics, the guard's `Drop` implementation will poison the future.
+            let _guard = PanicGuard { future: &sender_future };
+            sender_future.complete(f())
+        });
+
+        receiver_future
+    }
+}
+
+/// A "guard" to detect if this thread panics, and poison the `Future` in that case.
+#[derive(Debug)]
+struct PanicGuard<'a, T> {
+    future: &'a Future<T>,
+}
+
+impl<'a, T> Drop for PanicGuard<'a, T> {
+    fn drop(&mut self) {
+        if thread::panicking() {
+            // Unlikely, but if the future was already set, just ignore the error from
+            // `try_poison()` and carry on
+            let _ = self.future.try_poison();
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::executor::{PoisonError, Wait};
+    use std::panic;
+
+    #[test]
+    fn defer() {
+        let pool = builder().build().expect("building thread pool failed");
+        let future = pool.defer(|| 123);
+        assert_eq!(future.wait(), &123);
+    }
+
+    #[test]
+    fn poisoning() {
+        let pool = builder().build().expect("building thread pool failed");
+        let future = pool.defer(|| panic!("something went wrong"));
+        panic::catch_unwind(|| future.wait()).expect_err("wait() was expected to panic");
+        assert_eq!(future.try_get(), Some(Err(PoisonError)));
+    }
+}

src/executor/traits.rs

@@ -0,0 +1,33 @@
+use crate::executor::Future;
+
+/// Trait for objects that can be awaited until they reach their final state.
+///
+/// The main structure that implements this trait is [`Future`], but also any structure that wraps
+/// a `Future` may implement this trait.
+pub trait Wait {
+    type Output;
+
+    /// Blocks execution until the final state is reached.
+    fn wait(&self) -> &Self::Output;
+}
+
+/// Trait for objects that can run functions concurrently.
+pub trait Executor {
+    /// Runs the given closure `f` in a separate thread, and returns a [`Future`] that can be used
+    /// to obtain the result of `f` once its execution is complete.
+    fn defer<F, T>(&self, f: F) -> Future<T>
+    where
+        F: FnOnce() -> T + Send + 'static,
+        T: Send + Sync + 'static;
+}
+
+impl<E: Executor> Executor for &E {
+    #[inline]
+    fn defer<F, T>(&self, f: F) -> Future<T>
+    where
+        F: FnOnce() -> T + Send + 'static,
+        T: Send + Sync + 'static,
+    {
+        (**self).defer(f)
+    }
+}