Merge pull request #260 from George-Miao/feat/dispatcher-blocking

Author: George-Miao

compio-dispatcher/src/lib.rs

@@ -11,14 +11,14 @@ use std::{
     thread::{available_parallelism, JoinHandle},
 };
 
-use compio_driver::{AsyncifyPool, ProactorBuilder};
+use compio_driver::{AsyncifyPool, DispatchError, Dispatchable, ProactorBuilder};
 use compio_runtime::{event::Event, JoinHandle as CompioJoinHandle, Runtime};
-use flume::{unbounded, SendError, Sender};
+use flume::{unbounded, Sender};
 use futures_channel::oneshot;
 
-type Dispatching = Box<dyn Dispatchable + Send>;
+type Spawning = Box<dyn Spawnable + Send>;
 
-trait Dispatchable {
+trait Spawnable {
     fn spawn(self: Box<Self>, handle: &Runtime) -> CompioJoinHandle<()>;
 }
 
@@ -28,7 +28,14 @@ struct Concrete<F, R> {
     func: F,
 }
 
-impl<F, Fut, R> Dispatchable for Concrete<F, R>
+impl<F, R> Concrete<F, R> {
+    pub fn new(func: F) -> (Self, oneshot::Receiver<R>) {
+        let (tx, rx) = oneshot::channel();
+        (Self { callback: tx, func }, rx)
+    }
+}
+
+impl<F, Fut, R> Spawnable for Concrete<F, R>
 where
     F: FnOnce() -> Fut + Send + 'static,
     Fut: Future<Output = R>,
@@ -43,10 +50,22 @@ where
     }
 }
 
-#[derive(Debug)]
+impl<F, R> Dispatchable for Concrete<F, R>
+where
+    F: FnOnce() -> R + Send + 'static,
+    R: Send + 'static,
+{
+    fn run(self: Box<Self>) {
+        let Concrete { callback, func } = *self;
+        let res = func();
+        callback.send(res).ok();
+    }
+}
+
 /// The dispatcher. It manages the threads and dispatches the tasks.
+#[derive(Debug)]
 pub struct Dispatcher {
-    sender: Sender<Dispatching>,
+    sender: Sender<Spawning>,
     threads: Vec<JoinHandle<()>>,
     pool: AsyncifyPool,
 }
