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_language_api::Language;
25
26pub type SharedLanguageRegistry = Arc<std::sync::Mutex<HashMap<String, Arc<dyn Language>>>>;
27
28const QUEUE_DEPTH_SAMPLE_INTERVAL: Duration = Duration::from_millis(250);
32
33pub const CAMEL_AGGREGATOR_PENDING: &str = "CamelAggregatorPending";
34pub const CAMEL_AGGREGATED_SIZE: &str = "CamelAggregatedSize";
35pub const CAMEL_AGGREGATED_KEY: &str = "CamelAggregatedKey";
36pub const CAMEL_AGGREGATED_COMPLETION_REASON: &str = "CamelAggregatedCompletionReason";
37
38struct Bucket {
40 exchanges: Vec<Exchange>,
41 last_updated: Instant,
42}
43
44impl Bucket {
45 fn new() -> Self {
46 Self {
47 exchanges: Vec::new(),
48 last_updated: Instant::now(),
49 }
50 }
51
52 fn push(&mut self, exchange: Exchange) {
53 self.exchanges.push(exchange);
54 self.last_updated = Instant::now();
55 }
56
57 fn is_expired(&self, ttl: Duration) -> bool {
58 Instant::now().duration_since(self.last_updated) >= ttl
59 }
60}
61
62#[derive(Clone)]
63pub struct AggregatorService {
64 config: AggregatorConfig,
65 buckets: Arc<Mutex<HashMap<String, Bucket>>>,
66 timeout_tasks: Arc<Mutex<HashMap<String, CancellationToken>>>,
67 timeout_handles: Arc<Mutex<HashMap<String, JoinHandle<()>>>>,
68 late_tx: mpsc::Sender<Exchange>,
69 language_registry: SharedLanguageRegistry,
70 sweep_cancel: Arc<Mutex<CancellationToken>>,
77 sweep_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
85 queue_metrics: Option<(Arc<dyn MetricsCollector>, String)>,
90 clone_guard: Arc<()>,
94}
95
96impl std::fmt::Debug for AggregatorService {
97 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
98 f.debug_struct("AggregatorService").finish_non_exhaustive()
99 }
100}
101
102impl Drop for AggregatorService {
103 fn drop(&mut self) {
104 if Arc::strong_count(&self.clone_guard) > 1 {
114 return;
115 }
116 if let Some(handle) = self
117 .sweep_handle
118 .lock()
119 .unwrap_or_else(|e| e.into_inner())
120 .take()
121 {
122 handle.abort();
123 }
124 }
125}
126
127impl AggregatorService {
128 pub fn new(
139 config: AggregatorConfig,
140 late_tx: mpsc::Sender<Exchange>,
141 language_registry: SharedLanguageRegistry,
142 route_cancel: CancellationToken,
143 ) -> Self {
144 config.validate().expect(
146 "AggregatorService::new: config failed validation \
149 (need max_buckets, completionTimeout, or bucket_ttl)",
150 );
151
152 let has_timeout = has_timeout_condition(&config.completion);
156 if !has_timeout && config.bucket_ttl.is_none() {
157 tracing::warn!(
158 "Aggregator configured with Size/Predicate completion and no \
159 bucket_ttl: buckets accumulate until max_buckets is reached"
160 );
161 }
162
163 let buckets: Arc<Mutex<HashMap<String, Bucket>>> = Arc::new(Mutex::new(HashMap::new()));
166
167 Self {
176 config,
177 buckets,
178 timeout_tasks: Arc::new(Mutex::new(HashMap::new())),
179 timeout_handles: Arc::new(Mutex::new(HashMap::new())),
180 late_tx,
181 language_registry,
182 sweep_cancel: Arc::new(Mutex::new(route_cancel)),
183 sweep_handle: Arc::new(Mutex::new(None)),
184 queue_metrics: None,
185 clone_guard: Arc::new(()),
186 }
187 }
188
189 pub fn with_queue_metrics(
194 mut self,
195 metrics: Arc<dyn MetricsCollector>,
196 label: impl Into<String>,
197 ) -> Self {
198 self.queue_metrics = Some((metrics, label.into()));
199 self
200 }
201
202 pub fn config(&self) -> &AggregatorConfig {
203 &self.config
204 }
205
206 pub fn has_timeout(&self) -> bool {
207 has_timeout_condition(&self.config.completion)
208 }
209
210 pub fn force_complete_all(&self) {
211 let mut buckets_guard = self.buckets.lock().unwrap_or_else(|e| e.into_inner());
212 let keys: Vec<String> = buckets_guard.keys().cloned().collect();
213
214 for key in keys {
215 if let Some(bucket) = buckets_guard.remove(&key) {
216 if self.config.force_completion_on_stop {
217 cancel_timeout_task_with_handle(
218 &key,
219 &self.timeout_tasks,
220 &self.timeout_handles,
221 );
222 match aggregate(bucket.exchanges, &self.config.strategy) {
223 Ok(mut result) => {
224 result.set_property(
225 CAMEL_AGGREGATED_COMPLETION_REASON,
226 serde_json::json!(CompletionReason::Stop.as_str()),
227 );
228 if self.late_tx.try_send(result).is_err() {
229 tracing::warn!(
230 key = %key,
231 "aggregator force-complete emit dropped: late channel full"
232 );
233 }
234 }
235 Err(e) => {
236 tracing::warn!(
238 key = %key,
239 error = %e,
240 "aggregation failed in force_complete_all"
241 );
242 }
243 }
244 } else {
245 cancel_timeout_task_with_handle(
246 &key,
247 &self.timeout_tasks,
248 &self.timeout_handles,
249 );
250 }
251 }
252 }
253 }
254
255 pub(crate) async fn shutdown_inner(&self) {
258 {
260 let mut guard = self.timeout_tasks.lock().unwrap_or_else(|e| e.into_inner());
261 for token in guard.values() {
262 token.cancel();
263 }
264 guard.clear();
265 };
266
267 let handles: Vec<JoinHandle<()>> = {
269 let mut guard = self
270 .timeout_handles
271 .lock()
272 .unwrap_or_else(|e| e.into_inner());
273 guard.drain().map(|(_, handle)| handle).collect()
274 };
275
276 if handles.is_empty() {
277 return;
278 }
279
280 let _ = tokio::time::timeout(Duration::from_secs(5), async {
282 for handle in handles {
283 let _ = handle.await;
284 }
285 })
286 .await;
287 }
288}
289
290#[async_trait]
291impl StepLifecycle for AggregatorService {
292 fn name(&self) -> &'static str {
293 "aggregator"
294 }
295
296 async fn shutdown(&self, reason: StepShutdownReason) -> Result<(), CamelError> {
297 tracing::debug!(reason = ?reason, "Aggregator shutdown via StepLifecycle");
298
299 self.sweep_cancel
302 .lock()
303 .unwrap_or_else(|e| e.into_inner())
304 .cancel();
305
306 if let Some(handle) = self
309 .sweep_handle
310 .lock()
311 .unwrap_or_else(|e| e.into_inner())
312 .take()
313 {
314 handle.abort();
315 }
316
317 self.shutdown_inner().await;
318 Ok(())
319 }
320
321 async fn start(&self) -> Result<(), CamelError> {
325 *self.sweep_cancel.lock().unwrap_or_else(|e| e.into_inner()) = CancellationToken::new();
326
327 if let Some(handle) = self
328 .sweep_handle
329 .lock()
330 .unwrap_or_else(|e| e.into_inner())
331 .take()
332 {
333 handle.abort();
334 }
335
336 Ok(())
337 }
338}
339
340pub fn has_timeout_condition(mode: &CompletionMode) -> bool {
341 match mode {
342 CompletionMode::Single(CompletionCondition::Timeout(_)) => true,
343 CompletionMode::Any(conditions) => conditions
344 .iter()
345 .any(|c| matches!(c, CompletionCondition::Timeout(_))),
346 _ => false,
347 }
348}
349
350impl Service<Exchange> for AggregatorService {
351 type Response = Exchange;
352 type Error = CamelError;
353 type Future = Pin<Box<dyn Future<Output = Result<Exchange, CamelError>> + Send>>;
354
355 fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), CamelError>> {
356 let ttl = self.config.bucket_ttl;
365 if ttl.is_none() && self.queue_metrics.is_none() {
366 return Poll::Ready(Ok(()));
367 }
368 let mut g = self.sweep_handle.lock().unwrap_or_else(|e| e.into_inner());
369 if g.is_none() {
370 let interval = match ttl {
371 Some(ttl) => std::cmp::max(ttl / 2, Duration::from_millis(50)),
372 None => QUEUE_DEPTH_SAMPLE_INTERVAL,
373 };
374 let buckets = Arc::clone(&self.buckets);
375 let cancel = self
376 .sweep_cancel
377 .lock()
378 .unwrap_or_else(|e| e.into_inner())
379 .clone();
380 let queue_metrics = self.queue_metrics.clone();
381 *g = Some(tokio::spawn(async move {
382 loop {
383 tokio::select! {
384 _ = cancel.cancelled() => break,
385 _ = tokio::time::sleep(interval) => {
386 let mut guard =
387 buckets.lock().unwrap_or_else(|e| e.into_inner());
388 if let Some(ttl) = ttl {
389 guard.retain(|_, b| !b.is_expired(ttl));
390 }
391 if let Some((metrics, label)) = queue_metrics.as_ref() {
394 metrics.set_queue_depth(label, guard.len());
395 }
396 }
397 }
398 }
399 }));
400 }
401 Poll::Ready(Ok(()))
402 }
403
404 fn call(&mut self, exchange: Exchange) -> Self::Future {
405 let config = self.config.clone();
406 let buckets = Arc::clone(&self.buckets);
407 let timeout_tasks = Arc::clone(&self.timeout_tasks);
408 let timeout_handles = Arc::clone(&self.timeout_handles);
409 let late_tx = self.late_tx.clone();
410 let language_registry = Arc::clone(&self.language_registry);
411
412 Box::pin(async move {
413 let key_value =
414 extract_correlation_key(&exchange, &config.correlation, &language_registry).await?;
415
416 let key_str = serde_json::to_string(&key_value)
417 .map_err(|e| CamelError::ProcessorError(e.to_string()))?;
418
419 let predicate_satisfied =
423 evaluate_completion_predicate(&config.completion, &exchange, &language_registry)
424 .await?;
425
426 let completed_bucket = {
427 let mut guard = buckets.lock().unwrap_or_else(|e| e.into_inner());
428
429 if let Some(ttl) = config.bucket_ttl {
430 guard.retain(|_, bucket| !bucket.is_expired(ttl));
431 }
432
433 if let Some(max) = config.max_buckets
434 && !guard.contains_key(&key_str)
435 && guard.len() >= max
436 {
437 tracing::warn!(
438 max_buckets = max,
439 correlation_key = %key_str,
440 "Aggregator reached max buckets limit, rejecting new correlation key"
441 );
442 return Err(CamelError::ProcessorError(format!(
443 "Aggregator reached maximum {} buckets",
444 max
445 )));
446 }
447
448 let bucket = guard.entry(key_str.clone()).or_insert_with(Bucket::new);
449 bucket.push(exchange);
450
451 let (is_complete, reason) = check_sync_completion(
452 &config.completion,
453 &bucket.exchanges,
454 predicate_satisfied,
455 );
456
457 if is_complete {
458 let exchanges = guard.remove(&key_str).map(|b| b.exchanges);
459 (exchanges, reason)
460 } else {
461 (None, CompletionReason::Size) }
463 };
464
465 if completed_bucket.0.is_none() && has_timeout_condition(&config.completion) {
466 let timeout_dur = extract_timeout_duration(&config.completion);
467 if let Some(timeout) = timeout_dur {
468 let live_count = timeout_handles
474 .lock()
475 .unwrap_or_else(|e| e.into_inner())
476 .len();
477 if live_count >= config.max_timeout_tasks {
478 tracing::warn!(
479 live_timeout_tasks = live_count,
480 max_timeout_tasks = config.max_timeout_tasks,
481 correlation_key = %key_str,
482 "Aggregator timeout-task cap reached; bucket will rely on \
483 bucket_ttl eviction instead of a dedicated timeout task"
484 );
485 } else {
486 {
488 let tt_guard = timeout_tasks.lock().unwrap_or_else(|e| e.into_inner());
489 if let Some(existing) = tt_guard.get(&key_str) {
490 existing.cancel();
491 }
492 }
493 {
495 let mut hh = timeout_handles.lock().unwrap_or_else(|e| e.into_inner());
496 if let Some(old) = hh.remove(&key_str) {
497 old.abort();
498 }
499 }
500 let cancel = CancellationToken::new();
501 timeout_tasks
502 .lock()
503 .unwrap_or_else(|e| e.into_inner())
504 .insert(key_str.clone(), cancel.clone());
505 let handle = spawn_timeout_task(
506 key_str.clone(),
507 timeout,
508 cancel,
509 buckets.clone(),
510 timeout_tasks.clone(),
511 timeout_handles.clone(),
512 late_tx,
513 config.strategy.clone(),
514 config.discard_on_timeout,
515 );
516 timeout_handles
517 .lock()
518 .unwrap_or_else(|e| e.into_inner())
519 .insert(key_str.clone(), handle);
520 }
521 }
522 }
523
524 if let Some(exchanges) = completed_bucket.0 {
525 cancel_timeout_task_with_handle(&key_str, &timeout_tasks, &timeout_handles);
526 let reason = completed_bucket.1;
527 let size = exchanges.len();
528 let mut result = aggregate(exchanges, &config.strategy)?;
529 result.set_property(CAMEL_AGGREGATED_SIZE, serde_json::json!(size as u64));
530 result.set_property(CAMEL_AGGREGATED_KEY, key_value);
531 result.set_property(
532 CAMEL_AGGREGATED_COMPLETION_REASON,
533 serde_json::json!(reason.as_str()),
534 );
535 Ok(result)
536 } else {
537 let mut pending = Exchange::new(Message {
538 headers: Default::default(),
539 body: Body::Empty,
540 });
541 pending.set_property(CAMEL_AGGREGATOR_PENDING, serde_json::json!(true));
542 Ok(pending)
543 }
544 })
545 }
546}
547
548async fn extract_correlation_key(
549 exchange: &Exchange,
550 strategy: &CorrelationStrategy,
551 registry: &SharedLanguageRegistry,
552) -> Result<serde_json::Value, CamelError> {
553 match strategy {
554 CorrelationStrategy::HeaderName(h) => {
555 exchange.input.headers.get(h).cloned().ok_or_else(|| {
556 CamelError::ProcessorError(format!(
557 "Aggregator: missing correlation key header '{}'",
558 h
559 ))
560 })
561 }
562 CorrelationStrategy::Expression { expr, language } => {
563 let expression = {
564 let reg = registry.lock().unwrap_or_else(|e| e.into_inner());
565 let lang = reg.get(language).ok_or_else(|| {
566 CamelError::ProcessorError(format!(
567 "Aggregator: language '{}' not found in registry",
568 language
569 ))
570 })?;
571 lang.create_expression(expr)
572 .map_err(|e| CamelError::ProcessorError(e.to_string()))?
573 };
574 let value = expression
575 .evaluate(exchange)
576 .await
577 .map_err(|e| CamelError::ProcessorError(e.to_string()))?;
578 if value.is_null() {
579 return Err(CamelError::ProcessorError(format!(
580 "Aggregator: correlation expression '{}' evaluated to null",
581 expr
582 )));
583 }
584 Ok(value)
585 }
586 CorrelationStrategy::Fn(f) => f(exchange).map(serde_json::Value::String).ok_or_else(|| {
587 CamelError::ProcessorError("Aggregator: correlation function returned None".to_string())
588 }),
589 _ => Err(CamelError::ProcessorError(
591 "Aggregator: unsupported correlation strategy".to_string(),
592 )),
593 }
594}
595
596fn iter_predicate_exprs(mode: &CompletionMode) -> Vec<(&String, &String)> {
598 let mut out = Vec::new();
599 match mode {
600 CompletionMode::Single(CompletionCondition::PredicateExpr { expr, language }) => {
601 out.push((expr, language));
602 }
603 CompletionMode::Any(conds) => {
604 for c in conds {
605 if let CompletionCondition::PredicateExpr { expr, language } = c {
606 out.push((expr, language));
607 }
608 }
609 }
610 _ => {}
611 }
612 out
613}
614
615async fn evaluate_completion_predicate(
620 mode: &CompletionMode,
621 incoming: &Exchange,
622 registry: &SharedLanguageRegistry,
623) -> Result<bool, CamelError> {
624 let mut satisfied = false;
625 for (expr, language) in iter_predicate_exprs(mode) {
626 let expression = {
627 let reg = registry.lock().unwrap_or_else(|e| e.into_inner());
628 let lang = reg.get(language).ok_or_else(|| {
629 CamelError::ProcessorError(format!(
630 "Aggregator: language '{}' not found in registry",
631 language
632 ))
633 })?;
634 lang.create_expression(expr)
635 .map_err(|e| CamelError::ProcessorError(e.to_string()))?
636 }; let value = expression
638 .evaluate(incoming)
639 .await
640 .map_err(|e| CamelError::ProcessorError(e.to_string()))?;
641 if value.as_bool().unwrap_or(false) {
642 satisfied = true;
643 }
644 }
645 Ok(satisfied)
646}
647
648fn check_sync_completion(
649 mode: &CompletionMode,
650 exchanges: &[Exchange],
651 predicate_satisfied: bool,
652) -> (bool, CompletionReason) {
653 match mode {
654 CompletionMode::Single(cond) => check_single(cond, exchanges, predicate_satisfied),
655 CompletionMode::Any(conditions) => {
656 for cond in conditions {
657 if let CompletionCondition::Timeout(_) = cond {
658 continue;
659 }
660 let (done, reason) = check_single(cond, exchanges, predicate_satisfied);
661 if done {
662 return (true, reason);
663 }
664 }
665 (false, CompletionReason::Size)
666 }
667 _ => (false, CompletionReason::Size),
669 }
670}
671
672fn check_single(
673 cond: &CompletionCondition,
674 exchanges: &[Exchange],
675 predicate_satisfied: bool,
676) -> (bool, CompletionReason) {
677 match cond {
678 CompletionCondition::Size(n) => (exchanges.len() >= *n, CompletionReason::Size),
679 CompletionCondition::Predicate(pred) => (pred(exchanges), CompletionReason::Predicate),
680 CompletionCondition::PredicateExpr { .. } => {
681 (predicate_satisfied, CompletionReason::Predicate)
682 }
683 CompletionCondition::Timeout(_) => (false, CompletionReason::Timeout),
684 _ => (false, CompletionReason::Size),
686 }
687}
688
689fn extract_timeout_duration(mode: &CompletionMode) -> Option<Duration> {
690 match mode {
691 CompletionMode::Single(CompletionCondition::Timeout(d)) => Some(*d),
692 CompletionMode::Any(conditions) => conditions.iter().find_map(|c| {
693 if let CompletionCondition::Timeout(d) = c {
694 Some(*d)
695 } else {
696 None
697 }
698 }),
699 _ => None,
700 }
701}
702
703fn cancel_timeout_task(key: &str, timeout_tasks: &Arc<Mutex<HashMap<String, CancellationToken>>>) {
704 let mut guard = timeout_tasks.lock().unwrap_or_else(|e| e.into_inner());
705 if let Some(token) = guard.remove(key) {
706 token.cancel();
707 }
708}
709
710fn cancel_timeout_task_with_handle(
712 key: &str,
713 timeout_tasks: &Arc<Mutex<HashMap<String, CancellationToken>>>,
714 timeout_handles: &Arc<Mutex<HashMap<String, JoinHandle<()>>>>,
715) {
716 cancel_timeout_task(key, timeout_tasks);
717 let mut guard = timeout_handles.lock().unwrap_or_else(|e| e.into_inner());
718 guard.remove(key);
719}
720
721#[allow(clippy::too_many_arguments)]
722fn spawn_timeout_task(
723 key: String,
724 timeout: Duration,
725 cancel: CancellationToken,
726 buckets: Arc<Mutex<HashMap<String, Bucket>>>,
727 timeout_tasks: Arc<Mutex<HashMap<String, CancellationToken>>>,
728 timeout_handles: Arc<Mutex<HashMap<String, JoinHandle<()>>>>,
729 late_tx: mpsc::Sender<Exchange>,
730 strategy: AggregationStrategy,
731 discard: bool,
732) -> JoinHandle<()> {
733 let cancel_clone = cancel.clone();
734 tokio::spawn(async move {
735 tokio::select! {
736 _ = tokio::time::sleep(timeout) => {
737 let should_proceed = {
738 let mut tt_guard = timeout_tasks.lock().unwrap_or_else(|e| e.into_inner());
739 if cancel_clone.is_cancelled() {
740 false
741 } else {
742 tt_guard.remove(&key);
743 true
744 }
745 };
746 if !should_proceed {
747 return;
748 }
749 {
752 let mut hh = timeout_handles.lock().unwrap_or_else(|e| e.into_inner());
753 hh.remove(&key);
754 }
755 let bucket_exchanges = {
756 let mut guard = buckets.lock().unwrap_or_else(|e| e.into_inner());
757 guard.remove(&key).map(|b| b.exchanges)
758 };
759 if let Some(exchanges) = bucket_exchanges
760 && !discard
761 {
762 match aggregate(exchanges, &strategy) {
763 Ok(mut result) => {
764 result.set_property(
765 CAMEL_AGGREGATED_COMPLETION_REASON,
766 serde_json::json!(CompletionReason::Timeout.as_str()),
767 );
768 if late_tx.try_send(result).is_err() {
769 tracing::warn!(
770 key = %key,
771 "aggregator timeout emit dropped: late channel full"
772 );
773 }
774 }
775 Err(e) => {
776 tracing::warn!(
778 key = %key,
779 error = %e,
780 "aggregation failed in timeout task"
781 );
782 }
783 }
784 }
785 }
786 _ = cancel_clone.cancelled() => {}
787 }
788 })
789}
790
791fn aggregate(
792 exchanges: Vec<Exchange>,
793 strategy: &AggregationStrategy,
794) -> Result<Exchange, CamelError> {
795 match strategy {
796 AggregationStrategy::CollectAll => {
797 let bodies: Vec<serde_json::Value> = exchanges
798 .into_iter()
799 .map(|e| match e.input.body {
800 Body::Json(v) => v,
801 Body::Text(s) => serde_json::Value::String(s),
802 Body::Xml(s) => serde_json::Value::String(s),
803 Body::Bytes(b) => {
804 serde_json::Value::String(String::from_utf8_lossy(&b).into_owned())
805 }
806 Body::Stream(s) => serde_json::json!({
807 "_stream": {
808 "origin": s.metadata.origin,
809 "placeholder": true,
810 "hint": "Materialize exchange body with .into_bytes() before aggregation if content needed"
811 }
812 }),
813 _ => serde_json::Value::Null,
815 })
816 .collect();
817 Ok(Exchange::new(Message {
818 headers: Default::default(),
819 body: Body::Json(serde_json::Value::Array(bodies)),
820 }))
821 }
822 AggregationStrategy::Custom(f) => {
823 let mut iter = exchanges.into_iter();
824 let first = iter.next().ok_or_else(|| {
825 CamelError::ProcessorError("Aggregator: empty bucket".to_string())
826 })?;
827 Ok(iter.fold(first, |acc, next| f(acc, next)))
828 }
829 _ => Err(CamelError::ProcessorError(
831 "Aggregator: unsupported aggregation strategy".to_string(),
832 )),
833 }
834}
835
836#[cfg(test)]
837mod tests {
838 use super::*;
839 use std::collections::HashMap;
840
841 use camel_api::{
842 StepLifecycle, StepShutdownReason,
843 aggregator::{AggregationStrategy, AggregatorConfig},
844 body::Body,
845 exchange::Exchange,
846 message::Message,
847 };
848 use tokio::sync::mpsc;
849 use tokio_util::sync::CancellationToken;
850 use tower::ServiceExt;
851
852 fn make_exchange(header: &str, value: &str, body: &str) -> Exchange {
853 let mut msg = Message {
854 headers: Default::default(),
855 body: Body::Text(body.to_string()),
856 };
857 msg.headers
858 .insert(header.to_string(), serde_json::json!(value));
859 Exchange::new(msg)
860 }
861
862 fn config_size(n: usize) -> AggregatorConfig {
863 AggregatorConfig::correlate_by("orderId")
864 .complete_when_size(n)
865 .build()
866 .unwrap()
867 }
868
869 fn new_test_svc(config: AggregatorConfig) -> AggregatorService {
870 let (tx, _rx) = mpsc::channel(256);
871 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
872 let cancel = CancellationToken::new();
873 AggregatorService::new(config, tx, registry, cancel)
874 }
875
876 fn new_test_svc_with_registry(
880 config: AggregatorConfig,
881 registry: SharedLanguageRegistry,
882 ) -> AggregatorService {
883 let (tx, _rx) = mpsc::channel(256);
884 let cancel = CancellationToken::new();
885 AggregatorService::new(config, tx, registry, cancel)
886 }
887
888 #[tokio::test]
889 async fn test_pending_exchange_not_yet_complete() {
890 let mut svc = new_test_svc(config_size(3));
891 let ex = make_exchange("orderId", "A", "first");
892 let result = svc.ready().await.unwrap().call(ex).await.unwrap();
893 assert!(matches!(result.input.body, Body::Empty));
894 assert_eq!(
895 result.property(CAMEL_AGGREGATOR_PENDING),
896 Some(&serde_json::json!(true))
897 );
898 }
899
900 #[tokio::test]
901 async fn test_completes_on_size() {
902 let mut svc = new_test_svc(config_size(3));
903 for _ in 0..2 {
904 let ex = make_exchange("orderId", "A", "item");
905 let r = svc.ready().await.unwrap().call(ex).await.unwrap();
906 assert!(matches!(r.input.body, Body::Empty));
907 }
908 let ex = make_exchange("orderId", "A", "last");
909 let result = svc.ready().await.unwrap().call(ex).await.unwrap();
910 assert!(result.property(CAMEL_AGGREGATOR_PENDING).is_none());
911 assert_eq!(
912 result.property(CAMEL_AGGREGATED_SIZE),
913 Some(&serde_json::json!(3u64))
914 );
915 }
916
917 #[tokio::test]
918 async fn test_collect_all_produces_json_array() {
919 let mut svc = new_test_svc(config_size(2));
920 svc.ready()
921 .await
922 .unwrap()
923 .call(make_exchange("orderId", "A", "alpha"))
924 .await
925 .unwrap();
926 let result = svc
927 .ready()
928 .await
929 .unwrap()
930 .call(make_exchange("orderId", "A", "beta"))
931 .await
932 .unwrap();
933 let Body::Json(v) = &result.input.body else {
934 panic!("expected Body::Json")
935 };
936 let arr = v.as_array().unwrap();
937 assert_eq!(arr.len(), 2);
938 assert_eq!(arr[0], serde_json::json!("alpha"));
939 assert_eq!(arr[1], serde_json::json!("beta"));
940 }
941
942 #[tokio::test]
943 async fn test_two_keys_independent_buckets() {
944 let mut svc = new_test_svc(config_size(3));
946 svc.ready()
947 .await
948 .unwrap()
949 .call(make_exchange("orderId", "A", "a1"))
950 .await
951 .unwrap();
952 svc.ready()
953 .await
954 .unwrap()
955 .call(make_exchange("orderId", "B", "b1"))
956 .await
957 .unwrap();
958 svc.ready()
959 .await
960 .unwrap()
961 .call(make_exchange("orderId", "A", "a2"))
962 .await
963 .unwrap();
964 let ra = svc
966 .ready()
967 .await
968 .unwrap()
969 .call(make_exchange("orderId", "A", "a3"))
970 .await
971 .unwrap();
972 assert!(matches!(ra.input.body, Body::Json(_)));
974 let rb = svc
976 .ready()
977 .await
978 .unwrap()
979 .call(make_exchange("orderId", "B", "b_check"))
980 .await
981 .unwrap();
982 assert!(matches!(rb.input.body, Body::Empty));
983 }
984
985 #[tokio::test]
986 async fn test_bucket_resets_after_completion() {
987 let mut svc = new_test_svc(config_size(2));
988 svc.ready()
989 .await
990 .unwrap()
991 .call(make_exchange("orderId", "A", "x"))
992 .await
993 .unwrap();
994 svc.ready()
995 .await
996 .unwrap()
997 .call(make_exchange("orderId", "A", "x"))
998 .await
999 .unwrap(); let r = svc
1002 .ready()
1003 .await
1004 .unwrap()
1005 .call(make_exchange("orderId", "A", "new"))
1006 .await
1007 .unwrap();
1008 assert!(matches!(r.input.body, Body::Empty)); }
1010
1011 #[tokio::test]
1012 async fn test_completion_size_1_emits_immediately() {
1013 let mut svc = new_test_svc(config_size(1));
1014 let ex = make_exchange("orderId", "A", "solo");
1015 let result = svc.ready().await.unwrap().call(ex).await.unwrap();
1016 assert!(result.property(CAMEL_AGGREGATOR_PENDING).is_none());
1017 }
1018
1019 #[tokio::test]
1020 async fn test_custom_aggregation_strategy() {
1021 use camel_api::aggregator::AggregationFn;
1022 use std::sync::Arc;
1023
1024 let f: AggregationFn = Arc::new(|mut acc: Exchange, next: Exchange| {
1025 let combined = format!(
1026 "{}+{}",
1027 acc.input.body.as_text().unwrap_or(""),
1028 next.input.body.as_text().unwrap_or("")
1029 );
1030 acc.input.body = Body::Text(combined);
1031 acc
1032 });
1033 let config = AggregatorConfig::correlate_by("key")
1034 .complete_when_size(2)
1035 .strategy(AggregationStrategy::Custom(f))
1036 .build()
1037 .unwrap();
1038 let mut svc = new_test_svc(config);
1039 svc.ready()
1040 .await
1041 .unwrap()
1042 .call(make_exchange("key", "X", "hello"))
1043 .await
1044 .unwrap();
1045 let result = svc
1046 .ready()
1047 .await
1048 .unwrap()
1049 .call(make_exchange("key", "X", "world"))
1050 .await
1051 .unwrap();
1052 assert_eq!(result.input.body.as_text(), Some("hello+world"));
1053 }
1054
1055 #[tokio::test]
1056 async fn test_completion_predicate() {
1057 let config = AggregatorConfig::correlate_by("key")
1058 .complete_when(|bucket| {
1059 bucket
1060 .iter()
1061 .any(|e| e.input.body.as_text() == Some("DONE"))
1062 })
1063 .build()
1064 .unwrap();
1065 let mut svc = new_test_svc(config);
1066 svc.ready()
1067 .await
1068 .unwrap()
1069 .call(make_exchange("key", "K", "first"))
1070 .await
1071 .unwrap();
1072 svc.ready()
1073 .await
1074 .unwrap()
1075 .call(make_exchange("key", "K", "second"))
1076 .await
1077 .unwrap();
1078 let result = svc
1079 .ready()
1080 .await
1081 .unwrap()
1082 .call(make_exchange("key", "K", "DONE"))
1083 .await
1084 .unwrap();
1085 assert!(result.property(CAMEL_AGGREGATOR_PENDING).is_none());
1086 }
1087
1088 #[tokio::test]
1089 async fn test_missing_header_returns_error() {
1090 let mut svc = new_test_svc(config_size(2));
1091 let msg = Message {
1092 headers: Default::default(),
1093 body: Body::Text("no key".into()),
1094 };
1095 let ex = Exchange::new(msg);
1096 let result = svc.ready().await.unwrap().call(ex).await;
1097 assert!(result.is_err());
1098 assert!(matches!(
1099 result.unwrap_err(),
1100 camel_api::CamelError::ProcessorError(_)
1101 ));
1102 }
1103
1104 #[tokio::test]
1105 async fn test_cloned_service_shares_state() {
1106 let svc1 = new_test_svc(config_size(2));
1107 let mut svc2 = svc1.clone();
1108 svc1.clone()
1110 .ready()
1111 .await
1112 .unwrap()
1113 .call(make_exchange("orderId", "A", "from-svc1"))
1114 .await
1115 .unwrap();
1116 let result = svc2
1118 .ready()
1119 .await
1120 .unwrap()
1121 .call(make_exchange("orderId", "A", "from-svc2"))
1122 .await
1123 .unwrap();
1124 assert!(result.property(CAMEL_AGGREGATOR_PENDING).is_none());
1125 }
1126
1127 #[tokio::test]
1128 async fn test_camel_aggregated_key_property_set() {
1129 let mut svc = new_test_svc(config_size(1));
1130 let ex = make_exchange("orderId", "ORDER-42", "body");
1131 let result = svc.ready().await.unwrap().call(ex).await.unwrap();
1132 assert_eq!(
1133 result.property(CAMEL_AGGREGATED_KEY),
1134 Some(&serde_json::json!("ORDER-42"))
1135 );
1136 }
1137
1138 #[tokio::test]
1139 async fn test_aggregator_enforces_max_buckets() {
1140 let config = AggregatorConfig::correlate_by("orderId")
1141 .complete_when_size(2)
1142 .max_buckets(3)
1143 .build()
1144 .unwrap();
1145
1146 let mut svc = new_test_svc(config);
1147
1148 for i in 0..3 {
1150 let ex = make_exchange("orderId", &format!("key-{}", i), "body");
1151 let _ = svc.ready().await.unwrap().call(ex).await.unwrap();
1152 }
1153
1154 let ex = make_exchange("orderId", "key-4", "body");
1156 let result = svc.ready().await.unwrap().call(ex).await;
1157
1158 assert!(result.is_err(), "Should reject when max buckets reached");
1159 let err = result.unwrap_err().to_string();
1160 assert!(
1161 err.contains("maximum"),
1162 "Error message should contain 'maximum': {}",
1163 err
1164 );
1165 }
1166
1167 #[tokio::test]
1168 async fn test_max_buckets_allows_existing_key() {
1169 let config = AggregatorConfig::correlate_by("orderId")
1170 .complete_when_size(5) .max_buckets(2)
1172 .build()
1173 .unwrap();
1174
1175 let mut svc = new_test_svc(config);
1176
1177 let ex1 = make_exchange("orderId", "key-A", "body1");
1179 let _ = svc.ready().await.unwrap().call(ex1).await.unwrap();
1180 let ex2 = make_exchange("orderId", "key-B", "body2");
1181 let _ = svc.ready().await.unwrap().call(ex2).await.unwrap();
1182
1183 let ex3 = make_exchange("orderId", "key-A", "body3");
1185 let result = svc.ready().await.unwrap().call(ex3).await;
1186 assert!(
1187 result.is_ok(),
1188 "Should allow adding to existing bucket even at max limit"
1189 );
1190 }
1191
1192 #[tokio::test]
1193 async fn test_bucket_ttl_eviction() {
1194 let config = AggregatorConfig::correlate_by("orderId")
1195 .complete_when_size(10) .bucket_ttl(Duration::from_millis(50))
1197 .build()
1198 .unwrap();
1199
1200 let mut svc = new_test_svc(config);
1201
1202 let ex1 = make_exchange("orderId", "key-A", "body1");
1204 let _ = svc.ready().await.unwrap().call(ex1).await.unwrap();
1205
1206 tokio::time::sleep(Duration::from_millis(100)).await;
1208
1209 let ex2 = make_exchange("orderId", "key-B", "body2");
1211 let _ = svc.ready().await.unwrap().call(ex2).await.unwrap();
1212
1213 let ex3 = make_exchange("orderId", "key-A", "body3");
1216 let result = svc.ready().await.unwrap().call(ex3).await;
1217 assert!(result.is_ok(), "Should be able to recreate evicted bucket");
1218 }
1219
1220 #[tokio::test(start_paused = true)]
1221 async fn test_timeout_completes_bucket() {
1222 let config = AggregatorConfig::correlate_by("key")
1223 .complete_on_timeout(Duration::from_millis(100))
1224 .build()
1225 .unwrap();
1226 let mut svc = new_test_svc(config);
1227 let ex = make_exchange("key", "A", "data");
1228 let result = svc.ready().await.unwrap().call(ex).await.unwrap();
1229 assert!(result.property(CAMEL_AGGREGATOR_PENDING).is_some());
1230
1231 tokio::time::sleep(Duration::from_millis(200)).await;
1232
1233 assert_eq!(
1234 svc.buckets.lock().unwrap().len(),
1235 0,
1236 "bucket should be removed after timeout"
1237 );
1238 }
1239
1240 #[tokio::test(start_paused = true)]
1241 async fn test_timeout_resets_on_new_exchange() {
1242 let config = AggregatorConfig::correlate_by("key")
1243 .complete_on_timeout(Duration::from_millis(150))
1244 .build()
1245 .unwrap();
1246 let mut svc = new_test_svc(config);
1247
1248 let ex1 = make_exchange("key", "A", "first");
1249 let _ = svc.ready().await.unwrap().call(ex1).await.unwrap();
1250
1251 tokio::time::sleep(Duration::from_millis(100)).await;
1252
1253 let ex2 = make_exchange("key", "A", "second");
1254 let _ = svc.ready().await.unwrap().call(ex2).await.unwrap();
1255
1256 tokio::time::sleep(Duration::from_millis(100)).await;
1257
1258 assert_eq!(
1259 svc.buckets.lock().unwrap().len(),
1260 1,
1261 "bucket should still exist — timeout was reset"
1262 );
1263
1264 tokio::time::sleep(Duration::from_millis(100)).await;
1265
1266 assert_eq!(
1267 svc.buckets.lock().unwrap().len(),
1268 0,
1269 "bucket should be gone after timeout fires"
1270 );
1271 }
1272
1273 #[tokio::test]
1274 async fn test_composable_size_and_timeout() {
1275 let config = AggregatorConfig::correlate_by("key")
1276 .complete_on_size_or_timeout(2, Duration::from_millis(200))
1277 .build()
1278 .unwrap();
1279 let mut svc = new_test_svc(config);
1280
1281 let ex1 = make_exchange("key", "A", "first");
1282 let _ = svc.ready().await.unwrap().call(ex1).await.unwrap();
1283 assert!(svc.buckets.lock().unwrap().contains_key("\"A\""));
1284
1285 let ex2 = make_exchange("key", "A", "second");
1286 let result = svc.ready().await.unwrap().call(ex2).await.unwrap();
1287 assert!(result.property(CAMEL_AGGREGATOR_PENDING).is_none());
1288 assert_eq!(
1289 result.property(CAMEL_AGGREGATED_COMPLETION_REASON),
1290 Some(&serde_json::json!("size"))
1291 );
1292 }
1293
1294 #[tokio::test(start_paused = true)]
1295 async fn test_discard_on_timeout() {
1296 let config = AggregatorConfig::correlate_by("key")
1297 .complete_on_timeout(Duration::from_millis(50))
1298 .discard_on_timeout(true)
1299 .build()
1300 .unwrap();
1301 let (tx, mut rx) = mpsc::channel(256);
1302 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1303 let cancel = CancellationToken::new();
1304 let mut svc = AggregatorService::new(config, tx, registry, cancel);
1305
1306 let ex = make_exchange("key", "A", "data");
1307 let _ = svc.ready().await.unwrap().call(ex).await.unwrap();
1308
1309 tokio::time::sleep(Duration::from_millis(100)).await;
1310
1311 assert!(
1312 rx.try_recv().is_err(),
1313 "no emit expected with discard_on_timeout"
1314 );
1315 assert_eq!(svc.buckets.lock().unwrap().len(), 0);
1316 assert!(
1317 svc.timeout_tasks.lock().unwrap().is_empty(),
1318 "timeout task should be cleaned up"
1319 );
1320 }
1321
1322 #[tokio::test]
1323 async fn test_force_completion_on_stop() {
1324 let config = AggregatorConfig::correlate_by("key")
1325 .complete_when_size(10)
1326 .force_completion_on_stop(true)
1327 .build()
1328 .unwrap();
1329 let (tx, mut rx) = mpsc::channel(256);
1330 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1331 let cancel = CancellationToken::new();
1332 let svc = AggregatorService::new(config, tx, registry, cancel);
1333
1334 let mut call_svc = svc.clone();
1335 let ex = make_exchange("key", "A", "data");
1336 let _ = call_svc.ready().await.unwrap().call(ex).await.unwrap();
1337
1338 svc.force_complete_all();
1339
1340 let result = rx.try_recv().expect("should emit on force-complete");
1341 assert!(
1342 result.input.body.as_text().is_some() || matches!(result.input.body, Body::Json(_))
1343 );
1344 assert_eq!(
1345 result.property(CAMEL_AGGREGATED_COMPLETION_REASON),
1346 Some(&serde_json::json!("stop"))
1347 );
1348 }
1349
1350 #[tokio::test]
1351 async fn test_completion_reason_property_size() {
1352 let config = AggregatorConfig::correlate_by("key")
1353 .complete_when_size(1)
1354 .build()
1355 .unwrap();
1356 let mut svc = new_test_svc(config);
1357 let ex = make_exchange("key", "X", "body");
1358 let result = svc.ready().await.unwrap().call(ex).await.unwrap();
1359 assert_eq!(
1360 result.property(CAMEL_AGGREGATED_COMPLETION_REASON),
1361 Some(&serde_json::json!("size"))
1362 );
1363 }
1364
1365 #[tokio::test]
1366 async fn test_completion_reason_property_predicate() {
1367 let config = AggregatorConfig::correlate_by("key")
1368 .complete_when(|_| true)
1369 .build()
1370 .unwrap();
1371 let mut svc = new_test_svc(config);
1372 let ex = make_exchange("key", "X", "body");
1373 let result = svc.ready().await.unwrap().call(ex).await.unwrap();
1374 assert_eq!(
1375 result.property(CAMEL_AGGREGATED_COMPLETION_REASON),
1376 Some(&serde_json::json!("predicate"))
1377 );
1378 }
1379
1380 #[tokio::test(start_paused = true)]
1381 async fn test_size_completes_before_timeout() {
1382 let config = AggregatorConfig::correlate_by("key")
1383 .complete_on_size_or_timeout(2, Duration::from_millis(200))
1384 .build()
1385 .unwrap();
1386 let mut svc = new_test_svc(config);
1387
1388 let ex1 = make_exchange("key", "A", "first");
1389 let _ = svc.ready().await.unwrap().call(ex1).await.unwrap();
1390
1391 let ex2 = make_exchange("key", "A", "second");
1392 let result = svc.ready().await.unwrap().call(ex2).await.unwrap();
1393
1394 assert!(result.property(CAMEL_AGGREGATOR_PENDING).is_none());
1395 assert_eq!(
1396 result.property(CAMEL_AGGREGATED_COMPLETION_REASON),
1397 Some(&serde_json::json!("size"))
1398 );
1399 assert_eq!(svc.buckets.lock().unwrap().len(), 0);
1400
1401 tokio::time::sleep(Duration::from_millis(300)).await;
1402 assert_eq!(
1403 svc.buckets.lock().unwrap().len(),
1404 0,
1405 "no re-fire after timeout"
1406 );
1407 }
1408
1409 #[tokio::test(start_paused = true)]
1410 async fn test_concurrent_timeout_fire_and_new_exchange() {
1411 let config = AggregatorConfig::correlate_by("key")
1412 .complete_on_size_or_timeout(2, Duration::from_millis(100))
1413 .build()
1414 .unwrap();
1415 let (tx, mut rx) = mpsc::channel(256);
1416 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1417 let cancel = CancellationToken::new();
1418 let mut svc = AggregatorService::new(config, tx, registry, cancel);
1419
1420 let ex = make_exchange("key", "A", "data");
1421 let _ = svc.ready().await.unwrap().call(ex).await.unwrap();
1422
1423 tokio::time::sleep(Duration::from_millis(150)).await;
1425
1426 let ex2 = make_exchange("key", "A", "data2");
1428 let result = svc.ready().await.unwrap().call(ex2).await.unwrap();
1429 assert!(
1430 result.property(CAMEL_AGGREGATOR_PENDING).is_some(),
1431 "should be pending in new bucket"
1432 );
1433
1434 let mut late_count = 0;
1436 while rx.try_recv().is_ok() {
1437 late_count += 1;
1438 }
1439 assert_eq!(
1440 late_count, 1,
1441 "exactly 1 late emit from the timed-out bucket"
1442 );
1443 }
1444
1445 #[tokio::test(start_paused = true)]
1446 async fn test_late_channel_full_drops_with_warning() {
1447 let config = AggregatorConfig::correlate_by("key")
1448 .complete_on_timeout(Duration::from_millis(50))
1449 .build()
1450 .unwrap();
1451 let (tx, mut rx) = mpsc::channel(1);
1452 rx.close();
1453 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1454 let cancel = CancellationToken::new();
1455 let mut svc = AggregatorService::new(config, tx, registry, cancel);
1456
1457 let ex = make_exchange("key", "A", "data");
1458 let _ = svc.ready().await.unwrap().call(ex).await.unwrap();
1459
1460 tokio::time::sleep(Duration::from_millis(100)).await;
1461 assert_eq!(
1462 svc.buckets.lock().unwrap().len(),
1463 0,
1464 "bucket removed despite channel closed"
1465 );
1466 }
1467
1468 #[tokio::test]
1474 async fn test_da3_force_complete_all_drops_on_saturated_channel() {
1475 let config = AggregatorConfig::correlate_by("k")
1476 .complete_when_size(10)
1477 .force_completion_on_stop(true)
1478 .build()
1479 .unwrap();
1480 let (late_tx, mut late_rx) = mpsc::channel::<Exchange>(1);
1482 late_tx
1486 .try_send(make_exchange("k", "99", "dummy"))
1487 .expect("pre-fill succeeds");
1488
1489 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1490 let cancel = CancellationToken::new();
1491 let mut svc = AggregatorService::new(config, late_tx, registry, cancel);
1492
1493 for v in ["1", "2", "3"] {
1496 let ex = make_exchange("k", v, "body");
1497 let _ = svc.ready().await.unwrap().call(ex).await.unwrap();
1498 }
1499 assert_eq!(svc.buckets.lock().unwrap().len(), 3);
1500
1501 svc.force_complete_all();
1503
1504 let pre_fill = late_rx
1506 .try_recv()
1507 .expect("pre-fill should still be in channel");
1508 assert_eq!(
1509 pre_fill.input.headers.get("k"),
1510 Some(&serde_json::json!("99"))
1511 );
1512
1513 assert!(matches!(
1516 late_rx.try_recv(),
1517 Err(mpsc::error::TryRecvError::Empty)
1518 ));
1519
1520 assert!(svc.buckets.lock().unwrap().is_empty());
1522 }
1523
1524 #[tokio::test]
1525 async fn test_aggregate_stream_bodies_creates_valid_json() {
1526 use bytes::Bytes;
1527 use camel_api::{Body, StreamBody, StreamMetadata};
1528 use futures::stream;
1529 use tokio::sync::Mutex;
1530
1531 let chunks = vec![Ok(Bytes::from("test"))];
1532 let stream_body = StreamBody {
1533 stream: Arc::new(Mutex::new(Some(Box::pin(stream::iter(chunks))))),
1534 metadata: StreamMetadata {
1535 origin: Some("file:///test.txt".to_string()),
1536 ..Default::default()
1537 },
1538 };
1539
1540 let ex1 = Exchange::new(Message {
1541 headers: Default::default(),
1542 body: Body::Stream(stream_body),
1543 });
1544
1545 let exchanges = vec![ex1];
1546 let result = aggregate(exchanges, &AggregationStrategy::CollectAll);
1547
1548 let exchange = result.expect("Expected Ok result");
1549 assert!(
1550 matches!(exchange.input.body, Body::Json(_)),
1551 "Expected Json body"
1552 );
1553
1554 if let Body::Json(value) = exchange.input.body {
1555 let json_str = serde_json::to_string(&value).unwrap();
1556 let parsed: serde_json::Value = serde_json::from_str(&json_str).unwrap();
1557
1558 assert!(parsed.is_array(), "Result should be an array");
1559 let arr = parsed.as_array().unwrap();
1560 assert!(arr[0].is_object(), "First element should be an object");
1561 assert!(
1562 arr[0]["_stream"].is_object(),
1563 "Should contain _stream object"
1564 );
1565 assert_eq!(arr[0]["_stream"]["origin"], "file:///test.txt");
1566 assert_eq!(
1567 arr[0]["_stream"]["placeholder"], true,
1568 "placeholder flag should be true"
1569 );
1570 }
1571 }
1572
1573 #[tokio::test]
1574 async fn test_aggregate_stream_bodies_with_none_origin() {
1575 use bytes::Bytes;
1576 use camel_api::{Body, StreamBody, StreamMetadata};
1577 use futures::stream;
1578 use tokio::sync::Mutex;
1579
1580 let chunks = vec![Ok(Bytes::from("test"))];
1581 let stream_body = StreamBody {
1582 stream: Arc::new(Mutex::new(Some(Box::pin(stream::iter(chunks))))),
1583 metadata: StreamMetadata {
1584 origin: None,
1585 ..Default::default()
1586 },
1587 };
1588
1589 let ex1 = Exchange::new(Message {
1590 headers: Default::default(),
1591 body: Body::Stream(stream_body),
1592 });
1593
1594 let exchanges = vec![ex1];
1595 let result = aggregate(exchanges, &AggregationStrategy::CollectAll);
1596
1597 let exchange = result.expect("Expected Ok result");
1598 assert!(
1599 matches!(exchange.input.body, Body::Json(_)),
1600 "Expected Json body"
1601 );
1602
1603 if let Body::Json(value) = exchange.input.body {
1604 let json_str = serde_json::to_string(&value).unwrap();
1605 let parsed: serde_json::Value = serde_json::from_str(&json_str).unwrap();
1606
1607 assert!(parsed.is_array(), "Result should be an array");
1608 let arr = parsed.as_array().unwrap();
1609 assert!(arr[0].is_object(), "First element should be an object");
1610 assert!(
1611 arr[0]["_stream"].is_object(),
1612 "Should contain _stream object"
1613 );
1614 assert_eq!(
1615 arr[0]["_stream"]["origin"],
1616 serde_json::Value::Null,
1617 "origin should be null when None"
1618 );
1619 assert_eq!(
1620 arr[0]["_stream"]["placeholder"], true,
1621 "placeholder flag should be true"
1622 );
1623 }
1624 }
1625
1626 #[tokio::test]
1627 async fn timeout_completion_clears_handle_from_map() {
1628 let config = AggregatorConfig::correlate_by("key")
1633 .complete_on_timeout(Duration::from_millis(50))
1634 .build()
1635 .unwrap();
1636 let (tx, _rx) = mpsc::channel(256);
1637 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1638 let cancel = CancellationToken::new();
1639 let svc = AggregatorService::new(config, tx, registry, cancel);
1640
1641 let mut call_svc = svc.clone();
1643 let ex = make_exchange("key", "A", "data");
1644 let _ = call_svc.ready().await.unwrap().call(ex).await.unwrap();
1645 assert!(
1646 !svc.timeout_handles.lock().unwrap().is_empty(),
1647 "handle should exist while timeout pending"
1648 );
1649
1650 tokio::time::sleep(Duration::from_millis(200)).await;
1652
1653 assert!(
1654 svc.timeout_handles.lock().unwrap().is_empty(),
1655 "handle should be cleared from map after natural timeout completion (was leak)"
1656 );
1657 }
1658
1659 #[tokio::test]
1660 async fn aggregator_shutdown_via_trait_dispatch() {
1661 let config = AggregatorConfig::correlate_by("key")
1664 .complete_when_size(10)
1665 .build()
1666 .unwrap();
1667 let (tx, _rx) = mpsc::channel(256);
1668 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1669 let cancel = CancellationToken::new();
1670 let svc = AggregatorService::new(config, tx, registry, cancel);
1671
1672 let step: Arc<dyn StepLifecycle> = Arc::new(svc);
1673 step.shutdown(StepShutdownReason::RouteStop)
1674 .await
1675 .expect("first shutdown should succeed");
1676 step.shutdown(StepShutdownReason::RouteStop)
1677 .await
1678 .expect("second shutdown (idempotent) should succeed");
1679 }
1680
1681 #[tokio::test(start_paused = true)]
1682 async fn test_shutdown_awaits_timeout_handles() {
1683 let config = AggregatorConfig::correlate_by("key")
1684 .complete_on_timeout(Duration::from_millis(100))
1685 .build()
1686 .unwrap();
1687 let (tx, _rx) = mpsc::channel(256);
1688 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1689 let cancel = CancellationToken::new();
1690 let svc = AggregatorService::new(config, tx, registry, cancel);
1691
1692 let mut call_svc = svc.clone();
1694 let ex = make_exchange("key", "A", "data");
1695 let _ = call_svc.ready().await.unwrap().call(ex).await.unwrap();
1696
1697 assert!(
1699 !svc.timeout_handles.lock().unwrap().is_empty(),
1700 "should have a timeout handle"
1701 );
1702
1703 svc.shutdown_inner().await;
1706
1707 assert!(
1708 svc.timeout_handles.lock().unwrap().is_empty(),
1709 "all handles should be cleaned up after shutdown"
1710 );
1711 }
1712
1713 #[tokio::test]
1720 async fn test_unique_key_flood_stays_bounded_by_default() {
1721 let config = AggregatorConfig::correlate_by("orderId")
1723 .complete_when_size(1_000_000) .build()
1725 .unwrap();
1726 let mut svc = new_test_svc(config);
1727
1728 for i in 0..10_000usize {
1731 let ex = make_exchange("orderId", &format!("key-{i}"), "body");
1732 let result = svc.ready().await.unwrap().call(ex).await;
1733 assert!(result.is_ok(), "key {i} should be accepted under the cap");
1734 }
1735 let ex = make_exchange("orderId", "key-10001", "body");
1736 let result = svc.ready().await.unwrap().call(ex).await;
1737 assert!(
1738 result.is_err(),
1739 "10_001st unique key must be rejected by the max_buckets cap"
1740 );
1741 let err = result.unwrap_err().to_string();
1742 assert!(
1743 err.contains("maximum") || err.contains("max"),
1744 "error should mention cap: {err}"
1745 );
1746 }
1747
1748 #[tokio::test]
1754 async fn test_background_sweep_spawns_on_first_poll_not_construction() {
1755 let config = AggregatorConfig::correlate_by("key")
1756 .complete_when_size(10_000)
1757 .bucket_ttl(Duration::from_millis(50))
1758 .build()
1759 .unwrap();
1760 let cancel = CancellationToken::new();
1761 let (tx, _rx) = mpsc::channel(8);
1762 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1763 let mut svc = AggregatorService::new(config, tx, registry, cancel.clone());
1764
1765 assert!(
1767 svc.sweep_handle
1768 .lock()
1769 .unwrap_or_else(|e| e.into_inner())
1770 .is_none(),
1771 "sweep must NOT be spawned at construction (runtime-free new)"
1772 );
1773
1774 let _ = svc.ready().await.unwrap();
1777 let sweep_present = svc
1778 .sweep_handle
1779 .lock()
1780 .unwrap_or_else(|e| e.into_inner())
1781 .is_some();
1782 assert!(
1783 sweep_present,
1784 "sweep handle should be Some after first poll when bucket_ttl is set"
1785 );
1786
1787 cancel.cancel();
1789 tokio::time::sleep(Duration::from_millis(50)).await;
1791 }
1792
1793 struct QueueDepthRecorder(Mutex<Vec<(String, usize)>>);
1796 impl MetricsCollector for QueueDepthRecorder {
1797 fn record_exchange_duration(&self, _: &str, _: std::time::Duration) {}
1798 fn increment_errors(&self, _: &str, _: &str) {}
1799 fn increment_exchanges(&self, _: &str) {}
1800 fn set_queue_depth(&self, queue: &str, depth: usize) {
1801 self.0
1802 .lock()
1803 .unwrap_or_else(|e| e.into_inner())
1804 .push((queue.to_string(), depth));
1805 }
1806 fn record_circuit_breaker_change(&self, _: &str, _: &str, _: &str) {}
1807 }
1808
1809 async fn await_depth_sample(
1812 recorder: &QueueDepthRecorder,
1813 label: &str,
1814 pred: impl Fn(usize) -> bool,
1815 ) {
1816 let deadline = std::time::Instant::now() + Duration::from_secs(2);
1817 loop {
1818 let matched = recorder
1819 .0
1820 .lock()
1821 .unwrap_or_else(|e| e.into_inner())
1822 .iter()
1823 .any(|(q, d)| q == label && pred(*d));
1824 if matched {
1825 return;
1826 }
1827 assert!(
1828 std::time::Instant::now() < deadline,
1829 "no queue-depth sample for '{label}' matched within 2s"
1830 );
1831 tokio::time::sleep(Duration::from_millis(20)).await;
1832 }
1833 }
1834
1835 #[tokio::test]
1840 async fn test_sweep_reports_queue_depth() {
1841 let config = AggregatorConfig::correlate_by("orderId")
1842 .complete_when_size(3)
1843 .bucket_ttl(Duration::from_millis(100))
1847 .build()
1848 .unwrap();
1849 let cancel = CancellationToken::new();
1850 let (tx, _rx) = mpsc::channel(8);
1851 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1852 let recorder = Arc::new(QueueDepthRecorder(Mutex::new(Vec::new())));
1853 let mut svc = AggregatorService::new(config, tx, registry, cancel).with_queue_metrics(
1854 Arc::clone(&recorder) as Arc<dyn MetricsCollector>,
1855 "aggregator:t",
1856 );
1857
1858 let _ = svc.ready().await.unwrap();
1860
1861 let ex = make_exchange("orderId", "g1", "partial");
1863 let _ = svc.ready().await.unwrap().call(ex).await;
1864
1865 await_depth_sample(&recorder, "aggregator:t", |d| d > 0).await;
1866
1867 let ex = make_exchange("orderId", "g1", "b");
1869 let _ = svc.ready().await.unwrap().call(ex).await;
1870 let ex = make_exchange("orderId", "g1", "c");
1871 let _ = svc.ready().await.unwrap().call(ex).await;
1872
1873 await_depth_sample(&recorder, "aggregator:t", |d| d == 0).await;
1874
1875 svc.shutdown(StepShutdownReason::RouteStop).await.unwrap();
1876 }
1877
1878 #[tokio::test]
1883 async fn test_no_ttl_aggregator_reports_queue_depth() {
1884 use camel_api::aggregator::CorrelationStrategy;
1885
1886 let config = AggregatorConfig {
1887 header_name: "orderId".into(),
1888 completion: CompletionMode::Single(CompletionCondition::Size(3)),
1889 correlation: CorrelationStrategy::HeaderName("orderId".into()),
1890 strategy: AggregationStrategy::CollectAll,
1891 max_buckets: Some(100),
1892 bucket_ttl: None,
1893 force_completion_on_stop: false,
1894 discard_on_timeout: false,
1895 max_timeout_tasks: 64,
1896 };
1897 let cancel = CancellationToken::new();
1898 let (tx, _rx) = mpsc::channel(8);
1899 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1900 let recorder = Arc::new(QueueDepthRecorder(Mutex::new(Vec::new())));
1901 let mut svc = AggregatorService::new(config, tx, registry, cancel).with_queue_metrics(
1902 Arc::clone(&recorder) as Arc<dyn MetricsCollector>,
1903 "aggregator:nottl",
1904 );
1905
1906 let _ = svc.ready().await.unwrap();
1907 assert!(
1908 svc.sweep_handle
1909 .lock()
1910 .unwrap_or_else(|e| e.into_inner())
1911 .is_some(),
1912 "metrics-only config (bucket_ttl = None) must still spawn the sweep"
1913 );
1914
1915 let ex = make_exchange("orderId", "g1", "partial");
1917 let _ = svc.ready().await.unwrap().call(ex).await;
1918 await_depth_sample(&recorder, "aggregator:nottl", |d| d > 0).await;
1919
1920 let ex = make_exchange("orderId", "g1", "b");
1922 let _ = svc.ready().await.unwrap().call(ex).await;
1923 let ex = make_exchange("orderId", "g1", "c");
1924 let _ = svc.ready().await.unwrap().call(ex).await;
1925 await_depth_sample(&recorder, "aggregator:nottl", |d| d == 0).await;
1926
1927 svc.shutdown(StepShutdownReason::RouteStop).await.unwrap();
1928 }
1929
1930 #[tokio::test]
1935 async fn test_transient_clone_drop_keeps_sweep_sampling() {
1936 let config = AggregatorConfig::correlate_by("orderId")
1937 .complete_when_size(3)
1938 .bucket_ttl(Duration::from_millis(100))
1939 .build()
1940 .unwrap();
1941 let cancel = CancellationToken::new();
1942 let (tx, _rx) = mpsc::channel(8);
1943 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1944 let recorder = Arc::new(QueueDepthRecorder(Mutex::new(Vec::new())));
1945 let mut svc = AggregatorService::new(config, tx, registry, cancel).with_queue_metrics(
1946 Arc::clone(&recorder) as Arc<dyn MetricsCollector>,
1947 "aggregator:clone",
1948 );
1949
1950 let _ = svc.ready().await.unwrap();
1951
1952 drop(svc.clone());
1954 assert!(
1955 svc.sweep_handle
1956 .lock()
1957 .unwrap_or_else(|e| e.into_inner())
1958 .is_some(),
1959 "transient clone drop must not abort the shared sweep"
1960 );
1961
1962 let ex = make_exchange("orderId", "g1", "partial");
1964 let _ = svc.ready().await.unwrap().call(ex).await;
1965 await_depth_sample(&recorder, "aggregator:clone", |d| d > 0).await;
1966
1967 svc.shutdown(StepShutdownReason::RouteStop).await.unwrap();
1968 }
1969
1970 #[tokio::test]
1973 async fn test_aggregator_timeout_task_cap_no_panic_under_flood() {
1974 use camel_api::aggregator::CorrelationStrategy;
1977
1978 let config = AggregatorConfig {
1979 header_name: "k".into(),
1980 completion: CompletionMode::Any(vec![
1981 CompletionCondition::Size(999),
1982 CompletionCondition::Timeout(Duration::from_secs(30)),
1983 ]),
1984 correlation: CorrelationStrategy::HeaderName("k".into()),
1985 strategy: AggregationStrategy::CollectAll,
1986 max_buckets: Some(50),
1987 bucket_ttl: Some(Duration::from_secs(30)),
1988 force_completion_on_stop: false,
1989 discard_on_timeout: false,
1990 max_timeout_tasks: 2,
1991 };
1992 let (late_tx, mut late_rx) = mpsc::channel(64);
1993 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1994 let cancel = CancellationToken::new();
1995 let svc = AggregatorService::new(config, late_tx, registry, cancel);
1996
1997 for i in 0..20u64 {
1999 let mut ex = Exchange::new(Message {
2000 headers: HashMap::from([("k".to_string(), serde_json::json!(i))]),
2001 body: Body::Text(i.to_string()),
2002 });
2003 ex.input
2004 .headers
2005 .insert("k".to_string(), serde_json::json!(i));
2006 let outcome = tokio::time::timeout(Duration::from_secs(2), async {
2007 let mut s = svc.clone();
2008 use tower::ServiceExt;
2009 s.ready().await.unwrap().call(ex).await
2010 })
2011 .await;
2012 assert!(outcome.is_ok(), "call {} hung/panicked under task cap", i);
2013 let res = outcome.unwrap().unwrap();
2015 assert_eq!(
2016 res.properties
2017 .get(CAMEL_AGGREGATOR_PENDING)
2018 .and_then(|v| v.as_bool()),
2019 Some(true),
2020 "exchange {} should be pending",
2021 i
2022 );
2023 }
2024 let _ = late_rx.try_recv();
2026 }
2027
2028 #[tokio::test]
2029 async fn evaluate_completion_predicate_or_combines_any() {
2030 use camel_api::aggregator::{CompletionCondition, CompletionMode};
2031 use camel_language_api::Language;
2032 use std::collections::HashMap;
2033
2034 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
2036 registry.lock().unwrap().insert(
2037 "simple".to_string(),
2038 Arc::new(camel_language_simple::SimpleLanguage::new()) as Arc<dyn Language>,
2039 );
2040
2041 let incoming = make_exchange("k", "X", "hello");
2042
2043 let mode = CompletionMode::Any(vec![
2045 CompletionCondition::PredicateExpr {
2046 expr: "${body} == 'NOPE'".to_string(),
2047 language: "simple".to_string(),
2048 },
2049 CompletionCondition::PredicateExpr {
2050 expr: "${body} == 'hello'".to_string(),
2051 language: "simple".to_string(),
2052 },
2053 ]);
2054
2055 let satisfied = evaluate_completion_predicate(&mode, &incoming, ®istry)
2056 .await
2057 .expect("eval must succeed");
2058 assert!(satisfied, "second predicate matches → OR true");
2059 }
2060
2061 #[tokio::test]
2062 async fn evaluate_completion_predicate_no_predicate_skips_registry() {
2063 use camel_api::aggregator::{CompletionCondition, CompletionMode};
2064 use std::collections::HashMap;
2065
2066 let mode = CompletionMode::Single(CompletionCondition::Size(3));
2070 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
2071 let incoming = make_exchange("k", "X", "hello");
2072 let result = evaluate_completion_predicate(&mode, &incoming, ®istry).await;
2073 assert!(result.is_ok(), "fast path must not error: {:?}", result);
2074 assert!(!result.unwrap(), "no predicate → not satisfied");
2075 }
2076
2077 #[tokio::test]
2078 async fn evaluate_completion_predicate_unregistered_language_errors() {
2079 use camel_api::aggregator::{CompletionCondition, CompletionMode};
2080
2081 let mode = CompletionMode::Single(CompletionCondition::PredicateExpr {
2082 expr: "${body} == 'DONE'".to_string(),
2083 language: "nonexistent-lang".to_string(),
2084 });
2085 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
2086 let incoming = make_exchange("k", "X", "hello");
2087 let result = evaluate_completion_predicate(&mode, &incoming, ®istry).await;
2088 assert!(result.is_err(), "unregistered language must error");
2089 }
2090
2091 #[tokio::test]
2092 async fn evaluate_completion_predicate_all_miss_returns_false() {
2093 use camel_api::aggregator::{CompletionCondition, CompletionMode};
2094 use camel_language_api::Language;
2095 use std::collections::HashMap;
2096
2097 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
2098 registry.lock().unwrap().insert(
2099 "simple".to_string(),
2100 Arc::new(camel_language_simple::SimpleLanguage::new()) as Arc<dyn Language>,
2101 );
2102
2103 let mode = CompletionMode::Single(CompletionCondition::PredicateExpr {
2104 expr: "${body} == 'NOPE'".to_string(),
2105 language: "simple".to_string(),
2106 });
2107 let incoming = make_exchange("k", "X", "hello");
2108 let result = evaluate_completion_predicate(&mode, &incoming, ®istry).await;
2109 assert!(result.is_ok(), "eval must succeed: {:?}", result);
2110 assert!(!result.unwrap(), "predicate does not match");
2111 }
2112
2113 #[tokio::test]
2114 async fn completion_predicate_expr_completes_on_match() {
2115 use camel_api::aggregator::{CompletionCondition, CompletionMode};
2116 use camel_language_api::Language;
2117 use std::collections::HashMap;
2118
2119 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
2121 registry.lock().unwrap().insert(
2122 "simple".to_string(),
2123 Arc::new(camel_language_simple::SimpleLanguage::new()) as Arc<dyn Language>,
2124 );
2125
2126 let mut config = AggregatorConfig::correlate_by("key")
2131 .complete_when_size(999) .build()
2133 .expect("config builds");
2134 config.completion = CompletionMode::Single(CompletionCondition::PredicateExpr {
2135 expr: "${body} == 'DONE'".to_string(),
2136 language: "simple".to_string(),
2137 });
2138
2139 let mut svc = new_test_svc_with_registry(config, registry);
2140
2141 let r = svc
2143 .ready()
2144 .await
2145 .unwrap()
2146 .call(make_exchange("key", "K", "first"))
2147 .await
2148 .unwrap();
2149 assert!(r.property(CAMEL_AGGREGATOR_PENDING).is_some());
2150
2151 let r = svc
2153 .ready()
2154 .await
2155 .unwrap()
2156 .call(make_exchange("key", "K", "DONE"))
2157 .await
2158 .unwrap();
2159 assert!(r.property(CAMEL_AGGREGATOR_PENDING).is_none());
2160 assert_eq!(
2161 r.property(CAMEL_AGGREGATED_COMPLETION_REASON),
2162 Some(&serde_json::json!("predicate"))
2163 );
2164 }
2165
2166 #[tokio::test]
2178 async fn test_da1_strategy_receives_two_exchanges_first_message_preserved() {
2179 use camel_api::aggregator::{AggregationFn, AggregationStrategy};
2180 use std::sync::Arc;
2181
2182 let recorded: Arc<std::sync::Mutex<Vec<(String, String)>>> =
2183 Arc::new(std::sync::Mutex::new(Vec::new()));
2184 let recorded_for_closure = Arc::clone(&recorded);
2185
2186 let f: AggregationFn = Arc::new(move |old: Exchange, new: Exchange| {
2187 let old_body = old.input.body.as_text().unwrap_or("").to_string();
2188 let new_body = new.input.body.as_text().unwrap_or("").to_string();
2189 recorded_for_closure
2190 .lock()
2191 .expect("recorded mutex poisoned")
2192 .push((old_body, new_body));
2193 new
2194 });
2195
2196 let config = AggregatorConfig::correlate_by("k")
2197 .complete_when_size(2)
2198 .strategy(AggregationStrategy::Custom(f))
2199 .build()
2200 .unwrap();
2201 let mut svc = new_test_svc(config);
2202
2203 let first = svc
2205 .ready()
2206 .await
2207 .unwrap()
2208 .call(make_exchange("k", "1", "A"))
2209 .await
2210 .unwrap();
2211 assert!(
2212 first.property(CAMEL_AGGREGATOR_PENDING).is_some(),
2213 "first message must leave the bucket pending (size < 2)"
2214 );
2215 assert!(
2216 recorded.lock().expect("recorded mutex poisoned").is_empty(),
2217 "strategy must NOT be invoked on the first message of a bucket"
2218 );
2219
2220 let _completed = svc
2223 .ready()
2224 .await
2225 .unwrap()
2226 .call(make_exchange("k", "1", "B"))
2227 .await
2228 .unwrap();
2229
2230 let recorded = recorded.lock().expect("recorded mutex poisoned");
2231 assert_eq!(
2232 recorded.len(),
2233 1,
2234 "strategy must be invoked exactly once across the two-message bucket, got {recorded:?}"
2235 );
2236 assert_eq!(
2237 recorded[0],
2238 ("A".to_string(), "B".to_string()),
2239 "strategy must observe the first message as `old` and the second as `new`, \
2240 with both bodies preserved unchanged"
2241 );
2242 }
2243
2244 #[tokio::test]
2247 async fn sweep_shutdown_cancels_task() {
2248 let config = AggregatorConfig::correlate_by("key")
2249 .complete_when_size(10)
2250 .bucket_ttl(Duration::from_millis(100))
2251 .build()
2252 .unwrap();
2253 let (tx, _rx) = mpsc::channel(256);
2254 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
2255 let cancel = CancellationToken::new();
2256 let mut svc = AggregatorService::new(config, tx, registry, cancel);
2257
2258 let _ = svc.ready().await.unwrap();
2259
2260 assert!(
2261 svc.sweep_handle
2262 .lock()
2263 .unwrap_or_else(|e| e.into_inner())
2264 .is_some(),
2265 "sweep handle should be Some after poll_ready"
2266 );
2267
2268 svc.shutdown(StepShutdownReason::RouteStop).await.unwrap();
2269
2270 assert!(
2271 svc.sweep_handle
2272 .lock()
2273 .unwrap_or_else(|e| e.into_inner())
2274 .is_none(),
2275 "sweep handle should be None after shutdown (taken + aborted)"
2276 );
2277 assert!(
2278 svc.sweep_cancel
2279 .lock()
2280 .unwrap_or_else(|e| e.into_inner())
2281 .is_cancelled(),
2282 "sweep_cancel token should be cancelled after shutdown"
2283 );
2284
2285 tokio::time::sleep(Duration::from_millis(50)).await;
2286 }
2287
2288 #[tokio::test]
2289 async fn sweep_start_respawns_after_shutdown() {
2290 let config = AggregatorConfig::correlate_by("key")
2291 .complete_when_size(10)
2292 .bucket_ttl(Duration::from_millis(100))
2293 .build()
2294 .unwrap();
2295 let (tx, _rx) = mpsc::channel(256);
2296 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
2297 let cancel = CancellationToken::new();
2298 let mut svc = AggregatorService::new(config, tx, registry, cancel);
2299
2300 let _ = svc.ready().await.unwrap();
2301 svc.shutdown(StepShutdownReason::RouteStop).await.unwrap();
2302
2303 svc.start().await.unwrap();
2304 let _ = svc.ready().await.unwrap();
2305
2306 assert!(
2307 svc.sweep_handle
2308 .lock()
2309 .unwrap_or_else(|e| e.into_inner())
2310 .is_some(),
2311 "sweep handle should be Some after start + poll_ready"
2312 );
2313 assert!(
2314 !svc.sweep_cancel
2315 .lock()
2316 .unwrap_or_else(|e| e.into_inner())
2317 .is_cancelled(),
2318 "sweep_cancel should be a fresh uncancelled token after start"
2319 );
2320 }
2321
2322 #[tokio::test]
2323 async fn sweep_shutdown_hotswap_cancels_task() {
2324 let config = AggregatorConfig::correlate_by("key")
2325 .complete_when_size(10)
2326 .bucket_ttl(Duration::from_millis(100))
2327 .build()
2328 .unwrap();
2329 let (tx, _rx) = mpsc::channel(256);
2330 let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
2331 let cancel = CancellationToken::new();
2332 let mut svc = AggregatorService::new(config, tx, registry, cancel);
2333
2334 let _ = svc.ready().await.unwrap();
2335
2336 assert!(
2337 svc.sweep_handle
2338 .lock()
2339 .unwrap_or_else(|e| e.into_inner())
2340 .is_some(),
2341 "sweep handle should be Some after poll_ready"
2342 );
2343
2344 svc.shutdown(StepShutdownReason::HotSwap).await.unwrap();
2345
2346 assert!(
2347 svc.sweep_handle
2348 .lock()
2349 .unwrap_or_else(|e| e.into_inner())
2350 .is_none(),
2351 "sweep handle should be None after HotSwap shutdown"
2352 );
2353 assert!(
2354 svc.sweep_cancel
2355 .lock()
2356 .unwrap_or_else(|e| e.into_inner())
2357 .is_cancelled(),
2358 "sweep_cancel token should be cancelled after HotSwap shutdown"
2359 );
2360
2361 tokio::time::sleep(Duration::from_millis(50)).await;
2362 }
2363}