camel_api/

aggregator.rs

use std::sync::Arc;
use std::time::Duration;

use crate::error::{CamelError, ConfigValidationError};
use crate::exchange::Exchange;

/// Aggregation function — left-fold binary: (accumulated, next) -> merged.
pub type AggregationFn = Arc<dyn Fn(Exchange, Exchange) -> Exchange + Send + Sync>;

/// Strategy for correlating exchanges into aggregation buckets.
pub enum CorrelationStrategy {
    /// Correlate by the value of a named header.
    HeaderName(String),
    /// Correlate by evaluating an expression using a language registry.
    Expression { expr: String, language: String },
    /// Correlate using a custom function.
    #[allow(clippy::type_complexity)]
    Fn(Arc<dyn Fn(&Exchange) -> Option<String> + Send + Sync>),
}

impl std::fmt::Debug for CorrelationStrategy {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            CorrelationStrategy::HeaderName(h) => f.debug_tuple("HeaderName").field(h).finish(),
            CorrelationStrategy::Expression { expr, language } => f
                .debug_struct("Expression")
                .field("expr", expr)
                .field("language", language)
                .finish(),
            CorrelationStrategy::Fn(_) => f.write_str("Fn(..)"),
        }
    }
}

impl Clone for CorrelationStrategy {
    fn clone(&self) -> Self {
        match self {
            CorrelationStrategy::HeaderName(h) => CorrelationStrategy::HeaderName(h.clone()),
            CorrelationStrategy::Expression { expr, language } => CorrelationStrategy::Expression {
                expr: expr.clone(),
                language: language.clone(),
            },
            CorrelationStrategy::Fn(f) => CorrelationStrategy::Fn(Arc::clone(f)),
        }
    }
}

/// How to combine collected exchanges into one.
#[derive(Clone)]
pub enum AggregationStrategy {
    /// Collects all bodies into Body::Json([body1, body2, ...]).
    CollectAll,
    /// Left-fold: f(f(ex1, ex2), ex3), ...
    Custom(AggregationFn),
}

impl std::fmt::Debug for AggregationStrategy {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            AggregationStrategy::CollectAll => f.write_str("CollectAll"),
            AggregationStrategy::Custom(_) => f.write_str("Custom(..)"),
        }
    }
}

/// When the bucket is considered complete and should be emitted.
#[derive(Clone)]
pub enum CompletionCondition {
    /// Emit when bucket reaches exactly N exchanges.
    Size(usize),
    /// Emit when predicate returns true for current bucket.
    #[allow(clippy::type_complexity)]
    Predicate(Arc<dyn Fn(&[Exchange]) -> bool + Send + Sync>),
    /// Emit when the bucket has been inactive for the given duration.
    Timeout(Duration),
}

impl std::fmt::Debug for CompletionCondition {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            CompletionCondition::Size(n) => f.debug_tuple("Size").field(n).finish(),
            CompletionCondition::Predicate(_) => f.write_str("Predicate(..)"),
            CompletionCondition::Timeout(d) => f.debug_tuple("Timeout").field(d).finish(),
        }
    }
}

/// Determines how a bucket's completion is evaluated.
/// `Single` wraps one condition; `Any` completes when the first condition triggers.
#[derive(Clone)]
pub enum CompletionMode {
    Single(CompletionCondition),
    Any(Vec<CompletionCondition>),
}

impl std::fmt::Debug for CompletionMode {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            CompletionMode::Single(c) => f.debug_tuple("Single").field(c).finish(),
            CompletionMode::Any(conds) => f.debug_tuple("Any").field(conds).finish(),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum CompletionReason {
    Size,
    Predicate,
    Timeout,
    Stop,
}

impl CompletionReason {
    pub fn as_str(&self) -> &'static str {
        match self {
            CompletionReason::Size => "size",
            CompletionReason::Predicate => "predicate",
            CompletionReason::Timeout => "timeout",
            CompletionReason::Stop => "stop",
        }
    }
}

