apollo-opentelemetry 0.8.0

//! Span sampling configuration.

use apollo_configuration::ErrorCollector;
use apollo_configuration::Validate;
use apollo_configuration::configuration;

/// Sampler configuration for controlling which spans are recorded.
#[configuration]
#[derive(Default)]
pub(crate) enum Sampler {
    /// Always sample all spans.
    #[default]
    AlwaysOn,
    /// Never sample any spans.
    AlwaysOff,
    /// Sample a fixed ratio of traces.
    TraceIdRatioBased(TraceIdRatioBasedSamplerConfig),
    /// Make sampling decisions based on parent span.
    ParentBased(ParentBasedSamplerConfig),
}

/// Configuration for trace ID ratio-based sampling.
#[configuration]
pub(crate) struct TraceIdRatioBasedSamplerConfig {
    /// Sampling ratio between 0.0 and 1.0.
    #[config(required)]
    pub(crate) ratio: SamplingRatio,
}

/// A sampling ratio value between 0.0 and 1.0.
#[derive(Debug, Clone, Copy, serde::Deserialize, schemars::JsonSchema)]
#[serde(transparent)]
pub(crate) struct SamplingRatio(f64);

impl Validate for SamplingRatio {
    fn validate(&self, mut errors: ErrorCollector<'_>) {
        if !(0.0..=1.0).contains(&self.0) {
            errors.report_simple("sampling ratio must be between 0.0 and 1.0");
        }
    }
}

/// Configuration for parent-based sampling.
#[configuration]
pub(crate) struct ParentBasedSamplerConfig {
    /// Sampler for root spans (no parent).
    pub(crate) root: Option<Box<RootSampler>>,

    /// Sampler when parent is remote and sampled.
    pub(crate) remote_parent_sampled: Option<Box<RootSampler>>,

    /// Sampler when parent is remote and not sampled.
    pub(crate) remote_parent_not_sampled: Option<Box<RootSampler>>,

    /// Sampler when parent is local and sampled.
    pub(crate) local_parent_sampled: Option<Box<RootSampler>>,

    /// Sampler when parent is local and not sampled.
    pub(crate) local_parent_not_sampled: Option<Box<RootSampler>>,
}

/// Non-recursive sampler for use in parent-based configuration.
#[configuration]
pub(crate) enum RootSampler {
    /// Always sample.
    AlwaysOn,
    /// Never sample.
    AlwaysOff,
    /// Sample a fixed ratio.
    TraceIdRatioBased(TraceIdRatioBasedSamplerConfig),
}

// --- SDK Conversions ---

use opentelemetry_sdk::trace::Sampler as SdkSampler;

impl From<&Sampler> for SdkSampler {
    fn from(config: &Sampler) -> Self {
        match config {
            Sampler::AlwaysOn => SdkSampler::AlwaysOn,
            Sampler::AlwaysOff => SdkSampler::AlwaysOff,
            Sampler::TraceIdRatioBased(ratio_config) => {
                SdkSampler::TraceIdRatioBased(ratio_config.ratio.0)
            }
            Sampler::ParentBased(parent_config) => {
                // The SDK's ParentBased sampler uses the delegate for root spans.
                // Additional parent-based config options are not exposed in this SDK version.
                let root = parent_config
                    .root
                    .as_ref()
                    .map(|s| SdkSampler::from(s.as_ref()))
                    .unwrap_or(SdkSampler::AlwaysOn);

                SdkSampler::ParentBased(Box::new(root))
            }
        }
    }
}

impl From<&RootSampler> for SdkSampler {
    fn from(config: &RootSampler) -> Self {
        match config {
            RootSampler::AlwaysOn => SdkSampler::AlwaysOn,
            RootSampler::AlwaysOff => SdkSampler::AlwaysOff,
            RootSampler::TraceIdRatioBased(ratio_config) => {
                SdkSampler::TraceIdRatioBased(ratio_config.ratio.0)
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use apollo_configuration::parse_yaml;

    use super::*;
    use crate::config::OpenTelemetryConfig;

    #[test]
    fn always_on_conversion() {
        let config: OpenTelemetryConfig = parse_yaml(
            indoc::indoc! {"
                tracer_provider:
                  sampler: always_on
            "},
            &Default::default(),
        )
        .unwrap();

        let sampler: SdkSampler = (&config.tracer_provider.sampler).into();

        assert!(matches!(sampler, SdkSampler::AlwaysOn));
    }

    #[test]
    fn always_off_conversion() {
        let config: OpenTelemetryConfig = parse_yaml(
            indoc::indoc! {"
                tracer_provider:
                  sampler: always_off
            "},
            &Default::default(),
        )
        .unwrap();

        let sampler: SdkSampler = (&config.tracer_provider.sampler).into();

        assert!(matches!(sampler, SdkSampler::AlwaysOff));
    }

    #[test]
    fn trace_id_ratio_conversion() {
        let config: OpenTelemetryConfig = parse_yaml(
            indoc::indoc! {"
                tracer_provider:
                  sampler:
                    trace_id_ratio_based:
                      ratio: 0.5
            "},
            &Default::default(),
        )
        .unwrap();

        let sampler: SdkSampler = (&config.tracer_provider.sampler).into();

        let SdkSampler::TraceIdRatioBased(ratio) = sampler else {
            panic!("Expected TraceIdRatioBased sampler");
        };
        assert!((ratio - 0.5).abs() < f64::EPSILON);
    }

    #[test]
    fn parent_based_conversion() {
        let config: OpenTelemetryConfig = parse_yaml(
            indoc::indoc! {"
                tracer_provider:
                  sampler:
                    parent_based:
                      root:
                        trace_id_ratio_based:
                          ratio: 0.1
            "},
            &Default::default(),
        )
        .unwrap();

        let sampler: SdkSampler = (&config.tracer_provider.sampler).into();

        assert!(matches!(sampler, SdkSampler::ParentBased(_)));
    }

    #[test]
    fn parse_sampler_always_on() {
        let _config: OpenTelemetryConfig = parse_yaml(
            indoc::indoc! {"
                tracer_provider:
                  sampler: always_on
            "},
            &Default::default(),
        )
        .unwrap();
    }

    #[test]
    fn parse_sampler_ratio() {
        let _config: OpenTelemetryConfig = parse_yaml(
            indoc::indoc! {"
                tracer_provider:
                  sampler:
                    trace_id_ratio_based:
                      ratio: 0.5
            "},
            &Default::default(),
        )
        .unwrap();
    }

    #[test]
    fn parse_sampler_ratio_invalid() {
        let result: Result<OpenTelemetryConfig, _> = parse_yaml(
            indoc::indoc! {"
                tracer_provider:
                  sampler:
                    trace_id_ratio_based:
                      ratio: 1.5
            "},
            &Default::default(),
        );

        assert!(result.is_err());
        let err = format!("{:?}", result.unwrap_err());
        assert!(
            err.contains("sampling ratio must be between 0.0 and 1.0"),
            "unexpected error: {err}"
        );
    }
}