common_telemetry/
tracing_sampler.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::HashSet;
16
17use opentelemetry::KeyValue;
18use opentelemetry::trace::{Link, SpanKind, TraceContextExt, TraceId, TraceState};
19use opentelemetry_sdk::trace::{Sampler, SamplingDecision, SamplingResult, ShouldSample};
20use serde::{Deserialize, Serialize};
21
22#[derive(Clone, Debug, Serialize, Deserialize)]
23#[serde(default)]
24pub struct TracingSampleOptions {
25    pub default_ratio: f64,
26    pub rules: Vec<TracingSampleRule>,
27}
28
29impl Default for TracingSampleOptions {
30    fn default() -> Self {
31        Self {
32            default_ratio: 1.0,
33            rules: vec![],
34        }
35    }
36}
37
38/// Determine the sampling rate of a span according to the `rules` provided in `RuleSampler`.
39/// For spans that do not hit any `rules`, the `default_ratio` is used.
40#[derive(Clone, Default, Debug, Serialize, Deserialize)]
41#[serde(default)]
42pub struct TracingSampleRule {
43    pub protocol: String,
44    pub request_types: HashSet<String>,
45    pub ratio: f64,
46}
47
48impl TracingSampleRule {
49    pub fn match_rule(&self, protocol: &str, request_type: Option<&str>) -> Option<f64> {
50        if protocol == self.protocol {
51            if self.request_types.is_empty() {
52                Some(self.ratio)
53            } else if let Some(t) = request_type
54                && self.request_types.contains(t)
55            {
56                Some(self.ratio)
57            } else {
58                None
59            }
60        } else {
61            None
62        }
63    }
64}
65
66impl PartialEq for TracingSampleOptions {
67    fn eq(&self, other: &Self) -> bool {
68        self.default_ratio == other.default_ratio && self.rules == other.rules
69    }
70}
71impl PartialEq for TracingSampleRule {
72    fn eq(&self, other: &Self) -> bool {
73        self.protocol == other.protocol
74            && self.request_types == other.request_types
75            && self.ratio == other.ratio
76    }
77}
78
79impl Eq for TracingSampleOptions {}
80impl Eq for TracingSampleRule {}
81
82pub fn create_sampler(opt: &TracingSampleOptions) -> Box<dyn ShouldSample> {
83    if opt.rules.is_empty() {
84        Box::new(Sampler::TraceIdRatioBased(opt.default_ratio))
85    } else {
86        Box::new(opt.clone())
87    }
88}
89
90impl ShouldSample for TracingSampleOptions {
91    fn should_sample(
92        &self,
93        parent_context: Option<&opentelemetry::Context>,
94        trace_id: TraceId,
95        _name: &str,
96        _span_kind: &SpanKind,
97        attributes: &[KeyValue],
98        _links: &[Link],
99    ) -> SamplingResult {
100        let (mut protocol, mut request_type) = (None, None);
101        for kv in attributes {
102            match kv.key.as_str() {
103                "protocol" => protocol = Some(kv.value.as_str()),
104                "request_type" => request_type = Some(kv.value.as_str()),
105                _ => (),
106            }
107        }
108        let ratio = protocol
109            .and_then(|p| {
110                self.rules
111                    .iter()
112                    .find_map(|rule| rule.match_rule(p.as_ref(), request_type.as_deref()))
113            })
114            .unwrap_or(self.default_ratio);
115        SamplingResult {
116            decision: sample_based_on_probability(ratio, trace_id),
117            // No extra attributes ever set by the SDK samplers.
118            attributes: Vec::new(),
119            // all sampler in SDK will not modify trace state.
120            trace_state: match parent_context {
121                Some(ctx) => ctx.span().span_context().trace_state().clone(),
122                None => TraceState::default(),
123            },
124        }
125    }
126}
127
128/// The code here mainly refers to the relevant implementation of
129/// [opentelemetry](https://github.com/open-telemetry/opentelemetry-rust/blob/ef4701055cc39d3448d5e5392812ded00cdd4476/opentelemetry-sdk/src/trace/sampler.rs#L229),
130/// and determines whether the span needs to be collected based on the `TraceId` and sampling rate (i.e. `prob`).
131fn sample_based_on_probability(prob: f64, trace_id: TraceId) -> SamplingDecision {
132    if prob >= 1.0 {
133        SamplingDecision::RecordAndSample
134    } else {
135        let prob_upper_bound = (prob.max(0.0) * (1u64 << 63) as f64) as u64;
136        let bytes = trace_id.to_bytes();
137        let (_, low) = bytes.split_at(8);
138        let trace_id_low = u64::from_be_bytes(low.try_into().unwrap());
139        let rnd_from_trace_id = trace_id_low >> 1;
140
141        if rnd_from_trace_id < prob_upper_bound {
142            SamplingDecision::RecordAndSample
143        } else {
144            SamplingDecision::Drop
145        }
146    }
147}
148
149#[cfg(test)]
150mod test {
151    use std::collections::HashSet;
152
153    use crate::tracing_sampler::TracingSampleRule;
154
155    #[test]
156    fn test_rule() {
157        let rule = TracingSampleRule {
158            protocol: "http".to_string(),
159            request_types: HashSet::new(),
160            ratio: 1.0,
161        };
162        assert_eq!(rule.match_rule("not_http", None), None);
163        assert_eq!(rule.match_rule("http", None), Some(1.0));
164        assert_eq!(rule.match_rule("http", Some("abc")), Some(1.0));
165        let rule1 = TracingSampleRule {
166            protocol: "http".to_string(),
167            request_types: HashSet::from(["mysql".to_string()]),
168            ratio: 1.0,
169        };
170        assert_eq!(rule1.match_rule("http", None), None);
171        assert_eq!(rule1.match_rule("http", Some("abc")), None);
172        assert_eq!(rule1.match_rule("http", Some("mysql")), Some(1.0));
173    }
174}
common_telemetry/tracing_sampler.rs

common_telemetry/
tracing_sampler.rs
