tasks shared queue

src/builder.rs

@@ -4,28 +4,57 @@ use futures::Future;
 
 use crate::{Error, Message};
 
+#[derive(Clone, Copy, Debug)]
+pub struct Config {
+    pub queue_size: usize,
+    pub queue_per_task: bool,
+    pub ordered: bool,
+    pub task_count: u32,
+    pub lazy_task_creation: bool,
+}
+
+impl Default for Config {
+    fn default() -> Self {
+        Self {
+            queue_size: 4,
+            queue_per_task: false,
+            ordered: false,
+            task_count: 1,
+            lazy_task_creation: true,
+        }
+    }
+}
+
 pub trait Builder<M: Message>: Send + Sync + 'static {
     type Context: 'static;
     type BuildFut<'a>: Future<Output = Result<Self::Context, Error>> + Send + 'a
     where
         Self: 'a;
 
-    fn parallel(&self, _stream_id: u32) -> (u32, bool) {
+    fn config(&self, _stream_id: u32) -> Config {
-        (1, false)
+        Default::default()
-    }
-
-    fn queue_size(&self, _stream_id: u32, _task_id: u32) -> usize {
-        4
     }
 
     fn build(&self, stream_id: u32, _task_id: u32) -> Self::BuildFut<'_>;
 }
 
-pub struct DefaultBuilder<M: Message, H>(usize, PhantomData<(M, H)>);
+pub struct DefaultBuilder<M: Message, H> {
+    config: Config,
+    _m: PhantomData<(M, H)>,
+}
 
 impl<M: Message, H> DefaultBuilder<M, H> {
     pub fn new(queue_size: usize) -> Self {
-        Self(queue_size, Default::default())
+        Self {
+            config: Config {
+                queue_size,
+                queue_per_task: false,
+                ordered: false,
+                task_count: 1,
+                lazy_task_creation: true,
+            },
+            _m: PhantomData,
+        }
     }
 }
 
@@ -37,17 +66,15 @@ impl<M: Message, H: Sync + Send + Default + 'static> Builder<M> for DefaultBuild
         async move { Ok(<Self::Context as Default>::default()) }
     }
 
-    fn queue_size(&self, _stream_id: u32, _task_id: u32) -> usize {
+    fn config(&self, _stream_id: u32) -> Config {
-        self.0
+        self.config
     }
 }
 
 pub struct SharedBuilder<M, H, C, F> {
-    queue_size: usize,
+    config: Config,
-    parallel: u32,
     stream_handlers: dashmap::DashMap<u32, Arc<H>>,
     callback: C,
-    ordered: bool,
     _m: PhantomData<(M, F)>,
 }
 
@@ -58,13 +85,41 @@ where
     F: Sync + Send + Future<Output = Result<H, Error>> + 'static,
     C: Sync + Send + Fn(u32, u32) -> F + 'static,
 {
-    pub fn new(queue_size: usize, parallel: u32, ordered: bool, callback: C) -> Self {
+    pub fn new(queue_size: usize, task_count: u32, callback: C) -> Self {
         Self {
-            queue_size,
+            config: Config {
-            parallel,
+                queue_size,
+                queue_per_task: false,
+                ordered: false,
+                task_count,
+                lazy_task_creation: true,
+            },
             stream_handlers: Default::default(),
             callback,
-            ordered,
+            _m: PhantomData,
+        }
+    }
+
+    pub fn ordered(self) -> Self {
+        let mut config = self.config;
+        config.ordered = true;
+
+        Self {
+            config,
+            stream_handlers: self.stream_handlers,
+            callback: self.callback,
+            _m: PhantomData,
+        }
+    }
+
+    pub fn queue_per_task(self) -> Self {
+        let mut config = self.config;
+        config.queue_per_task = true;
+
+        Self {
+            config,
+            stream_handlers: self.stream_handlers,
+            callback: self.callback,
             _m: PhantomData,
         }
     }
@@ -92,11 +147,7 @@ where
         }
     }
 
-    fn queue_size(&self, _stream_id: u32, _task_id: u32) -> usize {
+    fn config(&self, _stream_id: u32) -> Config {
-        self.queue_size
+        self.config
-    }
-
-    fn parallel(&self, _stream_id: u32) -> (u32, bool) {
-        (self.parallel, self.ordered)
     }
 }

