Skip to main content

camel_processor/
aggregator.rs

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, 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
28pub const CAMEL_AGGREGATOR_PENDING: &str = "CamelAggregatorPending";
29pub const CAMEL_AGGREGATED_SIZE: &str = "CamelAggregatedSize";
30pub const CAMEL_AGGREGATED_KEY: &str = "CamelAggregatedKey";
31pub const CAMEL_AGGREGATED_COMPLETION_REASON: &str = "CamelAggregatedCompletionReason";
32
33/// Internal bucket structure with timestamp tracking for TTL eviction.
34struct Bucket {
35    exchanges: Vec<Exchange>,
36    last_updated: Instant,
37}
38
39impl Bucket {
40    fn new() -> Self {
41        Self {
42            exchanges: Vec::new(),
43            last_updated: Instant::now(),
44        }
45    }
46
47    fn push(&mut self, exchange: Exchange) {
48        self.exchanges.push(exchange);
49        self.last_updated = Instant::now();
50    }
51
52    fn is_expired(&self, ttl: Duration) -> bool {
53        Instant::now().duration_since(self.last_updated) >= ttl
54    }
55}
56
57#[derive(Clone)]
58pub struct AggregatorService {
59    config: AggregatorConfig,
60    buckets: Arc<Mutex<HashMap<String, Bucket>>>,
61    timeout_tasks: Arc<Mutex<HashMap<String, CancellationToken>>>,
62    timeout_handles: Arc<Mutex<HashMap<String, JoinHandle<()>>>>,
63    late_tx: mpsc::Sender<Exchange>,
64    language_registry: SharedLanguageRegistry,
65    route_cancel: CancellationToken,
66    /// Handle to the background TTL-sweep task. `None` when `config.bucket_ttl`
67    /// is `None` (no TTL → no sweep). When the TTL is set, this is populated
68    /// at construction by the auto-spawn in `new` — the caller does not need
69    /// to start the sweep explicitly. The task is bound to `route_cancel` —
70    /// cancelling the route token aborts the task. This is the fix for
71    /// R3-C1's "no background sweep" half: a flood of unique keys within the
72    /// inline-retain window can no longer grow unbounded between calls.
73    sweep_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
74}
75
76impl std::fmt::Debug for AggregatorService {
77    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
78        f.debug_struct("AggregatorService").finish_non_exhaustive()
79    }
80}
81
82impl Drop for AggregatorService {
83    fn drop(&mut self) {
84        // Defense-in-depth for the R3-C1 sweep: abort the background task when
85        // the service is dropped so it cannot leak even if the route owner
86        // forgets to cancel `route_cancel` (the primary shutdown path).
87        // `abort()` on an already-finished task is a no-op, so this is safe
88        // alongside the `select!` cancel path. This also gives `sweep_handle`
89        // a production reader so it is not dead-code.
90        if let Some(handle) = self
91            .sweep_handle
92            .lock()
93            .unwrap_or_else(|e| e.into_inner())
94            .take()
95        {
96            handle.abort();
97        }
98    }
99}
100
101impl AggregatorService {
102    /// Lifecycle invariant: `route_cancel` is owned by the route. Construction
103    /// is runtime-free (no `tokio::spawn` here) so callers that build an
104    /// `AggregatorService` outside a tokio runtime — e.g. route-spec tests —
105    /// do not panic. When `config.bucket_ttl` is `Some`, the TTL-sweep task is
106    /// spawned LAZILY on the first `poll_ready` (a runtime is guaranteed
107    /// there) and bound to `route_cancel` via `select!`. The route owner MUST
108    /// cancel it on shutdown; `Drop` also aborts it as defense-in-depth.
109    pub fn new(
110        config: AggregatorConfig,
111        late_tx: mpsc::Sender<Exchange>,
112        language_registry: SharedLanguageRegistry,
113        route_cancel: CancellationToken,
114    ) -> Self {
115        // Build the shared buckets map up front so the sweep task can
116        // share it via Arc::clone.
117        let buckets: Arc<Mutex<HashMap<String, Bucket>>> = Arc::new(Mutex::new(HashMap::new()));
118
119        // R3-C1 Batch 1: the TTL-sweep is NOT spawned here — construction
120        // must stay runtime-free so callers that build an AggregatorService
121        // outside a tokio runtime (route-spec tests) do not panic on
122        // `tokio::spawn`. The sweep is spawned lazily on the first
123        // `poll_ready`. The inline `guard.retain(...)` per `call` evicts
124        // expired buckets regardless, so the cap + TTL invariants hold
125        // whether or not the sweep has started yet.
126
127        Self {
128            config,
129            buckets,
130            timeout_tasks: Arc::new(Mutex::new(HashMap::new())),
131            timeout_handles: Arc::new(Mutex::new(HashMap::new())),
132            late_tx,
133            language_registry,
134            route_cancel,
135            sweep_handle: Arc::new(Mutex::new(None)),
136        }
137    }
138
139    pub fn config(&self) -> &AggregatorConfig {
140        &self.config
141    }
142
143    pub fn has_timeout(&self) -> bool {
144        has_timeout_condition(&self.config.completion)
145    }
146
147    pub fn force_complete_all(&self) {
148        let mut buckets_guard = self.buckets.lock().unwrap_or_else(|e| e.into_inner());
149        let keys: Vec<String> = buckets_guard.keys().cloned().collect();
150
151        for key in keys {
152            if let Some(bucket) = buckets_guard.remove(&key) {
153                if self.config.force_completion_on_stop {
154                    cancel_timeout_task_with_handle(
155                        &key,
156                        &self.timeout_tasks,
157                        &self.timeout_handles,
158                    );
159                    match aggregate(bucket.exchanges, &self.config.strategy) {
160                        Ok(mut result) => {
161                            result.set_property(
162                                CAMEL_AGGREGATED_COMPLETION_REASON,
163                                serde_json::json!(CompletionReason::Stop.as_str()),
164                            );
165                            if self.late_tx.try_send(result).is_err() {
166                                tracing::warn!(
167                                    key = %key,
168                                    "aggregator force-complete emit dropped: late channel full"
169                                );
170                            }
171                        }
172                        Err(e) => {
173                            // log-policy: handler-owned
174                            tracing::warn!(
175                                key = %key,
176                                error = %e,
177                                "aggregation failed in force_complete_all"
178                            );
179                        }
180                    }
181                } else {
182                    cancel_timeout_task_with_handle(
183                        &key,
184                        &self.timeout_tasks,
185                        &self.timeout_handles,
186                    );
187                }
188            }
189        }
190    }
191
192    /// Graceful shutdown: cancel all outstanding timeout tasks and await their
193    /// JoinHandles (with a 5s deadline) so that no tasks are leaked.
194    pub(crate) async fn shutdown_inner(&self) {
195        // Cancel all timeout cancellation tokens.
196        {
197            let mut guard = self.timeout_tasks.lock().unwrap_or_else(|e| e.into_inner());
198            for token in guard.values() {
199                token.cancel();
200            }
201            guard.clear();
202        };
203
204        // Remove and collect all JoinHandles.
205        let handles: Vec<JoinHandle<()>> = {
206            let mut guard = self
207                .timeout_handles
208                .lock()
209                .unwrap_or_else(|e| e.into_inner());
210            guard.drain().map(|(_, handle)| handle).collect()
211        };
212
213        if handles.is_empty() {
214            return;
215        }
216
217        // Await all handles with a deadline.
218        let _ = tokio::time::timeout(Duration::from_secs(5), async {
219            for handle in handles {
220                let _ = handle.await;
221            }
222        })
223        .await;
224    }
225}
226
227#[async_trait]
228impl StepLifecycle for AggregatorService {
229    fn name(&self) -> &'static str {
230        "aggregator"
231    }
232
233    async fn shutdown(&self, reason: StepShutdownReason) -> Result<(), CamelError> {
234        tracing::debug!(reason = ?reason, "Aggregator shutdown via StepLifecycle");
235        self.shutdown_inner().await;
236        Ok(())
237    }
238}
239
240pub fn has_timeout_condition(mode: &CompletionMode) -> bool {
241    match mode {
242        CompletionMode::Single(CompletionCondition::Timeout(_)) => true,
243        CompletionMode::Any(conditions) => conditions
244            .iter()
245            .any(|c| matches!(c, CompletionCondition::Timeout(_))),
246        _ => false,
247    }
248}
249
250impl Service<Exchange> for AggregatorService {
251    type Response = Exchange;
252    type Error = CamelError;
253    type Future = Pin<Box<dyn Future<Output = Result<Exchange, CamelError>> + Send>>;
254
255    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), CamelError>> {
256        // R3-C1: lazily spawn the TTL-sweep on the first readiness poll. A
257        // tokio runtime is guaranteed here (the runtime driving the service),
258        // unlike at construction. Single-spawn via the lock + is_none check.
259        if let Some(ttl) = self.config.bucket_ttl {
260            let mut g = self.sweep_handle.lock().unwrap_or_else(|e| e.into_inner());
261            if g.is_none() {
262                let interval = std::cmp::max(ttl / 2, Duration::from_millis(50));
263                let buckets = Arc::clone(&self.buckets);
264                let cancel = self.route_cancel.clone();
265                *g = Some(tokio::spawn(async move {
266                    loop {
267                        tokio::select! {
268                            _ = cancel.cancelled() => break,
269                            _ = tokio::time::sleep(interval) => {
270                                let mut guard =
271                                    buckets.lock().unwrap_or_else(|e| e.into_inner());
272                                guard.retain(|_, b| !b.is_expired(ttl));
273                            }
274                        }
275                    }
276                }));
277            }
278        }
279        Poll::Ready(Ok(()))
280    }
281
282    fn call(&mut self, exchange: Exchange) -> Self::Future {
283        let config = self.config.clone();
284        let buckets = Arc::clone(&self.buckets);
285        let timeout_tasks = Arc::clone(&self.timeout_tasks);
286        let timeout_handles = Arc::clone(&self.timeout_handles);
287        let late_tx = self.late_tx.clone();
288        let language_registry = Arc::clone(&self.language_registry);
289        let route_cancel = self.route_cancel.clone();
290
291        Box::pin(async move {
292            let key_value =
293                extract_correlation_key(&exchange, &config.correlation, &language_registry).await?;
294
295            let key_str = serde_json::to_string(&key_value)
296                .map_err(|e| CamelError::ProcessorError(e.to_string()))?;
297
298            let completed_bucket = {
299                let mut guard = buckets.lock().unwrap_or_else(|e| e.into_inner());
300
301                if let Some(ttl) = config.bucket_ttl {
302                    guard.retain(|_, bucket| !bucket.is_expired(ttl));
303                }
304
305                if let Some(max) = config.max_buckets
306                    && !guard.contains_key(&key_str)
307                    && guard.len() >= max
308                {
309                    tracing::warn!(
310                        max_buckets = max,
311                        correlation_key = %key_str,
312                        "Aggregator reached max buckets limit, rejecting new correlation key"
313                    );
314                    return Err(CamelError::ProcessorError(format!(
315                        "Aggregator reached maximum {} buckets",
316                        max
317                    )));
318                }
319
320                let bucket = guard.entry(key_str.clone()).or_insert_with(Bucket::new);
321                bucket.push(exchange);
322
323                let (is_complete, reason) =
324                    check_sync_completion(&config.completion, &bucket.exchanges);
325
326                if is_complete {
327                    let exchanges = guard.remove(&key_str).map(|b| b.exchanges);
328                    (exchanges, reason)
329                } else {
330                    (None, CompletionReason::Size) // placeholder; reason unused when None
331                }
332            };
333
334            if completed_bucket.0.is_none() && has_timeout_condition(&config.completion) {
335                let cancel = {
336                    let mut tt_guard = timeout_tasks.lock().unwrap_or_else(|e| e.into_inner());
337                    // Cancel and remove old handle for this key (if any).
338                    if let Some(existing) = tt_guard.get(&key_str) {
339                        existing.cancel();
340                    }
341                    let token = CancellationToken::new();
342                    tt_guard.insert(key_str.clone(), token.clone());
343                    token
344                };
345
346                let timeout_dur = extract_timeout_duration(&config.completion);
347                if let Some(timeout) = timeout_dur {
348                    // Remove old handle if present.
349                    {
350                        let mut hh = timeout_handles.lock().unwrap_or_else(|e| e.into_inner());
351                        if let Some(old) = hh.remove(&key_str) {
352                            old.abort();
353                        }
354                    }
355                    let handle = spawn_timeout_task(
356                        key_str.clone(),
357                        timeout,
358                        cancel,
359                        buckets.clone(),
360                        timeout_tasks.clone(),
361                        timeout_handles.clone(),
362                        late_tx,
363                        config.strategy.clone(),
364                        config.discard_on_timeout,
365                        route_cancel,
366                    );
367                    timeout_handles
368                        .lock()
369                        .unwrap_or_else(|e| e.into_inner())
370                        .insert(key_str.clone(), handle);
371                }
372            }
373
374            if let Some(exchanges) = completed_bucket.0 {
375                cancel_timeout_task_with_handle(&key_str, &timeout_tasks, &timeout_handles);
376                let reason = completed_bucket.1;
377                let size = exchanges.len();
378                let mut result = aggregate(exchanges, &config.strategy)?;
379                result.set_property(CAMEL_AGGREGATED_SIZE, serde_json::json!(size as u64));
380                result.set_property(CAMEL_AGGREGATED_KEY, key_value);
381                result.set_property(
382                    CAMEL_AGGREGATED_COMPLETION_REASON,
383                    serde_json::json!(reason.as_str()),
384                );
385                Ok(result)
386            } else {
387                let mut pending = Exchange::new(Message {
388                    headers: Default::default(),
389                    body: Body::Empty,
390                });
391                pending.set_property(CAMEL_AGGREGATOR_PENDING, serde_json::json!(true));
392                Ok(pending)
393            }
394        })
395    }
396}
397
398async fn extract_correlation_key(
399    exchange: &Exchange,
400    strategy: &CorrelationStrategy,
401    registry: &SharedLanguageRegistry,
402) -> Result<serde_json::Value, CamelError> {
403    match strategy {
404        CorrelationStrategy::HeaderName(h) => {
405            exchange.input.headers.get(h).cloned().ok_or_else(|| {
406                CamelError::ProcessorError(format!(
407                    "Aggregator: missing correlation key header '{}'",
408                    h
409                ))
410            })
411        }
412        CorrelationStrategy::Expression { expr, language } => {
413            let expression = {
414                let reg = registry.lock().unwrap_or_else(|e| e.into_inner());
415                let lang = reg.get(language).ok_or_else(|| {
416                    CamelError::ProcessorError(format!(
417                        "Aggregator: language '{}' not found in registry",
418                        language
419                    ))
420                })?;
421                lang.create_expression(expr)
422                    .map_err(|e| CamelError::ProcessorError(e.to_string()))?
423            };
424            let value = expression
425                .evaluate(exchange)
426                .await
427                .map_err(|e| CamelError::ProcessorError(e.to_string()))?;
428            if value.is_null() {
429                return Err(CamelError::ProcessorError(format!(
430                    "Aggregator: correlation expression '{}' evaluated to null",
431                    expr
432                )));
433            }
434            Ok(value)
435        }
436        CorrelationStrategy::Fn(f) => f(exchange).map(serde_json::Value::String).ok_or_else(|| {
437            CamelError::ProcessorError("Aggregator: correlation function returned None".to_string())
438        }),
439    }
440}
441
442fn check_sync_completion(
443    mode: &CompletionMode,
444    exchanges: &[Exchange],
445) -> (bool, CompletionReason) {
446    match mode {
447        CompletionMode::Single(cond) => check_single(cond, exchanges),
448        CompletionMode::Any(conditions) => {
449            for cond in conditions {
450                if let CompletionCondition::Timeout(_) = cond {
451                    continue;
452                }
453                let (done, reason) = check_single(cond, exchanges);
454                if done {
455                    return (true, reason);
456                }
457            }
458            (false, CompletionReason::Size)
459        }
460    }
461}
462
463fn check_single(cond: &CompletionCondition, exchanges: &[Exchange]) -> (bool, CompletionReason) {
464    match cond {
465        CompletionCondition::Size(n) => (exchanges.len() >= *n, CompletionReason::Size),
466        CompletionCondition::Predicate(pred) => (pred(exchanges), CompletionReason::Predicate),
467        CompletionCondition::Timeout(_) => (false, CompletionReason::Timeout),
468    }
469}
470
471fn extract_timeout_duration(mode: &CompletionMode) -> Option<Duration> {
472    match mode {
473        CompletionMode::Single(CompletionCondition::Timeout(d)) => Some(*d),
474        CompletionMode::Any(conditions) => conditions.iter().find_map(|c| {
475            if let CompletionCondition::Timeout(d) = c {
476                Some(*d)
477            } else {
478                None
479            }
480        }),
481        _ => None,
482    }
483}
484
485fn cancel_timeout_task(key: &str, timeout_tasks: &Arc<Mutex<HashMap<String, CancellationToken>>>) {
486    let mut guard = timeout_tasks.lock().unwrap_or_else(|e| e.into_inner());
487    if let Some(token) = guard.remove(key) {
488        token.cancel();
489    }
490}
491
492/// Also removes the stored JoinHandle for a cancelled/completed timeout task.
493fn cancel_timeout_task_with_handle(
494    key: &str,
495    timeout_tasks: &Arc<Mutex<HashMap<String, CancellationToken>>>,
496    timeout_handles: &Arc<Mutex<HashMap<String, JoinHandle<()>>>>,
497) {
498    cancel_timeout_task(key, timeout_tasks);
499    let mut guard = timeout_handles.lock().unwrap_or_else(|e| e.into_inner());
500    guard.remove(key);
501}
502
503#[allow(clippy::too_many_arguments)]
504fn spawn_timeout_task(
505    key: String,
506    timeout: Duration,
507    cancel: CancellationToken,
508    buckets: Arc<Mutex<HashMap<String, Bucket>>>,
509    timeout_tasks: Arc<Mutex<HashMap<String, CancellationToken>>>,
510    timeout_handles: Arc<Mutex<HashMap<String, JoinHandle<()>>>>,
511    late_tx: mpsc::Sender<Exchange>,
512    strategy: AggregationStrategy,
513    discard: bool,
514    _route_cancel: CancellationToken,
515) -> JoinHandle<()> {
516    let cancel_clone = cancel.clone();
517    tokio::spawn(async move {
518        tokio::select! {
519            _ = tokio::time::sleep(timeout) => {
520                let should_proceed = {
521                    let mut tt_guard = timeout_tasks.lock().unwrap_or_else(|e| e.into_inner());
522                    if cancel_clone.is_cancelled() {
523                        false
524                    } else {
525                        tt_guard.remove(&key);
526                        true
527                    }
528                };
529                if !should_proceed {
530                    return;
531                }
532                // Clean up our own JoinHandle from the map on natural completion.
533                // Without this, the handle entry leaks until route shutdown.
534                {
535                    let mut hh = timeout_handles.lock().unwrap_or_else(|e| e.into_inner());
536                    hh.remove(&key);
537                }
538                let bucket_exchanges = {
539                    let mut guard = buckets.lock().unwrap_or_else(|e| e.into_inner());
540                    guard.remove(&key).map(|b| b.exchanges)
541                };
542                if let Some(exchanges) = bucket_exchanges
543                    && !discard
544                {
545                    match aggregate(exchanges, &strategy) {
546                        Ok(mut result) => {
547                            result.set_property(
548                                CAMEL_AGGREGATED_COMPLETION_REASON,
549                                serde_json::json!(CompletionReason::Timeout.as_str()),
550                            );
551                            if late_tx.try_send(result).is_err() {
552                                tracing::warn!(
553                                    key = %key,
554                                    "aggregator timeout emit dropped: late channel full"
555                                );
556                            }
557                        }
558                        Err(e) => {
559                            // log-policy: handler-owned
560                            tracing::warn!(
561                                key = %key,
562                                error = %e,
563                                "aggregation failed in timeout task"
564                            );
565                        }
566                    }
567                }
568            }
569            _ = cancel_clone.cancelled() => {}
570        }
571    })
572}
573
574fn aggregate(
575    exchanges: Vec<Exchange>,
576    strategy: &AggregationStrategy,
577) -> Result<Exchange, CamelError> {
578    match strategy {
579        AggregationStrategy::CollectAll => {
580            let bodies: Vec<serde_json::Value> = exchanges
581                .into_iter()
582                .map(|e| match e.input.body {
583                    Body::Json(v) => v,
584                    Body::Text(s) => serde_json::Value::String(s),
585                    Body::Xml(s) => serde_json::Value::String(s),
586                    Body::Bytes(b) => {
587                        serde_json::Value::String(String::from_utf8_lossy(&b).into_owned())
588                    }
589                    Body::Empty => serde_json::Value::Null,
590                    Body::Stream(s) => serde_json::json!({
591                        "_stream": {
592                            "origin": s.metadata.origin,
593                            "placeholder": true,
594                            "hint": "Materialize exchange body with .into_bytes() before aggregation if content needed"
595                        }
596                    }),
597                })
598                .collect();
599            Ok(Exchange::new(Message {
600                headers: Default::default(),
601                body: Body::Json(serde_json::Value::Array(bodies)),
602            }))
603        }
604        AggregationStrategy::Custom(f) => {
605            let mut iter = exchanges.into_iter();
606            let first = iter.next().ok_or_else(|| {
607                CamelError::ProcessorError("Aggregator: empty bucket".to_string())
608            })?;
609            Ok(iter.fold(first, |acc, next| f(acc, next)))
610        }
611    }
612}
613
614#[cfg(test)]
615mod tests {
616    use super::*;
617    use std::collections::HashMap;
618
619    use camel_api::{
620        StepLifecycle, StepShutdownReason,
621        aggregator::{AggregationStrategy, AggregatorConfig},
622        body::Body,
623        exchange::Exchange,
624        message::Message,
625    };
626    use tokio::sync::mpsc;
627    use tokio_util::sync::CancellationToken;
628    use tower::ServiceExt;
629
630    fn make_exchange(header: &str, value: &str, body: &str) -> Exchange {
631        let mut msg = Message {
632            headers: Default::default(),
633            body: Body::Text(body.to_string()),
634        };
635        msg.headers
636            .insert(header.to_string(), serde_json::json!(value));
637        Exchange::new(msg)
638    }
639
640    fn config_size(n: usize) -> AggregatorConfig {
641        AggregatorConfig::correlate_by("orderId")
642            .complete_when_size(n)
643            .build()
644            .unwrap()
645    }
646
647    fn new_test_svc(config: AggregatorConfig) -> AggregatorService {
648        let (tx, _rx) = mpsc::channel(256);
649        let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
650        let cancel = CancellationToken::new();
651        AggregatorService::new(config, tx, registry, cancel)
652    }
653
654    #[tokio::test]
655    async fn test_pending_exchange_not_yet_complete() {
656        let mut svc = new_test_svc(config_size(3));
657        let ex = make_exchange("orderId", "A", "first");
658        let result = svc.ready().await.unwrap().call(ex).await.unwrap();
659        assert!(matches!(result.input.body, Body::Empty));
660        assert_eq!(
661            result.property(CAMEL_AGGREGATOR_PENDING),
662            Some(&serde_json::json!(true))
663        );
664    }
665
666    #[tokio::test]
667    async fn test_completes_on_size() {
668        let mut svc = new_test_svc(config_size(3));
669        for _ in 0..2 {
670            let ex = make_exchange("orderId", "A", "item");
671            let r = svc.ready().await.unwrap().call(ex).await.unwrap();
672            assert!(matches!(r.input.body, Body::Empty));
673        }
674        let ex = make_exchange("orderId", "A", "last");
675        let result = svc.ready().await.unwrap().call(ex).await.unwrap();
676        assert!(result.property(CAMEL_AGGREGATOR_PENDING).is_none());
677        assert_eq!(
678            result.property(CAMEL_AGGREGATED_SIZE),
679            Some(&serde_json::json!(3u64))
680        );
681    }
682
683    #[tokio::test]
684    async fn test_collect_all_produces_json_array() {
685        let mut svc = new_test_svc(config_size(2));
686        svc.ready()
687            .await
688            .unwrap()
689            .call(make_exchange("orderId", "A", "alpha"))
690            .await
691            .unwrap();
692        let result = svc
693            .ready()
694            .await
695            .unwrap()
696            .call(make_exchange("orderId", "A", "beta"))
697            .await
698            .unwrap();
699        let Body::Json(v) = &result.input.body else {
700            panic!("expected Body::Json")
701        };
702        let arr = v.as_array().unwrap();
703        assert_eq!(arr.len(), 2);
704        assert_eq!(arr[0], serde_json::json!("alpha"));
705        assert_eq!(arr[1], serde_json::json!("beta"));
706    }
707
708    #[tokio::test]
709    async fn test_two_keys_independent_buckets() {
710        // completionSize=3 so we can test that A and B accumulate independently.
711        let mut svc = new_test_svc(config_size(3));
712        svc.ready()
713            .await
714            .unwrap()
715            .call(make_exchange("orderId", "A", "a1"))
716            .await
717            .unwrap();
718        svc.ready()
719            .await
720            .unwrap()
721            .call(make_exchange("orderId", "B", "b1"))
722            .await
723            .unwrap();
724        svc.ready()
725            .await
726            .unwrap()
727            .call(make_exchange("orderId", "A", "a2"))
728            .await
729            .unwrap();
730        // A has 2 items, B has 1 item — neither complete yet
731        let ra = svc
732            .ready()
733            .await
734            .unwrap()
735            .call(make_exchange("orderId", "A", "a3"))
736            .await
737            .unwrap();
738        // A now has 3 → completes
739        assert!(matches!(ra.input.body, Body::Json(_)));
740        // B only has 1 → still pending
741        let rb = svc
742            .ready()
743            .await
744            .unwrap()
745            .call(make_exchange("orderId", "B", "b_check"))
746            .await
747            .unwrap();
748        assert!(matches!(rb.input.body, Body::Empty));
749    }
750
751    #[tokio::test]
752    async fn test_bucket_resets_after_completion() {
753        let mut svc = new_test_svc(config_size(2));
754        svc.ready()
755            .await
756            .unwrap()
757            .call(make_exchange("orderId", "A", "x"))
758            .await
759            .unwrap();
760        svc.ready()
761            .await
762            .unwrap()
763            .call(make_exchange("orderId", "A", "x"))
764            .await
765            .unwrap(); // completes
766        // New bucket starts
767        let r = svc
768            .ready()
769            .await
770            .unwrap()
771            .call(make_exchange("orderId", "A", "new"))
772            .await
773            .unwrap();
774        assert!(matches!(r.input.body, Body::Empty)); // pending again
775    }
776
777    #[tokio::test]
778    async fn test_completion_size_1_emits_immediately() {
779        let mut svc = new_test_svc(config_size(1));
780        let ex = make_exchange("orderId", "A", "solo");
781        let result = svc.ready().await.unwrap().call(ex).await.unwrap();
782        assert!(result.property(CAMEL_AGGREGATOR_PENDING).is_none());
783    }
784
785    #[tokio::test]
786    async fn test_custom_aggregation_strategy() {
787        use camel_api::aggregator::AggregationFn;
788        use std::sync::Arc;
789
790        let f: AggregationFn = Arc::new(|mut acc: Exchange, next: Exchange| {
791            let combined = format!(
792                "{}+{}",
793                acc.input.body.as_text().unwrap_or(""),
794                next.input.body.as_text().unwrap_or("")
795            );
796            acc.input.body = Body::Text(combined);
797            acc
798        });
799        let config = AggregatorConfig::correlate_by("key")
800            .complete_when_size(2)
801            .strategy(AggregationStrategy::Custom(f))
802            .build()
803            .unwrap();
804        let mut svc = new_test_svc(config);
805        svc.ready()
806            .await
807            .unwrap()
808            .call(make_exchange("key", "X", "hello"))
809            .await
810            .unwrap();
811        let result = svc
812            .ready()
813            .await
814            .unwrap()
815            .call(make_exchange("key", "X", "world"))
816            .await
817            .unwrap();
818        assert_eq!(result.input.body.as_text(), Some("hello+world"));
819    }
820
821    #[tokio::test]
822    async fn test_completion_predicate() {
823        let config = AggregatorConfig::correlate_by("key")
824            .complete_when(|bucket| {
825                bucket
826                    .iter()
827                    .any(|e| e.input.body.as_text() == Some("DONE"))
828            })
829            .build()
830            .unwrap();
831        let mut svc = new_test_svc(config);
832        svc.ready()
833            .await
834            .unwrap()
835            .call(make_exchange("key", "K", "first"))
836            .await
837            .unwrap();
838        svc.ready()
839            .await
840            .unwrap()
841            .call(make_exchange("key", "K", "second"))
842            .await
843            .unwrap();
844        let result = svc
845            .ready()
846            .await
847            .unwrap()
848            .call(make_exchange("key", "K", "DONE"))
849            .await
850            .unwrap();
851        assert!(result.property(CAMEL_AGGREGATOR_PENDING).is_none());
852    }
853
854    #[tokio::test]
855    async fn test_missing_header_returns_error() {
856        let mut svc = new_test_svc(config_size(2));
857        let msg = Message {
858            headers: Default::default(),
859            body: Body::Text("no key".into()),
860        };
861        let ex = Exchange::new(msg);
862        let result = svc.ready().await.unwrap().call(ex).await;
863        assert!(result.is_err());
864        assert!(matches!(
865            result.unwrap_err(),
866            camel_api::CamelError::ProcessorError(_)
867        ));
868    }
869
870    #[tokio::test]
871    async fn test_cloned_service_shares_state() {
872        let svc1 = new_test_svc(config_size(2));
873        let mut svc2 = svc1.clone();
874        // send first exchange via svc1
875        svc1.clone()
876            .ready()
877            .await
878            .unwrap()
879            .call(make_exchange("orderId", "A", "from-svc1"))
880            .await
881            .unwrap();
882        // send second exchange via svc2 — should complete because same Arc<Mutex>
883        let result = svc2
884            .ready()
885            .await
886            .unwrap()
887            .call(make_exchange("orderId", "A", "from-svc2"))
888            .await
889            .unwrap();
890        assert!(result.property(CAMEL_AGGREGATOR_PENDING).is_none());
891    }
892
893    #[tokio::test]
894    async fn test_camel_aggregated_key_property_set() {
895        let mut svc = new_test_svc(config_size(1));
896        let ex = make_exchange("orderId", "ORDER-42", "body");
897        let result = svc.ready().await.unwrap().call(ex).await.unwrap();
898        assert_eq!(
899            result.property(CAMEL_AGGREGATED_KEY),
900            Some(&serde_json::json!("ORDER-42"))
901        );
902    }
903
904    #[tokio::test]
905    async fn test_aggregator_enforces_max_buckets() {
906        let config = AggregatorConfig::correlate_by("orderId")
907            .complete_when_size(2)
908            .max_buckets(3)
909            .build()
910            .unwrap();
911
912        let mut svc = new_test_svc(config);
913
914        // Create 3 different correlation keys (fills limit)
915        for i in 0..3 {
916            let ex = make_exchange("orderId", &format!("key-{}", i), "body");
917            let _ = svc.ready().await.unwrap().call(ex).await.unwrap();
918        }
919
920        // 4th key should be rejected
921        let ex = make_exchange("orderId", "key-4", "body");
922        let result = svc.ready().await.unwrap().call(ex).await;
923
924        assert!(result.is_err(), "Should reject when max buckets reached");
925        let err = result.unwrap_err().to_string();
926        assert!(
927            err.contains("maximum"),
928            "Error message should contain 'maximum': {}",
929            err
930        );
931    }
932
933    #[tokio::test]
934    async fn test_max_buckets_allows_existing_key() {
935        let config = AggregatorConfig::correlate_by("orderId")
936            .complete_when_size(5) // Large size so bucket doesn't complete
937            .max_buckets(2)
938            .build()
939            .unwrap();
940
941        let mut svc = new_test_svc(config);
942
943        // Create 2 different correlation keys (fills limit)
944        let ex1 = make_exchange("orderId", "key-A", "body1");
945        let _ = svc.ready().await.unwrap().call(ex1).await.unwrap();
946        let ex2 = make_exchange("orderId", "key-B", "body2");
947        let _ = svc.ready().await.unwrap().call(ex2).await.unwrap();
948
949        // Should still allow adding to existing key
950        let ex3 = make_exchange("orderId", "key-A", "body3");
951        let result = svc.ready().await.unwrap().call(ex3).await;
952        assert!(
953            result.is_ok(),
954            "Should allow adding to existing bucket even at max limit"
955        );
956    }
957
958    #[tokio::test]
959    async fn test_bucket_ttl_eviction() {
960        let config = AggregatorConfig::correlate_by("orderId")
961            .complete_when_size(10) // Large size so bucket doesn't complete normally
962            .bucket_ttl(Duration::from_millis(50))
963            .build()
964            .unwrap();
965
966        let mut svc = new_test_svc(config);
967
968        // Create a bucket
969        let ex1 = make_exchange("orderId", "key-A", "body1");
970        let _ = svc.ready().await.unwrap().call(ex1).await.unwrap();
971
972        // Wait for TTL to expire
973        tokio::time::sleep(Duration::from_millis(100)).await;
974
975        // Create a new bucket - this should trigger eviction of the old one
976        let ex2 = make_exchange("orderId", "key-B", "body2");
977        let _ = svc.ready().await.unwrap().call(ex2).await.unwrap();
978
979        // The expired bucket should have been evicted, so we should be able to
980        // add a new key-A bucket again
981        let ex3 = make_exchange("orderId", "key-A", "body3");
982        let result = svc.ready().await.unwrap().call(ex3).await;
983        assert!(result.is_ok(), "Should be able to recreate evicted bucket");
984    }
985
986    #[tokio::test(start_paused = true)]
987    async fn test_timeout_completes_bucket() {
988        let config = AggregatorConfig::correlate_by("key")
989            .complete_on_timeout(Duration::from_millis(100))
990            .build()
991            .unwrap();
992        let mut svc = new_test_svc(config);
993        let ex = make_exchange("key", "A", "data");
994        let result = svc.ready().await.unwrap().call(ex).await.unwrap();
995        assert!(result.property(CAMEL_AGGREGATOR_PENDING).is_some());
996
997        tokio::time::sleep(Duration::from_millis(200)).await;
998
999        assert_eq!(
1000            svc.buckets.lock().unwrap().len(),
1001            0,
1002            "bucket should be removed after timeout"
1003        );
1004    }
1005
1006    #[tokio::test(start_paused = true)]
1007    async fn test_timeout_resets_on_new_exchange() {
1008        let config = AggregatorConfig::correlate_by("key")
1009            .complete_on_timeout(Duration::from_millis(150))
1010            .build()
1011            .unwrap();
1012        let mut svc = new_test_svc(config);
1013
1014        let ex1 = make_exchange("key", "A", "first");
1015        let _ = svc.ready().await.unwrap().call(ex1).await.unwrap();
1016
1017        tokio::time::sleep(Duration::from_millis(100)).await;
1018
1019        let ex2 = make_exchange("key", "A", "second");
1020        let _ = svc.ready().await.unwrap().call(ex2).await.unwrap();
1021
1022        tokio::time::sleep(Duration::from_millis(100)).await;
1023
1024        assert_eq!(
1025            svc.buckets.lock().unwrap().len(),
1026            1,
1027            "bucket should still exist — timeout was reset"
1028        );
1029
1030        tokio::time::sleep(Duration::from_millis(100)).await;
1031
1032        assert_eq!(
1033            svc.buckets.lock().unwrap().len(),
1034            0,
1035            "bucket should be gone after timeout fires"
1036        );
1037    }
1038
1039    #[tokio::test]
1040    async fn test_composable_size_and_timeout() {
1041        let config = AggregatorConfig::correlate_by("key")
1042            .complete_on_size_or_timeout(2, Duration::from_millis(200))
1043            .build()
1044            .unwrap();
1045        let mut svc = new_test_svc(config);
1046
1047        let ex1 = make_exchange("key", "A", "first");
1048        let _ = svc.ready().await.unwrap().call(ex1).await.unwrap();
1049        assert!(svc.buckets.lock().unwrap().contains_key("\"A\""));
1050
1051        let ex2 = make_exchange("key", "A", "second");
1052        let result = svc.ready().await.unwrap().call(ex2).await.unwrap();
1053        assert!(result.property(CAMEL_AGGREGATOR_PENDING).is_none());
1054        assert_eq!(
1055            result.property(CAMEL_AGGREGATED_COMPLETION_REASON),
1056            Some(&serde_json::json!("size"))
1057        );
1058    }
1059
1060    #[tokio::test(start_paused = true)]
1061    async fn test_discard_on_timeout() {
1062        let config = AggregatorConfig::correlate_by("key")
1063            .complete_on_timeout(Duration::from_millis(50))
1064            .discard_on_timeout(true)
1065            .build()
1066            .unwrap();
1067        let (tx, mut rx) = mpsc::channel(256);
1068        let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1069        let cancel = CancellationToken::new();
1070        let mut svc = AggregatorService::new(config, tx, registry, cancel);
1071
1072        let ex = make_exchange("key", "A", "data");
1073        let _ = svc.ready().await.unwrap().call(ex).await.unwrap();
1074
1075        tokio::time::sleep(Duration::from_millis(100)).await;
1076
1077        assert!(
1078            rx.try_recv().is_err(),
1079            "no emit expected with discard_on_timeout"
1080        );
1081        assert_eq!(svc.buckets.lock().unwrap().len(), 0);
1082        assert!(
1083            svc.timeout_tasks.lock().unwrap().is_empty(),
1084            "timeout task should be cleaned up"
1085        );
1086    }
1087
1088    #[tokio::test]
1089    async fn test_force_completion_on_stop() {
1090        let config = AggregatorConfig::correlate_by("key")
1091            .complete_when_size(10)
1092            .force_completion_on_stop(true)
1093            .build()
1094            .unwrap();
1095        let (tx, mut rx) = mpsc::channel(256);
1096        let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1097        let cancel = CancellationToken::new();
1098        let svc = AggregatorService::new(config, tx, registry, cancel);
1099
1100        let mut call_svc = svc.clone();
1101        let ex = make_exchange("key", "A", "data");
1102        let _ = call_svc.ready().await.unwrap().call(ex).await.unwrap();
1103
1104        svc.force_complete_all();
1105
1106        let result = rx.try_recv().expect("should emit on force-complete");
1107        assert!(
1108            result.input.body.as_text().is_some() || matches!(result.input.body, Body::Json(_))
1109        );
1110        assert_eq!(
1111            result.property(CAMEL_AGGREGATED_COMPLETION_REASON),
1112            Some(&serde_json::json!("stop"))
1113        );
1114    }
1115
1116    #[tokio::test]
1117    async fn test_completion_reason_property_size() {
1118        let config = AggregatorConfig::correlate_by("key")
1119            .complete_when_size(1)
1120            .build()
1121            .unwrap();
1122        let mut svc = new_test_svc(config);
1123        let ex = make_exchange("key", "X", "body");
1124        let result = svc.ready().await.unwrap().call(ex).await.unwrap();
1125        assert_eq!(
1126            result.property(CAMEL_AGGREGATED_COMPLETION_REASON),
1127            Some(&serde_json::json!("size"))
1128        );
1129    }
1130
1131    #[tokio::test]
1132    async fn test_completion_reason_property_predicate() {
1133        let config = AggregatorConfig::correlate_by("key")
1134            .complete_when(|_| true)
1135            .build()
1136            .unwrap();
1137        let mut svc = new_test_svc(config);
1138        let ex = make_exchange("key", "X", "body");
1139        let result = svc.ready().await.unwrap().call(ex).await.unwrap();
1140        assert_eq!(
1141            result.property(CAMEL_AGGREGATED_COMPLETION_REASON),
1142            Some(&serde_json::json!("predicate"))
1143        );
1144    }
1145
1146    #[tokio::test(start_paused = true)]
1147    async fn test_size_completes_before_timeout() {
1148        let config = AggregatorConfig::correlate_by("key")
1149            .complete_on_size_or_timeout(2, Duration::from_millis(200))
1150            .build()
1151            .unwrap();
1152        let mut svc = new_test_svc(config);
1153
1154        let ex1 = make_exchange("key", "A", "first");
1155        let _ = svc.ready().await.unwrap().call(ex1).await.unwrap();
1156
1157        let ex2 = make_exchange("key", "A", "second");
1158        let result = svc.ready().await.unwrap().call(ex2).await.unwrap();
1159
1160        assert!(result.property(CAMEL_AGGREGATOR_PENDING).is_none());
1161        assert_eq!(
1162            result.property(CAMEL_AGGREGATED_COMPLETION_REASON),
1163            Some(&serde_json::json!("size"))
1164        );
1165        assert_eq!(svc.buckets.lock().unwrap().len(), 0);
1166
1167        tokio::time::sleep(Duration::from_millis(300)).await;
1168        assert_eq!(
1169            svc.buckets.lock().unwrap().len(),
1170            0,
1171            "no re-fire after timeout"
1172        );
1173    }
1174
1175    #[tokio::test(start_paused = true)]
1176    async fn test_concurrent_timeout_fire_and_new_exchange() {
1177        let config = AggregatorConfig::correlate_by("key")
1178            .complete_on_size_or_timeout(2, Duration::from_millis(100))
1179            .build()
1180            .unwrap();
1181        let (tx, mut rx) = mpsc::channel(256);
1182        let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1183        let cancel = CancellationToken::new();
1184        let mut svc = AggregatorService::new(config, tx, registry, cancel);
1185
1186        let ex = make_exchange("key", "A", "data");
1187        let _ = svc.ready().await.unwrap().call(ex).await.unwrap();
1188
1189        // Advance time past timeout — timeout task fires and removes bucket
1190        tokio::time::sleep(Duration::from_millis(150)).await;
1191
1192        // New exchange arrives after timeout — starts a fresh bucket
1193        let ex2 = make_exchange("key", "A", "data2");
1194        let result = svc.ready().await.unwrap().call(ex2).await.unwrap();
1195        assert!(
1196            result.property(CAMEL_AGGREGATOR_PENDING).is_some(),
1197            "should be pending in new bucket"
1198        );
1199
1200        // Drain late emits from timeout
1201        let mut late_count = 0;
1202        while rx.try_recv().is_ok() {
1203            late_count += 1;
1204        }
1205        assert_eq!(
1206            late_count, 1,
1207            "exactly 1 late emit from the timed-out bucket"
1208        );
1209    }
1210
1211    #[tokio::test(start_paused = true)]
1212    async fn test_late_channel_full_drops_with_warning() {
1213        let config = AggregatorConfig::correlate_by("key")
1214            .complete_on_timeout(Duration::from_millis(50))
1215            .build()
1216            .unwrap();
1217        let (tx, mut rx) = mpsc::channel(1);
1218        rx.close();
1219        let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1220        let cancel = CancellationToken::new();
1221        let mut svc = AggregatorService::new(config, tx, registry, cancel);
1222
1223        let ex = make_exchange("key", "A", "data");
1224        let _ = svc.ready().await.unwrap().call(ex).await.unwrap();
1225
1226        tokio::time::sleep(Duration::from_millis(100)).await;
1227        assert_eq!(
1228            svc.buckets.lock().unwrap().len(),
1229            0,
1230            "bucket removed despite channel closed"
1231        );
1232    }
1233
1234    #[tokio::test]
1235    async fn test_aggregate_stream_bodies_creates_valid_json() {
1236        use bytes::Bytes;
1237        use camel_api::{Body, StreamBody, StreamMetadata};
1238        use futures::stream;
1239        use tokio::sync::Mutex;
1240
1241        let chunks = vec![Ok(Bytes::from("test"))];
1242        let stream_body = StreamBody {
1243            stream: Arc::new(Mutex::new(Some(Box::pin(stream::iter(chunks))))),
1244            metadata: StreamMetadata {
1245                origin: Some("file:///test.txt".to_string()),
1246                ..Default::default()
1247            },
1248        };
1249
1250        let ex1 = Exchange::new(Message {
1251            headers: Default::default(),
1252            body: Body::Stream(stream_body),
1253        });
1254
1255        let exchanges = vec![ex1];
1256        let result = aggregate(exchanges, &AggregationStrategy::CollectAll);
1257
1258        let exchange = result.expect("Expected Ok result");
1259        assert!(
1260            matches!(exchange.input.body, Body::Json(_)),
1261            "Expected Json body"
1262        );
1263
1264        if let Body::Json(value) = exchange.input.body {
1265            let json_str = serde_json::to_string(&value).unwrap();
1266            let parsed: serde_json::Value = serde_json::from_str(&json_str).unwrap();
1267
1268            assert!(parsed.is_array(), "Result should be an array");
1269            let arr = parsed.as_array().unwrap();
1270            assert!(arr[0].is_object(), "First element should be an object");
1271            assert!(
1272                arr[0]["_stream"].is_object(),
1273                "Should contain _stream object"
1274            );
1275            assert_eq!(arr[0]["_stream"]["origin"], "file:///test.txt");
1276            assert_eq!(
1277                arr[0]["_stream"]["placeholder"], true,
1278                "placeholder flag should be true"
1279            );
1280        }
1281    }
1282
1283    #[tokio::test]
1284    async fn test_aggregate_stream_bodies_with_none_origin() {
1285        use bytes::Bytes;
1286        use camel_api::{Body, StreamBody, StreamMetadata};
1287        use futures::stream;
1288        use tokio::sync::Mutex;
1289
1290        let chunks = vec![Ok(Bytes::from("test"))];
1291        let stream_body = StreamBody {
1292            stream: Arc::new(Mutex::new(Some(Box::pin(stream::iter(chunks))))),
1293            metadata: StreamMetadata {
1294                origin: None,
1295                ..Default::default()
1296            },
1297        };
1298
1299        let ex1 = Exchange::new(Message {
1300            headers: Default::default(),
1301            body: Body::Stream(stream_body),
1302        });
1303
1304        let exchanges = vec![ex1];
1305        let result = aggregate(exchanges, &AggregationStrategy::CollectAll);
1306
1307        let exchange = result.expect("Expected Ok result");
1308        assert!(
1309            matches!(exchange.input.body, Body::Json(_)),
1310            "Expected Json body"
1311        );
1312
1313        if let Body::Json(value) = exchange.input.body {
1314            let json_str = serde_json::to_string(&value).unwrap();
1315            let parsed: serde_json::Value = serde_json::from_str(&json_str).unwrap();
1316
1317            assert!(parsed.is_array(), "Result should be an array");
1318            let arr = parsed.as_array().unwrap();
1319            assert!(arr[0].is_object(), "First element should be an object");
1320            assert!(
1321                arr[0]["_stream"].is_object(),
1322                "Should contain _stream object"
1323            );
1324            assert_eq!(
1325                arr[0]["_stream"]["origin"],
1326                serde_json::Value::Null,
1327                "origin should be null when None"
1328            );
1329            assert_eq!(
1330                arr[0]["_stream"]["placeholder"], true,
1331                "placeholder flag should be true"
1332            );
1333        }
1334    }
1335
1336    #[tokio::test]
1337    async fn timeout_completion_clears_handle_from_map() {
1338        // Regression: Oracle audit found that natural timeout completion removed
1339        // the bucket but left the JoinHandle in `timeout_handles`, leaking the
1340        // entry until route shutdown. After fix, the timeout task itself cleans
1341        // its handle from the map on natural completion.
1342        let config = AggregatorConfig::correlate_by("key")
1343            .complete_on_timeout(Duration::from_millis(50))
1344            .build()
1345            .unwrap();
1346        let (tx, _rx) = mpsc::channel(256);
1347        let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1348        let cancel = CancellationToken::new();
1349        let svc = AggregatorService::new(config, tx, registry, cancel);
1350
1351        // Send an exchange to create a pending bucket with a timeout task.
1352        let mut call_svc = svc.clone();
1353        let ex = make_exchange("key", "A", "data");
1354        let _ = call_svc.ready().await.unwrap().call(ex).await.unwrap();
1355        assert!(
1356            !svc.timeout_handles.lock().unwrap().is_empty(),
1357            "handle should exist while timeout pending"
1358        );
1359
1360        // Wait real time for the 50ms timeout to fire + spawned task to complete.
1361        tokio::time::sleep(Duration::from_millis(200)).await;
1362
1363        assert!(
1364            svc.timeout_handles.lock().unwrap().is_empty(),
1365            "handle should be cleared from map after natural timeout completion (was leak)"
1366        );
1367    }
1368
1369    #[tokio::test]
1370    async fn aggregator_shutdown_via_trait_dispatch() {
1371        // RED: builds an AggregatorService, dispatches through Arc<dyn StepLifecycle>,
1372        // and asserts idempotent shutdown works.
1373        let config = AggregatorConfig::correlate_by("key")
1374            .complete_when_size(10)
1375            .build()
1376            .unwrap();
1377        let (tx, _rx) = mpsc::channel(256);
1378        let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1379        let cancel = CancellationToken::new();
1380        let svc = AggregatorService::new(config, tx, registry, cancel);
1381
1382        let step: Arc<dyn StepLifecycle> = Arc::new(svc);
1383        step.shutdown(StepShutdownReason::RouteStop)
1384            .await
1385            .expect("first shutdown should succeed");
1386        step.shutdown(StepShutdownReason::RouteStop)
1387            .await
1388            .expect("second shutdown (idempotent) should succeed");
1389    }
1390
1391    #[tokio::test(start_paused = true)]
1392    async fn test_shutdown_awaits_timeout_handles() {
1393        let config = AggregatorConfig::correlate_by("key")
1394            .complete_on_timeout(Duration::from_millis(100))
1395            .build()
1396            .unwrap();
1397        let (tx, _rx) = mpsc::channel(256);
1398        let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1399        let cancel = CancellationToken::new();
1400        let svc = AggregatorService::new(config, tx, registry, cancel);
1401
1402        // Send an exchange to create a pending bucket with a timeout task.
1403        let mut call_svc = svc.clone();
1404        let ex = make_exchange("key", "A", "data");
1405        let _ = call_svc.ready().await.unwrap().call(ex).await.unwrap();
1406
1407        // Verify timeout handle exists.
1408        assert!(
1409            !svc.timeout_handles.lock().unwrap().is_empty(),
1410            "should have a timeout handle"
1411        );
1412
1413        // Shutdown should complete within the 5s deadline (the timeout task
1414        // gets cancelled so it won't wait for the full 100ms sleep).
1415        svc.shutdown_inner().await;
1416
1417        assert!(
1418            svc.timeout_handles.lock().unwrap().is_empty(),
1419            "all handles should be cleaned up after shutdown"
1420        );
1421    }
1422
1423    // ── R3-C1 Batch 1: DoS cap + background sweep ───────────────────
1424
1425    /// R3-C1: a flood of unique correlation keys must stay bounded.
1426    /// Default `max_buckets` is 10_000; the 10_001st unique key is rejected
1427    /// with `Aggregator reached maximum N buckets` (or its updated equivalent
1428    /// after the fix). The unique-key flood does NOT OOM the process.
1429    #[tokio::test]
1430    async fn test_unique_key_flood_stays_bounded_by_default() {
1431        // Builder defaults to max_buckets = 10_000, bucket_ttl = 300s.
1432        let config = AggregatorConfig::correlate_by("orderId")
1433            .complete_when_size(1_000_000) // never completes normally
1434            .build()
1435            .unwrap();
1436        let mut svc = new_test_svc(config);
1437
1438        // Send 10_001 unique keys. The first 10_000 should be accepted
1439        // (pending in their buckets); the 10_001st MUST be rejected.
1440        for i in 0..10_000usize {
1441            let ex = make_exchange("orderId", &format!("key-{i}"), "body");
1442            let result = svc.ready().await.unwrap().call(ex).await;
1443            assert!(result.is_ok(), "key {i} should be accepted under the cap");
1444        }
1445        let ex = make_exchange("orderId", "key-10001", "body");
1446        let result = svc.ready().await.unwrap().call(ex).await;
1447        assert!(
1448            result.is_err(),
1449            "10_001st unique key must be rejected by the max_buckets cap"
1450        );
1451        let err = result.unwrap_err().to_string();
1452        assert!(
1453            err.contains("maximum") || err.contains("max"),
1454            "error should mention cap: {err}"
1455        );
1456    }
1457
1458    /// The `AggregatorService` exposes a `sweep_handle` for the background
1459    /// sweep task. When `config.bucket_ttl` is `Some`, `AggregatorService::new`
1460    /// automatically spawns the sweep task and stores the handle, so the
1461    /// caller never sees `None` for a TTL-configured service. Cancelling the
1462    /// route token (via `shutdown`) aborts the sweep.
1463    #[tokio::test]
1464    async fn test_background_sweep_spawns_on_first_poll_not_construction() {
1465        let config = AggregatorConfig::correlate_by("key")
1466            .complete_when_size(10_000)
1467            .bucket_ttl(Duration::from_millis(50))
1468            .build()
1469            .unwrap();
1470        let cancel = CancellationToken::new();
1471        let (tx, _rx) = mpsc::channel(8);
1472        let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1473        let mut svc = AggregatorService::new(config, tx, registry, cancel.clone());
1474
1475        // Construction is runtime-free: no sweep spawned yet.
1476        assert!(
1477            svc.sweep_handle
1478                .lock()
1479                .unwrap_or_else(|e| e.into_inner())
1480                .is_none(),
1481            "sweep must NOT be spawned at construction (runtime-free new)"
1482        );
1483
1484        // First readiness poll lazily spawns the sweep (a runtime is present
1485        // here). `ready()` drives `poll_ready` until Ready.
1486        let _ = svc.ready().await.unwrap();
1487        let sweep_present = svc
1488            .sweep_handle
1489            .lock()
1490            .unwrap_or_else(|e| e.into_inner())
1491            .is_some();
1492        assert!(
1493            sweep_present,
1494            "sweep handle should be Some after first poll when bucket_ttl is set"
1495        );
1496
1497        // Cancel the route token; the sweep task observes it and exits.
1498        cancel.cancel();
1499        // Give the task a moment to observe the cancel.
1500        tokio::time::sleep(Duration::from_millis(50)).await;
1501    }
1502}