xray_tracing/
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, Debug, Clone, Copy)]
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    Unknown,
25}
26
27impl<'a> From<&'a str> for SamplingDecision {
28    fn from(value: &'a str) -> Self {
29        match value {
30            "Sampled=1" => SamplingDecision::Sampled,
31            "Sampled=0" => SamplingDecision::NotSampled,
32            "Sampled=?" => SamplingDecision::Requested,
33            _ => SamplingDecision::Unknown,
34        }
35    }
36}
37
38impl Display for SamplingDecision {
39    fn fmt(
40        &self,
41        f: &mut fmt::Formatter,
42    ) -> fmt::Result {
43        write!(
44            f,
45            "{}",
46            match self {
47                SamplingDecision::Sampled => "Sampled=1",
48                SamplingDecision::NotSampled => "Sampled=0",
49                SamplingDecision::Requested => "Sampled=?",
50                _ => "",
51            }
52        )?;
53        Ok(())
54    }
55}
56
57impl Default for SamplingDecision {
58    fn default() -> Self {
59        SamplingDecision::Unknown
60    }
61}
62
63/// Parsed representation of `X-Amzn-Trace-Id` request header
64#[derive(PartialEq, Debug, Default, Clone)]
65pub struct Header {
66    pub(crate) trace_id: TraceId,
67    pub(crate) parent_id: Option<SegmentId>,
68    pub(crate) sampling_decision: SamplingDecision,
69    additional_data: HashMap<String, String>,
70}
71
72impl Header {
73    /// HTTP header name associated with X-Ray trace data
74    ///
75    /// HTTP header values should be the Display serialization of Header structs
76    pub const NAME: &'static str = "X-Amzn-Trace-Id";
77
78    pub fn new(trace_id: TraceId) -> Self {
79        Header {
80            trace_id,
81            ..Header::default()
82        }
83    }
84
85    pub fn with_parent_id(
86        &mut self,
87        parent_id: SegmentId,
88    ) -> &mut Self {
89        self.parent_id = Some(parent_id);
90        self
91    }
92
93    pub fn with_sampling_decision(
94        &mut self,
95        decision: SamplingDecision,
96    ) -> &mut Self {
97        self.sampling_decision = decision;
98        self
99    }
100
101    pub fn with_data<K, V>(
102        &mut self,
103        key: K,
104        value: V,
105    ) -> &mut Self
106    where
107        K: Into<String>,
108        V: Into<String>,
109    {
110        self.additional_data.insert(key.into(), value.into());
111        self
112    }
113}
114
115impl FromStr for Header {
116    type Err = String;
117    fn from_str(s: &str) -> Result<Self, Self::Err> {
118        s.split(';')
119            .try_fold(Header::default(), |mut header, line| {
120                if line.starts_with("Root=") {
121                    header.trace_id = TraceId::Rendered(line[5..].into())
122                } else if line.starts_with("Parent=") {
123                    header.parent_id = Some(SegmentId::Rendered(line[7..].into()))
124                } else if line.starts_with("Sampled=") {
125                    header.sampling_decision = line.into();
126                } else if !line.starts_with("Self=") {
127                    let pos = line
128                        .find('=')
129                        .ok_or_else(|| format!("invalid key=value: no `=` found in `{}`", s))?;
130                    let (key, value) = (&line[..pos], &line[pos + 1..]);
131                    header.additional_data.insert(key.into(), value.into());
132                }
133                Ok(header)
134            })
135    }
136}
137
138impl Display for Header {
139    fn fmt(
140        &self,
141        f: &mut fmt::Formatter,
142    ) -> fmt::Result {
143        write!(f, "Root={}", self.trace_id)?;
144        if let Some(parent) = &self.parent_id {
145            write!(f, ";Parent={}", parent)?;
146        }
147        if self.sampling_decision != SamplingDecision::Unknown {
148            write!(f, ";{}", self.sampling_decision)?;
149        }
150        for (k, v) in &self.additional_data {
151            write!(f, ";{}={}", k, v)?;
152        }
153        Ok(())
154    }
155}
156
157#[cfg(test)]
158mod tests {
159    use super::*;
160    #[test]
161    fn parse_with_parent_from_str() {
162        assert_eq!(
163            "Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1"
164                .parse::<Header>(),
165            Ok(Header {
166                trace_id: TraceId::Rendered("1-5759e988-bd862e3fe1be46a994272793".into()),
167                parent_id: Some(SegmentId::Rendered("53995c3f42cd8ad8".into())),
168                sampling_decision: SamplingDecision::Sampled,
169                ..Header::default()
170            })
171        )
172    }
173    #[test]
174    fn parse_no_parent_from_str() {
175        assert_eq!(
176            "Root=1-5759e988-bd862e3fe1be46a994272793;Sampled=1".parse::<Header>(),
177            Ok(Header {
178                trace_id: TraceId::Rendered("1-5759e988-bd862e3fe1be46a994272793".into()),
179                parent_id: None,
180                sampling_decision: SamplingDecision::Sampled,
181                ..Header::default()
182            })
183        )
184    }
185
186    #[test]
187    fn displays_as_header() {
188        let header = Header {
189            trace_id: TraceId::Rendered("1-5759e988-bd862e3fe1be46a994272793".into()),
190            ..Header::default()
191        };
192        assert_eq!(
193            format!("{}", header),
194            "Root=1-5759e988-bd862e3fe1be46a994272793"
195        );
196    }
197}
xray_tracing/header.rs

xray_tracing/
header.rs
