saluki_components/transforms/apm_stats/

mod.rs

//! APM Stats transform.
//!
//! Aggregates traces into time-bucketed statistics, producing `TraceStats` events.

use std::{
    sync::Arc,
    time::{Duration, SystemTime, UNIX_EPOCH},
};

use agent_data_plane_config::domains;
use async_trait::async_trait;
use datadog_agent_apm::semantics::SemanticRegistryProvider;
use saluki_core::{
    accounting::{MemoryBounds, MemoryBoundsBuilder},
    components::{transforms::*, BuildContext},
    data_model::{
        event::{
            trace::{AttributeValue, Trace},
            trace_stats::{ClientStatsPayload, TraceStats},
            Event, EventType,
        },
        origin::OriginTagCardinality,
        tags::TagSet,
    },
    topology::OutputDefinition,
};
use saluki_env::{
    host::providers::BoxedHostProvider, workload::EntityId, EnvironmentProvider, HostProvider, WorkloadProvider,
};
use saluki_error::{ErrorContext as _, GenericError};
use stringtheory::MetaString;
use tokio::{select, time::interval};
use tracing::{debug, error};

use crate::common::otlp::util::extract_container_tags_from_attributes_map;
use crate::config::AgentEnv;

mod aggregation;
pub(crate) use self::aggregation::{process_tags_hash, PayloadAggregationKey};

mod peer_ip_quantize;

mod peer_tags;

mod span_concentrator;
pub(crate) use self::span_concentrator::{InfraTags, SpanConcentrator};

mod statsraw;

mod weight;
use self::weight::weight;

/// Default flush interval for the APM stats transform.
const DEFAULT_FLUSH_INTERVAL: Duration = Duration::from_secs(10);

/// Tag key for process tags in span meta.
const TAG_PROCESS_TAGS: &str = "_dd.tags.process";

/// Maximum number of `ClientGroupedStats` entries per `TraceStats` event.
const MAX_STATS_GROUPS_PER_EVENT: usize = 4000;

/// APM Stats transform configuration.
///
/// Aggregates incoming `Trace` events into time-bucketed statistics, emitting
/// `TraceStats` events.
pub struct ApmStatsTransformConfiguration {
    compute_stats_by_span_kind: bool,
    peer_tags_aggregation: bool,
    peer_tags: Vec<MetaString>,
    agent_env: AgentEnv,
    default_hostname: Option<String>,
    workload_provider: Option<Arc<dyn WorkloadProvider + Send + Sync>>,
    semantic_registry: SemanticRegistryProvider,
}

impl ApmStatsTransformConfiguration {
    /// Creates a new `ApmStatsTransformConfiguration` from the resolved traces configuration.
    ///
    /// Traces whose spans and payload carry no environment are aggregated under `agent_env`.
    pub fn from_configuration(config: &domains::traces::Domain, agent_env: AgentEnv) -> Self {
        Self {
            compute_stats_by_span_kind: config.compute_stats_by_span_kind,
            peer_tags_aggregation: config.peer_tags_aggregation,
            peer_tags: config.peer_tags.iter().cloned().map(MetaString::from).collect(),
            agent_env,
            default_hostname: None,
            workload_provider: None,
            semantic_registry: SemanticRegistryProvider::default(),
        }
    }

    /// Sets the default hostname using the environment provider.
    pub async fn with_environment_provider<E>(mut self, env_provider: E) -> Result<Self, GenericError>
    where
        E: EnvironmentProvider<Host = BoxedHostProvider>,
    {
        let hostname = env_provider.host().get_hostname().await?;
        self.default_hostname = Some(hostname);
        Ok(self)
    }

    /// Sets the workload provider.
    ///
    /// Defaults to unset.
    pub fn with_workload_provider<W>(mut self, workload_provider: W) -> Self
    where
        W: WorkloadProvider + Send + Sync + 'static,
    {
        self.workload_provider = Some(Arc::new(workload_provider));
        self
    }

    /// Sets the provider of semantic mappings used to choose peer tag names for stats grouping.
    ///
    /// Peer tags describe the remote service or resource a span calls. The transform combines their names from the
    /// registry with configured custom peer tags. With a subscribed provider, it checks for changed mappings before
    /// each input event buffer; all spans in that buffer use the same key set.
    ///
    /// Defaults to embedded mappings without remote updates.
    pub fn with_semantic_registry(mut self, semantic_registry: SemanticRegistryProvider) -> Self {
        self.semantic_registry = semantic_registry;
        self
    }
}

#[async_trait]
impl TransformBuilder for ApmStatsTransformConfiguration {
    async fn build(&self, _context: BuildContext) -> Result<Box<dyn Transform + Send>, GenericError> {
        let agent_hostname = MetaString::from(self.default_hostname.clone().unwrap_or_default());
        let concentrator = SpanConcentrator::new(
            self.compute_stats_by_span_kind,
            self.peer_tags_aggregation,
            &self.peer_tags,
            self.semantic_registry.clone(),
            now_nanos(),
        );

        Ok(Box::new(ApmStats {
            concentrator,
            flush_interval: DEFAULT_FLUSH_INTERVAL,
            agent_env: self.agent_env.clone().into(),
            agent_hostname,
            workload_provider: self.workload_provider.clone(),
        }))
    }

    fn input_event_type(&self) -> EventType {
        EventType::Trace
    }

    fn outputs(&self) -> &[OutputDefinition<EventType>] {
        static OUTPUTS: &[OutputDefinition<EventType>] = &[OutputDefinition::default_output(EventType::TraceStats)];
        OUTPUTS
    }
}

impl MemoryBounds for ApmStatsTransformConfiguration {
    fn specify_bounds(&self, builder: &mut MemoryBoundsBuilder) {
        builder.minimum().with_single_value::<ApmStats>("component struct");
        // TODO: Think about everything we need to account for here.
    }
}

struct ApmStats {
    concentrator: SpanConcentrator,
    flush_interval: Duration,
    agent_env: MetaString,
    agent_hostname: MetaString,
    workload_provider: Option<Arc<dyn WorkloadProvider + Send + Sync>>,
}

impl ApmStats {
    fn process_trace(&mut self, trace: &Trace) {
        let root_span = trace
            .spans()
            .iter()
            .find(|s| s.parent_id() == 0)
            .or_else(|| trace.spans().first());

        let trace_weight = root_span.map(weight).unwrap_or(1.0);

        let process_tags = extract_process_tags(trace);

        let payload_key = self.build_payload_key(trace, &process_tags);
        let infra_tags = self.build_infra_tags(trace, &process_tags);

        let origin = trace
            .spans()
            .first()
            .and_then(|s| s.attributes.get("_dd.origin").and_then(AttributeValue::as_string))
            .map(|s| s.as_ref())
            .unwrap_or("");

        for span in trace.spans() {
            if let Some(stat_span) = self.concentrator.new_stat_span_from_span(span) {
                self.concentrator
                    .add_span(&stat_span, trace_weight, &payload_key, &infra_tags, origin);
            }
        }
    }

    fn build_infra_tags(&self, trace: &Trace, process_tags: &str) -> InfraTags {
        let container_id = trace.payload.container_id.clone();
        let mut container_tags = if container_id.is_empty() {
            TagSet::default()
        } else {
            let mut tags = TagSet::default();
            extract_container_tags_from_attributes_map(&trace.attributes, &mut tags);
            tags
        };

        if !container_id.is_empty() {
            if let Some(workload_provider) = &self.workload_provider {
                let entity_id = EntityId::Container(container_id.clone());
                if let Some(tags) = workload_provider.get_tags_for_entity(&entity_id, OriginTagCardinality::Low) {
                    container_tags.merge_shared(&tags);
                }
            }
        }

        InfraTags::new(container_id, container_tags, process_tags)
    }