src/chan.rs

@@ -58,6 +58,7 @@ impl<M: Message> BusSender<M> {
     }
 }
 
+#[derive(Clone)]
 pub(crate) struct Sender<T> {
     inner: kanal::AsyncSender<ChannelItem<T>>,
 }
@@ -83,6 +84,7 @@ impl<T> Sender<T> {
     }
 }
 
+#[derive(Clone)]
 pub(crate) struct Receiver<T> {
     inner: kanal::AsyncReceiver<ChannelItem<T>>,
 }

src/handler.rs

@@ -1,6 +1,7 @@
 use std::{
     any::{Any, TypeId},
     marker::PhantomData,
+    pin::Pin,
     sync::Arc,
 };
 
@@ -9,7 +10,7 @@ use tokio::sync::Notify;
 
 use crate::{
     builder::Builder,
-    chan::{channel, Sender},
+    chan::Receiver,
     message::Msg,
     task::{TaskCounter, TaskSpawner},
     BusInner, Error, IntoMessage, Message,
@@ -50,15 +51,15 @@ where
 {
     fn spawn_task(
         &self,
+        rx: Receiver<Msg<M>>,
         stream_id: u32,
         task_id: u32,
         _abort: Arc<Notify>,
         task_counter: Arc<TaskCounter>,
         bus: Arc<BusInner>,
-    ) -> Box<dyn Future<Output = Result<Sender<Msg<M>>, Error>> + Send + '_> {
+    ) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send + '_>> {
-        Box::new(async move {
+        Box::pin(async move {
             let bus = bus.clone();
-            let (tx, rx) = channel::<Msg<M>>(self.builder.queue_size(stream_id, task_id));
             let mut ctx = self.builder.build(stream_id, task_id).await?;
 
             let _handle = tokio::spawn(async move {
@@ -96,7 +97,7 @@ where
                     println!("TASK FINALIZE ERROR: {:?}", err);
                 }
             });
-            Ok(tx)
+            Ok(())
         })
     }
 
@@ -104,7 +105,7 @@ where
         false
     }
 
-    fn parallel(&self, stream_id: u32) -> (u32, bool) {
+    fn config(&self, stream_id: u32) -> crate::builder::Config {
-        self.builder.parallel(stream_id)
+        self.builder.config(stream_id)
     }
 }

src/lib.rs

@@ -19,7 +19,8 @@ use std::{
     },
 };
 
-use chan::{BusSender, BusSenderClose};
+use builder::Config;
+use chan::{channel, BusSender, BusSenderClose, Receiver, Sender};
 use dashmap::DashMap;
 use futures::Future;
 use message::Msg;
