xray_lite/
header.rs

1//! X-Ray [tracing header](https://docs.aws.amazon.com/xray/latest/devguide/xray-concepts.html?shortFooter=true#xray-concepts-tracingheader)
2//! parser
3
4use crate::{SegmentId, TraceId};
5use std::{
6    collections::HashMap,
7    fmt::{self, Display},
8    str::FromStr,
9};
10
11#[derive(PartialEq, Clone, Copy, Debug, Default)]
12pub enum SamplingDecision {
13    /// Sampled indicates the current segment has been
14    /// sampled and will be sent to the X-Ray daemon.
15    Sampled,
16    /// NotSampled indicates the current segment has
17    /// not been sampled.
18    NotSampled,
19    ///sampling decision will be
20    /// made by the downstream service and propagated
21    /// back upstream in the response.
22    Requested,
23    /// Unknown indicates no sampling decision will be made.
24    #[default]
25    Unknown,
26}
27
28impl<'a> From<&'a str> for SamplingDecision {
29    fn from(value: &'a str) -> Self {
30        match value {
31            "Sampled=1" => SamplingDecision::Sampled,
32            "Sampled=0" => SamplingDecision::NotSampled,
33            "Sampled=?" => SamplingDecision::Requested,
34            _ => SamplingDecision::Unknown,
35        }
36    }
37}
38
39impl Display for SamplingDecision {
40    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
41        write!(
42            f,
43            "{}",
44            match self {
45                SamplingDecision::Sampled => "Sampled=1",
46                SamplingDecision::NotSampled => "Sampled=0",
47                SamplingDecision::Requested => "Sampled=?",
48                _ => "",
49            }
50        )?;
51        Ok(())
52    }
53}
54
55/// Parsed representation of `X-Amzn-Trace-Id` request header
56#[derive(PartialEq, Clone, Debug, Default)]
57pub struct Header {
58    pub(crate) trace_id: TraceId,
59    pub(crate) parent_id: Option<SegmentId>,
60    pub(crate) sampling_decision: SamplingDecision,
61    additional_data: HashMap<String, String>,
62}
63
64impl Header {
65    /// HTTP header name associated with X-Ray trace data
66    ///
67    /// HTTP header values should be the Display serialization of Header structs
68    pub const NAME: &'static str = "X-Amzn-Trace-Id";
69
70    /// Creates a new Header with a given trace ID.
71    pub fn new(trace_id: TraceId) -> Self {
72        Header {
73            trace_id,
74            ..Header::default()
75        }
76    }
77
78    /// Creates a new Header with the parent ID replaced.
79    pub fn with_parent_id(&self, parent_id: SegmentId) -> Self {
80        Self {
81            trace_id: self.trace_id.clone(),
82            parent_id: Some(parent_id),
83            sampling_decision: self.sampling_decision,
84            additional_data: self.additional_data.clone(),
85        }
86    }
87
88    /// Creates a new Header with the sampling decision replaced.
89    pub fn with_sampling_decision(&self, decision: SamplingDecision) -> Self {
90        Self {
91            trace_id: self.trace_id.clone(),
92            parent_id: self.parent_id.clone(),
93            sampling_decision: decision,
94            additional_data: self.additional_data.clone(),
95        }
96    }
97
98    /// Inserts a key-value pair into the additional data map.
99    pub fn insert_data(&mut self, key: impl Into<String>, value: impl Into<String>) -> &mut Self {
100        self.additional_data.insert(key.into(), value.into());
101        self
102    }
103}
104
105impl FromStr for Header {
106    type Err = String;
107    fn from_str(s: &str) -> Result<Self, Self::Err> {
108        s.split(';')
109            .try_fold(Header::default(), |mut header, line| {
110                if let Some(trace_id) = line.strip_prefix("Root=") {
111                    header.trace_id = TraceId::Rendered(trace_id.into())
112                } else if let Some(parent_id) = line.strip_prefix("Parent=") {
113                    header.parent_id = Some(SegmentId::Rendered(parent_id.into()))
114                } else if line.starts_with("Sampled=") {
115                    header.sampling_decision = line.into();
116                } else if !line.starts_with("Self=") {
117                    let (key, value) = line
118                        .split_once('=')
119                        .ok_or_else(|| format!("invalid key=value: no `=` found in `{}`", s))?;
120                    header.additional_data.insert(key.into(), value.into());
121                }
122                Ok(header)
123            })
124    }
125}
126
127impl Display for Header {
128    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
129        write!(f, "Root={}", self.trace_id)?;
130        if let Some(parent) = &self.parent_id {
131            write!(f, ";Parent={}", parent)?;
132        }
133        if self.sampling_decision != SamplingDecision::Unknown {
134            write!(f, ";{}", self.sampling_decision)?;
135        }
136        for (k, v) in &self.additional_data {
137            write!(f, ";{}={}", k, v)?;
138        }
139        Ok(())
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146    #[test]
147    fn parse_with_parent_from_str() {
148        assert_eq!(
149            "Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1"
150                .parse::<Header>(),
151            Ok(Header {
152                trace_id: TraceId::Rendered("1-5759e988-bd862e3fe1be46a994272793".into()),
153                parent_id: Some(SegmentId::Rendered("53995c3f42cd8ad8".into())),
154                sampling_decision: SamplingDecision::Sampled,
155                ..Header::default()
156            })
157        )
158    }
159    #[test]
160    fn parse_no_parent_from_str() {
161        assert_eq!(
162            "Root=1-5759e988-bd862e3fe1be46a994272793;Sampled=1".parse::<Header>(),
163            Ok(Header {
164                trace_id: TraceId::Rendered("1-5759e988-bd862e3fe1be46a994272793".into()),
165                parent_id: None,
166                sampling_decision: SamplingDecision::Sampled,
167                ..Header::default()
168            })
169        )
170    }
171    #[test]
172    fn parse_with_additional_data_from_str() {
173        assert_eq!(
174            "Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1;Lineage=01234567:0;Unknown=unknown"
175                .parse::<Header>(),
176            Ok(Header {
177                trace_id: TraceId::Rendered("1-5759e988-bd862e3fe1be46a994272793".into()),
178                parent_id: Some(SegmentId::Rendered("53995c3f42cd8ad8".into())),
179                sampling_decision: SamplingDecision::Sampled,
180                additional_data: vec![
181                    ("Lineage".into(), "01234567:0".into()),
182                    ("Unknown".into(), "unknown".into()),
183                ].into_iter().collect()
184            })
185        )
186    }
187
188    #[test]
189    fn displays_as_header() {
190        let header = Header {
191            trace_id: TraceId::Rendered("1-5759e988-bd862e3fe1be46a994272793".into()),
192            ..Header::default()
193        };
194        assert_eq!(
195            format!("{}", header),
196            "Root=1-5759e988-bd862e3fe1be46a994272793"
197        );
198    }
199
200    #[test]
201    fn replace_parent_id() {
202        let header = Header {
203            trace_id: TraceId::Rendered("1-5759e988-bd862e3fe1be46a994272793".into()),
204            parent_id: Some(SegmentId::Rendered("53995c3f42cd8ad8".into())),
205            sampling_decision: SamplingDecision::Sampled,
206            ..Header::default()
207        };
208        assert_eq!(
209            header.with_parent_id(SegmentId::Rendered("35b167406b7746cf".into())),
210            Header {
211                trace_id: TraceId::Rendered("1-5759e988-bd862e3fe1be46a994272793".into()),
212                parent_id: Some(SegmentId::Rendered("35b167406b7746cf".into())),
213                sampling_decision: SamplingDecision::Sampled,
214                ..Header::default()
215            },
216        );
217    }
218
219    #[test]
220    fn replace_sampling_decision() {
221        let header = Header {
222            trace_id: TraceId::Rendered("1-5759e988-bd862e3fe1be46a994272793".into()),
223            parent_id: Some(SegmentId::Rendered("53995c3f42cd8ad8".into())),
224            sampling_decision: SamplingDecision::Sampled,
225            ..Header::default()
226        };
227        assert_eq!(
228            header.with_sampling_decision(SamplingDecision::NotSampled),
229            Header {
230                trace_id: TraceId::Rendered("1-5759e988-bd862e3fe1be46a994272793".into()),
231                parent_id: Some(SegmentId::Rendered("53995c3f42cd8ad8".into())),
232                sampling_decision: SamplingDecision::NotSampled,
233                ..Header::default()
234            },
235        );
236    }
237}
xray_lite/header.rs

xray_lite/
header.rs
