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: 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: RoutingKey::Value(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 {
250 self.routing_key = RoutingKey::Value(Some(v.into()));
251
252 self
253 }
254
255 pub fn retry(mut self, v: u8) -> Self {
259 self.retry = Some(v);
260
261 self
262 }
263
264 pub fn all(mut self) -> Self {
266 self.routing_key = RoutingKey::All;
267
268 self
269 }
270
271 pub fn aggregator<A: Aggregator>(mut self, id: impl Into<String>) -> Self {
273 self.aggregators
274 .insert(A::aggregator_type().to_owned(), id.into());
275
276 self
277 }
278
279 fn read_aggregators(&self) -> Vec<ReadAggregator> {
280 self.handlers
281 .values()
282 .map(|h| match self.aggregators.get(h.aggregator_type()) {
283 Some(id) => ReadAggregator {
284 aggregator_type: h.aggregator_type().to_owned(),
285 aggregator_id: Some(id.to_owned()),
286 name: if self.safety_disabled {
287 Some(h.event_name().to_owned())
288 } else {
289 None
290 },
291 },
292 _ => {
293 if self.safety_disabled {
294 ReadAggregator::event(h.aggregator_type(), h.event_name())
295 } else {
296 ReadAggregator::aggregator(h.aggregator_type())
297 }
298 }
299 })
300 .collect()
301 }
302
303 fn key(&self) -> String {
304 if let RoutingKey::Value(Some(ref key)) = self.routing_key {
305 return format!("{key}.{}", self.key);
306 }
307
308 self.key.to_owned()
309 }
310
311 #[tracing::instrument(
312 skip_all,
313 fields(
314 subscription = Empty,
315 aggregator_type = Empty,
316 aggregator_id = Empty,
317 event = Empty,
318 )
319 )]
320 async fn process(
321 &self,
322 executor: &E,
323 id: &Ulid,
324 aggregators: &[ReadAggregator],
325 ) -> anyhow::Result<bool> {
326 let mut interval = interval_at(
327 Instant::now() - Duration::from_millis(400),
328 Duration::from_millis(300),
329 );
330
331 tracing::Span::current().record("subscription", self.key());
332
333 loop {
334 interval.tick().await;
335
336 if !executor.is_subscriber_running(self.key(), *id).await? {
337 return Ok(false);
338 }
339
340 let cursor = executor.get_subscriber_cursor(self.key()).await?;
341
342 let res = executor
343 .read(
344 Some(aggregators.to_vec()),
345 Some(self.routing_key.to_owned()),
346 Args::forward(self.chunk_size, cursor),
347 )
348 .await?;
349
350 if res.edges.is_empty() {
351 return Ok(false);
352 }
353
354 let timestamp = executor
355 .latest_timestamp(
356 Some(aggregators.to_vec()),
357 Some(self.routing_key.to_owned()),
358 )
359 .await?;
360
361 let context = Context {
362 context: self.context.clone(),
363 executor,
364 };
365
366 for event in res.edges {
367 if let Some(ref rx) = self.shutdown_rx {
368 let mut rx = rx.lock().await;
369 if rx.try_recv().is_ok() {
370 tracing::info!(
371 key = self.key(),
372 "Subscription received shutdown signal, stopping gracefull"
373 );
374
375 return Ok(true);
376 }
377 drop(rx);
378 }
379
380 tracing::Span::current().record("aggregator_type", &event.node.aggregator_type);
381 tracing::Span::current().record("aggregator_id", &event.node.aggregator_id);
382 tracing::Span::current().record("event", &event.node.name);
383
384 let all_key = format!("{}_all", event.node.aggregator_type);
385 let key = format!("{}_{}", event.node.aggregator_type, event.node.name);
386 let Some(handler) = self.handlers.get(&all_key).or(self.handlers.get(&key)) else {
387 if !self.safety_disabled {
388 anyhow::bail!("no handler s={} k={key}", self.key());
389 }
390
391 continue;
392 };
393
394 if let Err(err) = handler.handle(&context, &event.node).await {
395 tracing::error!("failed");
396
397 return Err(err);
398 }
399
400 tracing::debug!("completed");
401
402 executor
403 .acknowledge(
404 self.key(),
405 event.cursor.to_owned(),
406 timestamp.saturating_sub(event.node.timestamp),
407 )
408 .await?;
409 }
410 }
411 }
412
413 pub async fn unretry_start(mut self, executor: &E) -> anyhow::Result<Subscription>
417 where
418 E: Clone,
419 {
420 self.retry = None;
421 self.start(executor).await
422 }
423
424 #[tracing::instrument(skip_all, fields(
429 subscription = self.key(),
430 aggregator_type = tracing::field::Empty,
431 aggregator_id = tracing::field::Empty,
432 event = tracing::field::Empty,
433 ))]
434 pub async fn start(mut self, executor: &E) -> anyhow::Result<Subscription>
435 where
436 E: Clone,
437 {
438 let executor = executor.clone();
439 let id = Ulid::new();
440 let subscription_id = id;
441 let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
442 self.shutdown_rx = Some(Mutex::new(shutdown_rx));
443
444 executor
445 .upsert_subscriber(self.key(), id.to_owned())
446 .await?;
447
448 let task_handle = tokio::spawn(async move {
449 let read_aggregators = self.read_aggregators();
450 let start = self
451 .delay
452 .map(|d| Instant::now() + d)
453 .unwrap_or_else(Instant::now);
454
455 let mut interval = interval_at(
456 start - Duration::from_millis(1200),
457 Duration::from_millis(1000),
458 );
459
460 loop {
461 interval.tick().await;
462
463 if let Some(ref rx) = self.shutdown_rx {
464 let mut rx = rx.lock().await;
465 if rx.try_recv().is_ok() {
466 tracing::info!(
467 key = self.key(),
468 "Subscription received shutdown signal, stopping gracefull"
469 );
470
471 break;
472 }
473 drop(rx);
474 }
475
476 let result = match self.retry {
477 Some(retry) => {
478 (|| async { self.process(&executor, &id, &read_aggregators).await })
479 .retry(ExponentialBuilder::default().with_max_times(retry.into()))
480 .sleep(tokio::time::sleep)
481 .notify(|err, dur| {
482 tracing::error!(
483 error = %err,
484 duration = ?dur,
485 "Failed to process event"
486 );
487 })
488 .await
489 }
490 _ => self.process(&executor, &id, &read_aggregators).await,
491 };
492
493 match result {
494 Ok(shutdown) => {
495 if shutdown {
496 break;
497 }
498 }
499 Err(err) => {
500 tracing::error!(error = %err, "Failed to process event");
501
502 if !self.is_accept_failure {
503 break;
504 }
505 }
506 };
507 }
508 });
509
510 Ok(Subscription {
511 id: subscription_id,
512 task_handle,
513 shutdown_tx,
514 })
515 }
516
517 pub async fn unretry_execute(mut self, executor: &E) -> anyhow::Result<()> {
521 self.retry = None;
522 self.execute(executor).await
523 }
524
525 #[tracing::instrument(skip_all, fields(
530 subscription = self.key(),
531 aggregator_type = tracing::field::Empty,
532 aggregator_id = tracing::field::Empty,
533 event = tracing::field::Empty,
534 ))]
535 pub async fn execute(&self, executor: &E) -> anyhow::Result<()> {
536 let id = Ulid::new();
537
538 executor
539 .upsert_subscriber(self.key(), id.to_owned())
540 .await?;
541
542 let read_aggregators = self.read_aggregators();
543
544 match self.retry {
545 Some(retry) => {
546 (|| async { self.process(executor, &id, &read_aggregators).await })
547 .retry(ExponentialBuilder::default().with_max_times(retry.into()))
548 .sleep(tokio::time::sleep)
549 .notify(|err, dur| {
550 tracing::error!(
551 error = %err,
552 duration = ?dur,
553 "Failed to process event"
554 );
555 })
556 .await
557 }
558 _ => self.process(executor, &id, &read_aggregators).await,
559 }?;
560
561 Ok(())
562 }
563}
564
565#[derive(Debug)]
584pub struct Subscription {
585 pub id: Ulid,
587 task_handle: tokio::task::JoinHandle<()>,
588 shutdown_tx: tokio::sync::oneshot::Sender<()>,
589}
590
591impl Subscription {
592 pub async fn shutdown(self) -> Result<(), tokio::task::JoinError> {
597 let _ = self.shutdown_tx.send(());
598
599 self.task_handle.await
600 }
601}
602
603struct SkipHandler<E: AggregatorEvent>(PhantomData<E>);
604
605impl<E: Executor, EV: AggregatorEvent + Send + Sync> Handler<E> for SkipHandler<EV> {
606 fn handle<'a>(
607 &'a self,
608 _context: &'a Context<'a, E>,
609 _event: &'a crate::Event,
610 ) -> Pin<Box<dyn Future<Output = anyhow::Result<()>> + Send + 'a>> {
611 Box::pin(async { Ok(()) })
612 }
613
614 fn aggregator_type(&self) -> &'static str {
615 EV::aggregator_type()
616 }
617
618 fn event_name(&self) -> &'static str {
619 EV::event_name()
620 }
621}