1use futures::TryStreamExt as _;
12use simploxide_api_types::events::{Event, EventData};
13#[cfg(feature = "cancellation")]
14use tokio_util::sync::CancellationToken;
15
16use std::{future::Future, pin::Pin, sync::Arc};
17
18use crate::{EventParser, EventStream, StreamEvents};
19
20pub struct DispatchChain<P, Ctx> {
22 events: EventStream<P>,
23 ctx: Ctx,
24}
25
26impl<P, Ctx> DispatchChain<P, Ctx>
27where
28 P: EventParser,
29{
30 pub fn with_ctx(events: EventStream<P>, ctx: Ctx) -> Self {
31 Self { ctx, events }
32 }
33
34 pub fn seq_fallback<E, F>(mut self, f: F) -> Dispatcher<P, Ctx, Fallback<F>>
36 where
37 F: AsyncFnMut(Event, &mut Ctx) -> Result<StreamEvents, E>,
38 {
39 self.events.accept_all();
40 Dispatcher {
41 events: self.events,
42 ctx: self.ctx,
43 chain: Fallback { f },
44 }
45 }
46
47 pub fn fallback<E, F, Fut>(mut self, f: F) -> Dispatcher<P, Ctx, Fallback<F>>
52 where
53 Ctx: 'static + Send,
54 E: 'static + Send + From<P::Error>,
55 F: Fn(Event, Ctx) -> Fut,
56 Fut: 'static + Send + Future<Output = Result<StreamEvents, E>>,
57 {
58 self.events.accept_all();
59 Dispatcher {
60 events: self.events,
61 ctx: self.ctx,
62 chain: Fallback { f },
63 }
64 }
65
66 pub fn seq<Ev, E, F>(mut self, f: F) -> Dispatcher<P, Ctx, Match<Ev, F>>
70 where
71 Ev: EventData,
72 F: AsyncFnMut(Arc<Ev>, &mut Ctx) -> Result<StreamEvents, E>,
73 {
74 self.events.reject_all();
75 self.events.accept(Ev::KIND);
76 Dispatcher {
77 events: self.events,
78 ctx: self.ctx,
79 chain: Match {
80 f,
81 _phantom: std::marker::PhantomData,
82 },
83 }
84 }
85
86 pub fn on<Ev, E, F, Fut>(mut self, f: F) -> Dispatcher<P, Ctx, Match<Ev, F>>
92 where
93 Ctx: 'static + Send,
94 E: 'static + Send,
95 Ev: 'static + EventData,
96 F: Fn(Arc<Ev>, Ctx) -> Fut,
97 Fut: 'static + Send + Future<Output = Result<StreamEvents, E>>,
98 {
99 self.events.reject_all();
100 self.events.accept(Ev::KIND);
101 Dispatcher {
102 events: self.events,
103 ctx: self.ctx,
104 chain: Match {
105 f,
106 _phantom: std::marker::PhantomData,
107 },
108 }
109 }
110}
111
112pub struct Dispatcher<P, Ctx, D> {
118 events: EventStream<P>,
119 ctx: Ctx,
120 chain: D,
121}
122
123impl<P, Ctx, D> Dispatcher<P, Ctx, D>
124where
125 D: DispatchEvent<Ctx>,
126{
127 pub fn seq<Ev, F>(mut self, f: F) -> Dispatcher<P, Ctx, Intercept<Match<Ev, F>, D>>
176 where
177 Ev: EventData,
178 F: AsyncFnMut(Arc<Ev>, &mut Ctx) -> Result<StreamEvents, D::Error>,
179 {
180 self.events.accept(Ev::KIND);
181 Dispatcher {
182 events: self.events,
183 ctx: self.ctx,
184 chain: Intercept {
185 d1: Match {
186 f,
187 _phantom: std::marker::PhantomData,
188 },
189 d2: self.chain,
190 },
191 }
192 }
193
194 pub async fn sequential_dispatch(self) -> Result<(EventStream<P>, Ctx), D::Error>
201 where
202 P: EventParser,
203 D::Error: From<P::Error>,
204 {
205 let Self {
206 ctx,
207 events,
208 mut chain,
209 } = self;
210
211 events.stream_events_with_ctx_mut(async move |ev, ctx| {
212 let Ok(handler) = chain.dispatch_event(ev, ctx) else {
213 unreachable!("EventStream filters set by seq/fallback_seq drop events without handlers during parsing");
214 };
215
216 handler.await
217
218 }, ctx).await
219 }
220
221 #[cfg(feature = "cancellation")]
224 pub async fn sequential_dispatch_with_cancellation(
225 self,
226 token: CancellationToken,
227 ) -> Result<(EventStream<P>, Ctx), D::Error>
228 where
229 P: EventParser,
230 D::Error: From<P::Error>,
231 {
232 let Self {
233 mut ctx,
234 mut events,
235 mut chain,
236 } = self;
237
238 loop {
239 tokio::select! {
240 biased;
241 _ = token.cancelled() => break,
242 res = events.try_next() => match res {
243 Ok(Some(ev)) => {
244 let Ok(handler) = chain.dispatch_event(ev, &mut ctx) else {
245 unreachable!("EventStream filters set by seq/fallback_seq drop events without handlers during parsing");
246 };
247 if let StreamEvents::Break = handler.await? {
248 break;
249 }
250 }
251 Ok(None) => break,
252 Err(e) => return Err(e.into()),
253 }
254 }
255 }
256
257 Ok((events, ctx))
258 }
259}
260
261impl<P, Ctx, D> Dispatcher<P, Ctx, D>
262where
263 P: 'static + EventParser,
264 Ctx: 'static + Send + Clone,
265 D: ConcurrentDispatchEvent<Ctx>,
266 D::Error: From<P::Error>,
267{
268 pub fn on<Ev, F, Fut>(mut self, f: F) -> Dispatcher<P, Ctx, Intercept<Match<Ev, F>, D>>
316 where
317 Ev: 'static + EventData,
318 F: Fn(Arc<Ev>, Ctx) -> Fut,
319 Fut: 'static + Send + Future<Output = Result<StreamEvents, D::Error>>,
320 {
321 self.events.accept(Ev::KIND);
322 Dispatcher {
323 events: self.events,
324 ctx: self.ctx,
325 chain: Intercept {
326 d1: Match {
327 f,
328 _phantom: std::marker::PhantomData,
329 },
330 d2: self.chain,
331 },
332 }
333 }
334
335 pub async fn dispatch(self) -> Result<(EventStream<P>, Ctx, Vec<Event>), D::Error> {
345 let chain = self.chain;
346 let ctx = self.ctx;
347 let mut events = self.events;
348 let (event_buffer, result) =
349 run_concurrent_dispatch(&chain, &ctx, &mut events, std::future::pending::<()>()).await;
350 match result {
351 Ok(inner) => inner.map(move |_| (events, ctx, event_buffer)),
352 Err(e) => std::panic::resume_unwind(e.into_panic()),
353 }
354 }
355
356 #[cfg(feature = "cancellation")]
359 pub async fn dispatch_with_cancellation(
360 self,
361 token: CancellationToken,
362 ) -> Result<(EventStream<P>, Ctx, Vec<Event>), D::Error> {
363 let chain = self.chain;
364 let ctx = self.ctx;
365 let mut events = self.events;
366 let (event_buffer, result) =
367 run_concurrent_dispatch(&chain, &ctx, &mut events, token.cancelled()).await;
368 match result {
369 Ok(inner) => inner.map(move |_| (events, ctx, event_buffer)),
370 Err(e) => std::panic::resume_unwind(e.into_panic()),
371 }
372 }
373
374 pub async fn dispatch_sequentially(self) -> Result<(EventStream<P>, Ctx), D::Error> {
380 let ctx = self.ctx;
381 let events = self.events;
382 let chain = self.chain;
383
384 events.stream_events_with_ctx_cloned(async move |ev, ctx| {
385 let Ok(handler) = chain.concurrent_dispatch_event(ev, ctx) else {
386 unreachable!("EventStream filters set by on/fallback drop events without handlers during parsing");
387 };
388 handler.await
389 }, ctx).await
390 }
391
392 #[cfg(feature = "cancellation")]
394 pub async fn dispatch_sequentially_with_cancellation(
395 self,
396 token: CancellationToken,
397 ) -> Result<(EventStream<P>, Ctx), D::Error> {
398 let Self {
399 ctx,
400 mut events,
401 chain,
402 } = self;
403
404 loop {
405 tokio::select! {
406 biased;
407 _ = token.cancelled() => break,
408 res = events.try_next() => match res {
409 Ok(Some(ev)) => {
410 let Ok(handler) = chain.concurrent_dispatch_event(ev, ctx.clone()) else {
411 unreachable!("EventStream filters set by on/fallback drop events without handlers during parsing");
412 };
413 if let StreamEvents::Break = handler.await? {
414 break;
415 }
416 }
417 Ok(None) => break,
418 Err(e) => return Err(e.into()),
419 }
420 }
421 }
422
423 Ok((events, ctx))
424 }
425}
426
427async fn run_concurrent_dispatch<P, Ctx, D, Fut>(
433 chain: &D,
434 ctx: &Ctx,
435 events: &mut EventStream<P>,
436 stop: Fut,
437) -> (
438 Vec<Event>,
439 Result<Result<StreamEvents, D::Error>, tokio::task::JoinError>,
440)
441where
442 P: 'static + EventParser,
443 Ctx: 'static + Send + Clone,
444 D: ConcurrentDispatchEvent<Ctx>,
445 D::Error: From<P::Error>,
446 Fut: Future<Output = ()>,
447{
448 let mut join_set: tokio::task::JoinSet<Result<StreamEvents, D::Error>> =
449 tokio::task::JoinSet::new();
450
451 let mut stop = std::pin::pin!(stop);
452
453 let mut result = loop {
454 tokio::select! {
455 _ = stop.as_mut() => break Ok(Ok(StreamEvents::Break)),
456 result = events.try_next() => match result {
457 Ok(Some(event)) => {
458 let Ok(handler) = chain.concurrent_dispatch_event(event, ctx.clone()) else {
459 unreachable!(
460 "EventStream filtering set by on and fallback methods drops events without handlers before parsing them"
461 );
462 };
463 join_set.spawn(handler);
464 }
465 Ok(None) => break Ok(Ok(StreamEvents::Break)),
466 Err(e) => break Ok(Err(e.into())),
467 },
468 result = join_set.join_next(), if !join_set.is_empty() => match result {
469 Some(Ok(Ok(StreamEvents::Continue))) => continue,
470 Some(Ok(Ok(StreamEvents::Break))) => break Ok(Ok(StreamEvents::Break)),
471 Some(err) => break err,
472 None => unreachable!("Dummy task must be running during the whole tokio select! loop"),
473 }
474 }
475 };
476
477 let mut event_buffer = Vec::new();
478
479 loop {
480 tokio::select! {
481 joined = join_set.join_next() => match joined {
482 Some(next) => {
483 if matches!(result, Ok(Ok(_))) {
484 result = next;
485 }
486 }
487 None => break,
488 },
489 event = events.try_next() => match event {
490 Ok(Some(ev)) => event_buffer.push(ev),
491 Ok(None) => (),
492 Err(e) => {
493 result = Ok(Err(e.into()));
494 break;
495 }
496 }
497 }
498 }
499
500 (event_buffer, result)
501}
502
503pub trait DispatchEvent<Ctx> {
504 type Error;
505 type Future<'s>: Future<Output = Result<StreamEvents, Self::Error>>
506 where
507 Self: 's,
508 Ctx: 's;
509
510 fn dispatch_event<'s>(
511 &'s mut self,
512 ev: Event,
513 ctx: &'s mut Ctx,
514 ) -> Result<Self::Future<'s>, (Event, &'s mut Ctx)>;
515}
516
517pub trait ConcurrentDispatchEvent<Ctx>
518where
519 Ctx: 'static + Send,
520{
521 type Error: 'static + Send;
522 type Future: 'static + Send + Future<Output = Result<StreamEvents, Self::Error>>;
523
524 fn concurrent_dispatch_event(&self, ev: Event, ctx: Ctx) -> Result<Self::Future, (Event, Ctx)>;
525}
526
527pub struct Fallback<F> {
530 f: F,
531}
532
533impl<Ctx, E, F> DispatchEvent<Ctx> for Fallback<F>
534where
535 F: AsyncFnMut(Event, &mut Ctx) -> Result<StreamEvents, E>,
536{
537 type Error = E;
538 type Future<'s>
540 = Pin<Box<dyn 's + Future<Output = Result<StreamEvents, E>>>>
541 where
542 Self: 's,
543 Ctx: 's;
544
545 fn dispatch_event<'s>(
546 &'s mut self,
547 ev: Event,
548 ctx: &'s mut Ctx,
549 ) -> Result<Self::Future<'s>, (Event, &'s mut Ctx)> {
550 Ok(Box::pin((self.f)(ev, ctx)))
551 }
552}
553
554impl<Ctx, E, F, Fut> ConcurrentDispatchEvent<Ctx> for Fallback<F>
555where
556 Ctx: 'static + Send,
557 E: 'static + Send,
558 F: Fn(Event, Ctx) -> Fut,
561 Fut: 'static + Send + Future<Output = Result<StreamEvents, E>>,
562{
563 type Error = E;
564 type Future = Fut;
565
566 fn concurrent_dispatch_event(&self, ev: Event, ctx: Ctx) -> Result<Self::Future, (Event, Ctx)> {
567 Ok((self.f)(ev, ctx))
568 }
569}
570
571pub struct Match<Ev, F> {
572 f: F,
573 _phantom: std::marker::PhantomData<Ev>,
574}
575
576impl<Ctx, Ev, E, F> DispatchEvent<Ctx> for Match<Ev, F>
577where
578 Ev: EventData,
579 F: AsyncFnMut(Arc<Ev>, &mut Ctx) -> Result<StreamEvents, E>,
580{
581 type Error = E;
582 type Future<'s>
584 = Pin<Box<dyn 's + Future<Output = Result<StreamEvents, E>>>>
585 where
586 Self: 's,
587 Ctx: 's;
588
589 fn dispatch_event<'s>(
590 &'s mut self,
591 ev: Event,
592 ctx: &'s mut Ctx,
593 ) -> Result<Self::Future<'s>, (Event, &'s mut Ctx)> {
594 match Ev::from_event(ev) {
595 Ok(ev) => Ok(Box::pin((self.f)(ev, ctx))),
596 Err(ev) => Err((ev, ctx)),
597 }
598 }
599}
600
601impl<Ctx, Ev, E, F, Fut> ConcurrentDispatchEvent<Ctx> for Match<Ev, F>
602where
603 Ctx: 'static + Send,
604 Ev: 'static + EventData,
605 E: 'static + Send,
606 F: Fn(Arc<Ev>, Ctx) -> Fut,
609 Fut: 'static + Send + Future<Output = Result<StreamEvents, E>>,
610{
611 type Error = E;
612 type Future = Fut;
613
614 fn concurrent_dispatch_event(&self, ev: Event, ctx: Ctx) -> Result<Self::Future, (Event, Ctx)> {
615 match Ev::from_event(ev) {
616 Ok(ev) => Ok((self.f)(ev, ctx)),
617 Err(ev) => Err((ev, ctx)),
618 }
619 }
620}
621
622pub struct Intercept<D1, D2> {
623 d1: D1,
624 d2: D2,
625}
626
627impl<Ctx, D1, D2> DispatchEvent<Ctx> for Intercept<D1, D2>
628where
629 D1: DispatchEvent<Ctx>,
630 D2: DispatchEvent<Ctx, Error = D1::Error>,
631{
632 type Error = D1::Error;
633 type Future<'s>
634 = futures::future::Either<D1::Future<'s>, D2::Future<'s>>
635 where
636 Self: 's,
637 Ctx: 's;
638
639 fn dispatch_event<'s>(
640 &'s mut self,
641 ev: Event,
642 ctx: &'s mut Ctx,
643 ) -> Result<Self::Future<'s>, (Event, &'s mut Ctx)> {
644 self.d1
645 .dispatch_event(ev, ctx)
646 .map(futures::future::Either::Left)
647 .or_else(|(ev, ctx)| {
648 self.d2
649 .dispatch_event(ev, ctx)
650 .map(futures::future::Either::Right)
651 })
652 }
653}
654
655impl<Ctx, D1, D2> ConcurrentDispatchEvent<Ctx> for Intercept<D1, D2>
656where
657 Ctx: 'static + Send,
658 D1: ConcurrentDispatchEvent<Ctx>,
659 D2: ConcurrentDispatchEvent<Ctx, Error = D1::Error>,
660{
661 type Error = D1::Error;
662 type Future = futures::future::Either<D1::Future, D2::Future>;
663
664 fn concurrent_dispatch_event(&self, ev: Event, ctx: Ctx) -> Result<Self::Future, (Event, Ctx)> {
665 self.d1
666 .concurrent_dispatch_event(ev, ctx)
667 .map(futures::future::Either::Left)
668 .or_else(|(ev, ctx)| {
669 self.d2
670 .concurrent_dispatch_event(ev, ctx)
671 .map(futures::future::Either::Right)
672 })
673 }
674}