Skip to content

Commit 779ffdd

Browse files
alocayswcollard
authored andcommitted
feat: adding config option for trace sampling (#370)
* feat: adding config option for trace sampling * chore: adding changeset entry
1 parent 198aeb9 commit 779ffdd

File tree

3 files changed

+119
-0
lines changed

3 files changed

+119
-0
lines changed
Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
### feat: adding config option for trace sampling - @alocay PR #366
2+
3+
Adding configuration option to sample traces. Can use the following options:
4+
1. Ratio based samples (ratio >= 1 is always sample)
5+
2. Always on
6+
3. Always off
7+
8+
Defaults to always on if not provided.

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
1+
mod sampler;
2+
13
use crate::runtime::Config;
24
use crate::runtime::filtering_exporter::FilteringExporter;
35
use crate::runtime::logging::Logging;
6+
use crate::runtime::telemetry::sampler::SamplerOption;
47
use apollo_mcp_server::generated::telemetry::TelemetryAttribute;
58
use opentelemetry::{Key, KeyValue, global, trace::TracerProvider as _};
69
use opentelemetry_otlp::WithExportConfig;
@@ -59,6 +62,7 @@ impl Default for OTLPMetricExporter {
5962
#[derive(Debug, Deserialize, JsonSchema)]
6063
pub struct TracingExporters {
6164
otlp: Option<OTLPTracingExporter>,
65+
sampler: Option<SamplerOption>,
6266
omitted_attributes: Option<HashSet<TelemetryAttribute>>,
6367
}
6468

@@ -190,6 +194,12 @@ fn init_tracer_provider(telemetry: &Telemetry) -> Result<SdkTracerProvider, anyh
190194
}
191195
};
192196

197+
let sampler: opentelemetry_sdk::trace::Sampler = tracer_exporters
198+
.as_ref()
199+
.and_then(|e| e.sampler.clone())
200+
.unwrap_or_default()
201+
.into();
202+
193203
let omitted_attributes: HashSet<Key> = tracer_exporters
194204
.and_then(|exporters| exporters.omitted_attributes.clone())
195205
.map(|set| set.iter().map(|a| a.to_key()).collect())
@@ -201,6 +211,7 @@ fn init_tracer_provider(telemetry: &Telemetry) -> Result<SdkTracerProvider, anyh
201211
.with_id_generator(RandomIdGenerator::default())
202212
.with_resource(resource(telemetry))
203213
.with_batch_exporter(filtering_exporter)
214+
.with_sampler(sampler)
204215
.build();
205216

206217
Ok(tracer_provider)
@@ -301,6 +312,7 @@ mod tests {
301312
}),
302313
Some(TracingExporters {
303314
otlp: Some(OTLPTracingExporter::default()),
315+
sampler: Default::default(),
304316
omitted_attributes: Some(ommitted),
305317
}),
306318
);
@@ -344,6 +356,7 @@ mod tests {
344356
protocol: "bogus".to_string(),
345357
endpoint: "http://localhost:4317".to_string(),
346358
}),
359+
sampler: Default::default(),
347360
omitted_attributes: None,
348361
}),
349362
);
Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
1+
use schemars::JsonSchema;
2+
use serde::Deserialize;
3+
4+
#[derive(Clone, Debug, Deserialize, JsonSchema)]
5+
#[serde(deny_unknown_fields, untagged)]
6+
pub(crate) enum SamplerOption {
7+
/// Sample a given fraction. Fractions >= 1 will always sample.
8+
RatioBased(f64),
9+
Always(Sampler),
10+
}
11+
12+
#[derive(Clone, Debug, Deserialize, JsonSchema)]
13+
#[serde(deny_unknown_fields, rename_all = "snake_case")]
14+
pub(crate) enum Sampler {
15+
/// Always sample
16+
AlwaysOn,
17+
/// Never sample
18+
AlwaysOff,
19+
}
20+
21+
impl From<Sampler> for opentelemetry_sdk::trace::Sampler {
22+
fn from(s: Sampler) -> Self {
23+
match s {
24+
Sampler::AlwaysOn => opentelemetry_sdk::trace::Sampler::AlwaysOn,
25+
Sampler::AlwaysOff => opentelemetry_sdk::trace::Sampler::AlwaysOff,
26+
}
27+
}
28+
}
29+
30+
impl From<SamplerOption> for opentelemetry_sdk::trace::Sampler {
31+
fn from(s: SamplerOption) -> Self {
32+
match s {
33+
SamplerOption::Always(s) => s.into(),
34+
SamplerOption::RatioBased(ratio) => {
35+
opentelemetry_sdk::trace::Sampler::TraceIdRatioBased(ratio)
36+
}
37+
}
38+
}
39+
}
40+
41+
impl Default for SamplerOption {
42+
fn default() -> Self {
43+
SamplerOption::Always(Sampler::AlwaysOn)
44+
}
45+
}
46+
47+
#[cfg(test)]
48+
mod tests {
49+
use super::*;
50+
51+
#[test]
52+
fn sampler_always_on_maps_to_otel_always_on() {
53+
assert!(matches!(
54+
Sampler::AlwaysOn.into(),
55+
opentelemetry_sdk::trace::Sampler::AlwaysOn
56+
));
57+
}
58+
59+
#[test]
60+
fn sampler_always_off_maps_to_otel_always_off() {
61+
assert!(matches!(
62+
Sampler::AlwaysOff.into(),
63+
opentelemetry_sdk::trace::Sampler::AlwaysOff
64+
));
65+
}
66+
67+
#[test]
68+
fn sampler_option_always_on_maps_to_otel_always_on() {
69+
assert!(matches!(
70+
SamplerOption::Always(Sampler::AlwaysOn).into(),
71+
opentelemetry_sdk::trace::Sampler::AlwaysOn
72+
));
73+
}
74+
75+
#[test]
76+
fn sampler_option_always_off_maps_to_otel_always_off() {
77+
assert!(matches!(
78+
SamplerOption::Always(Sampler::AlwaysOff).into(),
79+
opentelemetry_sdk::trace::Sampler::AlwaysOff
80+
));
81+
}
82+
83+
#[test]
84+
fn sampler_option_ratio_based_maps_to_otel_ratio_based_sampler() {
85+
assert!(matches!(
86+
SamplerOption::RatioBased(0.5).into(),
87+
opentelemetry_sdk::trace::Sampler::TraceIdRatioBased(0.5)
88+
));
89+
}
90+
91+
#[test]
92+
fn default_sampler_option_is_always_on() {
93+
assert!(matches!(
94+
SamplerOption::default(),
95+
SamplerOption::Always(Sampler::AlwaysOn)
96+
));
97+
}
98+
}

0 commit comments

Comments
 (0)