Skip to main content

dcontext_tracing/
field_mapping.rs

1use std::collections::{HashMap, HashSet};
2use std::fmt;
3
4use tracing_core::field::{Field, Visit};
5
6/// Extracted field values from a span, stored in span extensions.
7///
8/// Must be `Send + Sync` to satisfy span extension requirements.
9pub(crate) struct ExtractedFields {
10    pub string_values: HashMap<&'static str, String>,
11    pub u64_values: HashMap<&'static str, u64>,
12    pub i64_values: HashMap<&'static str, i64>,
13    pub bool_values: HashMap<&'static str, bool>,
14    /// Fields that were explicitly set by user code (not auto-recorded).
15    /// Auto-record should not overwrite these.
16    pub user_set_fields: HashSet<&'static str>,
17}
18
19impl ExtractedFields {
20    pub fn new() -> Self {
21        Self {
22            string_values: HashMap::new(),
23            u64_values: HashMap::new(),
24            i64_values: HashMap::new(),
25            bool_values: HashMap::new(),
26            user_set_fields: HashSet::new(),
27        }
28    }
29
30    pub fn is_empty(&self) -> bool {
31        self.string_values.is_empty()
32            && self.u64_values.is_empty()
33            && self.i64_values.is_empty()
34            && self.bool_values.is_empty()
35    }
36
37    /// Mark a field as user-set (not from auto-record).
38    pub fn mark_user_set(&mut self, field: &'static str) {
39        self.user_set_fields.insert(field);
40    }
41}
42
43/// Visitor that extracts field values from span attributes.
44pub(crate) struct FieldExtractor<'a> {
45    /// Set of field names we're interested in.
46    pub target_fields: &'a [&'static str],
47    pub extracted: ExtractedFields,
48}
49
50impl<'a> FieldExtractor<'a> {
51    pub fn new(target_fields: &'a [&'static str]) -> Self {
52        Self {
53            target_fields,
54            extracted: ExtractedFields::new(),
55        }
56    }
57
58    fn is_target(&self, field: &Field) -> Option<&'static str> {
59        self.target_fields
60            .iter()
61            .find(|&&f| f == field.name())
62            .copied()
63    }
64}
65
66impl<'a> Visit for FieldExtractor<'a> {
67    fn record_str(&mut self, field: &Field, value: &str) {
68        if let Some(key) = self.is_target(field) {
69            self.extracted.string_values.insert(key, value.to_string());
70        }
71    }
72
73    fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) {
74        if let Some(key) = self.is_target(field) {
75            // Fall back to Debug formatting for non-string types
76            if !self.extracted.string_values.contains_key(key)
77                && !self.extracted.u64_values.contains_key(key)
78                && !self.extracted.i64_values.contains_key(key)
79                && !self.extracted.bool_values.contains_key(key)
80            {
81                self.extracted
82                    .string_values
83                    .insert(key, format!("{:?}", value));
84            }
85        }
86    }
87
88    fn record_u64(&mut self, field: &Field, value: u64) {
89        if let Some(key) = self.is_target(field) {
90            self.extracted.u64_values.insert(key, value);
91        }
92    }
93
94    fn record_i64(&mut self, field: &Field, value: i64) {
95        if let Some(key) = self.is_target(field) {
96            self.extracted.i64_values.insert(key, value);
97        }
98    }
99
100    fn record_bool(&mut self, field: &Field, value: bool) {
101        if let Some(key) = self.is_target(field) {
102            self.extracted.bool_values.insert(key, value);
103        }
104    }
105}