@@ -47,13 +48,13 @@ struct BusInner {
 }
 
 impl BusInner {
-    fn get_task_id<M: Message>(&self, stream_id: u32, task_count: u32) -> u32 {
+    fn get_task_id<M: Message>(&self, stream_id: u32, config: &Config) -> u32 {
-        if task_count == 1 {
+        if !config.queue_per_task || config.task_count == 1 {
             return DEFAUL_TASK_ID;
         }
 
         let type_id = TypeId::of::<M>();
-        let (id1, id2) = self.rng.next_u32_pair(task_count);
+        let (id1, id2) = self.rng.next_u32_pair(config.task_count);
 
         let Some(l1) = self
             .senders
@@ -80,14 +81,40 @@ impl BusInner {
         }
     }
 
+    async fn spawn_task<M: Message>(
+        self: &Arc<Self>,
+        spawner: &TaskSpawnerWrapper<M>,
+        chan: (Sender<Msg<M>>, Receiver<Msg<M>>),
+        stream_id: u32,
+        task_id: u32,
+    ) -> Result<(), Error> {
+        let sender = spawner
+            .spawn_task(
+                chan,
+                stream_id,
+                task_id,
+                self.abort_notify.clone(),
+                self.task_counter.clone(),
+                self.clone(),
+            )
+            .await?;
+
+        self.senders.insert(
+            (stream_id, task_id, TypeId::of::<M>()),
+            Arc::new(sender) as _,
+        );
+
+        Ok(())
+    }
+
     async fn send_inner<M: Message>(
         self: &Arc<Self>,
         msg: Msg<M>,
-        task_count: u32,
+        config: Config,
     ) -> Result<(), Error> {
         let type_id = TypeId::of::<M>();
         let stream_id = msg.stream_id;
-        let task_id = self.get_task_id::<M>(stream_id, task_count);
+        let task_id = self.get_task_id::<M>(stream_id, &config);
 
         if !self.senders.contains_key(&(stream_id, task_id, type_id)) {
             let spawner = if let Some(spawner) = self.spawners.read().await.get(&type_id) {
@@ -99,36 +126,51 @@ impl BusInner {
                 return Err(Error::HandlerIsNotRegistered);
             };
 
-            let tx = spawner
+            if config.queue_per_task {
-                .spawn_task(
+                if config.lazy_task_creation {
-                    stream_id,
+                    self.spawn_task(
-                    task_id,
+                        &spawner,
-                    self.abort_notify.clone(),
+                        channel::<Msg<M>>(config.queue_size),
-                    self.task_counter.clone(),
+                        stream_id,
-                    self.clone(),
+                        task_id,
-                )
+                    )
-                .await?;
+                    .await?;
-
+                } else {
-            tx.send(msg).await.unwrap();
+                    for tid in 0..config.task_count {
-
+                        self.spawn_task(
-            self.senders
+                            &spawner,
-                .insert((stream_id, task_id, type_id), Arc::new(tx) as _);
+                            channel::<Msg<M>>(config.queue_size),
-        } else {
+                            stream_id,
-            let senders = self
+                            tid,
-                .senders
+                        )
-                .get(&(stream_id, task_id, type_id))
+                        .await?;
-                .unwrap()
+                    }
-                .clone();
+                }
-
+            } else {
-            senders
+                let (tx, rx) = channel::<Msg<M>>(config.queue_size);
-                .upcast()
+                for tid in 0..config.task_count - 1 {
-                .downcast_ref::<BusSender<M>>()
+                    self.spawn_task(&spawner, (tx.clone(), rx.clone()), stream_id, tid)
-                .unwrap()
+                        .await?;
-                .send(msg)
+                }
-                .await
+                self.spawn_task(&spawner, (tx, rx), stream_id, config.task_count - 1)
-                .unwrap();
+                    .await?;
+            };
         }
 
+        let senders = self
+            .senders
+            .get(&(stream_id, task_id, type_id))
+            .unwrap()
+            .clone();
+
+        senders
+            .upcast()
+            .downcast_ref::<BusSender<M>>()
+            .unwrap()
+            .send(msg)
+            .await
+            .unwrap();
+
         Ok(())
     }
 
@@ -136,19 +178,19 @@ impl BusInner {
         let type_id = TypeId::of::<M>();
         let stream_id = msg.stream_id;
 
-        let (task_count, ordered) = if let Some(spawner) = self
+        let config = if let Some(spawner) = self
             .spawners
             .read()
             .await
             .get(&type_id)
             .and_then(|x| x.downcast_ref::<TaskSpawnerWrapper<M>>())
         {
-            spawner.parallel(stream_id)
+            spawner.config(stream_id)
         } else {
-            (1, false)
+            Config::default()
         };
 
-        if ordered {
+        if config.ordered {
             // let queue = self
             //     .reordering_queue
             //     .get_or_insert(&(stream_id, type_id), task_count);
@@ -159,9 +201,9 @@ impl BusInner {
             // }
 
             // Ok(())
-            self.send_inner(msg, task_count).await
+            self.send_inner(msg, config).await
         } else {
-            self.send_inner(msg, task_count).await
+            self.send_inner(msg, config).await
         }
     }
 
@@ -419,7 +461,7 @@ mod tests {
     async fn test_tasks_shared() {
         let bus = Bus::default();
 
-        bus.register(SharedBuilder::new(2, 5, false, |sid, _tid| async move {
+        bus.register(SharedBuilder::new(2, 5, |sid, _tid| async move {
             Ok(TestConsumer(sid))
         }))
         .await;

src/producer.rs

@@ -1,11 +1,15 @@
-use std::{marker::PhantomData, pin::pin, sync::Arc};
+use std::{
+    marker::PhantomData,
+    pin::{pin, Pin},
+    sync::Arc,
+};
 
 use futures::{Future, Stream, StreamExt};
 use tokio::sync::Notify;
 
 use crate::{
     builder::Builder,
-    chan::{channel, Sender},
+    chan::Receiver,
     message::Msg,
     task::{TaskCounter, TaskSpawner},
     BusInner, Error, IntoMessage, Message,
@@ -46,14 +50,14 @@ where
 {
     fn spawn_task(
         &self,
+        rx: Receiver<Msg<M>>,
         stream_id: u32,
         task_id: u32,
         abort: Arc<Notify>,
         task_counter: Arc<TaskCounter>,
         bus: Arc<BusInner>,
-    ) -> Box<dyn Future<Output = Result<Sender<Msg<M>>, Error>> + Send + '_> {
+    ) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send + '_>> {
-        Box::new(async move {
+        Box::pin(async move {
-            let (tx, rx) = channel::<Msg<M>>(self.builder.queue_size(stream_id, task_id));
             let mut ctx = self.builder.build(stream_id, task_id).await?;
 
             let _handle = tokio::spawn(async move {
@@ -102,7 +106,7 @@ where
                 }
             });
 
-            Ok(tx)
+            Ok(())
         })
     }
 
@@ -110,7 +114,7 @@ where
         true
     }
 
-    fn parallel(&self, stream_id: u32) -> (u32, bool) {
+    fn config(&self, stream_id: u32) -> crate::builder::Config {
-        self.builder.parallel(stream_id)
+        self.builder.config(stream_id)
     }
 }

src/task.rs

@@ -1,13 +1,17 @@
-use std::sync::{
+use std::{
-    atomic::{AtomicUsize, Ordering},
+    pin::Pin,
-    Arc,
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc,
+    },
 };
 
 use futures::Future;
 use tokio::sync::Notify;
 
 use crate::{
-    chan::{BusSender, Sender},
+    builder::Config,
+    chan::{BusSender, Receiver, Sender},
     handler::HandlerSpawner,
     message::Msg,
     producer::ProducerSpawner,
@@ -41,16 +45,17 @@ impl TaskCounter {
 }
 
 pub(crate) trait TaskSpawner<M: Message>: Send + Sync {
-    fn parallel(&self, stream_id: u32) -> (u32, bool);
+    fn config(&self, stream_id: u32) -> Config;
     fn is_producer(&self) -> bool;
     fn spawn_task(
         &self,
+        rx: Receiver<Msg<M>>,
         stream_id: u32,
         task_id: u32,
         abort: Arc<Notify>,
         task_counter: Arc<TaskCounter>,
         bus: Arc<BusInner>,
-    ) -> Box<dyn Future<Output = Result<Sender<Msg<M>>, Error>> + Send + '_>;
+    ) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send + '_>>;
 }
 
 pub(crate) struct TaskSpawnerWrapper<M: Message> {
@@ -87,24 +92,22 @@ impl<M: Message> TaskSpawnerWrapper<M> {
     #[inline]
     pub async fn spawn_task(
         &self,
+        (tx, rx): (Sender<Msg<M>>, Receiver<Msg<M>>),
         stream_id: u32,
         task_id: u32,
         abort: Arc<Notify>,
         task_counter: Arc<TaskCounter>,
         bus: Arc<BusInner>,
     ) -> Result<BusSender<M>, Error> {
-        Ok(BusSender::new(
+        self.inner
-            self.inner.is_producer(),
+            .spawn_task(rx, stream_id, task_id, abort, task_counter, bus)
-            Box::into_pin(
+            .await?;
-                self.inner
+
-                    .spawn_task(stream_id, task_id, abort, task_counter, bus),
+        Ok(BusSender::new(self.inner.is_producer(), tx))
-            )
-            .await?,
-        ))
     }
 
     #[inline]
-    pub(crate) fn parallel(&self, stream_id: u32) -> (u32, bool) {
+    pub(crate) fn config(&self, stream_id: u32) -> Config {
-        self.inner.parallel(stream_id)
+        self.inner.config(stream_id)
     }
 }