1use std::collections::HashMap;
2use std::future::Future;
3use std::pin::Pin;
4use std::sync::{Arc, Mutex};
5use std::task::{Context, Poll};
6use std::time::{Duration, Instant};
7
8use tokio::sync::mpsc;
9use tokio::task::JoinHandle;
10use tokio_util::sync::CancellationToken;
11use tower::Service;
12
13use async_trait::async_trait;
14use camel_api::{
15 CamelError, MetricsCollector, StepLifecycle, StepShutdownReason,
16 aggregator::{
17 AggregationStrategy, AggregatorConfig, CompletionCondition, CompletionMode,
18 CompletionReason, CorrelationStrategy,
19 },
20 body::Body,
21 exchange::Exchange,
22 message::Message,
23};
24use camel_component_api::InFlightClaim;
25use camel_language_api::Language;
26
27pub type SharedLanguageRegistry = Arc<std::sync::Mutex<HashMap<String, Arc<dyn Language>>>>;
28
29pub struct AggregateEmission {
35 pub exchange: Exchange,
36 pub claims: Vec<Option<InFlightClaim>>,
37}
38
39pub struct AggregationReceipt {
47 pub reply: Result<Exchange, CamelError>,
50 pub claims: Vec<Option<InFlightClaim>>,
53}
54
55const QUEUE_DEPTH_SAMPLE_INTERVAL: Duration = Duration::from_millis(250);
59
60pub const CAMEL_AGGREGATOR_PENDING: &str = "CamelAggregatorPending";
61pub const CAMEL_AGGREGATED_SIZE: &str = "CamelAggregatedSize";
62pub const CAMEL_AGGREGATED_KEY: &str = "CamelAggregatedKey";
63pub const CAMEL_AGGREGATED_COMPLETION_REASON: &str = "CamelAggregatedCompletionReason";
64
65struct Bucket {
74 exchanges: Vec<Exchange>,
75 claims: Vec<Option<InFlightClaim>>,
76 last_updated: Instant,
77}
78
79impl Bucket {
80 fn new() -> Self {
81 Self {
82 exchanges: Vec::new(),
83 claims: Vec::new(),
84 last_updated: Instant::now(),
85 }
86 }
87
88 fn push(&mut self, exchange: Exchange, claim: Option<InFlightClaim>) {
89 self.exchanges.push(exchange);
90 self.claims.push(claim);
91 self.last_updated = Instant::now();
92 }
93
94 fn into_parts(self) -> (Vec<Exchange>, Vec<Option<InFlightClaim>>) {
97 (self.exchanges, self.claims)
98 }
99
100 fn is_expired(&self, ttl: Duration) -> bool {
101 Instant::now().duration_since(self.last_updated) >= ttl
102 }
103}
104
105#[derive(Clone)]
106pub struct AggregatorService {
107 config: AggregatorConfig,
108 buckets: Arc<Mutex<HashMap<String, Bucket>>>,
109 timeout_tasks: Arc<Mutex<HashMap<String, CancellationToken>>>,
110 timeout_handles: Arc<Mutex<HashMap<String, JoinHandle<()>>>>,
111 late_tx: mpsc::Sender<AggregateEmission>,
112 language_registry: SharedLanguageRegistry,
113 sweep_cancel: Arc<Mutex<CancellationToken>>,
120 sweep_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
128 queue_metrics: Option<(Arc<dyn MetricsCollector>, String)>,
133 clone_guard: Arc<()>,
137 cached_key: Arc<Mutex<Option<(serde_json::Value, String)>>>,
148 #[cfg(test)]
152 key_serializations: Arc<std::sync::atomic::AtomicUsize>,
153}
154
155impl std::fmt::Debug for AggregatorService {
156 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
157 f.debug_struct("AggregatorService").finish_non_exhaustive()
158 }
159}
160
161impl Drop for AggregatorService {
162 fn drop(&mut self) {
163 if Arc::strong_count(&self.clone_guard) > 1 {
173 return;
174 }
175 if let Some(handle) = self
176 .sweep_handle
177 .lock()
178 .unwrap_or_else(|e| e.into_inner())
179 .take()
180 {
181 handle.abort();
182 }
183 }
184}
185
186impl AggregatorService {
187 pub fn new(
198 config: AggregatorConfig,
199 late_tx: mpsc::Sender<AggregateEmission>,
200 language_registry: SharedLanguageRegistry,
201 route_cancel: CancellationToken,
202 ) -> Self {
203 config.validate().expect(
205 "AggregatorService::new: config failed validation \
208 (need max_buckets, completionTimeout, or bucket_ttl)",
209 );
210
211 let has_timeout = has_timeout_condition(&config.completion);
215 if !has_timeout && config.bucket_ttl.is_none() {
216 tracing::warn!(
217 "Aggregator configured with Size/Predicate completion and no \
218 bucket_ttl: buckets accumulate until max_buckets is reached"
219 );
220 }
221
222 let buckets: Arc<Mutex<HashMap<String, Bucket>>> = Arc::new(Mutex::new(HashMap::new()));
225
226 Self {
235 config,
236 buckets,
237 timeout_tasks: Arc::new(Mutex::new(HashMap::new())),
238 timeout_handles: Arc::new(Mutex::new(HashMap::new())),
239 late_tx,
240 language_registry,
241 sweep_cancel: Arc::new(Mutex::new(route_cancel)),
242 sweep_handle: Arc::new(Mutex::new(None)),
243 queue_metrics: None,
244 clone_guard: Arc::new(()),
245 cached_key: Arc::new(Mutex::new(None)),
246 #[cfg(test)]
247 key_serializations: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
248 }
249 }
250
251 pub fn with_queue_metrics(
256 mut self,
257 metrics: Arc<dyn MetricsCollector>,
258 label: impl Into<String>,
259 ) -> Self {
260 self.queue_metrics = Some((metrics, label.into()));
261 self
262 }
263
264 pub fn config(&self) -> &AggregatorConfig {
265 &self.config
266 }
267
268 pub fn has_timeout(&self) -> bool {
269 has_timeout_condition(&self.config.completion)
270 }
271
272 pub fn force_complete_all(&self) {
273 let mut buckets_guard = self.buckets.lock().unwrap_or_else(|e| e.into_inner());
274 let keys: Vec<String> = buckets_guard.keys().cloned().collect();
275
276 for key in keys {
277 if let Some(bucket) = buckets_guard.remove(&key) {
278 if self.config.force_completion_on_stop {
279 cancel_timeout_task_with_handle(
280 &key,
281 &self.timeout_tasks,
282 &self.timeout_handles,
283 );
284 let (exchanges, claims) = bucket.into_parts();
285 match aggregate(exchanges, &self.config.strategy) {
286 Ok(mut result) => {
287 result.set_property(
288 CAMEL_AGGREGATED_COMPLETION_REASON,
289 serde_json::json!(CompletionReason::Stop.as_str()),
290 );
291 if self
292 .late_tx
293 .try_send(AggregateEmission {
294 exchange: result,
295 claims,
296 })
297 .is_err()
298 {
299 tracing::warn!(
300 key = %key,
301 "aggregator force-complete emit dropped: late channel full"
302 );
303 }
304 }
305 Err(e) => {
306 tracing::warn!(
308 key = %key,
309 error = %e,
310 "aggregation failed in force_complete_all"
311 );
312 }
313 }
314 } else {
315 cancel_timeout_task_with_handle(
316 &key,
317 &self.timeout_tasks,
318 &self.timeout_handles,
319 );
320 }
321 }
322 }
323 }
324
325 pub fn release_unarmed_buckets(&self) {
342 let mut buckets_guard = self.buckets.lock().unwrap_or_else(|e| e.into_inner());
343 let timeout_guard = self.timeout_tasks.lock().unwrap_or_else(|e| e.into_inner());
344 let before = buckets_guard.len();
345 buckets_guard.retain(|key, _| timeout_guard.contains_key(key));
346 let released = before - buckets_guard.len();
347 if released > 0 {
348 tracing::debug!(
349 released,
350 armed = buckets_guard.len(),
351 "aggregator released unarmed buckets on consumer exit"
352 );
353 }
354 }
355
356 pub(crate) async fn shutdown_inner(&self) {
359 {
361 let mut guard = self.timeout_tasks.lock().unwrap_or_else(|e| e.into_inner());
362 for token in guard.values() {
363 token.cancel();
364 }
365 guard.clear();
366 };
367
368 let handles: Vec<JoinHandle<()>> = {
370 let mut guard = self
371 .timeout_handles
372 .lock()
373 .unwrap_or_else(|e| e.into_inner());
374 guard.drain().map(|(_, handle)| handle).collect()
375 };
376
377 if handles.is_empty() {
378 return;
379 }
380
381 let _ = tokio::time::timeout(Duration::from_secs(5), async {
383 for handle in handles {
384 let _ = handle.await;
385 }
386 })
387 .await;
388 }
389}
390
391#[async_trait]
392impl StepLifecycle for AggregatorService {
393 fn name(&self) -> &'static str {
394 "aggregator"
395 }
396
397 async fn shutdown(&self, reason: StepShutdownReason) -> Result<(), CamelError> {
398 tracing::debug!(reason = ?reason, "Aggregator shutdown via StepLifecycle");
399
400 self.sweep_cancel
403 .lock()
404 .unwrap_or_else(|e| e.into_inner())
405 .cancel();
406
407 if let Some(handle) = self
410 .sweep_handle
411 .lock()
412 .unwrap_or_else(|e| e.into_inner())
413 .take()
414 {
415 handle.abort();
416 }
417
418 self.shutdown_inner().await;
419 Ok(())
420 }
421
422 async fn start(&self) -> Result<(), CamelError> {
426 *self.sweep_cancel.lock().unwrap_or_else(|e| e.into_inner()) = CancellationToken::new();
427
428 if let Some(handle) = self
429 .sweep_handle
430 .lock()
431 .unwrap_or_else(|e| e.into_inner())
432 .take()
433 {
434 handle.abort();
435 }
436
437 Ok(())
438 }
439}
440
441pub fn has_timeout_condition(mode: &CompletionMode) -> bool {
442 match mode {
443 CompletionMode::Single(CompletionCondition::Timeout(_)) => true,
444 CompletionMode::Any(conditions) => conditions
445 .iter()
446 .any(|c| matches!(c, CompletionCondition::Timeout(_))),
447 _ => false,
448 }
449}
450
451impl Service<Exchange> for AggregatorService {
452 type Response = Exchange;
453 type Error = CamelError;
454 type Future = Pin<Box<dyn Future<Output = Result<Exchange, CamelError>> + Send>>;
455
456 fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), CamelError>> {
457 let ttl = self.config.bucket_ttl;
466 if ttl.is_none() && self.queue_metrics.is_none() {
467 return Poll::Ready(Ok(()));
468 }
469 let mut g = self.sweep_handle.lock().unwrap_or_else(|e| e.into_inner());
470 if g.is_none() {
471 let interval = match ttl {
472 Some(ttl) => std::cmp::max(ttl / 2, Duration::from_millis(50)),
473 None => QUEUE_DEPTH_SAMPLE_INTERVAL,
474 };
475 let buckets = Arc::clone(&self.buckets);
476 let cancel = self
477 .sweep_cancel
478 .lock()
479 .unwrap_or_else(|e| e.into_inner())
480 .clone();
481 let queue_metrics = self.queue_metrics.clone();
482 *g = Some(tokio::spawn(async move {
483 loop {
484 tokio::select! {
485 _ = cancel.cancelled() => break,
486 _ = tokio::time::sleep(interval) => {
487 let mut guard =
488 buckets.lock().unwrap_or_else(|e| e.into_inner());
489 if let Some(ttl) = ttl {
490 guard.retain(|_, b| !b.is_expired(ttl));
491 }
492 if let Some((metrics, label)) = queue_metrics.as_ref() {
495 metrics.set_queue_depth(label, guard.len());
496 }
497 }
498 }
499 }
500 }));
501 }
502 Poll::Ready(Ok(()))
503 }
504
505 fn call(&mut self, exchange: Exchange) -> Self::Future {
506 let svc = self.clone();
507 Box::pin(async move {
508 let mut exchange = exchange;
520 let claim = exchange.in_flight_claim.take();
521 let AggregationReceipt { reply, claims } = svc.submit_with_claim(exchange, claim).await;
522 match reply {
523 Ok(mut result) => {
524 result.in_flight_claim = claims.into_iter().flatten().next();
525 Ok(result)
526 }
527 Err(e) => Err(e),
528 }
529 })
530 }
531}
532
533impl AggregatorService {
534 pub async fn submit_with_claim(
545 &self,
546 exchange: Exchange,
547 claim: Option<InFlightClaim>,
548 ) -> AggregationReceipt {
549 let config = self.config.clone();
550 let buckets = Arc::clone(&self.buckets);
551 let timeout_tasks = Arc::clone(&self.timeout_tasks);
552 let timeout_handles = Arc::clone(&self.timeout_handles);
553 let late_tx = self.late_tx.clone();
554 let language_registry = Arc::clone(&self.language_registry);
555 let cached_key = Arc::clone(&self.cached_key);
556 #[cfg(test)]
557 let key_serializations = Arc::clone(&self.key_serializations);
558
559 let inner = async move {
560 let key_value =
561 extract_correlation_key(&exchange, &config.correlation, &language_registry).await?;
562
563 let key_is_memoizable = matches!(key_value, serde_json::Value::String(_));
574 let key_str = if key_is_memoizable {
575 let mut memo = cached_key.lock().unwrap_or_else(|e| e.into_inner());
576 if let Some((v, s)) = memo.as_ref()
577 && *v == key_value
578 {
579 s.clone()
580 } else {
581 #[cfg(test)]
582 key_serializations.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
583 let s = serde_json::to_string(&key_value)
584 .map_err(|e| CamelError::ProcessorError(e.to_string()))?;
585 *memo = Some((key_value.clone(), s.clone()));
586 s
587 }
588 } else {
589 #[cfg(test)]
590 key_serializations.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
591 serde_json::to_string(&key_value)
592 .map_err(|e| CamelError::ProcessorError(e.to_string()))?
593 };
594
595 let predicate_satisfied =
599 evaluate_completion_predicate(&config.completion, &exchange, &language_registry)
600 .await?;
601
602 let completed_bucket = {
603 let mut guard = buckets.lock().unwrap_or_else(|e| e.into_inner());
604
605 if let Some(ttl) = config.bucket_ttl {
606 guard.retain(|_, bucket| !bucket.is_expired(ttl));
607 }
608
609 if let Some(max) = config.max_buckets
610 && !guard.contains_key(&key_str)
611 && guard.len() >= max
612 {
613 tracing::warn!(
614 max_buckets = max,
615 correlation_key = %key_str,
616 "Aggregator reached max buckets limit, rejecting new correlation key"
617 );
618 return Err(CamelError::ProcessorError(format!(
619 "Aggregator reached maximum {} buckets",
620 max
621 )));
622 }
623
624 if let Some(max_size) = config.max_bucket_size
628 && let Some(existing) = guard.get(&key_str)
629 && existing.exchanges.len() >= max_size
630 {
631 tracing::warn!(
632 max_bucket_size = max_size,
633 correlation_key = %key_str,
634 "Aggregator reached per-bucket size limit, rejecting exchange"
635 );
636 return Err(CamelError::ProcessorError(format!(
637 "Aggregator bucket reached maximum {max_size} exchanges"
638 )));
639 }
640
641 let bucket = match guard.get_mut(&key_str) {
642 Some(b) => b,
643 None => guard.entry(key_str.clone()).or_insert_with(Bucket::new),
644 };
645 bucket.push(exchange, claim);
646
647 let (is_complete, reason) = check_sync_completion(
648 &config.completion,
649 &bucket.exchanges,
650 predicate_satisfied,
651 );
652
653 let bucket_parts = if is_complete {
654 guard.remove(&key_str).map(Bucket::into_parts)
655 } else {
656 None
657 };
658 (bucket_parts, reason)
659 };
660
661 if completed_bucket.0.is_none() && has_timeout_condition(&config.completion) {
662 let timeout_dur = extract_timeout_duration(&config.completion);
663 if let Some(timeout) = timeout_dur {
664 let live_count = timeout_handles
670 .lock()
671 .unwrap_or_else(|e| e.into_inner())
672 .len();
673 if live_count >= config.max_timeout_tasks {
674 tracing::warn!(
675 live_timeout_tasks = live_count,
676 max_timeout_tasks = config.max_timeout_tasks,
677 correlation_key = %key_str,
678 "Aggregator timeout-task cap reached; bucket will rely on \
679 bucket_ttl eviction instead of a dedicated timeout task"
680 );
681 } else {
682 {
684 let tt_guard = timeout_tasks.lock().unwrap_or_else(|e| e.into_inner());
685 if let Some(existing) = tt_guard.get(&key_str) {
686 existing.cancel();
687 }
688 }
689 {
691 let mut hh = timeout_handles.lock().unwrap_or_else(|e| e.into_inner());
692 if let Some(old) = hh.remove(&key_str) {
693 old.abort();
694 }
695 }
696 let cancel = CancellationToken::new();
697 timeout_tasks
698 .lock()
699 .unwrap_or_else(|e| e.into_inner())
700 .insert(key_str.clone(), cancel.clone());
701 let handle = spawn_timeout_task(
702 key_str.clone(),
703 timeout,
704 cancel,
705 buckets.clone(),
706 timeout_tasks.clone(),
707 timeout_handles.clone(),
708 late_tx,
709 config.strategy.clone(),
710 config.discard_on_timeout,
711 );
712 timeout_handles
713 .lock()
714 .unwrap_or_else(|e| e.into_inner())
715 .insert(key_str.clone(), handle);
716 }
717 }
718 }
719
720 if let Some((exchanges, claims)) = completed_bucket.0 {
721 cancel_timeout_task_with_handle(&key_str, &timeout_tasks, &timeout_handles);
722 let reason = completed_bucket.1;
723 let size = exchanges.len();
724 let mut result = aggregate(exchanges, &config.strategy)?;
725 result.set_property(CAMEL_AGGREGATED_SIZE, serde_json::json!(size as u64));
726 result.set_property(CAMEL_AGGREGATED_KEY, key_value);
727 result.set_property(
728 CAMEL_AGGREGATED_COMPLETION_REASON,
729 serde_json::json!(reason.as_str()),
730 );
731 Ok((result, claims))
732 } else {
733 let mut pending = Exchange::new(Message {
734 headers: Default::default(),
735 body: Body::Empty,
736 });
737 pending.set_property(CAMEL_AGGREGATOR_PENDING, serde_json::json!(true));
738 Ok((pending, Vec::new()))
740 }
741 }
742 .await;
743
744 match inner {
745 Ok((exchange, claims)) => AggregationReceipt {
746 reply: Ok(exchange),
747 claims,
748 },
749 Err(e) => AggregationReceipt {
752 reply: Err(e),
753 claims: Vec::new(),
754 },
755 }
756 }
757}
758
759async fn extract_correlation_key(
760 exchange: &Exchange,
761 strategy: &CorrelationStrategy,
762 registry: &SharedLanguageRegistry,
763) -> Result<serde_json::Value, CamelError> {
764 match strategy {
765 CorrelationStrategy::HeaderName(h) => {
766 exchange.input.headers.get(h).cloned().ok_or_else(|| {
767 CamelError::ProcessorError(format!(
768 "Aggregator: missing correlation key header '{}'",
769 h
770 ))
771 })
772 }
773 CorrelationStrategy::Expression { expr, language } => {
774 let expression = {
775 let reg = registry.lock().unwrap_or_else(|e| e.into_inner());
776 let lang = reg.get(language).ok_or_else(|| {
777 CamelError::ProcessorError(format!(
778 "Aggregator: language '{}' not found in registry",
779 language
780 ))
781 })?;
782 lang.create_expression(expr)
783 .map_err(|e| CamelError::ProcessorError(e.to_string()))?
784 };
785 let value = expression
786 .evaluate(exchange)
787 .await
788 .map_err(|e| CamelError::ProcessorError(e.to_string()))?;
789 if value.is_null() {
790 return Err(CamelError::ProcessorError(format!(
791 "Aggregator: correlation expression '{}' evaluated to null",
792 expr
793 )));
794 }
795 Ok(value)
796 }
797 CorrelationStrategy::Fn(f) => f(exchange).map(serde_json::Value::String).ok_or_else(|| {
798 CamelError::ProcessorError("Aggregator: correlation function returned None".to_string())
799 }),
800 _ => Err(CamelError::ProcessorError(
802 "Aggregator: unsupported correlation strategy".to_string(),
803 )),
804 }
805}
806
807fn iter_predicate_exprs(mode: &CompletionMode) -> Vec<(&String, &String)> {
809 let mut out = Vec::new();
810 match mode {
811 CompletionMode::Single(CompletionCondition::PredicateExpr { expr, language }) => {
812 out.push((expr, language));
813 }
814 CompletionMode::Any(conds) => {
815 for c in conds {
816 if let CompletionCondition::PredicateExpr { expr, language } = c {
817 out.push((expr, language));
818 }
819 }
820 }
821 _ => {}
822 }
823 out
824}
825
826async fn evaluate_completion_predicate(
831 mode: &CompletionMode,
832 incoming: &Exchange,
833 registry: &SharedLanguageRegistry,
834) -> Result<bool, CamelError> {
835 let mut satisfied = false;
836 for (expr, language) in iter_predicate_exprs(mode) {
837 let expression = {
838 let reg = registry.lock().unwrap_or_else(|e| e.into_inner());
839 let lang = reg.get(language).ok_or_else(|| {
840 CamelError::ProcessorError(format!(
841 "Aggregator: language '{}' not found in registry",
842 language
843 ))
844 })?;
845 lang.create_expression(expr)
846 .map_err(|e| CamelError::ProcessorError(e.to_string()))?
847 }; let value = expression
849 .evaluate(incoming)
850 .await
851 .map_err(|e| CamelError::ProcessorError(e.to_string()))?;
852 if value.as_bool().unwrap_or(false) {
853 satisfied = true;
854 }
855 }
856 Ok(satisfied)
857}
858
859fn check_sync_completion(
860 mode: &CompletionMode,
861 exchanges: &[Exchange],
862 predicate_satisfied: bool,
863) -> (bool, CompletionReason) {
864 match mode {
865 CompletionMode::Single(cond) => check_single(cond, exchanges, predicate_satisfied),
866 CompletionMode::Any(conditions) => {
867 for cond in conditions {
868 if let CompletionCondition::Timeout(_) = cond {
869 continue;
870 }
871 let (done, reason) = check_single(cond, exchanges, predicate_satisfied);
872 if done {
873 return (true, reason);
874 }
875 }
876 (false, CompletionReason::Size)
877 }
878 _ => (false, CompletionReason::Size),
880 }
881}
882
883fn check_single(
884 cond: &CompletionCondition,
885 exchanges: &[Exchange],
886 predicate_satisfied: bool,
887) -> (bool, CompletionReason) {
888 match cond {
889 CompletionCondition::Size(n) => (exchanges.len() >= *n, CompletionReason::Size),
890 CompletionCondition::Predicate(pred) => (pred(exchanges), CompletionReason::Predicate),
891 CompletionCondition::PredicateExpr { .. } => {
892 (predicate_satisfied, CompletionReason::Predicate)
893 }
894 CompletionCondition::Timeout(_) => (false, CompletionReason::Timeout),
895 _ => (false, CompletionReason::Size),
897 }
898}
899
900fn extract_timeout_duration(mode: &CompletionMode) -> Option<Duration> {
901 match mode {
902 CompletionMode::Single(CompletionCondition::Timeout(d)) => Some(*d),
903 CompletionMode::Any(conditions) => conditions.iter().find_map(|c| {
904 if let CompletionCondition::Timeout(d) = c {
905 Some(*d)
906 } else {
907 None
908 }
909 }),
910 _ => None,
911 }
912}
913
914fn cancel_timeout_task(key: &str, timeout_tasks: &Arc<Mutex<HashMap<String, CancellationToken>>>) {
915 let mut guard = timeout_tasks.lock().unwrap_or_else(|e| e.into_inner());
916 if let Some(token) = guard.remove(key) {
917 token.cancel();
918 }
919}
920
921fn cancel_timeout_task_with_handle(
923 key: &str,
924 timeout_tasks: &Arc<Mutex<HashMap<String, CancellationToken>>>,
925 timeout_handles: &Arc<Mutex<HashMap<String, JoinHandle<()>>>>,
926) {
927 cancel_timeout_task(key, timeout_tasks);
928 let mut guard = timeout_handles.lock().unwrap_or_else(|e| e.into_inner());
929 guard.remove(key);
930}
931
932#[allow(clippy::too_many_arguments)]
933fn spawn_timeout_task(
934 key: String,
935 timeout: Duration,
936 cancel: CancellationToken,
937 buckets: Arc<Mutex<HashMap<String, Bucket>>>,
938 timeout_tasks: Arc<Mutex<HashMap<String, CancellationToken>>>,
939 timeout_handles: Arc<Mutex<HashMap<String, JoinHandle<()>>>>,
940 late_tx: mpsc::Sender<AggregateEmission>,
941 strategy: AggregationStrategy,
942 discard: bool,
943) -> JoinHandle<()> {
944 let cancel_clone = cancel.clone();
945 tokio::spawn(async move {
946 tokio::select! {
947 _ = tokio::time::sleep(timeout) => {
948 let should_proceed = {
949 let mut tt_guard = timeout_tasks.lock().unwrap_or_else(|e| e.into_inner());
950 if cancel_clone.is_cancelled() {
951 false
952 } else {
953 tt_guard.remove(&key);
954 true
955 }
956 };
957 if !should_proceed {
958 return;
959 }
960 {
963 let mut hh = timeout_handles.lock().unwrap_or_else(|e| e.into_inner());
964 hh.remove(&key);
965 }
966 let bucket_parts = {
967 let mut guard = buckets.lock().unwrap_or_else(|e| e.into_inner());
968 guard.remove(&key).map(Bucket::into_parts)
969 };
970 if let Some((exchanges, claims)) = bucket_parts
971 && !discard
972 {
973 match aggregate(exchanges, &strategy) {
974 Ok(mut result) => {
975 result.set_property(
976 CAMEL_AGGREGATED_COMPLETION_REASON,
977 serde_json::json!(CompletionReason::Timeout.as_str()),
978 );
979 if late_tx
980 .try_send(AggregateEmission {
981 exchange: result,
982 claims,
983 })
984 .is_err()
985 {
986 tracing::warn!(
987 key = %key,
988 "aggregator timeout emit dropped: late channel full"
989 );
990 }
991 }
992 Err(e) => {
993 tracing::warn!(
995 key = %key,
996 error = %e,
997 "aggregation failed in timeout task"
998 );
999 }
1000 }
1001 }
1002 }
1003 _ = cancel_clone.cancelled() => {}
1004 }
1005 })
1006}
1007
1008fn aggregate(
1009 exchanges: Vec<Exchange>,
1010 strategy: &AggregationStrategy,
1011) -> Result<Exchange, CamelError> {
1012 match strategy {
1013 AggregationStrategy::CollectAll => {
1014 let bodies: Vec<serde_json::Value> = exchanges
1015 .into_iter()
1016 .map(|e| match e.input.body {
1017 Body::Json(v) => v,
1018 Body::Text(s) => serde_json::Value::String(s),
1019 Body::Xml(s) => serde_json::Value::String(s),
1020 Body::Bytes(b) => {
1021 serde_json::Value::String(String::from_utf8_lossy(&b).into_owned())
1022 }
1023 Body::Stream(s) => serde_json::json!({
1024 "_stream": {
1025 "origin": s.metadata.origin,
1026 "placeholder": true,
1027 "hint": "Materialize exchange body with .into_bytes() before aggregation if content needed"
1028 }
1029 }),
1030 _ => serde_json::Value::Null,
1032 })
1033 .collect();
1034 Ok(Exchange::new(Message {
1035 headers: Default::default(),
1036 body: Body::Json(serde_json::Value::Array(bodies)),
1037 }))
1038 }
1039 AggregationStrategy::Custom(f) => {
1040 let mut iter = exchanges.into_iter();
1041 let first = iter.next().ok_or_else(|| {
1042 CamelError::ProcessorError("Aggregator: empty bucket".to_string())
1043 })?;
1044 Ok(iter.fold(first, |acc, next| f(acc, next)))
1045 }
1046 _ => Err(CamelError::ProcessorError(
1048 "Aggregator: unsupported aggregation strategy".to_string(),
1049 )),
1050 }
1051}
1052
1053#[cfg(test)]
1054mod tests {
1055 use super::*;
1056 use std::collections::HashMap;
1057
1058 use camel_api::{
1059 StepLifecycle, StepShutdownReason,
1060 aggregator::{AggregationStrategy, AggregatorConfig},
1061 body::Body,
1062 exchange::Exchange,
1063 message::Message,
1064 };
1065 use tokio::sync::mpsc;
1066 use tokio_util::sync::CancellationToken;
1067 use tower::ServiceExt;
1068
1069 fn make_exchange(header: &str, value: &str, body: &str) -> Exchange {
1070 let mut msg = Message {
1071 headers: Default::default(),
1072 body: Body::Text(body.to_string()),
1073 };
1074 msg.headers
1075 .insert(header.to_string(), serde_json::json!(value));
1076 Exchange::new(msg)
1077 }
1078
1079 fn config_size(n: usize) -> AggregatorConfig {
1080 AggregatorConfig::correlate_by("orderId")
1081 .complete_when_size(n)
1082 .build()
1083 .unwrap()
1084 }
1085
1086 fn new_test_svc(config: AggregatorConfig) -> AggregatorService {
1087 let (tx, _rx) = mpsc::channel(256);
1088 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1089 let cancel = CancellationToken::new();
1090 AggregatorService::new(config, tx, registry, cancel)
1091 }
1092
1093 fn new_test_svc_with_registry(
1097 config: AggregatorConfig,
1098 registry: SharedLanguageRegistry,
1099 ) -> AggregatorService {
1100 let (tx, _rx) = mpsc::channel(256);
1101 let cancel = CancellationToken::new();
1102 AggregatorService::new(config, tx, registry, cancel)
1103 }
1104
1105 #[tokio::test]
1106 async fn test_pending_exchange_not_yet_complete() {
1107 let mut svc = new_test_svc(config_size(3));
1108 let ex = make_exchange("orderId", "A", "first");
1109 let result = svc.ready().await.unwrap().call(ex).await.unwrap();
1110 assert!(matches!(result.input.body, Body::Empty));
1111 assert_eq!(
1112 result.property(CAMEL_AGGREGATOR_PENDING),
1113 Some(&serde_json::json!(true))
1114 );
1115 }
1116
1117 #[tokio::test]
1120 async fn embedded_stash_keeps_carried_claim_until_completion() {
1121 use std::sync::atomic::{AtomicU64, Ordering};
1122
1123 let counter = Arc::new(AtomicU64::new(0));
1124 let svc = new_test_svc(config_size(2));
1125
1126 let mut first = make_exchange("orderId", "A", "first");
1130 first.in_flight_claim = Some(InFlightClaim::attach(&counter));
1131 let pending = svc.clone().oneshot(first).await.unwrap();
1132 assert_eq!(
1133 pending.property(CAMEL_AGGREGATOR_PENDING),
1134 Some(&serde_json::json!(true)),
1135 "first exchange must be stashed"
1136 );
1137 assert!(
1138 pending.in_flight_claim.is_none(),
1139 "the pending marker carries no claim (it completes immediately)"
1140 );
1141 assert_eq!(
1142 counter.load(Ordering::SeqCst),
1143 1,
1144 "stashed exchange must stay counted"
1145 );
1146
1147 let mut second = make_exchange("orderId", "A", "second");
1151 second.in_flight_claim = Some(InFlightClaim::attach(&counter));
1152 let aggregated = svc.clone().oneshot(second).await.unwrap();
1153 assert!(
1154 aggregated.property(CAMEL_AGGREGATOR_PENDING).is_none(),
1155 "second exchange must complete the bucket"
1156 );
1157 assert!(
1158 aggregated.in_flight_claim.is_some(),
1159 "aggregated output carries one claim downstream"
1160 );
1161 assert_eq!(
1162 counter.load(Ordering::SeqCst),
1163 1,
1164 "consumed input's claim released; the output's stays held"
1165 );
1166 drop(aggregated);
1167 assert_eq!(
1168 counter.load(Ordering::SeqCst),
1169 0,
1170 "dropping the completed output releases the last claim"
1171 );
1172 }
1173
1174 #[tokio::test]
1175 async fn embedded_rejection_releases_carried_claim() {
1176 use std::sync::atomic::{AtomicU64, Ordering};
1177
1178 let counter = Arc::new(AtomicU64::new(0));
1179 let svc = new_test_svc(config_size(2));
1180
1181 let mut bad = Exchange::new(Message::new("no-header"));
1184 bad.in_flight_claim = Some(InFlightClaim::attach(&counter));
1185 let result = svc.clone().oneshot(bad).await;
1186 assert!(result.is_err(), "missing correlation key must reject");
1187 assert_eq!(
1188 counter.load(Ordering::SeqCst),
1189 0,
1190 "rejected submission releases its claim"
1191 );
1192 }
1193
1194 #[tokio::test]
1195 async fn test_completes_on_size() {
1196 let mut svc = new_test_svc(config_size(3));
1197 for _ in 0..2 {
1198 let ex = make_exchange("orderId", "A", "item");
1199 let r = svc.ready().await.unwrap().call(ex).await.unwrap();
1200 assert!(matches!(r.input.body, Body::Empty));
1201 }
1202 let ex = make_exchange("orderId", "A", "last");
1203 let result = svc.ready().await.unwrap().call(ex).await.unwrap();
1204 assert!(result.property(CAMEL_AGGREGATOR_PENDING).is_none());
1205 assert_eq!(
1206 result.property(CAMEL_AGGREGATED_SIZE),
1207 Some(&serde_json::json!(3u64))
1208 );
1209 }
1210
1211 #[tokio::test]
1212 async fn test_collect_all_produces_json_array() {
1213 let mut svc = new_test_svc(config_size(2));
1214 svc.ready()
1215 .await
1216 .unwrap()
1217 .call(make_exchange("orderId", "A", "alpha"))
1218 .await
1219 .unwrap();
1220 let result = svc
1221 .ready()
1222 .await
1223 .unwrap()
1224 .call(make_exchange("orderId", "A", "beta"))
1225 .await
1226 .unwrap();
1227 let Body::Json(v) = &result.input.body else {
1228 panic!("expected Body::Json")
1229 };
1230 let arr = v.as_array().unwrap();
1231 assert_eq!(arr.len(), 2);
1232 assert_eq!(arr[0], serde_json::json!("alpha"));
1233 assert_eq!(arr[1], serde_json::json!("beta"));
1234 }
1235
1236 #[tokio::test]
1237 async fn test_two_keys_independent_buckets() {
1238 let mut svc = new_test_svc(config_size(3));
1240 svc.ready()
1241 .await
1242 .unwrap()
1243 .call(make_exchange("orderId", "A", "a1"))
1244 .await
1245 .unwrap();
1246 svc.ready()
1247 .await
1248 .unwrap()
1249 .call(make_exchange("orderId", "B", "b1"))
1250 .await
1251 .unwrap();
1252 svc.ready()
1253 .await
1254 .unwrap()
1255 .call(make_exchange("orderId", "A", "a2"))
1256 .await
1257 .unwrap();
1258 let ra = svc
1260 .ready()
1261 .await
1262 .unwrap()
1263 .call(make_exchange("orderId", "A", "a3"))
1264 .await
1265 .unwrap();
1266 assert!(matches!(ra.input.body, Body::Json(_)));
1268 let rb = svc
1270 .ready()
1271 .await
1272 .unwrap()
1273 .call(make_exchange("orderId", "B", "b_check"))
1274 .await
1275 .unwrap();
1276 assert!(matches!(rb.input.body, Body::Empty));
1277 }
1278
1279 #[tokio::test]
1280 async fn test_bucket_resets_after_completion() {
1281 let mut svc = new_test_svc(config_size(2));
1282 svc.ready()
1283 .await
1284 .unwrap()
1285 .call(make_exchange("orderId", "A", "x"))
1286 .await
1287 .unwrap();
1288 svc.ready()
1289 .await
1290 .unwrap()
1291 .call(make_exchange("orderId", "A", "x"))
1292 .await
1293 .unwrap(); let r = svc
1296 .ready()
1297 .await
1298 .unwrap()
1299 .call(make_exchange("orderId", "A", "new"))
1300 .await
1301 .unwrap();
1302 assert!(matches!(r.input.body, Body::Empty)); }
1304
1305 #[tokio::test]
1306 async fn test_completion_size_1_emits_immediately() {
1307 let mut svc = new_test_svc(config_size(1));
1308 let ex = make_exchange("orderId", "A", "solo");
1309 let result = svc.ready().await.unwrap().call(ex).await.unwrap();
1310 assert!(result.property(CAMEL_AGGREGATOR_PENDING).is_none());
1311 }
1312
1313 #[tokio::test]
1314 async fn test_custom_aggregation_strategy() {
1315 use camel_api::aggregator::AggregationFn;
1316 use std::sync::Arc;
1317
1318 let f: AggregationFn = Arc::new(|mut acc: Exchange, next: Exchange| {
1319 let combined = format!(
1320 "{}+{}",
1321 acc.input.body.as_text().unwrap_or(""),
1322 next.input.body.as_text().unwrap_or("")
1323 );
1324 acc.input.body = Body::Text(combined);
1325 acc
1326 });
1327 let config = AggregatorConfig::correlate_by("key")
1328 .complete_when_size(2)
1329 .strategy(AggregationStrategy::Custom(f))
1330 .build()
1331 .unwrap();
1332 let mut svc = new_test_svc(config);
1333 svc.ready()
1334 .await
1335 .unwrap()
1336 .call(make_exchange("key", "X", "hello"))
1337 .await
1338 .unwrap();
1339 let result = svc
1340 .ready()
1341 .await
1342 .unwrap()
1343 .call(make_exchange("key", "X", "world"))
1344 .await
1345 .unwrap();
1346 assert_eq!(result.input.body.as_text(), Some("hello+world"));
1347 }
1348
1349 #[tokio::test]
1350 async fn test_completion_predicate() {
1351 let config = AggregatorConfig::correlate_by("key")
1352 .complete_when(|bucket| {
1353 bucket
1354 .iter()
1355 .any(|e| e.input.body.as_text() == Some("DONE"))
1356 })
1357 .build()
1358 .unwrap();
1359 let mut svc = new_test_svc(config);
1360 svc.ready()
1361 .await
1362 .unwrap()
1363 .call(make_exchange("key", "K", "first"))
1364 .await
1365 .unwrap();
1366 svc.ready()
1367 .await
1368 .unwrap()
1369 .call(make_exchange("key", "K", "second"))
1370 .await
1371 .unwrap();
1372 let result = svc
1373 .ready()
1374 .await
1375 .unwrap()
1376 .call(make_exchange("key", "K", "DONE"))
1377 .await
1378 .unwrap();
1379 assert!(result.property(CAMEL_AGGREGATOR_PENDING).is_none());
1380 }
1381
1382 #[tokio::test]
1383 async fn test_missing_header_returns_error() {
1384 let mut svc = new_test_svc(config_size(2));
1385 let msg = Message {
1386 headers: Default::default(),
1387 body: Body::Text("no key".into()),
1388 };
1389 let ex = Exchange::new(msg);
1390 let result = svc.ready().await.unwrap().call(ex).await;
1391 assert!(result.is_err());
1392 assert!(matches!(
1393 result.unwrap_err(),
1394 camel_api::CamelError::ProcessorError(_)
1395 ));
1396 }
1397
1398 #[tokio::test]
1399 async fn test_cloned_service_shares_state() {
1400 let svc1 = new_test_svc(config_size(2));
1401 let mut svc2 = svc1.clone();
1402 svc1.clone()
1404 .ready()
1405 .await
1406 .unwrap()
1407 .call(make_exchange("orderId", "A", "from-svc1"))
1408 .await
1409 .unwrap();
1410 let result = svc2
1412 .ready()
1413 .await
1414 .unwrap()
1415 .call(make_exchange("orderId", "A", "from-svc2"))
1416 .await
1417 .unwrap();
1418 assert!(result.property(CAMEL_AGGREGATOR_PENDING).is_none());
1419 }
1420
1421 #[tokio::test]
1422 async fn test_camel_aggregated_key_property_set() {
1423 let mut svc = new_test_svc(config_size(1));
1424 let ex = make_exchange("orderId", "ORDER-42", "body");
1425 let result = svc.ready().await.unwrap().call(ex).await.unwrap();
1426 assert_eq!(
1427 result.property(CAMEL_AGGREGATED_KEY),
1428 Some(&serde_json::json!("ORDER-42"))
1429 );
1430 }
1431
1432 #[tokio::test]
1433 async fn test_aggregator_enforces_max_buckets() {
1434 let config = AggregatorConfig::correlate_by("orderId")
1435 .complete_when_size(2)
1436 .max_buckets(3)
1437 .build()
1438 .unwrap();
1439
1440 let mut svc = new_test_svc(config);
1441
1442 for i in 0..3 {
1444 let ex = make_exchange("orderId", &format!("key-{}", i), "body");
1445 let _ = svc.ready().await.unwrap().call(ex).await.unwrap();
1446 }
1447
1448 let ex = make_exchange("orderId", "key-4", "body");
1450 let result = svc.ready().await.unwrap().call(ex).await;
1451
1452 assert!(result.is_err(), "Should reject when max buckets reached");
1453 let err = result.unwrap_err().to_string();
1454 assert!(
1455 err.contains("maximum"),
1456 "Error message should contain 'maximum': {}",
1457 err
1458 );
1459 }
1460
1461 #[tokio::test]
1462 async fn test_max_buckets_allows_existing_key() {
1463 let config = AggregatorConfig::correlate_by("orderId")
1464 .complete_when_size(5) .max_buckets(2)
1466 .build()
1467 .unwrap();
1468
1469 let mut svc = new_test_svc(config);
1470
1471 let ex1 = make_exchange("orderId", "key-A", "body1");
1473 let _ = svc.ready().await.unwrap().call(ex1).await.unwrap();
1474 let ex2 = make_exchange("orderId", "key-B", "body2");
1475 let _ = svc.ready().await.unwrap().call(ex2).await.unwrap();
1476
1477 let ex3 = make_exchange("orderId", "key-A", "body3");
1479 let result = svc.ready().await.unwrap().call(ex3).await;
1480 assert!(
1481 result.is_ok(),
1482 "Should allow adding to existing bucket even at max limit"
1483 );
1484 }
1485
1486 #[tokio::test]
1490 async fn test_aggregator_enforces_max_bucket_size() {
1491 let config = AggregatorConfig::correlate_by("orderId")
1492 .complete_when(|_| false)
1494 .max_bucket_size(3)
1495 .build()
1496 .unwrap();
1497
1498 let mut svc = new_test_svc(config);
1499
1500 for _ in 0..3 {
1501 let ex = make_exchange("orderId", "hot-key", "body");
1502 let r = svc.ready().await.unwrap().call(ex).await;
1503 assert!(r.is_ok(), "first 3 exchanges accepted: {r:?}");
1504 }
1505
1506 let ex = make_exchange("orderId", "hot-key", "body");
1507 let result = svc.ready().await.unwrap().call(ex).await;
1508 assert!(result.is_err(), "4th exchange on hot key must be rejected");
1509 let err = result.unwrap_err().to_string();
1510 assert!(
1511 err.contains("maximum"),
1512 "error should mention the per-bucket maximum: {err}"
1513 );
1514
1515 let ex = make_exchange("orderId", "other-key", "body");
1517 let result = svc.ready().await.unwrap().call(ex).await;
1518 assert!(result.is_ok(), "other keys still accepted: {result:?}");
1519 }
1520
1521 #[tokio::test]
1522 async fn test_bucket_ttl_eviction() {
1523 let config = AggregatorConfig::correlate_by("orderId")
1524 .complete_when_size(10) .bucket_ttl(Duration::from_millis(50))
1526 .build()
1527 .unwrap();
1528
1529 let mut svc = new_test_svc(config);
1530
1531 let ex1 = make_exchange("orderId", "key-A", "body1");
1533 let _ = svc.ready().await.unwrap().call(ex1).await.unwrap();
1534
1535 tokio::time::sleep(Duration::from_millis(100)).await;
1537
1538 let ex2 = make_exchange("orderId", "key-B", "body2");
1540 let _ = svc.ready().await.unwrap().call(ex2).await.unwrap();
1541
1542 let ex3 = make_exchange("orderId", "key-A", "body3");
1545 let result = svc.ready().await.unwrap().call(ex3).await;
1546 assert!(result.is_ok(), "Should be able to recreate evicted bucket");
1547 }
1548
1549 #[tokio::test(start_paused = true)]
1550 async fn test_timeout_completes_bucket() {
1551 let config = AggregatorConfig::correlate_by("key")
1552 .complete_on_timeout(Duration::from_millis(100))
1553 .build()
1554 .unwrap();
1555 let mut svc = new_test_svc(config);
1556 let ex = make_exchange("key", "A", "data");
1557 let result = svc.ready().await.unwrap().call(ex).await.unwrap();
1558 assert!(result.property(CAMEL_AGGREGATOR_PENDING).is_some());
1559
1560 tokio::time::sleep(Duration::from_millis(200)).await;
1561
1562 assert_eq!(
1563 svc.buckets.lock().unwrap().len(),
1564 0,
1565 "bucket should be removed after timeout"
1566 );
1567 }
1568
1569 #[tokio::test(start_paused = true)]
1570 async fn test_timeout_resets_on_new_exchange() {
1571 let config = AggregatorConfig::correlate_by("key")
1572 .complete_on_timeout(Duration::from_millis(150))
1573 .build()
1574 .unwrap();
1575 let mut svc = new_test_svc(config);
1576
1577 let ex1 = make_exchange("key", "A", "first");
1578 let _ = svc.ready().await.unwrap().call(ex1).await.unwrap();
1579
1580 tokio::time::sleep(Duration::from_millis(100)).await;
1581
1582 let ex2 = make_exchange("key", "A", "second");
1583 let _ = svc.ready().await.unwrap().call(ex2).await.unwrap();
1584
1585 tokio::time::sleep(Duration::from_millis(100)).await;
1586
1587 assert_eq!(
1588 svc.buckets.lock().unwrap().len(),
1589 1,
1590 "bucket should still exist — timeout was reset"
1591 );
1592
1593 tokio::time::sleep(Duration::from_millis(100)).await;
1594
1595 assert_eq!(
1596 svc.buckets.lock().unwrap().len(),
1597 0,
1598 "bucket should be gone after timeout fires"
1599 );
1600 }
1601
1602 #[tokio::test]
1603 async fn test_composable_size_and_timeout() {
1604 let config = AggregatorConfig::correlate_by("key")
1605 .complete_on_size_or_timeout(2, Duration::from_millis(200))
1606 .build()
1607 .unwrap();
1608 let mut svc = new_test_svc(config);
1609
1610 let ex1 = make_exchange("key", "A", "first");
1611 let _ = svc.ready().await.unwrap().call(ex1).await.unwrap();
1612 assert!(svc.buckets.lock().unwrap().contains_key("\"A\""));
1613
1614 let ex2 = make_exchange("key", "A", "second");
1615 let result = svc.ready().await.unwrap().call(ex2).await.unwrap();
1616 assert!(result.property(CAMEL_AGGREGATOR_PENDING).is_none());
1617 assert_eq!(
1618 result.property(CAMEL_AGGREGATED_COMPLETION_REASON),
1619 Some(&serde_json::json!("size"))
1620 );
1621 }
1622
1623 #[tokio::test(start_paused = true)]
1624 async fn test_discard_on_timeout() {
1625 let config = AggregatorConfig::correlate_by("key")
1626 .complete_on_timeout(Duration::from_millis(50))
1627 .discard_on_timeout(true)
1628 .build()
1629 .unwrap();
1630 let (tx, mut rx) = mpsc::channel(256);
1631 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1632 let cancel = CancellationToken::new();
1633 let mut svc = AggregatorService::new(config, tx, registry, cancel);
1634
1635 let ex = make_exchange("key", "A", "data");
1636 let _ = svc.ready().await.unwrap().call(ex).await.unwrap();
1637
1638 tokio::time::sleep(Duration::from_millis(100)).await;
1639
1640 assert!(
1641 rx.try_recv().is_err(),
1642 "no emit expected with discard_on_timeout"
1643 );
1644 assert_eq!(svc.buckets.lock().unwrap().len(), 0);
1645 assert!(
1646 svc.timeout_tasks.lock().unwrap().is_empty(),
1647 "timeout task should be cleaned up"
1648 );
1649 }
1650
1651 #[tokio::test]
1652 async fn test_force_completion_on_stop() {
1653 let config = AggregatorConfig::correlate_by("key")
1654 .complete_when_size(10)
1655 .force_completion_on_stop(true)
1656 .build()
1657 .unwrap();
1658 let (tx, mut rx) = mpsc::channel(256);
1659 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1660 let cancel = CancellationToken::new();
1661 let svc = AggregatorService::new(config, tx, registry, cancel);
1662
1663 let mut call_svc = svc.clone();
1664 let ex = make_exchange("key", "A", "data");
1665 let _ = call_svc.ready().await.unwrap().call(ex).await.unwrap();
1666
1667 svc.force_complete_all();
1668
1669 let result = rx.try_recv().expect("should emit on force-complete");
1670 let result = result.exchange;
1671 assert!(
1672 result.input.body.as_text().is_some() || matches!(result.input.body, Body::Json(_))
1673 );
1674 assert_eq!(
1675 result.property(CAMEL_AGGREGATED_COMPLETION_REASON),
1676 Some(&serde_json::json!("stop"))
1677 );
1678 }
1679
1680 #[tokio::test]
1681 async fn test_completion_reason_property_size() {
1682 let config = AggregatorConfig::correlate_by("key")
1683 .complete_when_size(1)
1684 .build()
1685 .unwrap();
1686 let mut svc = new_test_svc(config);
1687 let ex = make_exchange("key", "X", "body");
1688 let result = svc.ready().await.unwrap().call(ex).await.unwrap();
1689 assert_eq!(
1690 result.property(CAMEL_AGGREGATED_COMPLETION_REASON),
1691 Some(&serde_json::json!("size"))
1692 );
1693 }
1694
1695 #[tokio::test]
1696 async fn test_completion_reason_property_predicate() {
1697 let config = AggregatorConfig::correlate_by("key")
1698 .complete_when(|_| true)
1699 .build()
1700 .unwrap();
1701 let mut svc = new_test_svc(config);
1702 let ex = make_exchange("key", "X", "body");
1703 let result = svc.ready().await.unwrap().call(ex).await.unwrap();
1704 assert_eq!(
1705 result.property(CAMEL_AGGREGATED_COMPLETION_REASON),
1706 Some(&serde_json::json!("predicate"))
1707 );
1708 }
1709
1710 #[tokio::test(start_paused = true)]
1711 async fn test_size_completes_before_timeout() {
1712 let config = AggregatorConfig::correlate_by("key")
1713 .complete_on_size_or_timeout(2, Duration::from_millis(200))
1714 .build()
1715 .unwrap();
1716 let mut svc = new_test_svc(config);
1717
1718 let ex1 = make_exchange("key", "A", "first");
1719 let _ = svc.ready().await.unwrap().call(ex1).await.unwrap();
1720
1721 let ex2 = make_exchange("key", "A", "second");
1722 let result = svc.ready().await.unwrap().call(ex2).await.unwrap();
1723
1724 assert!(result.property(CAMEL_AGGREGATOR_PENDING).is_none());
1725 assert_eq!(
1726 result.property(CAMEL_AGGREGATED_COMPLETION_REASON),
1727 Some(&serde_json::json!("size"))
1728 );
1729 assert_eq!(svc.buckets.lock().unwrap().len(), 0);
1730
1731 tokio::time::sleep(Duration::from_millis(300)).await;
1732 assert_eq!(
1733 svc.buckets.lock().unwrap().len(),
1734 0,
1735 "no re-fire after timeout"
1736 );
1737 }
1738
1739 #[tokio::test(start_paused = true)]
1740 async fn test_concurrent_timeout_fire_and_new_exchange() {
1741 let config = AggregatorConfig::correlate_by("key")
1742 .complete_on_size_or_timeout(2, Duration::from_millis(100))
1743 .build()
1744 .unwrap();
1745 let (tx, mut rx) = mpsc::channel(256);
1746 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1747 let cancel = CancellationToken::new();
1748 let mut svc = AggregatorService::new(config, tx, registry, cancel);
1749
1750 let ex = make_exchange("key", "A", "data");
1751 let _ = svc.ready().await.unwrap().call(ex).await.unwrap();
1752
1753 tokio::time::sleep(Duration::from_millis(150)).await;
1755
1756 let ex2 = make_exchange("key", "A", "data2");
1758 let result = svc.ready().await.unwrap().call(ex2).await.unwrap();
1759 assert!(
1760 result.property(CAMEL_AGGREGATOR_PENDING).is_some(),
1761 "should be pending in new bucket"
1762 );
1763
1764 let mut late_count = 0;
1766 while rx.try_recv().is_ok() {
1767 late_count += 1;
1768 }
1769 assert_eq!(
1770 late_count, 1,
1771 "exactly 1 late emit from the timed-out bucket"
1772 );
1773 }
1774
1775 #[tokio::test(start_paused = true)]
1776 async fn test_late_channel_full_drops_with_warning() {
1777 let config = AggregatorConfig::correlate_by("key")
1778 .complete_on_timeout(Duration::from_millis(50))
1779 .build()
1780 .unwrap();
1781 let (tx, mut rx) = mpsc::channel(1);
1782 rx.close();
1783 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1784 let cancel = CancellationToken::new();
1785 let mut svc = AggregatorService::new(config, tx, registry, cancel);
1786
1787 let ex = make_exchange("key", "A", "data");
1788 let _ = svc.ready().await.unwrap().call(ex).await.unwrap();
1789
1790 tokio::time::sleep(Duration::from_millis(100)).await;
1791 assert_eq!(
1792 svc.buckets.lock().unwrap().len(),
1793 0,
1794 "bucket removed despite channel closed"
1795 );
1796 }
1797
1798 #[tokio::test]
1804 async fn test_da3_force_complete_all_drops_on_saturated_channel() {
1805 let config = AggregatorConfig::correlate_by("k")
1806 .complete_when_size(10)
1807 .force_completion_on_stop(true)
1808 .build()
1809 .unwrap();
1810 let (late_tx, mut late_rx) = mpsc::channel::<AggregateEmission>(1);
1812 late_tx
1816 .try_send(AggregateEmission {
1817 exchange: make_exchange("k", "99", "dummy"),
1818 claims: Vec::new(),
1819 })
1820 .expect("pre-fill succeeds");
1821
1822 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1823 let cancel = CancellationToken::new();
1824 let mut svc = AggregatorService::new(config, late_tx, registry, cancel);
1825
1826 for v in ["1", "2", "3"] {
1829 let ex = make_exchange("k", v, "body");
1830 let _ = svc.ready().await.unwrap().call(ex).await.unwrap();
1831 }
1832 assert_eq!(svc.buckets.lock().unwrap().len(), 3);
1833
1834 svc.force_complete_all();
1836
1837 let pre_fill = late_rx
1839 .try_recv()
1840 .expect("pre-fill should still be in channel");
1841 assert_eq!(
1842 pre_fill.exchange.input.headers.get("k"),
1843 Some(&serde_json::json!("99"))
1844 );
1845
1846 assert!(matches!(
1849 late_rx.try_recv(),
1850 Err(mpsc::error::TryRecvError::Empty)
1851 ));
1852
1853 assert!(svc.buckets.lock().unwrap().is_empty());
1855 }
1856
1857 #[tokio::test]
1858 async fn test_aggregate_stream_bodies_creates_valid_json() {
1859 use bytes::Bytes;
1860 use camel_api::{Body, StreamBody, StreamMetadata};
1861 use futures::stream;
1862 use tokio::sync::Mutex;
1863
1864 let chunks = vec![Ok(Bytes::from("test"))];
1865 let stream_body = StreamBody {
1866 stream: Arc::new(Mutex::new(Some(Box::pin(stream::iter(chunks))))),
1867 metadata: StreamMetadata {
1868 origin: Some("file:///test.txt".to_string()),
1869 ..Default::default()
1870 },
1871 };
1872
1873 let ex1 = Exchange::new(Message {
1874 headers: Default::default(),
1875 body: Body::Stream(stream_body),
1876 });
1877
1878 let exchanges = vec![ex1];
1879 let result = aggregate(exchanges, &AggregationStrategy::CollectAll);
1880
1881 let exchange = result.expect("Expected Ok result");
1882 assert!(
1883 matches!(exchange.input.body, Body::Json(_)),
1884 "Expected Json body"
1885 );
1886
1887 if let Body::Json(value) = exchange.input.body {
1888 let json_str = serde_json::to_string(&value).unwrap();
1889 let parsed: serde_json::Value = serde_json::from_str(&json_str).unwrap();
1890
1891 assert!(parsed.is_array(), "Result should be an array");
1892 let arr = parsed.as_array().unwrap();
1893 assert!(arr[0].is_object(), "First element should be an object");
1894 assert!(
1895 arr[0]["_stream"].is_object(),
1896 "Should contain _stream object"
1897 );
1898 assert_eq!(arr[0]["_stream"]["origin"], "file:///test.txt");
1899 assert_eq!(
1900 arr[0]["_stream"]["placeholder"], true,
1901 "placeholder flag should be true"
1902 );
1903 }
1904 }
1905
1906 #[tokio::test]
1907 async fn test_aggregate_stream_bodies_with_none_origin() {
1908 use bytes::Bytes;
1909 use camel_api::{Body, StreamBody, StreamMetadata};
1910 use futures::stream;
1911 use tokio::sync::Mutex;
1912
1913 let chunks = vec![Ok(Bytes::from("test"))];
1914 let stream_body = StreamBody {
1915 stream: Arc::new(Mutex::new(Some(Box::pin(stream::iter(chunks))))),
1916 metadata: StreamMetadata {
1917 origin: None,
1918 ..Default::default()
1919 },
1920 };
1921
1922 let ex1 = Exchange::new(Message {
1923 headers: Default::default(),
1924 body: Body::Stream(stream_body),
1925 });
1926
1927 let exchanges = vec![ex1];
1928 let result = aggregate(exchanges, &AggregationStrategy::CollectAll);
1929
1930 let exchange = result.expect("Expected Ok result");
1931 assert!(
1932 matches!(exchange.input.body, Body::Json(_)),
1933 "Expected Json body"
1934 );
1935
1936 if let Body::Json(value) = exchange.input.body {
1937 let json_str = serde_json::to_string(&value).unwrap();
1938 let parsed: serde_json::Value = serde_json::from_str(&json_str).unwrap();
1939
1940 assert!(parsed.is_array(), "Result should be an array");
1941 let arr = parsed.as_array().unwrap();
1942 assert!(arr[0].is_object(), "First element should be an object");
1943 assert!(
1944 arr[0]["_stream"].is_object(),
1945 "Should contain _stream object"
1946 );
1947 assert_eq!(
1948 arr[0]["_stream"]["origin"],
1949 serde_json::Value::Null,
1950 "origin should be null when None"
1951 );
1952 assert_eq!(
1953 arr[0]["_stream"]["placeholder"], true,
1954 "placeholder flag should be true"
1955 );
1956 }
1957 }
1958
1959 #[tokio::test]
1960 async fn timeout_completion_clears_handle_from_map() {
1961 let config = AggregatorConfig::correlate_by("key")
1966 .complete_on_timeout(Duration::from_millis(50))
1967 .build()
1968 .unwrap();
1969 let (tx, _rx) = mpsc::channel(256);
1970 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1971 let cancel = CancellationToken::new();
1972 let svc = AggregatorService::new(config, tx, registry, cancel);
1973
1974 let mut call_svc = svc.clone();
1976 let ex = make_exchange("key", "A", "data");
1977 let _ = call_svc.ready().await.unwrap().call(ex).await.unwrap();
1978 assert!(
1979 !svc.timeout_handles.lock().unwrap().is_empty(),
1980 "handle should exist while timeout pending"
1981 );
1982
1983 tokio::time::sleep(Duration::from_millis(200)).await;
1985
1986 assert!(
1987 svc.timeout_handles.lock().unwrap().is_empty(),
1988 "handle should be cleared from map after natural timeout completion (was leak)"
1989 );
1990 }
1991
1992 #[tokio::test]
1993 async fn aggregator_shutdown_via_trait_dispatch() {
1994 let config = AggregatorConfig::correlate_by("key")
1997 .complete_when_size(10)
1998 .build()
1999 .unwrap();
2000 let (tx, _rx) = mpsc::channel(256);
2001 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
2002 let cancel = CancellationToken::new();
2003 let svc = AggregatorService::new(config, tx, registry, cancel);
2004
2005 let step: Arc<dyn StepLifecycle> = Arc::new(svc);
2006 step.shutdown(StepShutdownReason::RouteStop)
2007 .await
2008 .expect("first shutdown should succeed");
2009 step.shutdown(StepShutdownReason::RouteStop)
2010 .await
2011 .expect("second shutdown (idempotent) should succeed");
2012 }
2013
2014 #[tokio::test(start_paused = true)]
2015 async fn test_shutdown_awaits_timeout_handles() {
2016 let config = AggregatorConfig::correlate_by("key")
2017 .complete_on_timeout(Duration::from_millis(100))
2018 .build()
2019 .unwrap();
2020 let (tx, _rx) = mpsc::channel(256);
2021 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
2022 let cancel = CancellationToken::new();
2023 let svc = AggregatorService::new(config, tx, registry, cancel);
2024
2025 let mut call_svc = svc.clone();
2027 let ex = make_exchange("key", "A", "data");
2028 let _ = call_svc.ready().await.unwrap().call(ex).await.unwrap();
2029
2030 assert!(
2032 !svc.timeout_handles.lock().unwrap().is_empty(),
2033 "should have a timeout handle"
2034 );
2035
2036 svc.shutdown_inner().await;
2039
2040 assert!(
2041 svc.timeout_handles.lock().unwrap().is_empty(),
2042 "all handles should be cleaned up after shutdown"
2043 );
2044 }
2045
2046 #[tokio::test]
2053 async fn test_unique_key_flood_stays_bounded_by_default() {
2054 let config = AggregatorConfig::correlate_by("orderId")
2056 .complete_when_size(1_000_000) .build()
2058 .unwrap();
2059 let mut svc = new_test_svc(config);
2060
2061 for i in 0..10_000usize {
2064 let ex = make_exchange("orderId", &format!("key-{i}"), "body");
2065 let result = svc.ready().await.unwrap().call(ex).await;
2066 assert!(result.is_ok(), "key {i} should be accepted under the cap");
2067 }
2068 let ex = make_exchange("orderId", "key-10001", "body");
2069 let result = svc.ready().await.unwrap().call(ex).await;
2070 assert!(
2071 result.is_err(),
2072 "10_001st unique key must be rejected by the max_buckets cap"
2073 );
2074 let err = result.unwrap_err().to_string();
2075 assert!(
2076 err.contains("maximum") || err.contains("max"),
2077 "error should mention cap: {err}"
2078 );
2079 }
2080
2081 #[tokio::test]
2087 async fn test_background_sweep_spawns_on_first_poll_not_construction() {
2088 let config = AggregatorConfig::correlate_by("key")
2089 .complete_when_size(10_000)
2090 .bucket_ttl(Duration::from_millis(50))
2091 .build()
2092 .unwrap();
2093 let cancel = CancellationToken::new();
2094 let (tx, _rx) = mpsc::channel(8);
2095 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
2096 let mut svc = AggregatorService::new(config, tx, registry, cancel.clone());
2097
2098 assert!(
2100 svc.sweep_handle
2101 .lock()
2102 .unwrap_or_else(|e| e.into_inner())
2103 .is_none(),
2104 "sweep must NOT be spawned at construction (runtime-free new)"
2105 );
2106
2107 let _ = svc.ready().await.unwrap();
2110 let sweep_present = svc
2111 .sweep_handle
2112 .lock()
2113 .unwrap_or_else(|e| e.into_inner())
2114 .is_some();
2115 assert!(
2116 sweep_present,
2117 "sweep handle should be Some after first poll when bucket_ttl is set"
2118 );
2119
2120 cancel.cancel();
2122 tokio::time::sleep(Duration::from_millis(50)).await;
2124 }
2125
2126 struct QueueDepthRecorder(Mutex<Vec<(String, usize)>>);
2129 impl MetricsCollector for QueueDepthRecorder {
2130 fn record_exchange_duration(&self, _: &str, _: std::time::Duration) {}
2131 fn increment_errors(&self, _: &str, _: &str) {}
2132 fn increment_exchanges(&self, _: &str) {}
2133 fn set_queue_depth(&self, queue: &str, depth: usize) {
2134 self.0
2135 .lock()
2136 .unwrap_or_else(|e| e.into_inner())
2137 .push((queue.to_string(), depth));
2138 }
2139 fn record_circuit_breaker_change(&self, _: &str, _: &str, _: &str) {}
2140 }
2141
2142 async fn await_depth_sample(
2145 recorder: &QueueDepthRecorder,
2146 label: &str,
2147 pred: impl Fn(usize) -> bool,
2148 ) {
2149 let deadline = std::time::Instant::now() + Duration::from_secs(2);
2150 loop {
2151 let matched = recorder
2152 .0
2153 .lock()
2154 .unwrap_or_else(|e| e.into_inner())
2155 .iter()
2156 .any(|(q, d)| q == label && pred(*d));
2157 if matched {
2158 return;
2159 }
2160 assert!(
2161 std::time::Instant::now() < deadline,
2162 "no queue-depth sample for '{label}' matched within 2s"
2163 );
2164 tokio::time::sleep(Duration::from_millis(20)).await;
2165 }
2166 }
2167
2168 #[tokio::test]
2173 async fn test_sweep_reports_queue_depth() {
2174 let config = AggregatorConfig::correlate_by("orderId")
2175 .complete_when_size(3)
2176 .bucket_ttl(Duration::from_millis(100))
2180 .build()
2181 .unwrap();
2182 let cancel = CancellationToken::new();
2183 let (tx, _rx) = mpsc::channel(8);
2184 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
2185 let recorder = Arc::new(QueueDepthRecorder(Mutex::new(Vec::new())));
2186 let mut svc = AggregatorService::new(config, tx, registry, cancel).with_queue_metrics(
2187 Arc::clone(&recorder) as Arc<dyn MetricsCollector>,
2188 "aggregator:t",
2189 );
2190
2191 let _ = svc.ready().await.unwrap();
2193
2194 let ex = make_exchange("orderId", "g1", "partial");
2196 let _ = svc.ready().await.unwrap().call(ex).await;
2197
2198 await_depth_sample(&recorder, "aggregator:t", |d| d > 0).await;
2199
2200 let ex = make_exchange("orderId", "g1", "b");
2202 let _ = svc.ready().await.unwrap().call(ex).await;
2203 let ex = make_exchange("orderId", "g1", "c");
2204 let _ = svc.ready().await.unwrap().call(ex).await;
2205
2206 await_depth_sample(&recorder, "aggregator:t", |d| d == 0).await;
2207
2208 svc.shutdown(StepShutdownReason::RouteStop).await.unwrap();
2209 }
2210
2211 #[tokio::test]
2216 async fn test_no_ttl_aggregator_reports_queue_depth() {
2217 use camel_api::aggregator::CorrelationStrategy;
2218
2219 let config = AggregatorConfig {
2220 header_name: "orderId".into(),
2221 completion: CompletionMode::Single(CompletionCondition::Size(3)),
2222 correlation: CorrelationStrategy::HeaderName("orderId".into()),
2223 strategy: AggregationStrategy::CollectAll,
2224 max_buckets: Some(100),
2225 max_bucket_size: Some(100),
2226 bucket_ttl: None,
2227 force_completion_on_stop: false,
2228 discard_on_timeout: false,
2229 max_timeout_tasks: 64,
2230 };
2231 let cancel = CancellationToken::new();
2232 let (tx, _rx) = mpsc::channel(8);
2233 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
2234 let recorder = Arc::new(QueueDepthRecorder(Mutex::new(Vec::new())));
2235 let mut svc = AggregatorService::new(config, tx, registry, cancel).with_queue_metrics(
2236 Arc::clone(&recorder) as Arc<dyn MetricsCollector>,
2237 "aggregator:nottl",
2238 );
2239
2240 let _ = svc.ready().await.unwrap();
2241 assert!(
2242 svc.sweep_handle
2243 .lock()
2244 .unwrap_or_else(|e| e.into_inner())
2245 .is_some(),
2246 "metrics-only config (bucket_ttl = None) must still spawn the sweep"
2247 );
2248
2249 let ex = make_exchange("orderId", "g1", "partial");
2251 let _ = svc.ready().await.unwrap().call(ex).await;
2252 await_depth_sample(&recorder, "aggregator:nottl", |d| d > 0).await;
2253
2254 let ex = make_exchange("orderId", "g1", "b");
2256 let _ = svc.ready().await.unwrap().call(ex).await;
2257 let ex = make_exchange("orderId", "g1", "c");
2258 let _ = svc.ready().await.unwrap().call(ex).await;
2259 await_depth_sample(&recorder, "aggregator:nottl", |d| d == 0).await;
2260
2261 svc.shutdown(StepShutdownReason::RouteStop).await.unwrap();
2262 }
2263
2264 #[tokio::test]
2269 async fn test_transient_clone_drop_keeps_sweep_sampling() {
2270 let config = AggregatorConfig::correlate_by("orderId")
2271 .complete_when_size(3)
2272 .bucket_ttl(Duration::from_millis(100))
2273 .build()
2274 .unwrap();
2275 let cancel = CancellationToken::new();
2276 let (tx, _rx) = mpsc::channel(8);
2277 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
2278 let recorder = Arc::new(QueueDepthRecorder(Mutex::new(Vec::new())));
2279 let mut svc = AggregatorService::new(config, tx, registry, cancel).with_queue_metrics(
2280 Arc::clone(&recorder) as Arc<dyn MetricsCollector>,
2281 "aggregator:clone",
2282 );
2283
2284 let _ = svc.ready().await.unwrap();
2285
2286 drop(svc.clone());
2288 assert!(
2289 svc.sweep_handle
2290 .lock()
2291 .unwrap_or_else(|e| e.into_inner())
2292 .is_some(),
2293 "transient clone drop must not abort the shared sweep"
2294 );
2295
2296 let ex = make_exchange("orderId", "g1", "partial");
2298 let _ = svc.ready().await.unwrap().call(ex).await;
2299 await_depth_sample(&recorder, "aggregator:clone", |d| d > 0).await;
2300
2301 svc.shutdown(StepShutdownReason::RouteStop).await.unwrap();
2302 }
2303
2304 #[tokio::test]
2307 async fn test_aggregator_timeout_task_cap_no_panic_under_flood() {
2308 use camel_api::aggregator::CorrelationStrategy;
2311
2312 let config = AggregatorConfig {
2313 header_name: "k".into(),
2314 completion: CompletionMode::Any(vec![
2315 CompletionCondition::Size(999),
2316 CompletionCondition::Timeout(Duration::from_secs(30)),
2317 ]),
2318 correlation: CorrelationStrategy::HeaderName("k".into()),
2319 strategy: AggregationStrategy::CollectAll,
2320 max_buckets: Some(50),
2321 max_bucket_size: Some(50),
2322 bucket_ttl: Some(Duration::from_secs(30)),
2323 force_completion_on_stop: false,
2324 discard_on_timeout: false,
2325 max_timeout_tasks: 2,
2326 };
2327 let (late_tx, mut late_rx) = mpsc::channel(64);
2328 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
2329 let cancel = CancellationToken::new();
2330 let svc = AggregatorService::new(config, late_tx, registry, cancel);
2331
2332 for i in 0..20u64 {
2334 let mut ex = Exchange::new(Message {
2335 headers: HashMap::from([("k".to_string(), serde_json::json!(i))]),
2336 body: Body::Text(i.to_string()),
2337 });
2338 ex.input
2339 .headers
2340 .insert("k".to_string(), serde_json::json!(i));
2341 let outcome = tokio::time::timeout(Duration::from_secs(2), async {
2342 let mut s = svc.clone();
2343 use tower::ServiceExt;
2344 s.ready().await.unwrap().call(ex).await
2345 })
2346 .await;
2347 assert!(outcome.is_ok(), "call {} hung/panicked under task cap", i);
2348 let res = outcome.unwrap().unwrap();
2350 assert_eq!(
2351 res.properties
2352 .get(CAMEL_AGGREGATOR_PENDING)
2353 .and_then(|v| v.as_bool()),
2354 Some(true),
2355 "exchange {} should be pending",
2356 i
2357 );
2358 }
2359 let _ = late_rx.try_recv();
2361 }
2362
2363 #[tokio::test]
2371 async fn test_release_unarmed_buckets_discards_cap_exceeded_keeps_armed() {
2372 let config = AggregatorConfig::correlate_by("k")
2373 .complete_on_timeout(Duration::from_millis(400))
2374 .max_timeout_tasks(1)
2375 .build()
2376 .unwrap();
2377 let (late_tx, mut late_rx) = mpsc::channel(8);
2378 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
2379 let cancel = CancellationToken::new();
2380 let svc = AggregatorService::new(config, late_tx, registry, cancel);
2381
2382 let a = make_exchange("k", "a", "body-a");
2385 let mut sa = svc.clone();
2386 let _ = sa.ready().await.unwrap().call(a).await.unwrap();
2387 let b = make_exchange("k", "b", "body-b");
2388 let mut sb = svc.clone();
2389 let _ = sb.ready().await.unwrap().call(b).await.unwrap();
2390
2391 {
2392 let armed = svc.timeout_tasks.lock().unwrap_or_else(|e| e.into_inner());
2393 assert_eq!(armed.len(), 1, "cap=1 must arm exactly one timeout task");
2394 assert!(
2395 armed.contains_key("\"a\""),
2396 "first key must hold the armed slot, got {armed:?}"
2397 );
2398 }
2399
2400 svc.release_unarmed_buckets();
2402 {
2403 let buckets = svc.buckets.lock().unwrap_or_else(|e| e.into_inner());
2404 assert!(
2405 buckets.contains_key("\"a\""),
2406 "armed bucket must survive the release"
2407 );
2408 assert!(
2409 !buckets.contains_key("\"b\""),
2410 "cap-exceeded (unarmed) bucket must be released, not orphaned"
2411 );
2412 }
2413
2414 let emitted = tokio::time::timeout(Duration::from_secs(2), late_rx.recv())
2417 .await
2418 .expect("armed bucket must emit on its timeout")
2419 .expect("late channel must stay open");
2420 assert_eq!(
2421 emitted
2422 .exchange
2423 .properties
2424 .get(CAMEL_AGGREGATED_COMPLETION_REASON),
2425 Some(&serde_json::json!("timeout"))
2426 );
2427
2428 assert!(
2432 tokio::time::timeout(Duration::from_millis(600), late_rx.recv())
2433 .await
2434 .is_err(),
2435 "released unarmed bucket must not emit"
2436 );
2437
2438 svc.shutdown(StepShutdownReason::RouteStop).await.unwrap();
2439 }
2440
2441 #[tokio::test]
2442 async fn evaluate_completion_predicate_or_combines_any() {
2443 use camel_api::aggregator::{CompletionCondition, CompletionMode};
2444 use camel_language_api::Language;
2445 use std::collections::HashMap;
2446
2447 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
2449 registry.lock().unwrap().insert(
2450 "simple".to_string(),
2451 Arc::new(camel_language_simple::SimpleLanguage::new()) as Arc<dyn Language>,
2452 );
2453
2454 let incoming = make_exchange("k", "X", "hello");
2455
2456 let mode = CompletionMode::Any(vec![
2458 CompletionCondition::PredicateExpr {
2459 expr: "${body} == 'NOPE'".to_string(),
2460 language: "simple".to_string(),
2461 },
2462 CompletionCondition::PredicateExpr {
2463 expr: "${body} == 'hello'".to_string(),
2464 language: "simple".to_string(),
2465 },
2466 ]);
2467
2468 let satisfied = evaluate_completion_predicate(&mode, &incoming, ®istry)
2469 .await
2470 .expect("eval must succeed");
2471 assert!(satisfied, "second predicate matches → OR true");
2472 }
2473
2474 #[tokio::test]
2475 async fn evaluate_completion_predicate_no_predicate_skips_registry() {
2476 use camel_api::aggregator::{CompletionCondition, CompletionMode};
2477 use std::collections::HashMap;
2478
2479 let mode = CompletionMode::Single(CompletionCondition::Size(3));
2483 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
2484 let incoming = make_exchange("k", "X", "hello");
2485 let result = evaluate_completion_predicate(&mode, &incoming, ®istry).await;
2486 assert!(result.is_ok(), "fast path must not error: {:?}", result);
2487 assert!(!result.unwrap(), "no predicate → not satisfied");
2488 }
2489
2490 #[tokio::test]
2491 async fn evaluate_completion_predicate_unregistered_language_errors() {
2492 use camel_api::aggregator::{CompletionCondition, CompletionMode};
2493
2494 let mode = CompletionMode::Single(CompletionCondition::PredicateExpr {
2495 expr: "${body} == 'DONE'".to_string(),
2496 language: "nonexistent-lang".to_string(),
2497 });
2498 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
2499 let incoming = make_exchange("k", "X", "hello");
2500 let result = evaluate_completion_predicate(&mode, &incoming, ®istry).await;
2501 assert!(result.is_err(), "unregistered language must error");
2502 }
2503
2504 #[tokio::test]
2505 async fn evaluate_completion_predicate_all_miss_returns_false() {
2506 use camel_api::aggregator::{CompletionCondition, CompletionMode};
2507 use camel_language_api::Language;
2508 use std::collections::HashMap;
2509
2510 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
2511 registry.lock().unwrap().insert(
2512 "simple".to_string(),
2513 Arc::new(camel_language_simple::SimpleLanguage::new()) as Arc<dyn Language>,
2514 );
2515
2516 let mode = CompletionMode::Single(CompletionCondition::PredicateExpr {
2517 expr: "${body} == 'NOPE'".to_string(),
2518 language: "simple".to_string(),
2519 });
2520 let incoming = make_exchange("k", "X", "hello");
2521 let result = evaluate_completion_predicate(&mode, &incoming, ®istry).await;
2522 assert!(result.is_ok(), "eval must succeed: {:?}", result);
2523 assert!(!result.unwrap(), "predicate does not match");
2524 }
2525
2526 #[tokio::test]
2527 async fn completion_predicate_expr_completes_on_match() {
2528 use camel_api::aggregator::{CompletionCondition, CompletionMode};
2529 use camel_language_api::Language;
2530 use std::collections::HashMap;
2531
2532 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
2534 registry.lock().unwrap().insert(
2535 "simple".to_string(),
2536 Arc::new(camel_language_simple::SimpleLanguage::new()) as Arc<dyn Language>,
2537 );
2538
2539 let mut config = AggregatorConfig::correlate_by("key")
2544 .complete_when_size(999) .build()
2546 .expect("config builds");
2547 config.completion = CompletionMode::Single(CompletionCondition::PredicateExpr {
2548 expr: "${body} == 'DONE'".to_string(),
2549 language: "simple".to_string(),
2550 });
2551
2552 let mut svc = new_test_svc_with_registry(config, registry);
2553
2554 let r = svc
2556 .ready()
2557 .await
2558 .unwrap()
2559 .call(make_exchange("key", "K", "first"))
2560 .await
2561 .unwrap();
2562 assert!(r.property(CAMEL_AGGREGATOR_PENDING).is_some());
2563
2564 let r = svc
2566 .ready()
2567 .await
2568 .unwrap()
2569 .call(make_exchange("key", "K", "DONE"))
2570 .await
2571 .unwrap();
2572 assert!(r.property(CAMEL_AGGREGATOR_PENDING).is_none());
2573 assert_eq!(
2574 r.property(CAMEL_AGGREGATED_COMPLETION_REASON),
2575 Some(&serde_json::json!("predicate"))
2576 );
2577 }
2578
2579 #[tokio::test]
2591 async fn test_da1_strategy_receives_two_exchanges_first_message_preserved() {
2592 use camel_api::aggregator::{AggregationFn, AggregationStrategy};
2593 use std::sync::Arc;
2594
2595 let recorded: Arc<std::sync::Mutex<Vec<(String, String)>>> =
2596 Arc::new(std::sync::Mutex::new(Vec::new()));
2597 let recorded_for_closure = Arc::clone(&recorded);
2598
2599 let f: AggregationFn = Arc::new(move |old: Exchange, new: Exchange| {
2600 let old_body = old.input.body.as_text().unwrap_or("").to_string();
2601 let new_body = new.input.body.as_text().unwrap_or("").to_string();
2602 recorded_for_closure
2603 .lock()
2604 .expect("recorded mutex poisoned")
2605 .push((old_body, new_body));
2606 new
2607 });
2608
2609 let config = AggregatorConfig::correlate_by("k")
2610 .complete_when_size(2)
2611 .strategy(AggregationStrategy::Custom(f))
2612 .build()
2613 .unwrap();
2614 let mut svc = new_test_svc(config);
2615
2616 let first = svc
2618 .ready()
2619 .await
2620 .unwrap()
2621 .call(make_exchange("k", "1", "A"))
2622 .await
2623 .unwrap();
2624 assert!(
2625 first.property(CAMEL_AGGREGATOR_PENDING).is_some(),
2626 "first message must leave the bucket pending (size < 2)"
2627 );
2628 assert!(
2629 recorded.lock().expect("recorded mutex poisoned").is_empty(),
2630 "strategy must NOT be invoked on the first message of a bucket"
2631 );
2632
2633 let _completed = svc
2636 .ready()
2637 .await
2638 .unwrap()
2639 .call(make_exchange("k", "1", "B"))
2640 .await
2641 .unwrap();
2642
2643 let recorded = recorded.lock().expect("recorded mutex poisoned");
2644 assert_eq!(
2645 recorded.len(),
2646 1,
2647 "strategy must be invoked exactly once across the two-message bucket, got {recorded:?}"
2648 );
2649 assert_eq!(
2650 recorded[0],
2651 ("A".to_string(), "B".to_string()),
2652 "strategy must observe the first message as `old` and the second as `new`, \
2653 with both bodies preserved unchanged"
2654 );
2655 }
2656
2657 #[tokio::test]
2660 async fn sweep_shutdown_cancels_task() {
2661 let config = AggregatorConfig::correlate_by("key")
2662 .complete_when_size(10)
2663 .bucket_ttl(Duration::from_millis(100))
2664 .build()
2665 .unwrap();
2666 let (tx, _rx) = mpsc::channel(256);
2667 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
2668 let cancel = CancellationToken::new();
2669 let mut svc = AggregatorService::new(config, tx, registry, cancel);
2670
2671 let _ = svc.ready().await.unwrap();
2672
2673 assert!(
2674 svc.sweep_handle
2675 .lock()
2676 .unwrap_or_else(|e| e.into_inner())
2677 .is_some(),
2678 "sweep handle should be Some after poll_ready"
2679 );
2680
2681 svc.shutdown(StepShutdownReason::RouteStop).await.unwrap();
2682
2683 assert!(
2684 svc.sweep_handle
2685 .lock()
2686 .unwrap_or_else(|e| e.into_inner())
2687 .is_none(),
2688 "sweep handle should be None after shutdown (taken + aborted)"
2689 );
2690 assert!(
2691 svc.sweep_cancel
2692 .lock()
2693 .unwrap_or_else(|e| e.into_inner())
2694 .is_cancelled(),
2695 "sweep_cancel token should be cancelled after shutdown"
2696 );
2697
2698 tokio::time::sleep(Duration::from_millis(50)).await;
2699 }
2700
2701 #[tokio::test]
2702 async fn sweep_start_respawns_after_shutdown() {
2703 let config = AggregatorConfig::correlate_by("key")
2704 .complete_when_size(10)
2705 .bucket_ttl(Duration::from_millis(100))
2706 .build()
2707 .unwrap();
2708 let (tx, _rx) = mpsc::channel(256);
2709 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
2710 let cancel = CancellationToken::new();
2711 let mut svc = AggregatorService::new(config, tx, registry, cancel);
2712
2713 let _ = svc.ready().await.unwrap();
2714 svc.shutdown(StepShutdownReason::RouteStop).await.unwrap();
2715
2716 svc.start().await.unwrap();
2717 let _ = svc.ready().await.unwrap();
2718
2719 assert!(
2720 svc.sweep_handle
2721 .lock()
2722 .unwrap_or_else(|e| e.into_inner())
2723 .is_some(),
2724 "sweep handle should be Some after start + poll_ready"
2725 );
2726 assert!(
2727 !svc.sweep_cancel
2728 .lock()
2729 .unwrap_or_else(|e| e.into_inner())
2730 .is_cancelled(),
2731 "sweep_cancel should be a fresh uncancelled token after start"
2732 );
2733 }
2734
2735 #[tokio::test]
2736 async fn sweep_shutdown_hotswap_cancels_task() {
2737 let config = AggregatorConfig::correlate_by("key")
2738 .complete_when_size(10)
2739 .bucket_ttl(Duration::from_millis(100))
2740 .build()
2741 .unwrap();
2742 let (tx, _rx) = mpsc::channel(256);
2743 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
2744 let cancel = CancellationToken::new();
2745 let mut svc = AggregatorService::new(config, tx, registry, cancel);
2746
2747 let _ = svc.ready().await.unwrap();
2748
2749 assert!(
2750 svc.sweep_handle
2751 .lock()
2752 .unwrap_or_else(|e| e.into_inner())
2753 .is_some(),
2754 "sweep handle should be Some after poll_ready"
2755 );
2756
2757 svc.shutdown(StepShutdownReason::HotSwap).await.unwrap();
2758
2759 assert!(
2760 svc.sweep_handle
2761 .lock()
2762 .unwrap_or_else(|e| e.into_inner())
2763 .is_none(),
2764 "sweep handle should be None after HotSwap shutdown"
2765 );
2766 assert!(
2767 svc.sweep_cancel
2768 .lock()
2769 .unwrap_or_else(|e| e.into_inner())
2770 .is_cancelled(),
2771 "sweep_cancel token should be cancelled after HotSwap shutdown"
2772 );
2773
2774 tokio::time::sleep(Duration::from_millis(50)).await;
2775 }
2776
2777 #[tokio::test]
2785 async fn constant_key_skips_reserialization() {
2786 let mut svc = new_test_svc(config_size(3));
2787 let mut result = None;
2788 for body in ["first", "second", "third"] {
2789 result = Some(
2790 svc.ready()
2791 .await
2792 .unwrap()
2793 .call(make_exchange("orderId", "A", body))
2794 .await
2795 .unwrap(),
2796 );
2797 }
2798 let result = result.unwrap();
2799 assert_eq!(
2800 result.property(CAMEL_AGGREGATED_SIZE),
2801 Some(&serde_json::json!(3u64)),
2802 "all 3 fragments must aggregate out of a single bucket"
2803 );
2804 assert!(svc.buckets.lock().unwrap().is_empty());
2805 assert_eq!(
2806 svc.key_serializations
2807 .load(std::sync::atomic::Ordering::Relaxed),
2808 1,
2809 "constant scalar key must serialize exactly once"
2810 );
2811 }
2812
2813 #[tokio::test]
2817 async fn divergent_keys_keep_serde_semantics() {
2818 let mut svc = new_test_svc(config_size(10)); for key in ["k1", "k2", "k1"] {
2820 svc.ready()
2821 .await
2822 .unwrap()
2823 .call(make_exchange("orderId", key, "body"))
2824 .await
2825 .unwrap();
2826 }
2827 let guard = svc.buckets.lock().unwrap();
2828 assert_eq!(guard.len(), 2, "k1/k2/k1 → exactly two buckets");
2829 for key in ["k1", "k2"] {
2830 let expected = serde_json::to_string(&serde_json::json!(key)).unwrap();
2831 assert!(
2832 guard.contains_key(expected.as_str()),
2833 "bucket name must be byte-identical to serde_json::to_string: \
2834 expected {expected}, have {:?}",
2835 guard.keys().collect::<Vec<_>>()
2836 );
2837 }
2838 }
2839
2840 #[tokio::test]
2848 async fn object_keys_bypass_cache() {
2849 let mut svc = new_test_svc(config_size(10)); let obj_a = serde_json::json!({"a": 1, "b": 2});
2851 let obj_b = serde_json::json!({"b": 2, "a": 1});
2852 for obj in [&obj_a, &obj_b, &obj_a] {
2853 let mut msg = Message {
2854 headers: Default::default(),
2855 body: Body::Text("body".into()),
2856 };
2857 msg.headers.insert("orderId".to_string(), obj.clone());
2858 svc.ready()
2859 .await
2860 .unwrap()
2861 .call(Exchange::new(msg))
2862 .await
2863 .unwrap();
2864 }
2865 let expected_keys: std::collections::HashSet<String> = [&obj_a, &obj_b]
2866 .into_iter()
2867 .map(|o| serde_json::to_string(o).unwrap())
2868 .collect();
2869 let guard = svc.buckets.lock().unwrap();
2870 let actual_keys: std::collections::HashSet<String> = guard.keys().cloned().collect();
2871 assert_eq!(
2872 actual_keys, expected_keys,
2873 "object keys must bucket exactly per serde_json::to_string (cache bypassed)"
2874 );
2875 assert_eq!(
2876 svc.key_serializations
2877 .load(std::sync::atomic::Ordering::Relaxed),
2878 3,
2879 "each object fragment must serialize (memo never consulted for objects)"
2880 );
2881 }
2882
2883 #[tokio::test]
2889 async fn float_zero_sign_keys_stay_distinct() {
2890 let mut svc = new_test_svc(config_size(10)); for key in [0.0_f64, -0.0, 0.0] {
2892 let mut msg = Message {
2893 headers: Default::default(),
2894 body: Body::Text("body".into()),
2895 };
2896 msg.headers
2897 .insert("orderId".to_string(), serde_json::json!(key));
2898 svc.ready()
2899 .await
2900 .unwrap()
2901 .call(Exchange::new(msg))
2902 .await
2903 .unwrap();
2904 }
2905 let guard = svc.buckets.lock().unwrap();
2906 assert_eq!(guard.len(), 2, "±0.0 must stay in distinct buckets");
2907 for key in [0.0_f64, -0.0] {
2908 let expected = serde_json::to_string(&serde_json::json!(key)).unwrap();
2909 assert!(
2910 guard.contains_key(expected.as_str()),
2911 "bucket name must be byte-identical to serde_json::to_string: \
2912 expected {expected}, have {:?}",
2913 guard.keys().collect::<Vec<_>>()
2914 );
2915 }
2916 assert_eq!(
2917 svc.key_serializations
2918 .load(std::sync::atomic::Ordering::Relaxed),
2919 3,
2920 "number keys serialize per fragment (never memoized)"
2921 );
2922 }
2923}