diff --git a/.changesets/feat_trace_sampling_option_alocay.md b/.changesets/feat_trace_sampling_option_alocay.md new file mode 100644 index 00000000..38bec085 --- /dev/null +++ b/.changesets/feat_trace_sampling_option_alocay.md @@ -0,0 +1,8 @@ +### feat: adding config option for trace sampling - @alocay PR #366 + +Adding configuration option to sample traces. Can use the following options: +1. Ratio based samples (ratio >= 1 is always sample) +2. Always on +3. Always off + +Defaults to always on if not provided. diff --git a/crates/apollo-mcp-server/src/runtime/telemetry.rs b/crates/apollo-mcp-server/src/runtime/telemetry.rs index 11deee20..bd392702 100644 --- a/crates/apollo-mcp-server/src/runtime/telemetry.rs +++ b/crates/apollo-mcp-server/src/runtime/telemetry.rs @@ -1,6 +1,9 @@ +mod sampler; + use crate::runtime::Config; use crate::runtime::filtering_exporter::FilteringExporter; use crate::runtime::logging::Logging; +use crate::runtime::telemetry::sampler::SamplerOption; use apollo_mcp_server::generated::telemetry::TelemetryAttribute; use opentelemetry::{Key, KeyValue, global, trace::TracerProvider as _}; use opentelemetry_otlp::WithExportConfig; @@ -59,6 +62,7 @@ impl Default for OTLPMetricExporter { #[derive(Debug, Deserialize, JsonSchema)] pub struct TracingExporters { otlp: Option, + sampler: Option, omitted_attributes: Option>, } @@ -190,6 +194,12 @@ fn init_tracer_provider(telemetry: &Telemetry) -> Result = tracer_exporters .and_then(|exporters| exporters.omitted_attributes.clone()) .map(|set| set.iter().map(|a| a.to_key()).collect()) @@ -201,6 +211,7 @@ fn init_tracer_provider(telemetry: &Telemetry) -> Result= 1 will always sample. + RatioBased(f64), + Always(Sampler), +} + +#[derive(Clone, Debug, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields, rename_all = "snake_case")] +pub(crate) enum Sampler { + /// Always sample + AlwaysOn, + /// Never sample + AlwaysOff, +} + +impl From for opentelemetry_sdk::trace::Sampler { + fn from(s: Sampler) -> Self { + match s { + Sampler::AlwaysOn => opentelemetry_sdk::trace::Sampler::AlwaysOn, + Sampler::AlwaysOff => opentelemetry_sdk::trace::Sampler::AlwaysOff, + } + } +} + +impl From for opentelemetry_sdk::trace::Sampler { + fn from(s: SamplerOption) -> Self { + match s { + SamplerOption::Always(s) => s.into(), + SamplerOption::RatioBased(ratio) => { + opentelemetry_sdk::trace::Sampler::TraceIdRatioBased(ratio) + } + } + } +} + +impl Default for SamplerOption { + fn default() -> Self { + SamplerOption::Always(Sampler::AlwaysOn) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn sampler_always_on_maps_to_otel_always_on() { + assert!(matches!( + Sampler::AlwaysOn.into(), + opentelemetry_sdk::trace::Sampler::AlwaysOn + )); + } + + #[test] + fn sampler_always_off_maps_to_otel_always_off() { + assert!(matches!( + Sampler::AlwaysOff.into(), + opentelemetry_sdk::trace::Sampler::AlwaysOff + )); + } + + #[test] + fn sampler_option_always_on_maps_to_otel_always_on() { + assert!(matches!( + SamplerOption::Always(Sampler::AlwaysOn).into(), + opentelemetry_sdk::trace::Sampler::AlwaysOn + )); + } + + #[test] + fn sampler_option_always_off_maps_to_otel_always_off() { + assert!(matches!( + SamplerOption::Always(Sampler::AlwaysOff).into(), + opentelemetry_sdk::trace::Sampler::AlwaysOff + )); + } + + #[test] + fn sampler_option_ratio_based_maps_to_otel_ratio_based_sampler() { + assert!(matches!( + SamplerOption::RatioBased(0.5).into(), + opentelemetry_sdk::trace::Sampler::TraceIdRatioBased(0.5) + )); + } + + #[test] + fn default_sampler_option_is_always_on() { + assert!(matches!( + SamplerOption::default(), + SamplerOption::Always(Sampler::AlwaysOn) + )); + } +}