    fn build_payload_key(&self, trace: &Trace, process_tags: &str) -> PayloadAggregationKey {
        let root_span = trace
            .spans()
            .iter()
            .find(|s| s.parent_id() == 0)
            .or_else(|| trace.spans().first());

        // env fallback order mirrors get_trace_env (Go agent pkg/trace/traceutil/trace.go:GetEnv):
        // root span attrs → all span attrs → trace.payload.env (ADP extension) → agent_env.
        let env = root_span
            .and_then(|s| {
                s.attributes
                    .get("env")
                    .and_then(AttributeValue::as_string)
                    .filter(|s| !s.is_empty())
            })
            .cloned()
            .or_else(|| {
                trace.spans().iter().find_map(|s| {
                    s.attributes
                        .get("env")
                        .and_then(AttributeValue::as_string)
                        .filter(|s| !s.is_empty())
                        .cloned()
                })
            })
            .unwrap_or_else(|| {
                if !trace.payload.env.is_empty() {
                    trace.payload.env.clone()
                } else {
                    self.agent_env.clone()
                }
            });

        let hostname = root_span
            .and_then(|s| {
                s.attributes
                    .get("_dd.hostname")
                    .and_then(AttributeValue::as_string)
                    .filter(|s| !s.is_empty())
            })
            .cloned()
            .unwrap_or_else(|| {
                if !trace.payload.hostname.is_empty() {
                    trace.payload.hostname.clone()
                } else {
                    self.agent_hostname.clone()
                }
            });

        // Version resolution mirrors Go agent pkg/trace/version/version.go:GetAppVersionFromTrace:
        // root span meta["version"] (set by otel_span_to_dd_span with span-over-resource precedence)
        // takes priority, falling back to the resource-level payload.app_version.
        let version = root_span
            .and_then(|s| {
                s.attributes
                    .get("version")
                    .and_then(AttributeValue::as_string)
                    .filter(|s| !s.is_empty())
            })
            .cloned()
            .unwrap_or_else(|| trace.payload.app_version.clone());

        let container_id = if !trace.payload.container_id.is_empty() {
            trace.payload.container_id.clone()
        } else {
            root_span
                .and_then(|s| s.attributes.get("_dd.container_id").and_then(AttributeValue::as_string))
                .cloned()
                .unwrap_or_default()
        };

        let git_commit_sha = root_span
            .and_then(|s| {
                s.attributes
                    .get("_dd.git.commit.sha")
                    .and_then(AttributeValue::as_string)
                    .filter(|s| !s.is_empty())
            })
            .cloned()
            .unwrap_or_default();

        let image_tag = root_span
            .and_then(|s| {
                s.attributes
                    .get("_dd.image_tag")
                    .and_then(AttributeValue::as_string)
                    .filter(|s| !s.is_empty())
            })
            .cloned()
            .unwrap_or_default();

        let lang = if !trace.payload.language_name.is_empty() {
            trace.payload.language_name.clone()
        } else {
            root_span
                .and_then(|s| s.attributes.get("language").and_then(AttributeValue::as_string))
                .cloned()
                .unwrap_or_default()
        };

        PayloadAggregationKey {
            env,
            hostname,
            version,
            container_id,
            git_commit_sha,
            image_tag,
            lang,
            process_tags_hash: process_tags_hash(process_tags),
        }
    }
}

/// Splits stats payloads into multiple `TraceStats` events, each containing at most `max_entries_per_event` grouped
/// entries.
///
/// This function won't attempt to collapse/pack payloads to optimize the number of events generated: there will always
/// be at least as many output client payloads as there are input client payloads.
fn split_into_trace_stats(client_payloads: Vec<ClientStatsPayload>, max_entries_per_event: usize) -> Vec<TraceStats> {
    if client_payloads.is_empty() {
        return Vec::new();
    }

    // If the total number of grouped stats entries is below the specified threshold, then we don't do any splitting or
    // collapsing or anything.
    let total_grouped_entries = client_payloads
        .iter()
        .map(|p| p.stats().iter().map(|b| b.stats().len()).sum::<usize>())
        .sum::<usize>();
    if total_grouped_entries <= max_entries_per_event {
        return vec![TraceStats::new(client_payloads)];
    }

    let mut events = Vec::new();
    let mut current_client_payloads = Vec::new();
    let mut current_event_len = 0;

    for mut client_payload in client_payloads {
        // If the current payload can fit entirely in the current event, then collect it as-is.
        let client_payload_len = client_payload.stats().iter().map(|b| b.stats().len()).sum::<usize>();
        if current_event_len + client_payload_len <= max_entries_per_event {
            current_client_payloads.push(client_payload);
            current_event_len += client_payload_len;
            continue;
        }

        // Consume all the stats buckets from the current client payload, and set ourselves up to go through each
        // bucket, splitting them as necessary.
        //
        // We basically iterate until we find a bucket where adding its entries would cause us to exceed our threshold,
        // and then we split that bucket, creating a new client payload in the process, and keep repeating that until
        // we've exhausted all buckets for our starting client payload.
        let mut current_client_stats_buckets = Vec::new();
        for mut client_stats_bucket in client_payload.take_stats() {
            let bucket_len = client_stats_bucket.stats().len();
            // If this bucket can fit in the current event, then collect it as-is.
            if current_event_len + bucket_len <= max_entries_per_event {
                current_client_stats_buckets.push(client_stats_bucket);
                current_event_len += bucket_len;
                continue;
            }

            // We have to split this bucket. We take the grouped entries it has, and we subdivide those into a new
            // client bucket, or buckets in order to ensure we don't exceed our threshold.
            let mut bucket_entries = client_stats_bucket.take_stats();
            while current_event_len + bucket_entries.len() > max_entries_per_event {
                // We calculate the split point this way because the returned vector from `split_off` is referenced
                // against the end of the vector, so if we have 100 items, and we only want 20, we need to split at
                // index 80 (100 - 20 == 80).
                let split_amount = max_entries_per_event - current_event_len;
                let split_point = bucket_entries.len() - split_amount;
                let split_entries = bucket_entries.split_off(split_point);

                // Create a new "split" bucket based on the current bucket (`client_stats_bucket`) but containing only
                // the split-off entries. This will feed into the current client payload (`client_payload`) which we'll
                // finalize and add to the current event before starting a new event.
                let split_bucket = client_stats_bucket.clone().with_stats(split_entries);
                current_client_stats_buckets.push(split_bucket);

                let split_client_payload = client_payload
                    .clone()
                    .with_stats(std::mem::take(&mut current_client_stats_buckets));
                current_client_payloads.push(split_client_payload);

                events.push(TraceStats::new(std::mem::take(&mut current_client_payloads)));
                current_event_len = 0;
            }

            // If we have any leftover entries from the bucket after splitting it up, put them back in the current bucket
            // and add that bucket to the current client payload.
            if !bucket_entries.is_empty() {
                current_event_len += bucket_entries.len();
                current_client_stats_buckets.push(client_stats_bucket.with_stats(bucket_entries));
            }
        }

        // If we have buckets for this client payload (whether we split or not), add them back to the current client payload.
        if !current_client_stats_buckets.is_empty() {
            current_client_payloads.push(client_payload.with_stats(current_client_stats_buckets));
        }
    }

    // Stick the remaining entries into a final `TraceStats` event, if any.
    if !current_client_payloads.is_empty() {
        events.push(TraceStats::new(current_client_payloads));
    }

    events
}

