1use backon::{ExponentialBuilder, Retryable};
35use std::{
36 collections::HashMap, future::Future, marker::PhantomData, ops::Deref, pin::Pin, time::Duration,
37};
38use tokio::{
39 sync::{oneshot::Receiver, Mutex},
40 time::{interval_at, Instant},
41};
42use tracing::field::Empty;
43use ulid::Ulid;
44
45use crate::{context, cursor::Args, Aggregator, AggregatorEvent, Executor, ReadAggregator};
46
47#[derive(Clone)]
52pub enum RoutingKey {
53 All,
55 Value(Option<String>),
57}
58
59#[derive(Clone)]
83pub struct Context<'a, E: Executor> {
84 context: context::RwContext,
85 pub executor: &'a E,
87}
88
89impl<'a, E: Executor> Deref for Context<'a, E> {
90 type Target = context::RwContext;
91
92 fn deref(&self) -> &Self::Target {
93 &self.context
94 }
95}
96
97pub trait Handler<E: Executor>: Sync + Send {
105 fn handle<'a>(
110 &'a self,
111 context: &'a Context<'a, E>,
112 event: &'a crate::Event,
113 ) -> Pin<Box<dyn Future<Output = anyhow::Result<()>> + Send + 'a>>;
114
115 fn aggregator_type(&self) -> &'static str;
117 fn event_name(&self) -> &'static str;
119}
120
121pub struct SubscriptionBuilder<E: Executor> {
143 key: String,
144 handlers: HashMap<String, Box<dyn Handler<E>>>,
145 context: context::RwContext,
146 routing_key: Option<RoutingKey>,
147 delay: Option<Duration>,
148 chunk_size: u16,
149 is_accept_failure: bool,
150 retry: Option<u8>,
151 aggregators: HashMap<String, String>,
152 safety_disabled: bool,
153 shutdown_rx: Option<Mutex<Receiver<()>>>,
154}
155
156impl<E: Executor + 'static> SubscriptionBuilder<E> {
157 pub fn new(key: impl Into<String>) -> Self {
161 Self {
162 key: key.into(),
163 handlers: HashMap::new(),
164 safety_disabled: true,
165 context: Default::default(),
166 delay: None,
167 retry: Some(30),
168 chunk_size: 300,
169 is_accept_failure: false,
170 routing_key: None,
171 aggregators: Default::default(),
172 shutdown_rx: None,
173 }
174 }
175
176 pub fn safety_check(mut self) -> Self {
180 self.safety_disabled = false;
181
182 self
183 }
184
185 pub fn handler<H: Handler<E> + 'static>(mut self, h: H) -> Self {
191 let key = format!("{}_{}", h.aggregator_type(), h.event_name());
192 if self.handlers.insert(key.to_owned(), Box::new(h)).is_some() {
193 panic!("Cannot register event handler: key {} already exists", key);
194 }
195 self
196 }
197
198 pub fn skip<EV: AggregatorEvent + Send + Sync + 'static>(self) -> Self {
206 self.handler(SkipHandler::<EV>(PhantomData))
207 }
208
209 pub fn data<D: Send + Sync + 'static>(self, v: D) -> Self {
213 self.context.insert(v);
214
215 self
216 }
217
218 pub fn accept_failure(mut self) -> Self {
223 self.is_accept_failure = true;
224
225 self
226 }
227
228 pub fn chunk_size(mut self, v: u16) -> Self {
232 self.chunk_size = v;
233
234 self
235 }
236
237 pub fn delay(mut self, v: Duration) -> Self {
241 self.delay = Some(v);
242
243 self
244 }
245
246 pub fn routing_key(mut self, v: impl Into<String>) -> Self {
251 self.routing_key = Some(RoutingKey::Value(Some(v.into())));
252
253 self
254 }
255
256 pub fn retry(mut self, v: u8) -> Self {
260 self.retry = Some(v);
261
262 self
263 }
264
265 pub fn all(mut self) -> Self {
269 self.routing_key = Some(RoutingKey::All);
270
271 self
272 }
273
274 pub fn aggregator<A: Aggregator>(mut self, id: impl Into<String>) -> Self {
276 self.aggregators
277 .insert(A::aggregator_type().to_owned(), id.into());
278
279 self
280 }
281
282 fn read_aggregators(&self) -> Vec<ReadAggregator> {
283 self.handlers
284 .values()
285 .map(|h| match self.aggregators.get(h.aggregator_type()) {
286 Some(id) => ReadAggregator {
287 aggregator_type: h.aggregator_type().to_owned(),
288 aggregator_id: Some(id.to_owned()),
289 name: if self.safety_disabled {
290 Some(h.event_name().to_owned())
291 } else {
292 None
293 },
294 },
295 _ => {
296 if self.safety_disabled {
297 ReadAggregator::event(h.aggregator_type(), h.event_name())
298 } else {
299 ReadAggregator::aggregator(h.aggregator_type())
300 }
301 }
302 })
303 .collect()
304 }
305
306 fn key(&self) -> String {
307 if let Some(RoutingKey::Value(Some(ref key))) = self.routing_key {
308 return format!("{key}.{}", self.key);
309 }
310
311 self.key.to_owned()
312 }
313
314 fn resolve_routing_key(&mut self, executor: &E) {
320 if self.routing_key.is_some() {
321 return;
322 }
323 self.routing_key = Some(match executor.default_routing_key() {
324 Some(k) => RoutingKey::Value(Some(k.to_owned())),
325 None => RoutingKey::Value(None),
326 });
327 }
328
329 fn effective_routing_key(&self) -> RoutingKey {
330 self.routing_key.clone().unwrap_or(RoutingKey::Value(None))
331 }
332
333 #[tracing::instrument(
334 skip_all,
335 fields(
336 subscription = Empty,
337 aggregator_type = Empty,
338 aggregator_id = Empty,
339 event = Empty,
340 )
341 )]
342 async fn process(
343 &self,
344 executor: &E,
345 id: &Ulid,
346 aggregators: &[ReadAggregator],
347 ) -> anyhow::Result<bool> {
348 let mut interval = interval_at(
349 Instant::now() - Duration::from_millis(400),
350 Duration::from_millis(300),
351 );
352
353 tracing::Span::current().record("subscription", self.key());
354
355 loop {
356 interval.tick().await;
357
358 if !executor.is_subscriber_running(self.key(), *id).await? {
359 return Ok(false);
360 }
361
362 let cursor = executor.get_subscriber_cursor(self.key()).await?;
363
364 let res = executor
365 .read(
366 Some(aggregators.to_vec()),
367 Some(self.effective_routing_key()),
368 Args::forward(self.chunk_size, cursor),
369 )
370 .await?;
371
372 if res.edges.is_empty() {
373 return Ok(false);
374 }
375
376 let timestamp = executor
377 .latest_timestamp(
378 Some(aggregators.to_vec()),
379 Some(self.effective_routing_key()),
380 )
381 .await?;
382
383 let context = Context {
384 context: self.context.clone(),
385 executor,
386 };
387
388 for event in res.edges {
389 if let Some(ref rx) = self.shutdown_rx {
390 let mut rx = rx.lock().await;
391 if rx.try_recv().is_ok() {
392 tracing::info!(
393 key = self.key(),
394 "Subscription received shutdown signal, stopping gracefull"
395 );
396
397 return Ok(true);
398 }
399 drop(rx);
400 }
401
402 tracing::Span::current().record("aggregator_type", &event.node.aggregator_type);
403 tracing::Span::current().record("aggregator_id", &event.node.aggregator_id);
404 tracing::Span::current().record("event", &event.node.name);
405
406 let all_key = format!("{}_all", event.node.aggregator_type);
407 let key = format!("{}_{}", event.node.aggregator_type, event.node.name);
408 let Some(handler) = self.handlers.get(&all_key).or(self.handlers.get(&key)) else {
409 if !self.safety_disabled {
410 anyhow::bail!("no handler s={} k={key}", self.key());
411 }
412
413 continue;
414 };
415
416 if let Err(err) = handler.handle(&context, &event.node).await {
417 tracing::error!("failed");
418
419 return Err(err);
420 }
421
422 tracing::debug!("completed");
423
424 executor
425 .acknowledge(
426 self.key(),
427 event.cursor.to_owned(),
428 timestamp.saturating_sub(event.node.timestamp),
429 )
430 .await?;
431 }
432 }
433 }
434
435 pub async fn unretry_start(mut self, executor: &E) -> anyhow::Result<Subscription>
439 where
440 E: Clone,
441 {
442 self.retry = None;
443 self.start(executor).await
444 }
445
446 #[tracing::instrument(skip_all, fields(
451 subscription = self.key(),
452 aggregator_type = tracing::field::Empty,
453 aggregator_id = tracing::field::Empty,
454 event = tracing::field::Empty,
455 ))]
456 pub async fn start(mut self, executor: &E) -> anyhow::Result<Subscription>
457 where
458 E: Clone,
459 {
460 self.resolve_routing_key(executor);
461 let executor = executor.clone();
462 let id = Ulid::new();
463 let subscription_id = id;
464 let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
465 self.shutdown_rx = Some(Mutex::new(shutdown_rx));
466
467 executor
468 .upsert_subscriber(self.key(), id.to_owned())
469 .await?;
470
471 let task_handle = tokio::spawn(async move {
472 let read_aggregators = self.read_aggregators();
473 let start = self
474 .delay
475 .map(|d| Instant::now() + d)
476 .unwrap_or_else(Instant::now);
477
478 let mut interval = interval_at(
479 start - Duration::from_millis(1200),
480 Duration::from_millis(1000),
481 );
482
483 loop {
484 interval.tick().await;
485
486 if let Some(ref rx) = self.shutdown_rx {
487 let mut rx = rx.lock().await;
488 if rx.try_recv().is_ok() {
489 tracing::info!(
490 key = self.key(),
491 "Subscription received shutdown signal, stopping gracefull"
492 );
493
494 break;
495 }
496 drop(rx);
497 }
498
499 let result = match self.retry {
500 Some(retry) => {
501 (|| async { self.process(&executor, &id, &read_aggregators).await })
502 .retry(ExponentialBuilder::default().with_max_times(retry.into()))
503 .sleep(tokio::time::sleep)
504 .notify(|err, dur| {
505 tracing::error!(
506 error = %err,
507 duration = ?dur,
508 "Failed to process event"
509 );
510 })
511 .await
512 }
513 _ => self.process(&executor, &id, &read_aggregators).await,
514 };
515
516 match result {
517 Ok(shutdown) => {
518 if shutdown {
519 break;
520 }
521 }
522 Err(err) => {
523 tracing::error!(error = %err, "Failed to process event");
524
525 if !self.is_accept_failure {
526 break;
527 }
528 }
529 };
530 }
531 });
532
533 Ok(Subscription {
534 id: subscription_id,
535 task_handle,
536 shutdown_tx,
537 })
538 }
539
540 pub async fn unretry_execute(mut self, executor: &E) -> anyhow::Result<()> {
544 self.retry = None;
545 self.execute(executor).await
546 }
547
548 #[tracing::instrument(skip_all, fields(
553 subscription = self.key(),
554 aggregator_type = tracing::field::Empty,
555 aggregator_id = tracing::field::Empty,
556 event = tracing::field::Empty,
557 ))]
558 pub async fn execute(&mut self, executor: &E) -> anyhow::Result<()> {
559 self.resolve_routing_key(executor);
560 let id = Ulid::new();
561
562 executor
563 .upsert_subscriber(self.key(), id.to_owned())
564 .await?;
565
566 let read_aggregators = self.read_aggregators();
567
568 match self.retry {
569 Some(retry) => {
570 (|| async { self.process(executor, &id, &read_aggregators).await })
571 .retry(ExponentialBuilder::default().with_max_times(retry.into()))
572 .sleep(tokio::time::sleep)
573 .notify(|err, dur| {
574 tracing::error!(
575 error = %err,
576 duration = ?dur,
577 "Failed to process event"
578 );
579 })
580 .await
581 }
582 _ => self.process(executor, &id, &read_aggregators).await,
583 }?;
584
585 Ok(())
586 }
587}
588
589#[derive(Debug)]
608pub struct Subscription {
609 pub id: Ulid,
611 task_handle: tokio::task::JoinHandle<()>,
612 shutdown_tx: tokio::sync::oneshot::Sender<()>,
613}
614
615impl Subscription {
616 pub async fn shutdown(self) -> Result<(), tokio::task::JoinError> {
621 let _ = self.shutdown_tx.send(());
622
623 self.task_handle.await
624 }
625}
626
627struct SkipHandler<E: AggregatorEvent>(PhantomData<E>);
628
629impl<E: Executor, EV: AggregatorEvent + Send + Sync> Handler<E> for SkipHandler<EV> {
630 fn handle<'a>(
631 &'a self,
632 _context: &'a Context<'a, E>,
633 _event: &'a crate::Event,
634 ) -> Pin<Box<dyn Future<Output = anyhow::Result<()>> + Send + 'a>> {
635 Box::pin(async { Ok(()) })
636 }
637
638 fn aggregator_type(&self) -> &'static str {
639 EV::aggregator_type()
640 }
641
642 fn event_name(&self) -> &'static str {
643 EV::event_name()
644 }
645}