@@ -57,7 +76,7 @@ impl Dispatcher {
         let mut proactor_builder = builder.proactor_builder;
         proactor_builder.force_reuse_thread_pool();
         let pool = proactor_builder.create_or_get_thread_pool();
-        let (sender, receiver) = unbounded::<Dispatching>();
+        let (sender, receiver) = unbounded::<Spawning>();
 
         let threads = (0..builder.nthreads)
             .map({
@@ -123,28 +142,51 @@ impl Dispatcher {
     ///
     /// If all threads have panicked, this method will return an error with the
     /// sent closure.
-    pub fn dispatch<Fut, Fn, R>(&self, f: Fn) -> Result<oneshot::Receiver<R>, SendError<Fn>>
+    pub fn dispatch<Fn, Fut, R>(&self, f: Fn) -> Result<oneshot::Receiver<R>, DispatchError<Fn>>
     where
         Fn: (FnOnce() -> Fut) + Send + 'static,
         Fut: Future<Output = R> + 'static,
         R: Send + 'static,
     {
-        let (tx, rx) = oneshot::channel();
-        let concrete: Concrete<Fn, R> = Concrete {
-            callback: tx,
-            func: f,
-        };
+        let (concrete, rx) = Concrete::new(f);
+
         match self.sender.send(Box::new(concrete)) {
             Ok(_) => Ok(rx),
             Err(err) => {
                 // SAFETY: We know the dispatchable we sent has type `Concrete<Fn, R>`
                 let recovered =
                     unsafe { Box::from_raw(Box::into_raw(err.0) as *mut Concrete<Fn, R>) };
-                Err(SendError(recovered.func))
+                Err(DispatchError(recovered.func))
             }
         }
     }
 
+    /// Dispatch a blocking task to the threads.
+    ///
+    /// Blocking pool of the dispatcher will be obtained from the proactor
+    /// builder. So any configuration of the proactor's blocking pool will be
+    /// applied to the dispatcher.
+    ///
+    /// # Error
+    ///
+    /// If all threads are busy and the thread pool is full, this method will
+    /// return an error with the original closure. The limit can be configured
+    /// with [`DispatcherBuilder::proactor_builder`] and
+    /// [`ProactorBuilder::thread_pool_limit`].
+    pub fn dispatch_blocking<Fn, R>(&self, f: Fn) -> Result<oneshot::Receiver<R>, DispatchError<Fn>>
+    where
+        Fn: FnOnce() -> R + Send + 'static,
+        R: Send + 'static,
+    {
+        let (concrete, rx) = Concrete::new(f);
+
+        self.pool
+            .dispatch(concrete)
+            .map_err(|e| DispatchError(e.0.func))?;
+
+        Ok(rx)
+    }
+
     /// Stop the dispatcher and wait for the threads to complete. If there is a
     /// thread panicked, this method will resume the panic.
     pub async fn join(self) -> io::Result<()> {
@@ -163,7 +205,7 @@ impl Dispatcher {
                 handle.notify();
             }
         }) {
-            std::thread::spawn(f);
+            std::thread::spawn(f.0);
         }
         event.wait().await;
         let mut guard = results.lock().unwrap();

compio-driver/src/asyncify.rs

@@ -1,4 +1,5 @@
 use std::{
+    fmt,
     sync::{
         atomic::{AtomicUsize, Ordering},
         Arc,
@@ -8,7 +9,51 @@ use std::{
 
 use crossbeam_channel::{bounded, Receiver, Sender, TrySendError};
 
-type BoxClosure = Box<dyn FnOnce() + Send>;
+/// An error that may be emitted when all worker threads are busy. It simply
+/// returns the dispatchable value with a convenient [`fmt::Debug`] and
+/// [`fmt::Display`] implementation.
+#[derive(Copy, Clone, PartialEq, Eq)]
+pub struct DispatchError<T>(pub T);
+
+impl<T> DispatchError<T> {
+    /// Consume the error, yielding the dispatchable that failed to be sent.
+    pub fn into_inner(self) -> T {
+        self.0
+    }
+}
+
+impl<T> fmt::Debug for DispatchError<T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        "DispatchError(..)".fmt(f)
+    }
+}
+
+impl<T> fmt::Display for DispatchError<T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        "all threads are busy".fmt(f)
+    }
+}
+
+impl<T> std::error::Error for DispatchError<T> {}
+
+type BoxedDispatchable = Box<dyn Dispatchable + Send>;
+
+/// A trait for dispatching a closure. It's implemented for all `FnOnce() + Send
+/// + 'static` but may also be implemented for any other types that are `Send`
+///   and `'static`.
+pub trait Dispatchable: Send + 'static {
+    /// Run the dispatchable
+    fn run(self: Box<Self>);
+}
+
+impl<F> Dispatchable for F
+where
+    F: FnOnce() + Send + 'static,
+{
+    fn run(self: Box<Self>) {
+        (*self)()
+    }
+}
 
 struct CounterGuard(Arc<AtomicUsize>);
 
@@ -19,24 +64,24 @@ impl Drop for CounterGuard {
 }
 
 fn worker(
-    receiver: Receiver<BoxClosure>,
+    receiver: Receiver<BoxedDispatchable>,
     counter: Arc<AtomicUsize>,
     timeout: Duration,
 ) -> impl FnOnce() {
     move || {
         counter.fetch_add(1, Ordering::AcqRel);
         let _guard = CounterGuard(counter);
         while let Ok(f) = receiver.recv_timeout(timeout) {
-            f();
+            f.run();
         }
     }
 }
 
 /// A thread pool to perform blocking operations in other threads.
 #[derive(Debug, Clone)]
 pub struct AsyncifyPool {
-    sender: Sender<BoxClosure>,
-    receiver: Receiver<BoxClosure>,
+    sender: Sender<BoxedDispatchable>,
+    receiver: Receiver<BoxedDispatchable>,
     counter: Arc<AtomicUsize>,
     thread_limit: usize,
     recv_timeout: Duration,
@@ -56,15 +101,20 @@ impl AsyncifyPool {
         }
     }
 
-    /// Send a closure to another thread. Usually the user should not use it.
-    pub fn dispatch<F: FnOnce() + Send + 'static>(&self, f: F) -> Result<(), F> {
-        match self.sender.try_send(Box::new(f) as BoxClosure) {
+    /// Send a dispatchable, usually a closure, to another thread. Usually the
+    /// user should not use it. When all threads are busy and thread number
+    /// limit has been reached, it will return an error with the original
+    /// dispatchable.
+    pub fn dispatch<D: Dispatchable>(&self, f: D) -> Result<(), DispatchError<D>> {
+        match self.sender.try_send(Box::new(f) as BoxedDispatchable) {
             Ok(_) => Ok(()),
             Err(e) => match e {
                 TrySendError::Full(f) => {
                     if self.counter.load(Ordering::Acquire) >= self.thread_limit {
                         // Safety: we can ensure the type
-                        Err(*unsafe { Box::from_raw(Box::into_raw(f).cast()) })
+                        Err(DispatchError(*unsafe {
+                            Box::from_raw(Box::into_raw(f).cast())
+                        }))
                     } else {
                         std::thread::spawn(worker(
                             self.receiver.clone(),