aleph_alpha_client/
tracing.rs

1use std::iter;
2
3/// Trace context that is propagated through HTTP headers to enable distributed tracing.
4///
5/// Currently still missing support for tracestate, otherwise compliant with
6/// https://www.w3.org/TR/trace-context-2/#design-overview, which standardizes how
7/// context information is sent and modified between services.
8#[derive(Clone, PartialEq, Eq, Hash)]
9pub struct TraceContext {
10    /// https://www.w3.org/TR/trace-context/#trace-id
11    trace_id: u128,
12    /// https://www.w3.org/TR/trace-context/#parent-id
13    span_id: u64,
14    /// https://www.w3.org/TR/trace-context/#sampled-flag
15    sampled: bool,
16    /// https://www.w3.org/TR/trace-context/#tracestate-header
17    state: Option<String>,
18}
19
20impl TraceContext {
21    /// Construct a new trace context.
22    ///
23    /// Version 0 of the trace context specification only supports the `sampled` flag,
24    /// so we can offer a bool input for this flag.
25    ///
26    /// [W3C TraceContext specification]: https://www.w3.org/TR/trace-context-2/#design-overview
27    pub fn new(trace_id: u128, span_id: u64, sampled: bool, state: Option<String>) -> Self {
28        Self {
29            trace_id,
30            span_id,
31            sampled,
32            state,
33        }
34    }
35
36    /// Construct a new trace context with the `sampled` flag set to true.
37    pub fn new_sampled(trace_id: u128, span_id: u64, state: Option<String>) -> Self {
38        Self::new(trace_id, span_id, true, state)
39    }
40
41    /// Construct a new trace context with the `sampled` flag set to false.
42    pub fn new_unsampled(trace_id: u128, span_id: u64, state: Option<String>) -> Self {
43        Self::new(trace_id, span_id, false, state)
44    }
45
46    /// Render the context as w3c trace context headers.
47    ///
48    /// <https://www.w3.org/TR/trace-context-2>
49    pub fn as_w3c_headers(&self) -> impl Iterator<Item = (&'static str, String)> {
50        let mut first = Some(("traceparent", self.traceparent()));
51        let mut second = if let Some(state) = &self.state {
52            // Vendors MUST accept empty tracestate headers but SHOULD avoid sending them.
53            if !state.is_empty() {
54                Some(("tracestate", state.clone()))
55            } else {
56                None
57            }
58        } else {
59            None
60        };
61
62        // By returning an iterator, we avoid allocating a Vec/HashMap.
63        let mut counter = 0;
64        iter::from_fn(move || {
65            counter += 1;
66            if counter == 1 {
67                first.take()
68            } else if counter == 2 {
69                second.take()
70            } else {
71                None
72            }
73        })
74    }
75
76    /// The version of the trace context specification that we support.
77    ///
78    /// https://www.w3.org/TR/trace-context-2/#version
79    const SUPPORTED_VERSION: u8 = 0;
80
81    fn traceparent(&self) -> String {
82        format!(
83            "{:02x}-{:032x}-{:016x}-{:02x}",
84            Self::SUPPORTED_VERSION,
85            self.trace_id,
86            self.span_id,
87            self.trace_flags()
88        )
89    }
90
91    /// The trace flags of this context.
92    ///
93    /// Version 0 of the trace context specification only supports the `sampled` flag.
94    /// [W3C TraceContext specification]: https://www.w3.org/TR/trace-context/#sampled-flag
95    fn trace_flags(&self) -> u8 {
96        if self.sampled {
97            0x01
98        } else {
99            0x00
100        }
101    }
102}
103
104#[cfg(test)]
105mod tests {
106    use super::TraceContext;
107
108    #[test]
109    fn trace_flags_if_sampled() {
110        let trace_id = 0x4bf92f3577b34da6a3ce929d0e0e4736;
111        let span_id = 0x00f067aa0ba902b7;
112        let trace_context = TraceContext::new_sampled(trace_id, span_id, None);
113        assert_eq!(trace_context.trace_flags(), 0x01);
114    }
115
116    #[test]
117    fn trace_flags_if_not_sampled() {
118        let trace_id = 0x4bf92f3577b34da6a3ce929d0e0e4736;
119        let span_id = 0x00f067aa0ba902b7;
120        let trace_context = TraceContext::new_unsampled(trace_id, span_id, None);
121        assert_eq!(trace_context.trace_flags(), 0x00);
122    }
123
124    #[test]
125    fn traceparent_generation_if_sampled() {
126        let trace_id = 0x4bf92f3577b34da6a3ce929d0e0e4736;
127        let span_id = 0x00f067aa0ba902b7;
128        let trace_context = TraceContext::new_sampled(trace_id, span_id, None);
129        assert_eq!(
130            trace_context.traceparent(),
131            "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"
132        );
133    }
134
135    #[test]
136    fn traceparent_generation_if_not_sampled() {
137        let trace_id = 0x4bf92f3577b34da6a3ce929d0e0e4736;
138        let span_id = 0x00f067aa0ba902b7;
139        let trace_context = TraceContext::new_unsampled(trace_id, span_id, None);
140        assert_eq!(
141            trace_context.traceparent(),
142            "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-00"
143        );
144    }
145
146    #[test]
147    fn headers_include_traceparent() {
148        let trace_id = 0x4bf92f3577b34da6a3ce929d0e0e4736;
149        let span_id = 0x00f067aa0ba902b7;
150        let trace_context = TraceContext::new_sampled(trace_id, span_id, None);
151        let mut headers = trace_context.as_w3c_headers();
152
153        // get first header
154        let header = headers.next().unwrap();
155        assert_eq!(header.0, "traceparent");
156        assert_eq!(header.1, trace_context.traceparent());
157        assert!(headers.next().is_none());
158    }
159
160    #[test]
161    fn non_empty_tracestate_is_included() {
162        let trace_id = 0x4bf92f3577b34da6a3ce929d0e0e4736;
163        let span_id = 0x00f067aa0ba902b7;
164        let trace_context =
165            TraceContext::new_sampled(trace_id, span_id, Some("foo=bar".to_string()));
166        let mut headers = trace_context.as_w3c_headers();
167
168        // get first header
169        let header = headers.next().unwrap();
170        assert_eq!(header.0, "traceparent");
171        assert_eq!(header.1, trace_context.traceparent());
172
173        // get second header
174        let header = headers.next().unwrap();
175        assert_eq!(header.0, "tracestate");
176        assert_eq!(header.1, "foo=bar");
177        assert!(headers.next().is_none());
178    }
179}
aleph_alpha_client/tracing.rs

aleph_alpha_client/
tracing.rs
