register_mapper

This commit is contained in:
Andrey Tkachenko 2023-11-23 09:58:39 +04:00
parent 33db3cb15d
commit d57dd8fed9
4 changed files with 155 additions and 32 deletions

View File

@ -39,13 +39,25 @@ pub trait Builder<M: Message>: Send + Sync + 'static {
) -> impl Future<Output = Result<Self::Context, Error>> + Send + '_;
}
pub struct DefaultBuilder<M: Message, H> {
pub struct DefaultBuilder<M: Send + Sync, H: Send + Sync, C: Send + Sync, F: Send> {
config: Config,
_m: PhantomData<(M, H)>,
callback: C,
_m: PhantomData<(M, H, F)>,
}
impl<M: Message, H> DefaultBuilder<M, H> {
pub fn new(queue_size: usize) -> Self {
unsafe impl<M: Send + Sync, H: Send + Sync, C: Send + Sync, F: Send> Sync
for DefaultBuilder<M, H, C, F>
{
}
impl<M, H, C, F> DefaultBuilder<M, H, C, F>
where
M: Message,
H: Sync + Send + 'static,
F: Send + Future<Output = Result<H, Error>> + 'static,
C: Sync + Send + Fn(u32, u32) -> F + 'static,
{
pub fn new(queue_size: usize, callback: C) -> Self {
Self {
config: Config {
queue_size,
@ -54,6 +66,7 @@ impl<M: Message, H> DefaultBuilder<M, H> {
task_count: 1,
lazy_task_creation: true,
},
callback,
_m: PhantomData,
}
}
@ -64,16 +77,34 @@ impl<M: Message, H> DefaultBuilder<M, H> {
Self {
config,
callback: self.callback,
_m: PhantomData,
}
}
pub fn tasks(self, tasks: u32) -> Self {
let mut config = self.config;
config.task_count = tasks;
Self {
config,
callback: self.callback,
_m: PhantomData,
}
}
}
impl<M: Message, H: Sync + Send + Default + 'static> Builder<M> for DefaultBuilder<M, H> {
impl<M, H, C, F> Builder<M> for DefaultBuilder<M, H, C, F>
where
M: Message,
H: Sync + Send + 'static,
F: Send + Future<Output = Result<H, Error>> + 'static,
C: Sync + Send + Fn(u32, u32) -> F + 'static,
{
type Context = H;
async fn build(&self, _stream_id: u32, _task_id: u32) -> Result<Self::Context, Error> {
Ok(<Self::Context as Default>::default())
async fn build(&self, stream_id: u32, task_id: u32) -> Result<Self::Context, Error> {
(self.callback)(stream_id, task_id).await
}
fn config(&self, _stream_id: u32) -> Config {
@ -88,11 +119,16 @@ pub struct SharedBuilder<M, H, C, F> {
_m: PhantomData<(M, F)>,
}
unsafe impl<M: Send + Sync, H: Send + Sync, C: Send + Sync, F: Send> Sync
for SharedBuilder<M, H, C, F>
{
}
impl<M, H, C, F> SharedBuilder<M, H, C, F>
where
M: Message,
H: Sync + Send + 'static,
F: Sync + Send + Future<Output = Result<H, Error>> + 'static,
F: Send + Future<Output = Result<H, Error>> + 'static,
C: Sync + Send + Fn(u32, u32) -> F + 'static,
{
pub fn new(queue_size: usize, task_count: u32, callback: C) -> Self {
@ -139,7 +175,7 @@ impl<M, H, C, F> Builder<M> for SharedBuilder<M, H, C, F>
where
M: Message,
H: Sync + Send + 'static,
F: Sync + Send + Future<Output = Result<H, Error>> + 'static,
F: Send + Future<Output = Result<H, Error>> + 'static,
C: Sync + Send + Fn(u32, u32) -> F + 'static,
{
type Context = Arc<H>;

View File

@ -14,7 +14,16 @@ pub enum Error {
HandlerError(Arc<dyn ErrorMessage>),
}
#[derive(Debug)]
pub enum VoidError {}
impl std::fmt::Display for VoidError {
fn fmt(&self, _f: &mut fmt::Formatter<'_>) -> fmt::Result {
Ok(())
}
}
impl std::error::Error for Error {}
impl ErrorMessage for VoidError {}
impl Eq for Error {}

View File

@ -14,6 +14,7 @@ mod task;
use std::{
any::{Any, TypeId},
collections::HashMap,
marker::PhantomData,
sync::{
atomic::{AtomicBool, AtomicU64, Ordering},
Arc,
@ -32,7 +33,7 @@ use tokio::sync::{Notify, RwLock};
pub use async_iter::*;
pub use builder::{Builder, DefaultBuilder, SharedBuilder};
pub use error::Error;
pub use error::{Error, VoidError};
pub use handler::Handler;
pub use message::{ErrorMessage, IntoMessages, Message};
@ -193,10 +194,12 @@ impl BusInner {
}
} else {
let (tx, rx) = channel::<Msg<M>>(config.queue_size);
for tid in 0..config.task_count - 1 {
self.spawn_task(&spawner, (tx.clone(), rx.clone()), stream_id, tid)
.await?;
}
self.spawn_task(&spawner, (tx, rx), stream_id, config.task_count - 1)
.await?;
};
@ -244,8 +247,6 @@ impl BusInner {
let mut queue = ReorderQueueInner::new(reorder_buff as _);
while let Ok(msg) = rx.recv().await {
println!(" > {:?}", msg);
if let Some(index) = queue.push(msg.index, msg.inner) {
bus.send_inner(
Some(Err::<M, _>(Error::ReorderingDropMessage(index))),
@ -372,6 +373,24 @@ impl Bus {
self
}
#[inline]
pub async fn register_mapper<
M: Message,
R: Message,
E: ErrorMessage,
C: Send + Clone + Sync + FnMut(u32, u32, M) -> Result<R, E> + 'static,
>(
&self,
cb: C,
) -> &Self {
let mapper = DefaultBuilder::new(0, move |_, _| {
let cb = cb.clone();
async move { Ok(Mapper { cb, m: PhantomData }) }
});
self.inner.clone().register(mapper).await;
self
}
#[inline]
pub async fn send<M: Message>(&self, inner: M) -> Result<(), Error> {
self.send_with_stream(DEFAUL_STREAM_ID, inner).await
@ -417,6 +436,46 @@ impl Bus {
}
}
struct Mapper<M, R, E, C> {
cb: C,
m: PhantomData<(M, R, E)>,
}
impl<M: Message, R: Message, E: ErrorMessage, C> Handler<M> for Mapper<M, R, E, C>
where
M: Message,
R: Message,
E: ErrorMessage,
C: Send + Sync + FnMut(u32, u32, M) -> Result<R, E> + 'static,
{
type Result = R;
type Error = E;
async fn handle(
&mut self,
_msg: M,
_stream_id: u32,
_task_id: u32,
_bus: Bus,
) -> Result<impl IntoMessages<Self::Result, Self::Error>, Self::Error> {
(self.cb)(_stream_id, _task_id, _msg).map(|x| [x])
}
async fn handle_error(
&mut self,
_err: Error,
_stream_id: u32,
_task_id: u32,
_bus: Bus,
) -> Result<impl IntoMessages<Self::Result, Self::Error>, Self::Error> {
Ok(None)
}
async fn finalize(self, _bus: Bus) -> Result<(), Self::Error> {
Ok(())
}
}
#[cfg(test)]
mod tests {
use std::{sync::Arc, time::Duration};
@ -631,10 +690,14 @@ mod tests {
async fn test_streams() {
let bus = Bus::default();
bus.register(DefaultBuilder::<u64, TestConsumer>::new(2))
bus.register(DefaultBuilder::<u64, _, _, _>::new(2, |_, _| async move {
Ok(TestConsumer::default())
}))
.await;
bus.register(DefaultBuilder::<u32, TestProducer>::new(2))
bus.register(DefaultBuilder::<u32, _, _, _>::new(2, |_, _| async move {
Ok(TestProducer)
}))
.await;
for start in 0u32..10 {
@ -655,7 +718,9 @@ mod tests {
}))
.await;
bus.register(DefaultBuilder::<u32, TestProducer>::new(2))
bus.register(DefaultBuilder::<u32, _, _, _>::new(2, |_, _| async move {
Ok(TestProducer)
}))
.await;
for start in 0u32..10 {
@ -680,7 +745,9 @@ mod tests {
)
.await;
bus.register(DefaultBuilder::<_, TestCollector>::new(4))
bus.register(DefaultBuilder::<_, _, _, _>::new(4, |_, _| async move {
Ok(TestCollector::default())
}))
.await;
for i in 0i16..1024 {

View File

@ -41,9 +41,17 @@ impl<M> ReorderQueueInner<M> {
}
pub fn push(&mut self, index: u64, inner: M) -> Option<u64> {
if let Some(ri) = self.recent_index {
if index <= ri {
return Some(index);
}
}
self.heap.push(Entry { inner, index });
if self.heap.len() == self.cap {
if self.heap.len() > self.cap {
let _ = self.heap.pop();
self.recent_index = self.recent_index.map(|x| x + 1);
self.recent_index
} else {
@ -52,13 +60,7 @@ impl<M> ReorderQueueInner<M> {
}
pub fn pop(&mut self) -> Option<(u64, M)> {
match self.recent_index {
None => {
let e = self.heap.pop()?;
self.recent_index = Some(e.index);
Some((e.index, e.inner))
}
Some(ri) => {
if let Some(ri) = self.recent_index {
let e = self.heap.peek()?;
if e.index == ri + 1 {
self.recent_index = Some(e.index);
@ -66,6 +68,15 @@ impl<M> ReorderQueueInner<M> {
} else {
None
}
} else {
let e = self.heap.peek()?;
if e.index == 0 {
let e = self.heap.pop()?;
self.recent_index = Some(e.index);
Some((e.index, e.inner))
} else {
None
}
}
}