vector/sinks/datadog/traces/apm_stats/
weight.rs

1use std::collections::BTreeMap;
2
3use crate::event::{ObjectMap, Value};
4
5const SAMPLING_RATE_KEY: &str = "_sample_rate";
6
7/// This extracts the relative weights from the top level span (i.e. the span that does not have a parent).
8pub(crate) fn extract_weight_from_root_span(spans: &[&ObjectMap]) -> f64 {
9    // Based on https://github.com/DataDog/datadog-agent/blob/cfa750c7412faa98e87a015f8ee670e5828bbe7f/pkg/trace/stats/weight.go#L17-L26.
10
11    // TODO this logic likely has a bug(s) that need to be root caused. The root span is not reliably found and defaults to "1.0"
12    // regularly for users even when sampling is disabled in the Agent.
13    // GH issue to track that: https://github.com/vectordotdev/vector/issues/14859
14
15    if spans.is_empty() {
16        return 1.0;
17    }
18
19    let mut trace_id: Option<usize> = None;
20
21    let mut parent_id_to_child_weight = BTreeMap::<i64, f64>::new();
22    let mut span_ids = Vec::<i64>::new();
23    for s in spans.iter() {
24        // TODO these need to change to u64 when the following issue is fixed:
25        // https://github.com/vectordotdev/vector/issues/14687
26        let parent_id = match s.get("parent_id") {
27            Some(Value::Integer(val)) => *val,
28            None => 0,
29            _ => panic!("`parent_id` should be an i64"),
30        };
31        let span_id = match s.get("span_id") {
32            Some(Value::Integer(val)) => *val,
33            None => 0,
34            _ => panic!("`span_id` should be an i64"),
35        };
36        if trace_id.is_none() {
37            trace_id = match s.get("trace_id") {
38                Some(Value::Integer(v)) => Some(*v as usize),
39                _ => panic!("`trace_id` should be an i64"),
40            }
41        }
42        let weight = s
43            .get("metrics")
44            .and_then(|m| m.as_object())
45            .map(|m| match m.get(SAMPLING_RATE_KEY) {
46                Some(Value::Float(v)) => {
47                    let sample_rate = v.into_inner();
48                    if sample_rate <= 0.0 || sample_rate > 1.0 {
49                        1.0
50                    } else {
51                        1.0 / sample_rate
52                    }
53                }
54                _ => 1.0,
55            })
56            .unwrap_or(1.0);
57
58        // found root
59        if parent_id == 0 {
60            return weight;
61        }
62
63        span_ids.push(span_id);
64
65        parent_id_to_child_weight.insert(parent_id, weight);
66    }
67
68    // Remove all spans that have a parent
69    span_ids.iter().for_each(|id| {
70        parent_id_to_child_weight.remove(id);
71    });
72
73    // There should be only one value remaining, the weight from the root span
74    if parent_id_to_child_weight.len() != 1 {
75        // TODO remove the debug print and emit the Error event as outlined in
76        // https://github.com/vectordotdev/vector/issues/14859
77        debug!(
78            "Didn't reliably find the root span for weight calculation of trace_id {:?}.",
79            trace_id
80        );
81    }
82
83    *parent_id_to_child_weight
84        .values()
85        .next()
86        .unwrap_or_else(|| {
87            // TODO remove the debug print and emit the Error event as outlined in
88            // https://github.com/vectordotdev/vector/issues/14859
89            debug!(
90                "Root span was not found. Defaulting to weight of 1.0 for trace_id {:?}.",
91                trace_id
92            );
93            &1.0
94        })
95}
vector/sinks/datadog/traces/apm_stats/weight.rs

vector/sinks/datadog/traces/apm_stats/
weight.rs
