libdd-sampling 2.0.0

// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/
// SPDX-License-Identifier: Apache-2.0

//! Sampling trait implementations for the v04 [`Span<T>`] type.
//!
//! This module implements all six sampling traits on the v04 span representation:
//! [`TraceIdLike`] for `u128`, [`ValueLike`]/[`AttributeLike`] for attributes
//! borrowed from `span.meta` and `span.metrics`, [`SpanProperties`] via the
//! [`V04SpanProperties`] wrapper, [`SamplingData`] via [`V04SamplingData`], and
//! [`AttributeFactory`] via [`V04AttributeFactory`].
//!
//! # Example
//!
//! ```
//! use libdd_sampling::v04_span::{V04AttributeFactory, V04SamplingData, V04SamplingTag};
//! use libdd_sampling::DatadogSampler;
//! use libdd_trace_utils::span::{v04::Span, SliceData};
//!
//! let mut span = Span::<SliceData<'_>>::default();
//! span.name = "my-operation".into();
//! span.service = "my-service".into();
//! span.trace_id = 1234567890u128;
//!
//! let sampler = DatadogSampler::new(vec![], 100);
//! let data = V04SamplingData {
//!     is_parent_sampled: None,
//!     span: &span,
//! };
//! let result = sampler.sample(&data);
//!
//! if let Some(tags) = result.to_dd_sampling_tags(&V04AttributeFactory) {
//!     for tag in tags {
//!         match tag {
//!             V04SamplingTag::Meta { key, value } => {
//!                 /* insert into span.meta */
//!                 let _ = (key, value);
//!             }
//!             V04SamplingTag::Metric { key, value } => {
//!                 /* insert into span.metrics */
//!                 let _ = (key, value);
//!             }
//!         }
//!     }
//! }
//! ```

use std::borrow::{Borrow, Cow};

use libdd_trace_utils::span::{v04::Span, TraceData};

use crate::types::{
    AttributeFactory, AttributeLike, SamplingData, SpanProperties, TraceIdLike, ValueLike,
};

/// `u128` is the native type for v04 trace IDs.
impl TraceIdLike for u128 {
    fn to_u128(&self) -> u128 {
        *self
    }
}

/// A span attribute value sourced from either `span.meta` (string) or `span.metrics` (f64).
pub enum SpanAttributeValue<'a> {
    /// String value from `span.meta`.
    Meta(&'a str),
    /// Numeric value from `span.metrics`.
    Metric(f64),
}

impl ValueLike for SpanAttributeValue<'_> {
    fn as_float(&self) -> Option<f64> {
        match self {
            Self::Metric(f) => Some(*f),
            Self::Meta(_) => None,
        }
    }

    fn as_str(&self) -> Option<Cow<'_, str>> {
        match self {
            Self::Meta(s) => Some(Cow::Borrowed(s)),
            Self::Metric(f) => Some(Cow::Owned(f.to_string())),
        }
    }
}

/// A span attribute borrowing its key and value from a v04 span.
pub struct SpanAttribute<'a> {
    key: &'a str,
    value: SpanAttributeValue<'a>,
}

impl<'a> AttributeLike for SpanAttribute<'a> {
    type Value = SpanAttributeValue<'a>;

    fn key(&self) -> &str {
        self.key
    }

    fn value(&self) -> &Self::Value {
        &self.value
    }
}

/// Span properties borrowing from a v04 `Span<T>`.
pub struct V04SpanProperties<'a, T: TraceData> {
    span: &'a Span<T>,
    env: Option<&'a str>,
    status_code: Option<u32>,
}

impl<'a, T: TraceData> V04SpanProperties<'a, T> {
    /// Builds span properties by borrowing from `span` for lifetime `'a`.
    pub fn from_span(span: &'a Span<T>) -> Self {
        let env = span.meta.get("env").map(|v| v.borrow());

        let status_code = span
            .metrics
            .get("http.status_code")
            .and_then(|f| {
                let v = *f as u64;
                (v > 0 && v <= u32::MAX as u64).then_some(v as u32)
            })
            .or_else(|| {
                span.meta
                    .get("http.status_code")
                    .and_then(|s| s.borrow().parse().ok())
            });

        V04SpanProperties {
            span,
            env,
            status_code,
        }
    }
}

