use std::{
any::{type_name, Any},
marker::PhantomData,
num::NonZeroU32,
pin::{pin, Pin},
sync::{
atomic::{AtomicU64, Ordering},
Arc,
},
};
use futures::Future;
use tokio::sync::Notify;
use crate::{
builder::Builder,
chan::Receiver,
message::{IntoMessages, Msg},
task::{TaskCounter, TaskSpawner},
AsyncIterator, Bus, BusInner, Error, ErrorMessage, Message,
};
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct Context {
pub stream_id: u32,
pub task_id: u32,
pub index: u64,
pub load: (u32, u32),
}
pub trait Handler<M: Message>: Send + 'static {
type Result: Message + Unpin;
type Error: ErrorMessage + Unpin;
fn handle(
&mut self,
msg: M,
ctx: Context,
bus: crate::Bus,
) -> impl Future<
Output = Result<impl IntoMessages<Self::Result, Self::Error> + Send + '_, Self::Error>,
> + Send
+ '_;
fn handle_error(
&mut self,
_err: Error,
_ctx: Context,
_bus: crate::Bus,
) -> impl Future<
Output = Result<impl IntoMessages<Self::Result, Self::Error> + Send + '_, Self::Error>,
> + Send
+ '_;
fn finalize(self, bus: crate::Bus) -> impl Future<Output = Result<(), Self::Error>> + Send;
}
pub(crate) struct HandlerSpawner<M, B> {
pub(crate) builder: B,
_m: PhantomData<M>,
}
impl<M, B> HandlerSpawner<M, B> {
pub(crate) fn new(builder: B) -> Self {
Self {
builder,
_m: PhantomData,
}
}
}
impl<M: Message, B: Builder<M>> TaskSpawner<M> for HandlerSpawner<M, B>
where
B::Context: Any + Handler<M>,
{
fn spawn_task(
&self,
rx: Receiver<Msg<M>>,
stream_id: u32,
task_id: u32,
_abort: Arc<Notify>,
task_counter: Arc<TaskCounter>,
spawn_counter: Arc<TaskCounter>,
index_counter: Arc<AtomicU64>,
bus: Arc<BusInner>,
) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send + '_>> {
Box::pin(async move {
let bus = Bus { inner: bus.clone() };
let config = self.builder.config(stream_id);
let mut handler = self.builder.build(stream_id, task_id).await?;
let _handle = tokio::spawn(async move {
let _test = spawn_counter.clone().lease_unit(|| true);
while let Some(msg) = rx.recv().await {
let _test = task_counter.clone().lease_unit(|| rx.is_empty());
let ctx = Context {
stream_id,
task_id,
index: msg.index,
load: rx.load(),
};
let res = match msg.inner {
Some(Ok(m)) => {
send_result(
&bus.inner,
&index_counter,
msg.index,
stream_id,
&config,
Some(
handler
.handle(m, ctx, bus.clone())
.await
.map(IntoMessages::into_messages),
),
)
.await
}
Some(Err(err)) => {
send_result(
&bus.inner,
&index_counter,
msg.index,
stream_id,
&config,
Some(
handler
.handle_error(err, ctx, bus.clone())
.await
.map(IntoMessages::into_messages),
),
)
.await
}
None => {
send_result::<
<B::Context as Handler<M>>::Result,
<B::Context as Handler<M>>::Error,
>(
&bus.inner,
&index_counter,
msg.index,
stream_id,
&config,
None::<Result<crate::Empty<_>, _>>,
)
.await
}
};
if let Err(err) = res {
println!(
"Messagebus Send Error: {}/{} {}: {}",
stream_id,
task_id,
type_name::<<B::Context as Handler<M>>::Result>(),
err,
);
}
}
println!(
"TASK #{} of type `{}` ENDED",
task_id,
std::any::type_name::<B>()
);
if let Err(err) = handler.finalize(bus.clone()).await {
println!("TASK FINALIZE ERROR: {:?}", err);
}
});
Ok(())
})
}
fn is_producer(&self) -> bool {
false
}
fn config(&self, stream_id: u32) -> crate::builder::Config {
self.builder.config(stream_id)
}
}
async fn send_result<'a, M: Message, E: ErrorMessage>(
bus: &Arc<BusInner>,
index_counter: &AtomicU64,
index: u64,
stream_id: u32,
config: &crate::builder::Config,
res: Option<Result<impl AsyncIterator<Item = Result<M, E>> + Send + 'a, E>>,
) -> Result<(), Error> {
let reorder_buff = if config.ordered && config.task_count > 1 {
config
.ordering_buffer_size
.map(NonZeroU32::get)
.unwrap_or(config.task_count)
} else {
0
};
let one = match res {
Some(Ok(iter)) => {
let hint = iter.size_hint();
let mut iter = pin!(iter);
match hint {
(_, Some(0)) => None,
(_, Some(1)) => iter.next().await,
_ => {
while let Some(item) = iter.as_mut().next().await {
let index = index_counter.fetch_add(1, Ordering::Relaxed);
let stream_id = if config.stream_per_message {
bus.next_stream_id()
} else {
stream_id
};
bus.send::<M>(
Some(item.map_err(Into::into)),
index,
stream_id,
reorder_buff,
)
.await?;
}
return Ok(());
}
}
}
Some(Err(err)) => Some(Err(err)),
None => None,
};
bus.send(
one.map(|x| x.map_err(Into::into)),
index,
stream_id,
reorder_buff,
)
.await?;
Ok(())
}