/// Configuration for the Aggregator EIP.
#[derive(Clone)]
pub struct AggregatorConfig {
    /// Name of the header used as correlation key.
    pub header_name: String,
    /// When to emit the aggregated exchange.
    pub completion: CompletionMode,
    /// Strategy for determining correlation keys.
    pub correlation: CorrelationStrategy,
    /// How to combine the bucket into one exchange.
    pub strategy: AggregationStrategy,
    /// Maximum number of correlation key buckets (memory protection).
    /// When limit is reached, new correlation keys are rejected.
    pub max_buckets: Option<usize>,
    /// Time-to-live for inactive buckets (memory protection).
    /// Buckets not updated for this duration are evicted.
    pub bucket_ttl: Option<Duration>,
    /// Force-complete all pending buckets when the route is stopped.
    pub force_completion_on_stop: bool,
    /// Discard bucket contents on timeout instead of emitting.
    pub discard_on_timeout: bool,
    /// Maximum number of concurrently-live per-bucket timeout tasks (DoS cap, R3-M3).
    /// When the cap is reached, new buckets skip the dedicated timeout spawn and
    /// rely on `bucket_ttl` eviction (graceful degradation under a key flood).
    /// Default 1024.
    pub max_timeout_tasks: usize,
}

impl std::fmt::Debug for AggregatorConfig {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("AggregatorConfig")
            .field("header_name", &self.header_name)
            .field("completion", &self.completion)
            .field("correlation", &self.correlation)
            .field("strategy", &self.strategy)
            .field("max_buckets", &self.max_buckets)
            .field("bucket_ttl", &self.bucket_ttl)
            .field("force_completion_on_stop", &self.force_completion_on_stop)
            .field("discard_on_timeout", &self.discard_on_timeout)
            .field("max_timeout_tasks", &self.max_timeout_tasks)
            .finish()
    }
}

impl AggregatorConfig {
    /// Start building config with correlation key extracted from the named header.
    pub fn correlate_by(header: impl Into<String>) -> AggregatorConfigBuilder {
        let header_name = header.into();
        AggregatorConfigBuilder {
            header_name: header_name.clone(),
            completion: None,
            correlation: CorrelationStrategy::HeaderName(header_name),
            strategy: AggregationStrategy::CollectAll,
            // R3-C1 Batch 1: bounded defaults — the correlation map MUST NOT
            // grow without a cap. A flood of unique correlation keys is the
            // remote-OOM vector; the default cap (10_000) is a sane production
            // ceiling the operator may still override. The TTL (5 minutes) is
            // the inline-retain + background-sweep eviction window.
            max_buckets: Some(10_000),
            bucket_ttl: Some(Duration::from_secs(300)),
            force_completion_on_stop: false,
            discard_on_timeout: false,
            // R3-M3: cap concurrently-live per-bucket timeout tasks.
            max_timeout_tasks: 1024,
        }
    }

    /// Validate that at least one memory-release bound is configured (R3-M2).
    ///
    /// At least one of `max_buckets`, a `Timeout` completion condition, or
    /// `bucket_ttl` MUST be set, otherwise a flood of unique correlation keys
    /// grows the bucket map without limit (remote-OOM vector).
    ///
    /// Additionally, when a `Timeout` completion condition is present,
    /// `bucket_ttl` MUST also be set. The R3-M3 timeout-task cap may skip
    /// spawning a dedicated timeout task under flood; without `bucket_ttl`
    /// there is no fallback eviction path and the bucket leaks until shutdown.
    /// Requiring `bucket_ttl` whenever Timeout is present makes the cap-skip
    /// degradation safe by construction.
    pub fn validate(&self) -> Result<(), CamelError> {
        let has_timeout = match &self.completion {
            CompletionMode::Single(CompletionCondition::Timeout(_)) => true,
            CompletionMode::Any(conds) => conds
                .iter()
                .any(|c| matches!(c, CompletionCondition::Timeout(_))),
            _ => false,
        };
        let has_bound = self.max_buckets.is_some() || has_timeout || self.bucket_ttl.is_some();
        if !has_bound {
            return Err(CamelError::from(
                ConfigValidationError::AggregatorMissingMemoryBound,
            ));
        }
        // R3-M3: Timeout completion requires bucket_ttl so the cap-skip
        // degradation always has an eviction path.
        if has_timeout && self.bucket_ttl.is_none() {
            return Err(CamelError::from(
                ConfigValidationError::AggregatorTimeoutRequiresTtl,
            ));
        }
        Ok(())
    }
}

