dcontext-tracing 0.8.0

Automatic dcontext scope management via tracing spans
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137

//! Shared logic for AsyncDcontextLayer and SyncDcontextLayer.
//!
//! These helpers handle field extraction into span extensions and span recording —
//! the parts that are identical between async and sync layers and do NOT access
//! the context store directly.
//!
//! Functions that read/write the context store (apply_field_extraction,
//! set_span_info, record_context_to_span) are defined separately in each layer
//! since they target different stores.

use std::cell::Cell;

use tracing_subscriber::registry::{LookupSpan, SpanRef};

use crate::field_mapping::{ExtractedFields, FieldExtractor};
use crate::tracing_field::get_tracing_fields;

// Flag to prevent on_record from processing values recorded by our own layer.
thread_local! {
    pub(crate) static SELF_RECORDING: Cell<bool> = const { Cell::new(false) };
}

/// Extract fields from span attributes and store in span extensions.
/// Called from `on_new_span` in both layers.
pub(crate) fn extract_span_fields<S>(
    attrs: &tracing_core::span::Attributes<'_>,
    span: &SpanRef<'_, S>,
) where
    S: for<'a> LookupSpan<'a>,
{
    let metadata_fields = get_tracing_fields();

    let extract_names: Vec<&'static str> = metadata_fields
        .iter()
        .filter(|e| e.extract.is_some())
        .map(|e| e.span_field)
        .collect();

    let record_fields: Vec<&'static str> = metadata_fields
        .iter()
        .filter(|e| e.span_fmt_fn.is_some())
        .map(|e| e.record_field)
        .collect();

    if extract_names.is_empty() && record_fields.is_empty() {
        return;
    }

    let all_target: Vec<&'static str> = extract_names
        .iter()
        .chain(record_fields.iter())
        .copied()
        .collect();
    let mut extractor = FieldExtractor::new(&all_target);
    attrs.record(&mut extractor);

    for &rf in &record_fields {
        if extractor.extracted.string_values.contains_key(rf)
            || extractor.extracted.u64_values.contains_key(rf)
            || extractor.extracted.i64_values.contains_key(rf)
            || extractor.extracted.bool_values.contains_key(rf)
        {
            extractor.extracted.mark_user_set(rf);
        }
    }

    if !extractor.extracted.is_empty() || !extractor.extracted.user_set_fields.is_empty() {
        let mut extensions = span.extensions_mut();
        extensions.insert(extractor.extracted);
    }
}

/// Merge late-recorded field values into span extensions.
/// Called from `on_record` in both layers.
pub(crate) fn merge_recorded_fields<S>(
    values: &tracing_core::span::Record<'_>,
    span: &SpanRef<'_, S>,
) where
    S: for<'a> LookupSpan<'a>,
{
    if SELF_RECORDING.with(|f| f.get()) {
        return;
    }

    let metadata_fields = get_tracing_fields();
    let extract_names: Vec<&'static str> = metadata_fields
        .iter()
        .filter(|e| e.extract.is_some())
        .map(|e| e.span_field)
        .collect();

    let record_fields: Vec<&'static str> = metadata_fields
        .iter()
        .filter(|e| e.span_fmt_fn.is_some())
        .map(|e| e.record_field)
        .collect();

    if extract_names.is_empty() && record_fields.is_empty() {
        return;
    }

    let all_target: Vec<&'static str> = extract_names
        .iter()
        .chain(record_fields.iter())
        .copied()
        .collect();
    let mut extractor = FieldExtractor::new(&all_target);
    values.record(&mut extractor);

    if !extractor.extracted.is_empty() {
        let mut extensions = span.extensions_mut();

        for &rf in &record_fields {
            if extractor.extracted.string_values.contains_key(rf)
                || extractor.extracted.u64_values.contains_key(rf)
                || extractor.extracted.i64_values.contains_key(rf)
                || extractor.extracted.bool_values.contains_key(rf)
            {
                extractor.extracted.mark_user_set(rf);
            }
        }

        if let Some(existing) = extensions.get_mut::<ExtractedFields>() {
            existing
                .string_values
                .extend(extractor.extracted.string_values);
            existing.u64_values.extend(extractor.extracted.u64_values);
            existing.i64_values.extend(extractor.extracted.i64_values);
            existing.bool_values.extend(extractor.extracted.bool_values);
            existing
                .user_set_fields
                .extend(extractor.extracted.user_set_fields);
        } else {
            extensions.insert(extractor.extracted);
        }
    }
}