#[async_trait]
impl Transform for ApmStats {
    async fn run(mut self: Box<Self>, mut context: TransformContext) -> Result<(), GenericError> {
        let mut health = context.take_health_handle();

        let mut flush_ticker = interval(self.flush_interval);
        flush_ticker.tick().await;

        let mut final_flush = false;

        health.mark_ready();
        debug!("APM Stats transform started.");

        loop {
            select! {
                _ = health.live() => continue,

                _ = flush_ticker.tick() => {
                    let stats_payloads = self.concentrator.flush(now_nanos(), final_flush);
                    if !stats_payloads.is_empty() {
                        debug!(stats_payloads = stats_payloads.len(), "Flushing APM stats.");

                        let events = split_into_trace_stats(stats_payloads, MAX_STATS_GROUPS_PER_EVENT);
                        let dispatcher = context.dispatcher().buffered()
                            .error_context("Default output should be available.")?;

                        if let Err(e) = dispatcher.send_all(events.into_iter().map(Event::TraceStats)).await {
                            error!(error = %e, "Failed to dispatch events.");
                        }
                    }

                    if final_flush {
                        debug!("Final APM stats flush complete.");
                        break;
                    }
                },

                maybe_events = context.events().next(), if !final_flush => {
                    match maybe_events {
                        Some(events) => {
                            // Use one peer tag key set for the whole buffer.
                            self.concentrator.refresh_registry();
                            for event in events {
                                if let Event::Trace(trace) = event {
                                    self.process_trace(&trace);
                                }
                            }
                        },
                        None => {
                            // We've reached the end of our input stream, so mark ourselves for a final flush and reset the
                            // interval so it ticks immediately on the next loop iteration.
                            final_flush = true;
                            flush_ticker.reset_immediately();
                            debug!("APM Stats transform stopping, triggering final flush...");
                        }
                    }
                },
            }
        }

        debug!("APM Stats transform stopped.");
        Ok(())
    }
}

/// Returns the current time as nanoseconds since Unix epoch.
fn now_nanos() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_nanos() as u64
}

/// Extracts process tags from trace, checking both span and trace attributes.
fn extract_process_tags(trace: &Trace) -> MetaString {
    let root_span = trace
        .spans()
        .iter()
        .find(|s| s.parent_id() == 0)
        .or_else(|| trace.spans().first());
    if let Some(span) = root_span {
        if let Some(tags) = span
            .attributes
            .get(TAG_PROCESS_TAGS)
            .and_then(AttributeValue::as_string)
            .filter(|s| !s.is_empty())
        {
            return tags.clone();
        }
    }
    if let Some(AttributeValue::String(tags)) = trace.attributes.get(TAG_PROCESS_TAGS) {
        if !tags.is_empty() {
            return tags.clone();
        }
    }
    MetaString::empty()
}

#[cfg(test)]
mod tests {
    use datadog_agent_apm::semantics::test_util::TestRegistryPublisher;
    use datadog_agent_apm::semantics::{test_util::registry_json, Registry};
    use proptest::prelude::*;
    use saluki_common::collections::FastHashMap;
    use saluki_core::components::test_util::TestComponentDriver;
    use saluki_core::data_model::event::trace::{AttributeValue, Span, TraceIngestInfo};
    use saluki_core::data_model::event::trace_stats::ClientGroupedStats;
    use saluki_core::data_model::event::trace_stats::ClientStatsBucket;

    use super::aggregation::BUCKET_DURATION_NS;
    use super::span_concentrator::METRIC_PARTIAL_VERSION;
    use super::*;

    /// Helper to align timestamp to bucket boundary
    fn align_ts(ts: u64, bsize: u64) -> u64 {
        ts - ts % bsize
    }

    /// Creates a test span with the given parameters.
    #[allow(clippy::too_many_arguments)]
    fn test_span(
        aligned_now: u64, span_id: u64, parent_id: u64, duration: u64, bucket_offset: u64, service: &str,
        resource: &str, error: i32, meta: Option<FastHashMap<MetaString, MetaString>>,
        metrics: Option<FastHashMap<MetaString, f64>>,
    ) -> Span {
        let bucket_start = aligned_now - bucket_offset * BUCKET_DURATION_NS;
        let start = bucket_start - duration;

        let mut attrs: FastHashMap<MetaString, AttributeValue> = FastHashMap::default();
        if let Some(m) = meta {
            attrs.extend(m.into_iter().map(|(k, v)| (k, AttributeValue::String(v))));
        }
        if let Some(m) = metrics {
            attrs.extend(m.into_iter().map(|(k, v)| (k, AttributeValue::Float(v))));
        }
        Span::new(
            service, "query", resource, "db", span_id, parent_id, start, duration, error,
        )
        .with_attributes(attrs)
    }

    /// Creates a simple measured span for basic tests
    fn make_test_span(service: &str, name: &str, resource: &str) -> Span {
        let mut attrs = FastHashMap::default();
        attrs.insert(MetaString::from("_dd.measured"), AttributeValue::Float(1.0));
        Span::new(service, name, resource, "web", 1, 0, 1000000000, 100000000, 0).with_attributes(attrs)
    }

    /// Creates a top-level span (`parent_id` = 0, has `_top_level` metric)
    fn make_top_level_span(
        aligned_now: u64, span_id: u64, duration: u64, bucket_offset: u64, service: &str, resource: &str, error: i32,
        meta: Option<FastHashMap<MetaString, MetaString>>,
    ) -> Span {
        let mut metrics = FastHashMap::default();
        metrics.insert(MetaString::from("_top_level"), 1.0);
        test_span(
            aligned_now,
            span_id,
            0,
            duration,
            bucket_offset,
            service,
            resource,
            error,
            meta,
            Some(metrics),
        )
    }