/// Builder for `AggregatorConfig`.
pub struct AggregatorConfigBuilder {
    header_name: String,
    completion: Option<CompletionMode>,
    correlation: CorrelationStrategy,
    strategy: AggregationStrategy,
    max_buckets: Option<usize>,
    bucket_ttl: Option<Duration>,
    force_completion_on_stop: bool,
    discard_on_timeout: bool,
    max_timeout_tasks: usize,
}

impl AggregatorConfigBuilder {
    /// Emit when bucket has N exchanges.
    pub fn complete_when_size(mut self, n: usize) -> Self {
        self.completion = Some(CompletionMode::Single(CompletionCondition::Size(n)));
        self
    }

    /// Emit when predicate returns true for the current bucket.
    pub fn complete_when<F>(mut self, predicate: F) -> Self
    where
        F: Fn(&[Exchange]) -> bool + Send + Sync + 'static,
    {
        self.completion = Some(CompletionMode::Single(CompletionCondition::Predicate(
            Arc::new(predicate),
        )));
        self
    }

    /// Emit when the bucket has been inactive for the given duration.
    pub fn complete_on_timeout(mut self, duration: Duration) -> Self {
        self.completion = Some(CompletionMode::Single(CompletionCondition::Timeout(
            duration,
        )));
        self
    }

    /// Emit when the bucket reaches `size` OR has been inactive for `timeout`.
    pub fn complete_on_size_or_timeout(mut self, size: usize, timeout: Duration) -> Self {
        self.completion = Some(CompletionMode::Any(vec![
            CompletionCondition::Size(size),
            CompletionCondition::Timeout(timeout),
        ]));
        self
    }

    /// Enable force-completion of pending buckets when the route is stopped.
    pub fn force_completion_on_stop(mut self, enabled: bool) -> Self {
        self.force_completion_on_stop = enabled;
        self
    }

    /// Discard bucket contents on timeout instead of emitting the aggregated exchange.
    pub fn discard_on_timeout(mut self, enabled: bool) -> Self {
        self.discard_on_timeout = enabled;
        self
    }

    /// Override the correlation strategy with a header-based key.
    pub fn correlate_by(mut self, header: impl Into<String>) -> Self {
        let header = header.into();
        self.header_name = header.clone();
        self.correlation = CorrelationStrategy::HeaderName(header);
        self
    }

    /// Override the default `CollectAll` aggregation strategy.
    pub fn strategy(mut self, strategy: AggregationStrategy) -> Self {
        self.strategy = strategy;
        self
    }

    /// Set the maximum number of correlation key buckets.
    /// When the limit is reached, new correlation keys are rejected with an error.
    pub fn max_buckets(mut self, max: usize) -> Self {
        self.max_buckets = Some(max);
        self
    }

    /// Set the time-to-live for inactive buckets.
    /// Buckets that haven't been updated for this duration will be evicted.
    pub fn bucket_ttl(mut self, ttl: Duration) -> Self {
        self.bucket_ttl = Some(ttl);
        self
    }

    /// Override the maximum number of concurrently-live per-bucket timeout tasks.
    pub fn max_timeout_tasks(mut self, max: usize) -> Self {
        self.max_timeout_tasks = max;
        self
    }