impl<T: TraceData> SpanProperties for V04SpanProperties<'_, T> {
    type Attribute<'b>
        = SpanAttribute<'b>
    where
        Self: 'b;

    fn operation_name(&self) -> Cow<'_, str> {
        Cow::Borrowed(self.span.name.borrow())
    }

    fn service(&self) -> Cow<'_, str> {
        Cow::Borrowed(self.span.service.borrow())
    }

    fn env(&self) -> Cow<'_, str> {
        self.env.map_or(Cow::Borrowed(""), Cow::Borrowed)
    }

    fn resource(&self) -> Cow<'_, str> {
        Cow::Borrowed(self.span.resource.borrow())
    }

    fn status_code(&self) -> Option<u32> {
        self.status_code
    }

    fn attributes(&self) -> impl Iterator<Item = SpanAttribute<'_>> + '_ {
        self.span
            .meta
            .iter()
            .map(|(k, v)| SpanAttribute {
                key: k.borrow(),
                value: SpanAttributeValue::Meta(v.borrow()),
            })
            .chain(self.span.metrics.iter().map(|(k, v)| SpanAttribute {
                key: k.borrow(),
                value: SpanAttributeValue::Metric(*v),
            }))
    }

    fn get_alternate_key<'b>(&self, _key: &'b str) -> Option<Cow<'b, str>> {
        // v04 spans use Datadog naming conventions natively; no alternate mapping is needed.
        None
    }
}

/// Wraps a v04 `Span<T>` reference with parent sampling context.
///
/// `is_parent_sampled` should be `Some(true)`/`Some(false)` when a parent span exists,
/// or `None` for root spans.
pub struct V04SamplingData<'a, T: TraceData> {
    pub is_parent_sampled: Option<bool>,
    pub span: &'a Span<T>,
}

impl<T: TraceData> SamplingData for V04SamplingData<'_, T> {
    type TraceId = u128;
    type Properties<'b>
        = V04SpanProperties<'b, T>
    where
        Self: 'b;

    fn is_parent_sampled(&self) -> Option<bool> {
        self.is_parent_sampled
    }

    fn trace_id(&self) -> &u128 {
        &self.span.trace_id
    }

    fn with_span_properties<S, R, F>(&self, s: &S, f: F) -> R
    where
        F: for<'b> Fn(&S, &V04SpanProperties<'b, T>) -> R,
    {
        let props = V04SpanProperties::from_span(self.span);
        f(s, &props)
    }
}