    #[tokio::test]
    async fn run_refreshes_peer_tag_keys_before_each_buffer() {
        use saluki_core::components::test_util::TestComponentDriver;

        let now = now_nanos();
        let (publisher, provider) = TestRegistryPublisher::new();
        let config = domains::traces::Domain {
            compute_stats_by_span_kind: true,
            peer_tags_aggregation: true,
            ..Default::default()
        };
        let builder =
            ApmStatsTransformConfiguration::from_configuration(&config, AgentEnv::resolve(None, None, &[], &[]))
                .with_semantic_registry(provider);
        let (control, mut outputs) = TestComponentDriver::transform(builder)
            .await
            .expect("transform should build");

        // The buffer refresh must see a registry published after construction.
        let remapped = registry_json(
            r#"{"peer.service": {"fallbacks": [
                {"name": "custom.remote.service", "provider": "otel", "type": "string"}]}}"#,
        );
        publisher.accept_registry(
            Registry::from_json(&remapped).expect("remapped registry should parse"),
            "remapped",
        );

        let mut meta = FastHashMap::default();
        meta.insert(MetaString::from("span.kind"), MetaString::from("client"));
        meta.insert(MetaString::from("custom.remote.service"), MetaString::from("billing"));
        let span = make_top_level_span(
            align_ts(now, BUCKET_DURATION_NS),
            1,
            100,
            0,
            "svc",
            "resource",
            0,
            Some(meta),
        );
        control
            .send_event(Event::Trace(Trace::new(vec![span], TraceIngestInfo::Datadog)))
            .await;

        // Closing the input makes the run loop flush and stop.
        control.shutdown().await.expect("run should succeed");

        let mut peer_tags = Vec::new();
        for event in outputs.default_output().collect_events().await {
            let Event::TraceStats(stats) = event else { continue };
            for payload in stats.stats() {
                for bucket in payload.stats() {
                    for grouped in bucket.stats() {
                        peer_tags.extend(grouped.peer_tags().iter().map(|tag| tag.to_string()));
                    }
                }
            }
        }
        assert_eq!(peer_tags, ["custom.remote.service:billing"]);
    }

    #[tokio::test]
    async fn traces_without_an_env_aggregate_under_the_agent_env() {
        let agent_env = AgentEnv::resolve(None, Some("Prod"), &[], &[]);
        let config = ApmStatsTransformConfiguration::from_configuration(&domains::traces::Domain::default(), agent_env);
        let (control, mut outputs) = TestComponentDriver::transform(config)
            .await
            .expect("the transform should start");

        let span = make_top_level_span(
            align_ts(now_nanos(), BUCKET_DURATION_NS),
            1,
            100,
            0,
            "svc",
            "resource",
            0,
            None,
        );
        control
            .send_event(Event::Trace(Trace::new(vec![span], TraceIngestInfo::Datadog)))
            .await;
        // Shutting down flushes every bucket.
        control.shutdown().await.expect("the transform should stop cleanly");

        let mut envs = Vec::new();
        for event in outputs.default_output().collect_events().await {
            let Event::TraceStats(stats) = event else {
                panic!("unexpected event: {event:?}");
            };
            envs.extend(stats.stats().iter().map(|payload| payload.env().to_string()));
        }
        assert_eq!(envs, ["prod"]);
    }

    #[test]
    fn test_process_trace_creates_stats() {
        let now = now_nanos();

        let concentrator = SpanConcentrator::new(true, true, &[], Default::default(), now);
        let mut transform = ApmStats {
            concentrator,
            flush_interval: DEFAULT_FLUSH_INTERVAL,
            agent_env: MetaString::from("none"),
            agent_hostname: MetaString::default(),
            workload_provider: None,
        };

        let span = make_test_span("test-service", "test-operation", "test-resource");
        let trace = Trace::new(vec![span], TraceIngestInfo::Datadog);

        transform.process_trace(&trace);

        // Flush and verify we got stats
        let stats = transform.concentrator.flush(now + BUCKET_DURATION_NS * 2, true);
        assert!(!stats.is_empty(), "Expected stats to be produced");
    }

    #[test]
    fn test_weight_applied_to_stats() {
        let now = now_nanos();

        let concentrator = SpanConcentrator::new(true, true, &[], Default::default(), now);
        let mut transform = ApmStats {
            concentrator,
            flush_interval: DEFAULT_FLUSH_INTERVAL,
            agent_env: MetaString::from("none"),
            agent_hostname: MetaString::default(),
            workload_provider: None,
        };

        // Create a span with 0.5 sample rate (weight = 2.0)
        let mut attrs = FastHashMap::default();
        attrs.insert(MetaString::from("_dd.measured"), AttributeValue::Float(1.0));
        attrs.insert(MetaString::from("_sample_rate"), AttributeValue::Float(0.5));

        let span = Span::new(
            "test-service",
            "test-op",
            "test-resource",
            "web",
            1,
            0,
            now,
            100000000,
            0,
        )
        .with_attributes(attrs);

        let trace = Trace::new(vec![span], TraceIngestInfo::Datadog);
        transform.process_trace(&trace);

        let stats = transform.concentrator.flush(now + BUCKET_DURATION_NS * 2, true);
        assert!(!stats.is_empty());

        // The hits should be weighted (approximately 2 due to 0.5 sample rate)
        let bucket = &stats[0].stats()[0];
        let grouped = &bucket.stats()[0];
        // With stochastic rounding, hits could be 1 or 2, but with weight 2.0 it should round to 2
        assert!(grouped.hits() >= 1, "Expected weighted hits");
    }

    #[test]
    fn test_force_flush() {
        let now = now_nanos();
        let aligned_now = align_ts(now, BUCKET_DURATION_NS);

        let mut concentrator = SpanConcentrator::new(true, true, &[], Default::default(), now);

        // Add a span
        let span = make_top_level_span(aligned_now, 1, 50, 5, "A1", "resource1", 0, None);
        let trace = Trace::new(vec![span], TraceIngestInfo::Datadog);

        let payload_key = PayloadAggregationKey {
            env: MetaString::from("test"),
            hostname: MetaString::from("host"),
            ..Default::default()
        };
        let infra_tags = InfraTags::default();

        for span in trace.spans() {
            if let Some(stat_span) = concentrator.new_stat_span_from_span(span) {
                concentrator.add_span(&stat_span, 1.0, &payload_key, &infra_tags, "");
            }
        }

        // ts=0 so that flush always considers buckets not old enough
        let ts: u64 = 0;

        // Without force flush, should skip the bucket
        let stats = concentrator.flush(ts, false);
        assert!(stats.is_empty(), "Non-force flush should return empty");

        // With force flush, should flush buckets regardless of age
        let stats = concentrator.flush(ts, true);
        assert!(!stats.is_empty(), "Force flush should return stats");
        assert_eq!(stats[0].stats().len(), 1, "Should have 1 bucket");
    }

    #[test]
    fn test_ignores_partial_spans() {
        let now = now_nanos();
        let aligned_now = align_ts(now, BUCKET_DURATION_NS);

        let mut concentrator = SpanConcentrator::new(true, true, &[], Default::default(), now);

        // Create a partial span (has _dd.partial_version metric)
        let mut metrics = FastHashMap::default();
        metrics.insert(MetaString::from("_top_level"), 1.0);
        metrics.insert(MetaString::from(METRIC_PARTIAL_VERSION), 830604.0);

        let span = test_span(aligned_now, 1, 0, 50, 5, "A1", "resource1", 0, None, Some(metrics));
        let trace = Trace::new(vec![span], TraceIngestInfo::Datadog);

        let payload_key = PayloadAggregationKey {
            env: MetaString::from("test"),
            hostname: MetaString::from("tracer-hostname"),
            ..Default::default()
        };
        let infra_tags = InfraTags::default();

        for span in trace.spans() {
            if let Some(stat_span) = concentrator.new_stat_span_from_span(span) {
                concentrator.add_span(&stat_span, 1.0, &payload_key, &infra_tags, "");
            }
        }

        // Partial spans should be ignored
        let stats = concentrator.flush(now + BUCKET_DURATION_NS * 3, true);
        assert!(stats.is_empty(), "Partial spans should be ignored");
    }

    #[test]
    fn test_concentrator_stats_totals() {
        let now = now_nanos();
        let aligned_now = align_ts(now, BUCKET_DURATION_NS);

        // Set oldestTs to allow old buckets
        let oldest_ts = aligned_now - 2 * BUCKET_DURATION_NS;
        let mut concentrator = SpanConcentrator::new(true, true, &[], Default::default(), oldest_ts);

        // Build spans spread over time windows
        let spans = vec![
            make_top_level_span(aligned_now, 1, 50, 5, "A1", "resource1", 0, None),
            make_top_level_span(aligned_now, 2, 40, 4, "A1", "resource1", 0, None),
            make_top_level_span(aligned_now, 3, 30, 3, "A1", "resource1", 0, None),
            make_top_level_span(aligned_now, 4, 20, 2, "A1", "resource1", 0, None),
            make_top_level_span(aligned_now, 5, 10, 1, "A1", "resource1", 0, None),
            make_top_level_span(aligned_now, 6, 1, 0, "A1", "resource1", 0, None),
        ];

        let payload_key = PayloadAggregationKey {
            env: MetaString::from("none"),
            ..Default::default()
        };
        let infra_tags = InfraTags::default();

        for span in &spans {
            if let Some(stat_span) = concentrator.new_stat_span_from_span(span) {
                concentrator.add_span(&stat_span, 1.0, &payload_key, &infra_tags, "");
            }
        }

        // Flush all and collect totals
        let all_stats = concentrator.flush(now + BUCKET_DURATION_NS * 10, true);

        let mut total_duration: u64 = 0;
        let mut total_hits: u64 = 0;
        let mut total_errors: u64 = 0;
        let mut total_top_level_hits: u64 = 0;

        for payload in &all_stats {
            for bucket in payload.stats() {
                for grouped in bucket.stats() {
                    total_duration += grouped.duration();
                    total_hits += grouped.hits();
                    total_errors += grouped.errors();
                    total_top_level_hits += grouped.top_level_hits();
                }
            }
        }

        assert_eq!(total_duration, 50 + 40 + 30 + 20 + 10 + 1, "Wrong total duration");
        assert_eq!(total_hits, 6, "Wrong total hits");
        assert_eq!(total_top_level_hits, 6, "Wrong total top level hits");
        assert_eq!(total_errors, 0, "Wrong total errors");
    }

    #[test]
    fn test_root_tag() {
        let now = now_nanos();
        let aligned_now = align_ts(now, BUCKET_DURATION_NS);

        let mut concentrator = SpanConcentrator::new(true, true, &[], Default::default(), now);

        // Root span (parent_id = 0, top_level)
        let mut root_metrics = FastHashMap::default();
        root_metrics.insert(MetaString::from("_top_level"), 1.0);
        let root_span = test_span(
            aligned_now,
            1,
            0,
            40,
            10,
            "A1",
            "resource1",
            0,
            None,
            Some(root_metrics),
        );

        // Non-root but top level span (has _top_level but parent_id != 0)
        let mut top_level_metrics = FastHashMap::default();
        top_level_metrics.insert(MetaString::from("_top_level"), 1.0);
        let top_level_span = test_span(
            aligned_now,
            4,
            1000,
            10,
            10,
            "A1",
            "resource1",
            0,
            None,
            Some(top_level_metrics),
        );

        // Client span (non-root, non-top level, but has span.kind = client)
        let mut client_meta = FastHashMap::default();
        client_meta.insert(MetaString::from("span.kind"), MetaString::from("client"));
        let client_span = test_span(aligned_now, 3, 2, 20, 10, "A1", "resource1", 0, Some(client_meta), None);

        let spans = vec![root_span, top_level_span, client_span];

        let payload_key = PayloadAggregationKey {
            env: MetaString::from("none"),
            ..Default::default()
        };
        let infra_tags = InfraTags::default();

        for span in &spans {
            if let Some(stat_span) = concentrator.new_stat_span_from_span(span) {
                concentrator.add_span(&stat_span, 1.0, &payload_key, &infra_tags, "");
            }
        }

        let stats = concentrator.flush(now + BUCKET_DURATION_NS * 20, true);
        assert!(!stats.is_empty(), "Should have stats");

        // Count grouped stats - should be split by IsTraceRoot
        let mut total_grouped = 0;
        let mut root_count = 0;
        let mut non_root_count = 0;

        for payload in &stats {
            for bucket in payload.stats() {
                for grouped in bucket.stats() {
                    total_grouped += 1;
                    match grouped.is_trace_root() {
                        Some(true) => root_count += 1,
                        Some(false) => non_root_count += 1,
                        None => {}
                    }
                }
            }
        }

        // We expect 3 grouped stats:
        // 1. Root span (is_trace_root = true)
        // 2. Non-root top-level span (is_trace_root = false)
        // 3. Client span (is_trace_root = false, span.kind = client)
        assert_eq!(total_grouped, 3, "Expected 3 grouped stats");
        assert_eq!(root_count, 1, "Expected 1 root span");
        assert_eq!(non_root_count, 2, "Expected 2 non-root spans");
    }

    #[test]
    fn test_compute_stats_through_span_kind_check() {
        let now = now_nanos();

        // Test with compute_stats_by_span_kind DISABLED
        {
            let mut concentrator = SpanConcentrator::new(false, true, &[], Default::default(), now);

            let mut attrs = FastHashMap::default();
            attrs.insert(MetaString::from("_top_level"), AttributeValue::Float(1.0));
            let span = Span::new("myservice", "query", "GET /users", "web", 1, 0, now, 500, 0).with_attributes(attrs);

            let payload_key = PayloadAggregationKey {
                env: MetaString::from("test"),
                ..Default::default()
            };
            let infra_tags = InfraTags::default();

            if let Some(stat_span) = concentrator.new_stat_span_from_span(&span) {
                concentrator.add_span(&stat_span, 1.0, &payload_key, &infra_tags, "");
            }

            // Client span with span.kind=client but no _top_level or _dd.measured
            // Should NOT produce stats when compute_stats_by_span_kind is disabled
            let mut client_attrs = FastHashMap::default();
            client_attrs.insert(
                MetaString::from("span.kind"),
                AttributeValue::String(MetaString::from("client")),
            );
            let client_span = Span::new("myservice", "postgres.query", "SELECT ...", "db", 2, 1, now, 75, 0)
                .with_attributes(client_attrs);

            if let Some(stat_span) = concentrator.new_stat_span_from_span(&client_span) {
                concentrator.add_span(&stat_span, 1.0, &payload_key, &infra_tags, "");
            }

            let stats = concentrator.flush(now + BUCKET_DURATION_NS * 3, true);

            let mut count = 0;
            for payload in &stats {
                for bucket in payload.stats() {
                    count += bucket.stats().len();
                }
            }

            // When disabled, only top_level span gets stats (client span has no top_level/measured)
            assert_eq!(count, 1, "Expected 1 stat when span kind check disabled");
        }

        // Test with compute_stats_by_span_kind ENABLED
        {
            let mut concentrator = SpanConcentrator::new(true, true, &[], Default::default(), now);

            let mut attrs = FastHashMap::default();
            attrs.insert(MetaString::from("_top_level"), AttributeValue::Float(1.0));
            let span = Span::new("myservice", "query", "GET /users", "web", 1, 0, now, 500, 0).with_attributes(attrs);

            let payload_key = PayloadAggregationKey {
                env: MetaString::from("test"),
                ..Default::default()
            };
            let infra_tags = InfraTags::default();

            if let Some(stat_span) = concentrator.new_stat_span_from_span(&span) {
                concentrator.add_span(&stat_span, 1.0, &payload_key, &infra_tags, "");
            }

            // Client span with span.kind=client
            // SHOULD produce stats when compute_stats_by_span_kind is enabled
            let mut client_attrs = FastHashMap::default();
            client_attrs.insert(
                MetaString::from("span.kind"),
                AttributeValue::String(MetaString::from("client")),
            );
            let client_span = Span::new("myservice", "postgres.query", "SELECT ...", "db", 2, 1, now, 75, 0)
                .with_attributes(client_attrs);

            if let Some(stat_span) = concentrator.new_stat_span_from_span(&client_span) {
                concentrator.add_span(&stat_span, 1.0, &payload_key, &infra_tags, "");
            }

            let stats = concentrator.flush(now + BUCKET_DURATION_NS * 3, true);

            let mut count = 0;
            for payload in &stats {
                for bucket in payload.stats() {
                    count += bucket.stats().len();
                }
            }

            // When enabled, both spans get stats
            assert_eq!(count, 2, "Expected 2 stats when span kind check enabled");
        }
    }

    #[test]
    fn test_peer_tags() {
        let now = now_nanos();

        // Test without peer tags aggregation enabled
        {
            let mut concentrator = SpanConcentrator::new(true, false, &[], Default::default(), now);

            let mut attrs = FastHashMap::default();
            attrs.insert(
                MetaString::from("span.kind"),
                AttributeValue::String(MetaString::from("client")),
            );
            attrs.insert(
                MetaString::from("db.instance"),
                AttributeValue::String(MetaString::from("i-1234")),
            );
            attrs.insert(
                MetaString::from("db.system"),
                AttributeValue::String(MetaString::from("postgres")),
            );
            attrs.insert(MetaString::from("_dd.measured"), AttributeValue::Float(1.0));
            let client_span =
                Span::new("myservice", "postgres.query", "SELECT ...", "db", 2, 1, now, 75, 0).with_attributes(attrs);

            let payload_key = PayloadAggregationKey {
                env: MetaString::from("test"),
                ..Default::default()
            };
            let infra_tags = InfraTags::default();

            if let Some(stat_span) = concentrator.new_stat_span_from_span(&client_span) {
                concentrator.add_span(&stat_span, 1.0, &payload_key, &infra_tags, "");
            }

            let stats = concentrator.flush(now + BUCKET_DURATION_NS * 3, true);

            // Without peer tags aggregation, peer_tags should be empty
            for payload in &stats {
                for bucket in payload.stats() {
                    for grouped in bucket.stats() {
                        assert!(
                            grouped.peer_tags().is_empty(),
                            "Peer tags should be empty when peer_tags_aggregation is false"
                        );
                    }
                }
            }
        }

        // Test with peer tags aggregation enabled
        {
            // Note: BASE_PEER_TAGS already includes db.instance and db.system
            let mut concentrator = SpanConcentrator::new(true, true, &[], Default::default(), now);

            let mut attrs = FastHashMap::default();
            attrs.insert(
                MetaString::from("span.kind"),
                AttributeValue::String(MetaString::from("client")),
            );
            attrs.insert(
                MetaString::from("db.instance"),
                AttributeValue::String(MetaString::from("i-1234")),
            );
            attrs.insert(
                MetaString::from("db.system"),
                AttributeValue::String(MetaString::from("postgres")),
            );
            attrs.insert(MetaString::from("_dd.measured"), AttributeValue::Float(1.0));
            let client_span =
                Span::new("myservice", "postgres.query", "SELECT ...", "db", 2, 1, now, 75, 0).with_attributes(attrs);

            let payload_key = PayloadAggregationKey {
                env: MetaString::from("test"),
                ..Default::default()
            };
            let infra_tags = InfraTags::default();

            if let Some(stat_span) = concentrator.new_stat_span_from_span(&client_span) {
                concentrator.add_span(&stat_span, 1.0, &payload_key, &infra_tags, "");
            }

            let stats = concentrator.flush(now + BUCKET_DURATION_NS * 3, true);

            // With peer tags aggregation, client span should have peer_tags
            let mut found_client_with_peer_tags = false;
            for payload in &stats {
                for bucket in payload.stats() {
                    for grouped in bucket.stats() {
                        if grouped.resource() == "SELECT ..." {
                            assert!(!grouped.peer_tags().is_empty(), "Client span should have peer tags");
                            // Check that peer tags contain db.instance and db.system
                            let peer_tags: Vec<&str> = grouped.peer_tags().iter().map(|s| s.as_ref()).collect();
                            assert!(
                                peer_tags.iter().any(|t| t.starts_with("db.instance:")),
                                "Should have db.instance peer tag"
                            );
                            assert!(
                                peer_tags.iter().any(|t| t.starts_with("db.system:")),
                                "Should have db.system peer tag"
                            );
                            found_client_with_peer_tags = true;
                        }
                    }
                }
            }
            assert!(
                found_client_with_peer_tags,
                "Should have found client span with peer tags"
            );
        }
    }

    #[test]
    fn test_concentrator_oldest_ts() {
        let now = now_nanos();
        let aligned_now = align_ts(now, BUCKET_DURATION_NS);

        // Test "cold" scenario - all spans in the past should end up in current bucket
        {
            // Start concentrator at current time (cold start)
            let mut concentrator = SpanConcentrator::new(true, true, &[], Default::default(), now);

            // Build spans spread over many time windows (all in the past)
            let spans = vec![
                make_top_level_span(aligned_now, 1, 50, 5, "A1", "resource1", 0, None),
                make_top_level_span(aligned_now, 2, 40, 4, "A1", "resource1", 0, None),
                make_top_level_span(aligned_now, 3, 30, 3, "A1", "resource1", 0, None),
                make_top_level_span(aligned_now, 4, 20, 2, "A1", "resource1", 0, None),
                make_top_level_span(aligned_now, 5, 10, 1, "A1", "resource1", 0, None),
                make_top_level_span(aligned_now, 6, 1, 0, "A1", "resource1", 0, None),
            ];

            let payload_key = PayloadAggregationKey {
                env: MetaString::from("none"),
                ..Default::default()
            };
            let infra_tags = InfraTags::default();

            for span in &spans {
                if let Some(stat_span) = concentrator.new_stat_span_from_span(span) {
                    concentrator.add_span(&stat_span, 1.0, &payload_key, &infra_tags, "");
                }
            }

            // Flush multiple times without force
            let mut flush_time = now;
            let buffer_len = 2; // DEFAULT_BUFFER_LEN

            for _ in 0..buffer_len {
                let stats = concentrator.flush(flush_time, false);
                assert!(stats.is_empty(), "Should not flush before buffer fills");
                flush_time += BUCKET_DURATION_NS;
            }

            // After buffer_len flushes, should get aggregated stats
            let stats = concentrator.flush(flush_time, false);
            assert!(!stats.is_empty(), "Should flush after buffer fills");

            // All spans should be aggregated into one bucket (oldest bucket aggregates old data)
            let mut total_hits: u64 = 0;
            let mut total_duration: u64 = 0;
            for payload in &stats {
                for bucket in payload.stats() {
                    for grouped in bucket.stats() {
                        total_hits += grouped.hits();
                        total_duration += grouped.duration();
                    }
                }
            }

            assert_eq!(total_hits, 6, "All 6 spans should be counted");
            assert_eq!(
                total_duration,
                50 + 40 + 30 + 20 + 10 + 1,
                "Total duration should match"
            );
        }
    }

    #[test]
    fn test_compute_stats_for_span_kind() {
        use super::span_concentrator::compute_stats_for_span_kind;

        // Valid span kinds (case insensitive)
        assert!(compute_stats_for_span_kind("server"));
        assert!(compute_stats_for_span_kind("consumer"));
        assert!(compute_stats_for_span_kind("client"));
        assert!(compute_stats_for_span_kind("producer"));

        // Uppercase
        assert!(compute_stats_for_span_kind("SERVER"));
        assert!(compute_stats_for_span_kind("CONSUMER"));
        assert!(compute_stats_for_span_kind("CLIENT"));
        assert!(compute_stats_for_span_kind("PRODUCER"));

        // Mixed case
        assert!(compute_stats_for_span_kind("SErVER"));
        assert!(compute_stats_for_span_kind("COnSUMER"));
        assert!(compute_stats_for_span_kind("CLiENT"));
        assert!(compute_stats_for_span_kind("PRoDUCER"));

        // Invalid span kinds
        assert!(!compute_stats_for_span_kind("internal"));
        assert!(!compute_stats_for_span_kind("INTERNAL"));
        assert!(!compute_stats_for_span_kind("INtERNAL"));
        assert!(!compute_stats_for_span_kind(""));
    }

    #[test]
    fn test_extract_process_tags() {
        // Test with no process tags
        {
            let span = Span::default();
            let trace = Trace::new(vec![span], TraceIngestInfo::Datadog);
            let process_tags = extract_process_tags(&trace);
            assert!(process_tags.is_empty(), "Should be empty when no _dd.tags.process");
        }

        // Test with process tags in first span meta
        {
            let mut attrs = FastHashMap::default();
            attrs.insert(
                MetaString::from(TAG_PROCESS_TAGS),
                AttributeValue::String(MetaString::from("a:1,b:2,c:3")),
            );
            let span = Span::default().with_attributes(attrs);
            let trace = Trace::new(vec![span], TraceIngestInfo::Datadog);
            let process_tags = extract_process_tags(&trace);
            assert_eq!(process_tags, "a:1,b:2,c:3");
        }

        // Test with empty process tags
        {
            let mut attrs = FastHashMap::default();
            attrs.insert(
                MetaString::from(TAG_PROCESS_TAGS),
                AttributeValue::String(MetaString::from("")),
            );
            let span = Span::default().with_attributes(attrs);
            let trace = Trace::new(vec![span], TraceIngestInfo::Datadog);
            let process_tags = extract_process_tags(&trace);
            assert!(
                process_tags.is_empty(),
                "Should be empty when _dd.tags.process is empty string"
            );
        }

        // Test with empty trace
        {
            let trace = Trace::new(vec![], TraceIngestInfo::Datadog);
            let process_tags = extract_process_tags(&trace);
            assert!(process_tags.is_empty(), "Should be empty when trace has no spans");
        }
    }

    #[test]
    fn test_process_tags_hash_computation() {
        use super::aggregation::process_tags_hash;

        // Empty string should return 0
        assert_eq!(process_tags_hash(""), 0);

        // Same tags should produce same hash
        let hash1 = process_tags_hash("a:1,b:2,c:3");
        let hash2 = process_tags_hash("a:1,b:2,c:3");
        assert_eq!(hash1, hash2);

        // Different tags should produce different hash
        let hash3 = process_tags_hash("a:1,b:2");
        assert_ne!(hash1, hash3);
    }

    // Helper to create a ClientGroupedStats for testing
    fn make_grouped_stats(service: &str, resource: &str) -> ClientGroupedStats {
        ClientGroupedStats::new(service, "operation", resource)
            .with_hits(1)
            .with_duration(100)
    }

    // Helper to create a ClientStatsBucket with N stats
    fn make_bucket_with_stats(n: usize) -> ClientStatsBucket {
        let stats: Vec<ClientGroupedStats> = (0..n)
            .map(|i| make_grouped_stats("service", &format!("resource-{}", i)))
            .collect();
        ClientStatsBucket::new(1000, 10_000_000_000, stats)
    }

    // Helper to create a ClientStatsPayload with specified buckets
    fn make_payload_with_buckets(hostname: &str, buckets: Vec<ClientStatsBucket>) -> ClientStatsPayload {
        ClientStatsPayload::new(hostname, "test-env", "1.0.0")
            .with_stats(buckets)
            .with_container_id("container-123")
            .with_lang("rust")
    }

    // Helper to count total ClientGroupedStats across all payloads in a TraceStats
    fn count_grouped_stats(trace_stats: &TraceStats) -> usize {
        trace_stats
            .stats()
            .iter()
            .flat_map(|p| p.stats())
            .map(|b| b.stats().len())
            .sum()
    }

    #[test]
    fn test_split_into_trace_stats_empty_input() {
        let result = split_into_trace_stats(vec![], 100);
        assert!(result.is_empty());
    }

    #[test]
    fn test_split_into_trace_stats_no_split_needed() {
        // 50 stats with max 100 - no split needed
        let bucket = make_bucket_with_stats(50);
        let payload = make_payload_with_buckets("host1", vec![bucket]);

        let result = split_into_trace_stats(vec![payload], 100);

        assert_eq!(result.len(), 1);
        assert_eq!(count_grouped_stats(&result[0]), 50);
    }

    #[test]
    fn test_split_into_trace_stats_exact_threshold() {
        // Exactly 100 stats with max 100 - no split needed
        let bucket = make_bucket_with_stats(100);
        let payload = make_payload_with_buckets("host1", vec![bucket]);

        let result = split_into_trace_stats(vec![payload], 100);

        assert_eq!(result.len(), 1);
        assert_eq!(count_grouped_stats(&result[0]), 100);
    }

    #[test]
    fn test_split_into_trace_stats_splits_single_bucket() {
        // 250 stats in one bucket with max 100 - should split into 3 TraceStats
        let bucket = make_bucket_with_stats(250);
        let payload = make_payload_with_buckets("host1", vec![bucket]);

        let result = split_into_trace_stats(vec![payload], 100);

        assert_eq!(result.len(), 3);
        assert_eq!(count_grouped_stats(&result[0]), 100);
        assert_eq!(count_grouped_stats(&result[1]), 100);
        assert_eq!(count_grouped_stats(&result[2]), 50);

        // Total should be preserved
        let total: usize = result.iter().map(count_grouped_stats).sum();
        assert_eq!(total, 250);
    }

    #[test]
    fn test_split_into_trace_stats_splits_across_payloads() {
        // Two payloads with 60 stats each, max 100 - should combine into 2 TraceStats
        let payload1 = make_payload_with_buckets("host1", vec![make_bucket_with_stats(60)]);
        let payload2 = make_payload_with_buckets("host2", vec![make_bucket_with_stats(60)]);

        let result = split_into_trace_stats(vec![payload1, payload2], 100);

        assert_eq!(result.len(), 2);
        // First TraceStats: 60 from payload1 + 40 from payload2 = 100
        assert_eq!(count_grouped_stats(&result[0]), 100);
        // Second TraceStats: remaining 20 from payload2
        assert_eq!(count_grouped_stats(&result[1]), 20);
    }

    #[test]
    fn test_split_into_trace_stats_splits_single_payload_multiple_buckets() {
        // One payload with two buckets (70 + 80 = 150 stats), max 100
        let bucket1 = make_bucket_with_stats(70);
        let bucket2 = make_bucket_with_stats(80);
        let payload = make_payload_with_buckets("host1", vec![bucket1, bucket2]);

        let result = split_into_trace_stats(vec![payload], 100);

        assert_eq!(result.len(), 2);
        let total: usize = result.iter().map(count_grouped_stats).sum();
        assert_eq!(total, 150);
    }

    #[test]
    fn test_split_into_trace_stats_preserves_metadata() {
        let bucket = make_bucket_with_stats(250);
        let payload = ClientStatsPayload::new("test-host", "prod", "2.0.0")
            .with_stats(vec![bucket])
            .with_container_id("container-abc")
            .with_lang("go")
            .with_git_commit_sha("abc123")
            .with_image_tag("v1.2.3")
            .with_process_tags_hash(12345)
            .with_process_tags("tag1,tag2");

        let result = split_into_trace_stats(vec![payload], 100);

        // All split payloads should have the same metadata
        for trace_stats in &result {
            for p in trace_stats.stats() {
                assert_eq!(p.hostname(), "test-host");
                assert_eq!(p.env(), "prod");
                assert_eq!(p.version(), "2.0.0");
                assert_eq!(p.container_id(), "container-abc");
                assert_eq!(p.lang(), "go");
                assert_eq!(p.git_commit_sha(), "abc123");
                assert_eq!(p.image_tag(), "v1.2.3");
                assert_eq!(p.process_tags_hash(), 12345);
                assert_eq!(p.process_tags(), "tag1,tag2");
            }
        }
    }

    #[test]
    fn test_split_into_trace_stats_preserves_bucket_metadata() {
        // Create a bucket with specific start/duration/time_shift
        let stats: Vec<ClientGroupedStats> = (0..150)
            .map(|i| make_grouped_stats("svc", &format!("res-{}", i)))
            .collect();
        let bucket = ClientStatsBucket::new(999_000_000, 10_000_000_000, stats).with_agent_time_shift(42);
        let payload = make_payload_with_buckets("host1", vec![bucket]);

        let result = split_into_trace_stats(vec![payload], 100);

        // All split buckets should have the same start/duration/time_shift
        for trace_stats in &result {
            for p in trace_stats.stats() {
                for b in p.stats() {
                    assert_eq!(b.start(), 999_000_000);
                    assert_eq!(b.duration(), 10_000_000_000);
                    assert_eq!(b.agent_time_shift(), 42);
                }
            }
        }
    }

    #[test]
    fn test_split_into_trace_stats_handles_empty_bucket() {
        let empty_bucket = ClientStatsBucket::new(1000, 10_000_000_000, vec![]);
        let payload = make_payload_with_buckets("host1", vec![empty_bucket]);

        let result = split_into_trace_stats(vec![payload], 100);

        assert_eq!(result.len(), 1);
        assert_eq!(count_grouped_stats(&result[0]), 0);
    }

    #[test]
    fn test_split_into_trace_stats_large_split() {
        // 10,000 stats with max 4000 - should produce 3 TraceStats
        let bucket = make_bucket_with_stats(10_000);
        let payload = make_payload_with_buckets("host1", vec![bucket]);

        let result = split_into_trace_stats(vec![payload], 4000);

        assert_eq!(result.len(), 3);
        assert_eq!(count_grouped_stats(&result[0]), 4000);
        assert_eq!(count_grouped_stats(&result[1]), 4000);
        assert_eq!(count_grouped_stats(&result[2]), 2000);
    }

    // Property test strategies

    /// Strategy to generate arbitrary ClientGroupedStats.
    fn arb_grouped_stats() -> impl Strategy<Value = ClientGroupedStats> {
        (0..100u64, 0..1000u64).prop_map(|(hits, duration)| {
            ClientGroupedStats::new("service", "operation", "resource")
                .with_hits(hits)
                .with_duration(duration)
        })
    }

    /// Strategy to generate a bucket with 0..=max_stats_per_bucket grouped stats.
    fn arb_bucket(max_stats_per_bucket: usize) -> impl Strategy<Value = ClientStatsBucket> {
        proptest::collection::vec(arb_grouped_stats(), 0..=max_stats_per_bucket)
            .prop_map(|stats| ClientStatsBucket::new(1000, 10_000_000_000, stats))
    }

    /// Strategy to generate a payload with 1..=max_buckets buckets.
    fn arb_payload(max_buckets: usize, max_stats_per_bucket: usize) -> impl Strategy<Value = ClientStatsPayload> {
        proptest::collection::vec(arb_bucket(max_stats_per_bucket), 1..=max_buckets)
            .prop_map(|buckets| ClientStatsPayload::new("host", "env", "1.0.0").with_stats(buckets))
    }

    /// Strategy to generate test inputs for the split function.
    ///
    /// Parameters:
    /// - `num_payloads`: 1..=10
    /// - `num_buckets_per_payload`: 1..=5
    /// - `num_stats_per_bucket`: 0..=500
    /// - `max_entries_per_event`: 1..=1000 (never zero)
    fn arb_split_inputs() -> impl Strategy<Value = (Vec<ClientStatsPayload>, usize)> {
        let payloads_strategy = proptest::collection::vec(arb_payload(5, 500), 1..=10);
        let max_entries_strategy = 1..=1000usize;

        (payloads_strategy, max_entries_strategy)
    }

    // Mirrors Go agent behavior: for OTLP traces the version on the root span's meta["version"]
    // (set with span-over-resource precedence by otel_span_to_dd_span) should win over the
    // resource-level version carried in payload.app_version.
    // See: pkg/trace/version/version.go:GetAppVersionFromTrace and pkg/trace/api/otlp.go.
    #[test]
    fn test_version_span_beats_resource_for_otlp() {
        let now = now_nanos();
        let concentrator = SpanConcentrator::new(true, true, &[], Default::default(), now);
        let transform = ApmStats {
            concentrator,
            flush_interval: DEFAULT_FLUSH_INTERVAL,
            agent_env: MetaString::default(),
            agent_hostname: MetaString::default(),
            workload_provider: None,
        };

        // Root span carries a span-level version attribute (set by otel_span_to_dd_span with
        // span-over-resource precedence).
        let mut attrs = FastHashMap::default();
        attrs.insert(
            MetaString::from("version"),
            AttributeValue::String(MetaString::from("span-v2")),
        );
        let root_span = Span::new("svc", "op", "res", "web", 1, 0, now, 1_000_000, 0).with_attributes(attrs);

        let mut trace = Trace::new(vec![root_span], TraceIngestInfo::Datadog);
        // Simulate OTLP resource extraction: payload.app_version comes from resource attrs only.
        trace.payload.app_version = MetaString::from("resource-v1");

        let key = transform.build_payload_key(&trace, "");

        // The Go agent resolves version from root_span.meta["version"] first, which carries
        // span-over-resource precedence. The span attribute ("span-v2") must win.
        assert_eq!(
            key.version.as_ref(),
            "span-v2",
            "span-level version must take precedence over resource-level payload.app_version for OTLP traces"
        );
    }

    proptest! {
        #[test]
        #[cfg_attr(miri, ignore)]
        fn property_test_split_respects_max_entries((payloads, max_entries_per_event) in arb_split_inputs()) {
            let input_total: usize = payloads.iter().flat_map(|p| p.stats()).map(|b| b.stats().len()).sum();

            let result = split_into_trace_stats(payloads, max_entries_per_event);

            // Property 1: No TraceStats should exceed max_entries_per_event
            for trace_stats in &result {
                let count = count_grouped_stats(trace_stats);
                prop_assert!(
                    count <= max_entries_per_event,
                    "TraceStats has {} grouped stats, exceeds max of {}",
                    count,
                    max_entries_per_event
                );
            }

            // Property 2: Total stats should be preserved
            let output_total: usize = result.iter().map(count_grouped_stats).sum();
            prop_assert_eq!(input_total, output_total, "Total stats count should be preserved");
        }
    }
}