    pub fn try_build(self) -> Result<AggregatorConfig, CamelError> {
        // R3-C1 Batch 1: a completion-bound is mandatory. A config with no
        // completion bound lets a bucket live forever — combined with a
        // unique-key flood, that is the remote-OOM vector. Typed
        // ConfigValidationError (ADR-0033) — operators can match on the
        // `AggregatorMissingCompletionBound` variant.
        let completion = self.completion.ok_or_else(|| {
            CamelError::from(ConfigValidationError::AggregatorMissingCompletionBound)
        })?;
        Ok(AggregatorConfig {
            header_name: self.header_name,
            completion,
            correlation: self.correlation,
            strategy: self.strategy,
            max_buckets: self.max_buckets,
            bucket_ttl: self.bucket_ttl,
            force_completion_on_stop: self.force_completion_on_stop,
            discard_on_timeout: self.discard_on_timeout,
            max_timeout_tasks: self.max_timeout_tasks,
        })
    }

    /// Build the config. Returns an error if no completion condition was set.
    pub fn build(self) -> Result<AggregatorConfig, CamelError> {
        self.try_build()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_aggregator_config_complete_when_size() {
        let config = AggregatorConfig::correlate_by("orderId")
            .complete_when_size(3)
            .build()
            .unwrap();
        assert_eq!(config.header_name, "orderId");
        assert!(matches!(
            config.completion,
            CompletionMode::Single(CompletionCondition::Size(3))
        ));
        assert!(matches!(config.strategy, AggregationStrategy::CollectAll));
    }

    #[test]
    fn test_aggregator_config_complete_when_predicate() {
        let config = AggregatorConfig::correlate_by("key")
            .complete_when(|bucket| bucket.len() >= 2)
            .build()
            .unwrap();
        assert!(matches!(
            config.completion,
            CompletionMode::Single(CompletionCondition::Predicate(_))
        ));
    }

    #[test]
    fn test_aggregator_config_custom_strategy() {
        use std::sync::Arc;
        let f: AggregationFn = Arc::new(|acc, _next| acc);
        let config = AggregatorConfig::correlate_by("key")
            .complete_when_size(1)
            .strategy(AggregationStrategy::Custom(f))
            .build()
            .unwrap();
        assert!(matches!(config.strategy, AggregationStrategy::Custom(_)));
    }

    #[test]
    fn test_aggregator_config_missing_completion_returns_err() {
        let result = AggregatorConfig::correlate_by("key").build();
        let err = match result {
            Err(e) => e,
            Ok(_) => panic!("expected error, got Ok"),
        };
        assert!(
            err.to_string().contains("completion"),
            "error message should mention 'completion': {err}"
        );
    }

    #[test]
    fn test_complete_on_size_or_timeout() {
        let config = AggregatorConfig::correlate_by("key")
            .complete_on_size_or_timeout(3, Duration::from_secs(5))
            .build()
            .unwrap();
        assert!(matches!(config.completion, CompletionMode::Any(v) if v.len() == 2));
    }

    #[test]
    fn test_force_completion_on_stop_default() {
        let config = AggregatorConfig::correlate_by("key")
            .complete_when_size(1)
            .build()
            .unwrap();
        assert!(!config.force_completion_on_stop);
        assert!(!config.discard_on_timeout);
    }

    #[test]
    fn test_builder_sets_timeout_and_flags_and_limits() {
        let config = AggregatorConfig::correlate_by("key")
            .complete_on_timeout(Duration::from_secs(2))
            .max_buckets(7)
            .bucket_ttl(Duration::from_secs(10))
            .force_completion_on_stop(true)
            .discard_on_timeout(true)
            .build()
            .unwrap();

        assert!(matches!(
            config.completion,
            CompletionMode::Single(CompletionCondition::Timeout(d)) if d == Duration::from_secs(2)
        ));
        assert_eq!(config.max_buckets, Some(7));
        assert_eq!(config.bucket_ttl, Some(Duration::from_secs(10)));
        assert!(config.force_completion_on_stop);
        assert!(config.discard_on_timeout);
    }

    #[test]
    fn test_builder_correlate_by_overrides_header_and_strategy() {
        let config = AggregatorConfig::correlate_by("original")
            .correlate_by("override")
            .complete_when_size(1)
            .build()
            .unwrap();

        assert_eq!(config.header_name, "override");
        assert!(matches!(
            config.correlation,
            CorrelationStrategy::HeaderName(ref h) if h == "override"
        ));
    }

    #[test]
    fn test_completion_reason_as_str_all_variants() {
        assert_eq!(CompletionReason::Size.as_str(), "size");
        assert_eq!(CompletionReason::Predicate.as_str(), "predicate");
        assert_eq!(CompletionReason::Timeout.as_str(), "timeout");
        assert_eq!(CompletionReason::Stop.as_str(), "stop");
    }

    #[test]
    fn test_correlation_strategy_clone_and_debug() {
        let strategy = CorrelationStrategy::Expression {
            expr: "${header.orderId}".to_string(),
            language: "simple".to_string(),
        };
        let cloned = strategy.clone();
        assert!(matches!(
            cloned,
            CorrelationStrategy::Expression { ref expr, ref language }
                if expr == "${header.orderId}" && language == "simple"
        ));

        let f = CorrelationStrategy::Fn(Arc::new(|_| Some("k".to_string())));
        assert_eq!(format!("{:?}", f), "Fn(..)");
    }

    #[test]
    fn test_complete_on_size_or_timeout_contains_both_conditions() {
        let config = AggregatorConfig::correlate_by("k")
            .complete_on_size_or_timeout(4, Duration::from_millis(250))
            .build()
            .unwrap();

        match config.completion {
            CompletionMode::Any(conditions) => {
                assert!(matches!(conditions[0], CompletionCondition::Size(4)));
                assert!(matches!(
                    conditions[1],
                    CompletionCondition::Timeout(d) if d == Duration::from_millis(250)
                ));
            }
            _ => panic!("expected CompletionMode::Any"),
        }
    }

    #[test]
    #[allow(clippy::type_complexity)]
    fn test_correlation_strategy_fn_clone_shares_same_arc() {
        let f: Arc<dyn Fn(&Exchange) -> Option<String> + Send + Sync> =
            Arc::new(|_| Some("shared".to_string()));
        let strategy = CorrelationStrategy::Fn(f.clone());
        let cloned = strategy.clone();

        match cloned {
            CorrelationStrategy::Fn(cloned_fn) => assert!(Arc::ptr_eq(&f, &cloned_fn)),
            _ => panic!("expected fn strategy"),
        }
    }

    #[test]
    fn test_builder_correlate_by_overrides_previous() {
        let config = AggregatorConfig::correlate_by("first")
            .correlate_by("second")
            .complete_when_size(2)
            .build()
            .unwrap();

        assert_eq!(config.header_name, "second");
        assert!(
            matches!(config.correlation, CorrelationStrategy::HeaderName(ref h) if h == "second")
        );
    }

    #[test]
    fn test_aggregator_try_build_missing_completion_returns_error() {
        let result = AggregatorConfig::correlate_by("key").try_build();
        assert!(result.is_err());
    }

    // ── R3-C1 Batch 1: DoS caps + completion-bound validation ────────

    /// The builder default for `max_buckets` is `Some(10_000)`. The spec fixes
    /// this as the bounded default; the operator may still override.
    #[test]
    fn test_default_max_buckets_is_10000() {
        let cfg = AggregatorConfig::correlate_by("k")
            .complete_when_size(1)
            .build()
            .unwrap();
        assert_eq!(cfg.max_buckets, Some(10_000));
    }

    /// The builder default for `bucket_ttl` is `Some(Duration::from_secs(300))`.
    /// Inline retain + background sweep both use this TTL.
    #[test]
    fn test_default_bucket_ttl_is_300s() {
        let cfg = AggregatorConfig::correlate_by("k")
            .complete_when_size(1)
            .build()
            .unwrap();
        assert_eq!(cfg.bucket_ttl, Some(Duration::from_secs(300)));
    }

    /// Configs that override `max_buckets(0)` still build (the operator chose it).
    /// The cap may be set to 1 by the operator; the bound itself is not validated.
    /// The completion-bound check below is what Batch 1 enforces.
    #[test]
    fn test_explicit_max_buckets_zero_is_accepted_at_build() {
        let cfg = AggregatorConfig::correlate_by("k")
            .complete_when_size(1)
            .max_buckets(0)
            .build()
            .unwrap();
        assert_eq!(cfg.max_buckets, Some(0));
    }

    /// A config with no completion bound — neither size, nor timeout, nor predicate —
    /// is rejected at `try_build` with `AggregatorMissingCompletionBound`. Spec §11
    /// RESOLVED: at least one completion bound is mandatory.
    #[test]
    fn test_aggregator_rejects_no_completion_bound() {
        // Builder has no `complete_*` call → try_build returns Err.
        // Use match (not unwrap_err) because AggregatorConfig is not Debug.
        let err = match AggregatorConfig::correlate_by("k").try_build() {
            Err(e) => e,
            Ok(_) => panic!("expected error, got Ok"),
        };
        assert!(
            matches!(
                err,
                CamelError::ConfigValidation(
                    ConfigValidationError::AggregatorMissingCompletionBound
                )
            ),
            "expected ConfigValidation(AggregatorMissingCompletionBound), got: {err}"
        );
    }

    // ── R3-M2: memory-bound validation ────────────────────────────────

    #[test]
    fn test_aggregator_config_rejects_no_memory_bound() {
        // Direct construction bypassing the builder defaults — simulate a config
        // with no max_buckets, no timeout, no ttl.
        let config = AggregatorConfig {
            header_name: "k".into(),
            completion: CompletionMode::Single(CompletionCondition::Size(2)),
            correlation: CorrelationStrategy::HeaderName("k".into()),
            strategy: AggregationStrategy::CollectAll,
            max_buckets: None,
            bucket_ttl: None,
            force_completion_on_stop: false,
            discard_on_timeout: false,
            max_timeout_tasks: 1024,
        };
        let err = config.validate().unwrap_err();
        assert!(
            err.to_string().contains("max_buckets")
                || err.to_string().contains("completionTimeout")
                || err.to_string().contains("bucket_ttl"),
            "error should explain the required bound: {err}"
        );
    }

    #[test]
    fn test_aggregator_config_accepts_size_only_with_max_buckets() {
        // Builder path defaults max_buckets + bucket_ttl — must validate OK.
        let config = AggregatorConfig::correlate_by("k")
            .complete_when_size(2)
            .build()
            .unwrap();
        assert!(config.validate().is_ok());
    }

    /// R3-M3: Timeout completion requires bucket_ttl so the cap-skip
    /// degradation always has an eviction path. A config with Timeout but
    /// no bucket_ttl is rejected by validate().
    #[test]
    fn test_aggregator_timeout_requires_bucket_ttl() {
        let config = AggregatorConfig {
            header_name: "k".into(),
            completion: CompletionMode::Single(CompletionCondition::Timeout(Duration::from_secs(
                5,
            ))),
            correlation: CorrelationStrategy::HeaderName("k".into()),
            strategy: AggregationStrategy::CollectAll,
            max_buckets: Some(100),
            bucket_ttl: None, // <-- missing ttl fallback
            force_completion_on_stop: false,
            discard_on_timeout: false,
            max_timeout_tasks: 1024,
        };
        let err = config.validate().unwrap_err();
        assert!(
            err.to_string().contains("bucket_ttl") || err.to_string().contains("Timeout"),
            "error should explain the timeout-requires-ttl invariant: {err}"
        );
    }
}