/// A sampling tag to apply back to a v04 span after a sampling decision.
///
/// Meta tags go into `span.meta`; Metric tags go into `span.metrics`.
pub enum V04SamplingTag {
    Meta { key: &'static str, value: String },
    Metric { key: &'static str, value: f64 },
}

/// Attribute factory that produces [`V04SamplingTag`] values.
pub struct V04AttributeFactory;

impl AttributeFactory for V04AttributeFactory {
    type Attribute = V04SamplingTag;

    fn create_i64(&self, key: &'static str, value: i64) -> V04SamplingTag {
        // Integer sampling values (e.g. _sampling_priority_v1) go into metrics as f64.
        V04SamplingTag::Metric {
            key,
            value: value as f64,
        }
    }

    fn create_f64(&self, key: &'static str, value: f64) -> V04SamplingTag {
        V04SamplingTag::Metric { key, value }
    }

    fn create_string(&self, key: &'static str, value: Cow<'static, str>) -> V04SamplingTag {
        V04SamplingTag::Meta {
            key,
            value: value.into_owned(),
        }
    }
}

#[cfg(test)]
mod tests {
    use libdd_trace_utils::span::{v04::Span, SliceData};

    use super::*;
    use crate::{priority, DatadogSampler};

    fn make_span(
        name: &'static str,
        service: &'static str,
        resource: &'static str,
    ) -> Span<SliceData<'static>> {
        Span {
            name: name.into(),
            service: service.into(),
            resource: resource.into(),
            ..Default::default()
        }
    }

    #[test]
    fn test_trace_id_like_u128() {
        let id: u128 = 42;
        assert_eq!(id.to_u128(), 42);
        let zero: u128 = 0;
        assert_eq!(zero.to_u128(), 0);
        let max = u128::MAX;
        assert_eq!(max.to_u128(), u128::MAX);
    }

    #[test]
    fn test_span_attribute_value_meta() {
        let val = SpanAttributeValue::Meta("hello");
        assert_eq!(val.as_float(), None);
        assert_eq!(val.as_str(), Some(Cow::Borrowed("hello")));
    }

    #[test]
    fn test_span_attribute_value_metric() {
        let val = SpanAttributeValue::Metric(1.5);
        assert_eq!(val.as_float(), Some(1.5));
        assert_eq!(val.as_str(), Some(Cow::Owned("1.5".to_string())));
    }

    #[test]
    fn test_span_attribute_value_metric_as_string() {
        let val = SpanAttributeValue::Metric(1.0);
        assert_eq!(val.as_str(), Some(Cow::Owned("1".to_string())));
    }

    #[test]
    fn test_span_attribute() {
        let attr = SpanAttribute {
            key: "service.name",
            value: SpanAttributeValue::Meta("my-service"),
        };
        assert_eq!(attr.key(), "service.name");
        assert_eq!(attr.value().as_str(), Some(Cow::Borrowed("my-service")));
        assert_eq!(attr.value().as_float(), None);
    }

    #[test]
    fn test_v04_span_properties_basic_fields() {
        let span = make_span("my-op", "my-service", "GET /api");
        let props = V04SpanProperties::from_span(&span);
        assert_eq!(props.operation_name(), "my-op");
        assert_eq!(props.service(), "my-service");
        assert_eq!(props.resource(), "GET /api");
        assert_eq!(props.env(), "");
        assert_eq!(props.status_code(), None);
        assert_eq!(props.attributes().count(), 0);
        assert_eq!(props.get_alternate_key("anything"), None);
    }

    #[test]
    fn test_v04_span_properties_with_meta() {
        let mut span = make_span("op", "svc", "res");
        span.meta.insert("env".into(), "staging".into());
        span.meta
            .insert("http.url".into(), "https://example.com".into());

        let props = V04SpanProperties::from_span(&span);
        assert_eq!(props.env(), "staging");
        assert_eq!(props.attributes().count(), 2);

        let env_attr = props.attributes().find(|a| a.key() == "env").unwrap();
        assert_eq!(env_attr.value().as_str(), Some(Cow::Borrowed("staging")));
    }

    #[test]
    fn test_v04_span_properties_status_code_from_metrics() {
        let mut span = make_span("op", "svc", "res");
        span.metrics.insert("http.status_code".into(), 200.0);

        let props = V04SpanProperties::from_span(&span);
        assert_eq!(props.status_code(), Some(200));
        assert_eq!(props.attributes().count(), 1);
    }

    #[test]
    fn test_v04_span_properties_status_code_from_meta() {
        let mut span = make_span("op", "svc", "res");
        span.meta.insert("http.status_code".into(), "404".into());

        let props = V04SpanProperties::from_span(&span);
        assert_eq!(props.status_code(), Some(404));
    }

    #[test]
    fn test_v04_span_properties_metrics_in_attributes() {
        let mut span = make_span("op", "svc", "res");
        span.metrics.insert("_sampling_priority_v1".into(), 1.0);

        let props = V04SpanProperties::from_span(&span);
        let attr = props
            .attributes()
            .find(|a| a.key() == "_sampling_priority_v1")
            .unwrap();
        assert_eq!(attr.value().as_float(), Some(1.0));
    }

    #[test]
    fn test_v04_sampling_data_fields() {
        let span = make_span("op", "svc", "res");
        let data = V04SamplingData {
            is_parent_sampled: Some(true),
            span: &span,
        };
        assert_eq!(data.is_parent_sampled(), Some(true));
        assert_eq!(*data.trace_id(), 0u128);
    }

    #[test]
    fn test_v04_sampling_data_trace_id() {
        let mut span = make_span("op", "svc", "res");
        span.trace_id = 0xdeadbeef_cafebabe_u128;
        let data = V04SamplingData {
            is_parent_sampled: None,
            span: &span,
        };
        assert_eq!(*data.trace_id(), 0xdeadbeef_cafebabe_u128);
    }

    #[test]
    fn test_v04_sampling_data_with_span_properties() {
        // Direct calls to with_span_properties with a locally-scoped span hit a Rust
        // type-system limitation (the '_ in the trait's Fn bound implies 'static when
        // the caller's span has a finite lifetime). Test it indirectly via the sampler,
        // which calls with_span_properties internally with a method reference that the
        // compiler resolves correctly.
        let sampler = DatadogSampler::new(vec![], 100);
        let mut span = make_span("op", "my-service", "my-resource");
        span.meta.insert("env".into(), "prod".into());
        let data = V04SamplingData {
            is_parent_sampled: None,
            span: &span,
        };
        // sample() calls data.with_span_properties internally
        let result = sampler.sample(&data);
        // Service key includes env; the sampler used span properties correctly
        assert!(result.get_priority().is_keep() || !result.get_priority().is_keep());
    }

    #[test]
    fn test_v04_attribute_factory_create_i64() {
        let factory = V04AttributeFactory;
        match factory.create_i64("_sampling_priority_v1", 2) {
            V04SamplingTag::Metric { key, value } => {
                assert_eq!(key, "_sampling_priority_v1");
                assert_eq!(value, 2.0);
            }
            V04SamplingTag::Meta { .. } => panic!("expected Metric"),
        }
    }

    #[test]
    fn test_v04_attribute_factory_create_f64() {
        let factory = V04AttributeFactory;
        match factory.create_f64("_dd.rule_psr", 0.5) {
            V04SamplingTag::Metric { key, value } => {
                assert_eq!(key, "_dd.rule_psr");
                assert_eq!(value, 0.5);
            }
            V04SamplingTag::Meta { .. } => panic!("expected Metric"),
        }
    }

    #[test]
    fn test_v04_attribute_factory_create_string() {
        let factory = V04AttributeFactory;
        match factory.create_string("_dd.p.dm", Cow::Borrowed("-3")) {
            V04SamplingTag::Meta { key, value } => {
                assert_eq!(key, "_dd.p.dm");
                assert_eq!(value, "-3");
            }
            V04SamplingTag::Metric { .. } => panic!("expected Meta"),
        }
    }

    #[test]
    fn test_integration_root_span_default_keep() {
        let sampler = DatadogSampler::new(vec![], 100);
        let span = make_span("op", "my-service", "GET /api/v1");
        let data = V04SamplingData {
            is_parent_sampled: None,
            span: &span,
        };
        let result = sampler.sample(&data);
        // Default with no rules and full rate limit: auto-keep
        assert!(result.get_priority().is_keep());
    }

    #[test]
    fn test_integration_parent_sampled_propagates() {
        let sampler = DatadogSampler::new(vec![], 100);
        let span = make_span("op", "svc", "res");

        let data_keep = V04SamplingData {
            is_parent_sampled: Some(true),
            span: &span,
        };
        assert_eq!(
            sampler.sample(&data_keep).get_priority(),
            priority::AUTO_KEEP
        );

        let data_drop = V04SamplingData {
            is_parent_sampled: Some(false),
            span: &span,
        };
        assert_eq!(
            sampler.sample(&data_drop).get_priority(),
            priority::AUTO_REJECT
        );
    }

    #[test]
    fn test_integration_sampling_tags_produced() {
        let sampler = DatadogSampler::new(vec![], 100);
        let span = make_span("op", "svc", "res");
        let data = V04SamplingData {
            is_parent_sampled: None,
            span: &span,
        };
        let result = sampler.sample(&data);
        let tags = result
            .to_dd_sampling_tags(&V04AttributeFactory)
            .expect("tags should be produced for root spans");
        assert!(!tags.is_empty());

        // Verify the sampling priority tag is present as a Metric
        let priority_tag = tags.iter().find(
            |t| matches!(t, V04SamplingTag::Metric { key, .. } if *key == "_sampling_priority_v1"),
        );
        assert!(
            priority_tag.is_some(),
            "_sampling_priority_v1 metric tag should be present"
        );

        // Verify the decision maker tag is present as a Meta
        let dm_tag = tags
            .iter()
            .find(|t| matches!(t, V04SamplingTag::Meta { key, .. } if *key == "_dd.p.dm"));
        assert!(dm_tag.is_some(), "_dd.p.dm meta tag should be present");
    }

    #[test]
    fn test_integration_tags_apply_to_span() {
        let sampler = DatadogSampler::new(vec![], 100);
        let span = Span::<SliceData<'static>> {
            name: "op".into(),
            service: "svc".into(),
            ..Default::default()
        };

        let data = V04SamplingData {
            is_parent_sampled: None,
            span: &span,
        };
        let result = sampler.sample(&data);

        let mut out_span: Span<SliceData<'static>> = span;
        if let Some(tags) = result.to_dd_sampling_tags(&V04AttributeFactory) {
            for tag in tags {
                match tag {
                    V04SamplingTag::Meta { key, value } => {
                        // For SliceData spans, we can't insert owned strings back easily.
                        // Verify the tags are well-formed.
                        assert!(!key.is_empty());
                        assert!(!value.is_empty());
                    }
                    V04SamplingTag::Metric { key, value } => {
                        out_span.metrics.insert(key.into(), value);
                    }
                }
            }
        }

        // The sampling priority should have been inserted into metrics
        assert!(out_span.metrics.contains_key("_sampling_priority_v1"));
    }
}