1use std::collections::{BTreeMap, HashMap};
5use std::sync::atomic::{AtomicBool, Ordering};
6use std::sync::{Arc, Mutex, Weak};
7use std::time::Duration;
8
9use async_trait::async_trait;
10use camel_api::exchange::Exchange;
11use camel_api::resequencer::{CapacityPolicy, GapPolicy};
12use camel_language_api::Expression;
13use tokio::sync::mpsc;
14use tokio::task::JoinHandle;
15use tokio_util::sync::CancellationToken;
16
17use super::ResequencePolicy;
18
19pub struct StreamPolicy {
26 sequence_expr: Arc<dyn Expression>,
27 capacity: usize,
28 gap_timeout: Duration,
29 on_gap: GapPolicy,
30 on_capacity_exceeded: CapacityPolicy,
31 dedup: bool,
32
33 weak_self: Weak<Self>,
35
36 queue: Mutex<BTreeMap<u64, Exchange>>,
38
39 next_expected: Mutex<u64>,
41
42 gap_tokens: Mutex<HashMap<u64, CancellationToken>>,
44
45 gap_handles: Mutex<HashMap<u64, JoinHandle<()>>>,
47
48 driver_tx: Mutex<Option<mpsc::Sender<Exchange>>>,
50
51 shutdown_started: AtomicBool,
53}
54
55impl StreamPolicy {
56 pub fn new_cyclic(
58 sequence_expr: Arc<dyn Expression>,
59 capacity: usize,
60 gap_timeout_ms: u64,
61 on_gap: GapPolicy,
62 on_capacity_exceeded: CapacityPolicy,
63 dedup: bool,
64 ) -> Arc<Self> {
65 Arc::new_cyclic(|weak| Self {
66 sequence_expr,
67 capacity,
68 gap_timeout: Duration::from_millis(gap_timeout_ms),
69 on_gap,
70 on_capacity_exceeded,
71 dedup,
72 weak_self: weak.clone(),
73 queue: Mutex::new(BTreeMap::new()),
74 next_expected: Mutex::new(1),
75 gap_tokens: Mutex::new(HashMap::new()),
76 gap_handles: Mutex::new(HashMap::new()),
77 driver_tx: Mutex::new(None),
78 shutdown_started: AtomicBool::new(false),
79 })
80 }
81
82 fn set_driver_tx(&self, tx: mpsc::Sender<Exchange>) {
84 let mut guard = self.driver_tx.lock().unwrap_or_else(|e| e.into_inner());
85 *guard = Some(tx);
86 }
87
88 async fn eval_seq(&self, exchange: &Exchange) -> Result<u64, String> {
90 let val = self
91 .sequence_expr
92 .evaluate(exchange)
93 .await
94 .map_err(|e| format!("sequence expression evaluation failed: {e}"))?;
95 match val {
96 serde_json::Value::Number(n) => n
97 .as_u64()
98 .ok_or_else(|| format!("sequence value must be a non-negative integer, got {n}")),
99 _ => Err(format!(
100 "sequence expression must evaluate to a number, got {}",
101 val
102 )),
103 }
104 }
105
106 fn next_expected(&self) -> u64 {
108 *self.next_expected.lock().unwrap_or_else(|e| e.into_inner())
109 }
110
111 fn set_next_expected(&self, v: u64) {
113 *self.next_expected.lock().unwrap_or_else(|e| e.into_inner()) = v;
114 }
115
116 fn drain_contiguous(&self) -> Vec<Exchange> {
119 let mut queue = self.queue.lock().unwrap_or_else(|e| e.into_inner());
120 let mut expected = *self.next_expected.lock().unwrap_or_else(|e| e.into_inner());
121 let mut emitted = Vec::new();
122
123 while let Some(ex) = queue.remove(&expected) {
124 self.cancel_gap_timer(expected);
126 emitted.push(ex);
127 expected += 1;
128 }
129
130 *self.next_expected.lock().unwrap_or_else(|e| e.into_inner()) = expected;
131 emitted
132 }
133
134 fn drain_all_with_max(&self) -> (Vec<Exchange>, u64) {
138 let mut queue = self.queue.lock().unwrap_or_else(|e| e.into_inner());
139 let keys: Vec<u64> = queue.keys().copied().collect();
140 let max_seq = keys.iter().max().copied().unwrap_or(0);
141 let mut held = Vec::new();
142 for k in keys {
143 if let Some(ex) = queue.remove(&k) {
144 self.cancel_gap_timer(k);
145 held.push(ex);
146 }
147 }
148 (held, max_seq)
149 }
150
151 fn has_gap_timer(&self, seq: u64) -> bool {
153 let tokens = self.gap_tokens.lock().unwrap_or_else(|e| e.into_inner());
154 tokens.contains_key(&seq)
155 }
156
157 fn cancel_gap_timer(&self, seq: u64) {
159 {
160 let mut tokens = self.gap_tokens.lock().unwrap_or_else(|e| e.into_inner());
161 if let Some(token) = tokens.remove(&seq) {
162 token.cancel();
163 }
164 }
165 {
166 let mut handles = self.gap_handles.lock().unwrap_or_else(|e| e.into_inner());
167 handles.remove(&seq);
168 }
169 }
170
171 fn cancel_all_gap_timers(&self) {
173 let tokens: HashMap<u64, CancellationToken> = {
174 let mut guard = self.gap_tokens.lock().unwrap_or_else(|e| e.into_inner());
175 std::mem::take(&mut *guard)
176 };
177 for (_, token) in tokens {
178 token.cancel();
179 }
180 {
181 let mut handles = self.gap_handles.lock().unwrap_or_else(|e| e.into_inner());
182 handles.clear();
183 }
184 }
185
186 fn spawn_gap_timer(&self, missing_seq: u64) {
188 let cancel = CancellationToken::new();
189 let cancel_clone = cancel.clone();
190
191 {
192 let mut tokens = self.gap_tokens.lock().unwrap_or_else(|e| e.into_inner());
193 tokens.insert(missing_seq, cancel);
194 }
195
196 let weak = self.weak_self.clone();
197 let gap_timeout = self.gap_timeout;
198 let on_gap = self.on_gap;
199 let driver_tx_opt = {
200 let guard = self.driver_tx.lock().unwrap_or_else(|e| e.into_inner());
201 guard.clone()
202 };
203
204 let handle = tokio::spawn(async move {
205 tokio::select! {
206 _ = tokio::time::sleep(gap_timeout) => {
207 if cancel_clone.is_cancelled() {
208 return;
209 }
210 }
211 _ = cancel_clone.cancelled() => {
212 return;
213 }
214 }
215
216 let Some(policy) = weak.upgrade() else {
217 return;
218 };
219
220 if policy.shutdown_started.load(Ordering::SeqCst) {
221 return;
222 }
223
224 match on_gap {
225 GapPolicy::EmitPartial => {
226 policy.set_next_expected(missing_seq + 1);
230 let emitted = policy.drain_contiguous();
231
232 if emitted.is_empty() {
233 return;
234 }
235
236 if let Some(tx) = &driver_tx_opt {
237 for ex in emitted {
238 if tx.send(ex).await.is_err() {
239 break;
240 }
241 }
242 }
243 }
244 _ => {
247 let (held, max_seq) = policy.drain_all_with_max();
248
249 if held.is_empty() {
250 return;
251 }
252
253 policy.set_next_expected(max_seq + 1);
254
255 for ex in &held {
256 tracing::warn!(
257 correlation_id = %ex.correlation_id(),
258 "stream resequencer: gap timeout — dropping held exchange (no dead-letter sink wired)"
259 );
260 }
261 let _ = held;
262 }
263 }
264
265 {
267 let mut handles = policy.gap_handles.lock().unwrap_or_else(|e| e.into_inner());
268 handles.remove(&missing_seq);
269 }
270 {
271 let mut tokens = policy.gap_tokens.lock().unwrap_or_else(|e| e.into_inner());
272 tokens.remove(&missing_seq);
273 }
274 });
275
276 {
277 let mut handles = self.gap_handles.lock().unwrap_or_else(|e| e.into_inner());
278 handles.insert(missing_seq, handle);
279 }
280 }
281}
282
283#[async_trait]
284impl ResequencePolicy for StreamPolicy {
285 async fn accept(&self, input: Exchange) -> Vec<Exchange> {
286 let seq = match self.eval_seq(&input).await {
287 Ok(s) => s,
288 Err(e) => {
289 tracing::warn!(
290 error = %e,
291 correlation_id = %input.correlation_id(),
292 "StreamPolicy: sequence expression failed, dropping exchange"
293 );
294 return vec![];
295 }
296 };
297
298 let expected = self.next_expected();
299
300 if seq == expected {
301 self.cancel_gap_timer(expected);
305 self.set_next_expected(seq + 1);
307 let mut emitted = vec![input];
308 emitted.append(&mut self.drain_contiguous());
309 emitted
310 } else if seq < expected {
311 if self.dedup {
313 tracing::debug!(
315 seq = seq,
316 expected = expected,
317 "StreamPolicy: ignoring duplicate/late sequence (dedup enabled)"
318 );
319 return vec![];
320 }
321 {
323 let mut queue = self.queue.lock().unwrap_or_else(|e| e.into_inner());
324 queue.insert(seq, input);
325 }
326 vec![]
327 } else {
328 {
330 let mut queue = self.queue.lock().unwrap_or_else(|e| e.into_inner());
331 let queue_len = queue.len();
332
333 if self.dedup && queue.contains_key(&seq) {
336 tracing::debug!(
337 seq = seq,
338 "StreamPolicy: ignoring redelivered held sequence (dedup enabled)"
339 );
340 return vec![];
341 }
342
343 if queue_len >= self.capacity {
345 match self.on_capacity_exceeded {
346 CapacityPolicy::DropOldest => {
347 let oldest_key = queue.keys().next().copied();
349 if let Some(oldest) = oldest_key {
350 let dropped = queue.remove(&oldest);
351 self.cancel_gap_timer(oldest);
352 tracing::debug!(
353 dropped_seq = oldest,
354 "StreamPolicy: capacity exceeded, dropped oldest exchange"
355 );
356 let _ = dropped;
357 }
358 }
359 _ => {
361 tracing::warn!(
362 seq = seq,
363 capacity = self.capacity,
364 "StreamPolicy: capacity exceeded, dropping incoming exchange"
365 );
366 return vec![];
367 }
368 }
369 }
370
371 queue.insert(seq, input);
372 }
373
374 if !self.has_gap_timer(expected) {
376 self.spawn_gap_timer(expected);
377 }
378
379 vec![]
380 }
381 }
382
383 async fn flush(&self) -> Vec<Exchange> {
384 self.shutdown_started.store(true, Ordering::SeqCst);
385 self.cancel_all_gap_timers();
386
387 let mut queue = self.queue.lock().unwrap_or_else(|e| e.into_inner());
388 let keys: Vec<u64> = queue.keys().copied().collect();
389 let mut held = Vec::new();
390 for k in keys {
391 if let Some(ex) = queue.remove(&k) {
392 held.push(ex);
393 }
394 }
395 held
396 }
397
398 fn name(&self) -> &'static str {
399 "stream-resequencer"
400 }
401
402 fn buffered(&self) -> usize {
403 let queue = self.queue.lock().unwrap_or_else(|e| e.into_inner());
404 queue.len()
405 }
406
407 fn set_timeout_tx(&self, tx: mpsc::Sender<Exchange>) {
408 self.set_driver_tx(tx);
409 }
410}
411
412#[cfg(test)]
415mod tests {
416 use super::*;
417 use camel_api::exchange::ExchangePattern;
418 use camel_api::message::Message;
419
420 struct PropExpr(String);
422
423 #[async_trait::async_trait]
424 impl Expression for PropExpr {
425 async fn evaluate(
426 &self,
427 exchange: &Exchange,
428 ) -> Result<serde_json::Value, camel_language_api::LanguageError> {
429 Ok(exchange
430 .property(&self.0)
431 .cloned()
432 .unwrap_or(serde_json::Value::Null))
433 }
434 }
435
436 fn mk_exchange(seq: u64) -> Exchange {
437 let mut ex = Exchange::new(Message::new(camel_api::body::Body::Text(format!(
438 "msg-{seq}"
439 ))));
440 ex.set_property("seq", serde_json::json!(seq));
441 ex.pattern = ExchangePattern::InOnly;
442 ex
443 }
444
445 fn default_policy() -> Arc<StreamPolicy> {
446 StreamPolicy::new_cyclic(
447 Arc::new(PropExpr("seq".into())),
448 100,
449 5000,
450 GapPolicy::EmitPartial,
451 CapacityPolicy::LogAndDrop,
452 false,
453 )
454 }
455
456 fn seq_of(ex: &Exchange) -> u64 {
457 ex.property("seq").and_then(|v| v.as_u64()).unwrap_or(0)
458 }
459
460 #[tokio::test]
463 async fn stream_emit_contiguous_run() {
464 let policy = default_policy();
465
466 let result2 = policy.accept(mk_exchange(2)).await;
468 assert!(
469 result2.is_empty(),
470 "seq 2 should be held (not yet contiguous)"
471 );
472
473 let result1 = policy.accept(mk_exchange(1)).await;
475 assert_eq!(result1.len(), 2, "should emit 1 then drain 2");
476 let seqs1: Vec<u64> = result1.iter().map(seq_of).collect();
477 assert_eq!(seqs1, vec![1, 2]);
478
479 let result4 = policy.accept(mk_exchange(4)).await;
481 assert!(result4.is_empty(), "seq 4 should be held");
482
483 let result3 = policy.accept(mk_exchange(3)).await;
485 assert_eq!(result3.len(), 2, "should emit 3 then drain 4");
486 let seqs3: Vec<u64> = result3.iter().map(seq_of).collect();
487 assert_eq!(seqs3, vec![3, 4]);
488 }
489
490 #[tokio::test]
493 async fn stream_gap_timeout_emit_partial() {
494 let policy = StreamPolicy::new_cyclic(
495 Arc::new(PropExpr("seq".into())),
496 100,
497 50, GapPolicy::EmitPartial,
499 CapacityPolicy::LogAndDrop,
500 false,
501 );
502
503 let (tx, mut rx) = mpsc::channel::<Exchange>(16);
504 policy.set_timeout_tx(tx);
505
506 assert!(policy.accept(mk_exchange(2)).await.is_empty());
508 assert!(policy.accept(mk_exchange(3)).await.is_empty());
509
510 let emitted: Vec<Exchange> = tokio::time::timeout(Duration::from_millis(500), async {
512 let mut out = Vec::new();
513 out.push(rx.recv().await.unwrap());
514 out.push(rx.recv().await.unwrap());
515 out
516 })
517 .await
518 .expect("gap timer should fire within 500ms");
519
520 assert_eq!(emitted.len(), 2, "should emit all held exchanges");
521 let seqs: Vec<u64> = emitted.iter().map(seq_of).collect();
522 assert_eq!(seqs, vec![2, 3], "should emit in sequence order");
523
524 assert_eq!(
526 policy.next_expected(),
527 4,
528 "next_expected should advance past gap to max(drained)+1"
529 );
530 }
531
532 #[tokio::test]
534 async fn stream_capacity_exceeded_log_and_drop() {
535 let policy = StreamPolicy::new_cyclic(
536 Arc::new(PropExpr("seq".into())),
537 2, 5000,
539 GapPolicy::EmitPartial,
540 CapacityPolicy::LogAndDrop,
541 false,
542 );
543
544 assert!(policy.accept(mk_exchange(3)).await.is_empty());
546 assert!(policy.accept(mk_exchange(4)).await.is_empty());
547
548 {
549 let queue = policy.queue.lock().unwrap();
550 assert_eq!(queue.len(), 2, "queue should be full");
551 }
552
553 let result = policy.accept(mk_exchange(5)).await;
555 assert!(
556 result.is_empty(),
557 "overflow exchange should be dead-lettered (empty result)"
558 );
559
560 {
561 let queue = policy.queue.lock().unwrap();
562 assert_eq!(queue.len(), 2, "queue should stay at capacity");
563 }
564 }
565
566 #[tokio::test]
568 async fn stream_capacity_exceeded_drop_oldest() {
569 let policy = StreamPolicy::new_cyclic(
570 Arc::new(PropExpr("seq".into())),
571 2,
572 5000,
573 GapPolicy::EmitPartial,
574 CapacityPolicy::DropOldest,
575 false,
576 );
577
578 assert!(policy.accept(mk_exchange(3)).await.is_empty());
580 assert!(policy.accept(mk_exchange(4)).await.is_empty());
581
582 let result = policy.accept(mk_exchange(5)).await;
584 assert!(
585 result.is_empty(),
586 "overflow with DropOldest should not emit"
587 );
588
589 {
590 let queue = policy.queue.lock().unwrap();
591 assert_eq!(queue.len(), 2, "queue should still be at capacity");
592 assert!(!queue.contains_key(&3), "oldest seq 3 should be dropped");
593 assert!(queue.contains_key(&4), "seq 4 should remain");
594 assert!(queue.contains_key(&5), "seq 5 should be inserted");
595 }
596 }
597
598 #[tokio::test]
600 async fn stream_dedup_on_ignores_duplicate() {
601 let policy = StreamPolicy::new_cyclic(
602 Arc::new(PropExpr("seq".into())),
603 100,
604 5000,
605 GapPolicy::EmitPartial,
606 CapacityPolicy::LogAndDrop,
607 true, );
609
610 let result1 = policy.accept(mk_exchange(1)).await;
612 assert_eq!(result1.len(), 1, "seq 1 should be emitted");
613
614 let result2 = policy.accept(mk_exchange(1)).await;
616 assert!(
617 result2.is_empty(),
618 "duplicate seq 1 should be ignored with dedup on"
619 );
620 }
621
622 #[tokio::test]
625 async fn stream_dedup_off_inserts_duplicate() {
626 let policy = StreamPolicy::new_cyclic(
627 Arc::new(PropExpr("seq".into())),
628 100,
629 5000,
630 GapPolicy::EmitPartial,
631 CapacityPolicy::LogAndDrop,
632 false, );
634
635 let result1 = policy.accept(mk_exchange(1)).await;
637 assert_eq!(result1.len(), 1, "seq 1 should be emitted");
638
639 let result2 = policy.accept(mk_exchange(1)).await;
641 assert!(
642 result2.is_empty(),
643 "duplicate seq 1 with dedup off should be inserted, not emitted"
644 );
645
646 {
647 let queue = policy.queue.lock().unwrap();
648 assert!(
649 queue.contains_key(&1),
650 "duplicate seq 1 should be in queue (dedup off)"
651 );
652 }
653 }
654
655 #[tokio::test]
657 async fn stream_flush_emits_remaining_sorted() {
658 let policy = default_policy();
659
660 assert!(!policy.accept(mk_exchange(1)).await.is_empty());
662
663 assert!(policy.accept(mk_exchange(5)).await.is_empty());
665 assert!(policy.accept(mk_exchange(3)).await.is_empty());
666
667 let flushed = policy.flush().await;
669 assert_eq!(flushed.len(), 2, "should emit all remaining held exchanges");
670 let seqs: Vec<u64> = flushed.iter().map(seq_of).collect();
671 assert_eq!(seqs, vec![3, 5], "should be in sequence order");
672 }
673
674 #[tokio::test]
676 async fn stream_late_sequence_after_advance() {
677 let policy_dedup = StreamPolicy::new_cyclic(
679 Arc::new(PropExpr("seq".into())),
680 100,
681 50,
682 GapPolicy::EmitPartial,
683 CapacityPolicy::LogAndDrop,
684 true,
685 );
686
687 {
689 let mut ne = policy_dedup.next_expected.lock().unwrap();
690 *ne = 5;
691 }
692
693 let result = policy_dedup.accept(mk_exchange(3)).await;
695 assert!(
696 result.is_empty(),
697 "late seq with dedup=true should be ignored"
698 );
699
700 let policy_no_dedup = StreamPolicy::new_cyclic(
702 Arc::new(PropExpr("seq".into())),
703 100,
704 50,
705 GapPolicy::EmitPartial,
706 CapacityPolicy::LogAndDrop,
707 false,
708 );
709
710 {
711 let mut ne = policy_no_dedup.next_expected.lock().unwrap();
712 *ne = 5;
713 }
714
715 let result = policy_no_dedup.accept(mk_exchange(3)).await;
717 assert!(
718 result.is_empty(),
719 "late seq with dedup=false should be inserted"
720 );
721
722 {
723 let queue = policy_no_dedup.queue.lock().unwrap();
724 assert!(
725 queue.contains_key(&3),
726 "late seq should be in queue with dedup off"
727 );
728 }
729 }
730
731 #[tokio::test]
734 async fn stream_stale_gap_timer_cancelled_on_normal_arrival() {
735 let policy = StreamPolicy::new_cyclic(
736 Arc::new(PropExpr("seq".into())),
737 100,
738 50, GapPolicy::DropAndLog,
740 CapacityPolicy::LogAndDrop,
741 false,
742 );
743
744 let (tx, mut rx) = mpsc::channel::<Exchange>(16);
745 policy.set_timeout_tx(tx);
746
747 assert!(policy.accept(mk_exchange(2)).await.is_empty());
749 assert!(
750 policy.has_gap_timer(1),
751 "gap timer should be armed for seq 1"
752 );
753
754 let result = policy.accept(mk_exchange(1)).await;
756 assert_eq!(result.len(), 2, "should emit seq 1 and drained seq 2");
757 let seqs: Vec<u64> = result.iter().map(seq_of).collect();
758 assert_eq!(seqs, vec![1, 2]);
759
760 assert!(
762 !policy.has_gap_timer(1),
763 "gap timer for seq 1 should be cancelled after normal arrival"
764 );
765
766 tokio::time::sleep(Duration::from_millis(200)).await;
769 match rx.try_recv() {
770 Ok(ex) => {
771 panic!(
772 "stale gap timer fired and sent exchange with seq={} — corruption!",
773 seq_of(&ex)
774 );
775 }
776 Err(tokio::sync::mpsc::error::TryRecvError::Empty) => {
777 }
779 Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) => {}
780 }
781 }
782
783 #[tokio::test]
785 async fn stream_dedup_held_future_seq() {
786 let policy = StreamPolicy::new_cyclic(
787 Arc::new(PropExpr("seq".into())),
788 100,
789 5000,
790 GapPolicy::EmitPartial,
791 CapacityPolicy::LogAndDrop,
792 true, );
794
795 assert!(policy.accept(mk_exchange(5)).await.is_empty());
797
798 let result = policy.accept(mk_exchange(5)).await;
800 assert!(result.is_empty(), "redelivered held seq should be ignored");
801
802 {
804 let queue = policy.queue.lock().unwrap();
805 assert_eq!(queue.len(), 1, "queue should still have exactly one entry");
806 assert!(queue.contains_key(&5), "seq 5 should still be in queue");
807 }
808 }
809}