feat: adding config option for trace sampling

.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.

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<OTLPTracingExporter>,
+    sampler: Option<SamplerOption>,
     omitted_attributes: Option<HashSet<TelemetryAttribute>>,
 }
 
@@ -190,6 +194,12 @@ fn init_tracer_provider(telemetry: &Telemetry) -> Result<SdkTracerProvider, anyh
         }
     };
 
+    let sampler: opentelemetry_sdk::trace::Sampler = tracer_exporters
+        .as_ref()
+        .and_then(|e| e.sampler.clone())
+        .unwrap_or_default()
+        .into();
+
     let omitted_attributes: HashSet<Key> = 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<SdkTracerProvider, anyh
         .with_id_generator(RandomIdGenerator::default())
         .with_resource(resource(telemetry))
         .with_batch_exporter(filtering_exporter)
+        .with_sampler(sampler)
         .build();
 
     Ok(tracer_provider)
@@ -301,6 +312,7 @@ mod tests {
             }),
             Some(TracingExporters {
                 otlp: Some(OTLPTracingExporter::default()),
+                sampler: Default::default(),
                 omitted_attributes: Some(ommitted),
             }),
         );
@@ -344,6 +356,7 @@ mod tests {
                     protocol: "bogus".to_string(),
                     endpoint: "http://localhost:4317".to_string(),
                 }),
+                sampler: Default::default(),
                 omitted_attributes: None,
             }),
         );

crates/apollo-mcp-server/src/runtime/telemetry/sampler.rs

@@ -0,0 +1,98 @@
+use schemars::JsonSchema;
+use serde::Deserialize;
+
+#[derive(Clone, Debug, Deserialize, JsonSchema)]
+#[serde(deny_unknown_fields, untagged)]
+pub(crate) enum SamplerOption {
+    /// Sample a given fraction. Fractions >= 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<Sampler> 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<SamplerOption> 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)
+        ));
+    }
+}