Skip to main content

fakecloud_xray/
segment.rs

1//! X-Ray trace segment document parsing.
2//!
3//! `PutTraceSegments` accepts a list of segment documents, each a JSON string in
4//! the [X-Ray segment document
5//! format](https://docs.aws.amazon.com/xray/latest/devguide/xray-api-segmentdocuments.html).
6//! We parse the fields fakecloud actually models -- `trace_id`, `id`, `name`,
7//! `start_time`, `end_time`, `parent_id`, `origin`, the `http` request/response,
8//! the `error`/`fault`/`throttle` flags, and the nested `subsegments` (walking
9//! `namespace: "remote"` subsegments as downstream service calls) -- into a
10//! [`StoredSegment`] kept for the data-plane reads. The raw document string is
11//! retained verbatim so `BatchGetTraces` echoes exactly what was ingested.
12
13use serde::{Deserialize, Serialize};
14use serde_json::Value;
15
16/// A downstream service call discovered inside a segment's subsegments (a
17/// subsegment with `namespace: "remote"`). Drives the service-graph edges.
18#[derive(Debug, Clone, Serialize, Deserialize)]
19pub struct Downstream {
20    pub name: String,
21    pub start_time: f64,
22    pub end_time: Option<f64>,
23    pub fault: bool,
24    pub error: bool,
25    pub throttle: bool,
26    /// The subsegment's declared type, if any (e.g. `AWS::DynamoDB::Table`).
27    pub kind: Option<String>,
28}
29
30/// A parsed, stored X-Ray segment. Retains the raw `document` for verbatim
31/// re-emission and the extracted fields the data-plane reads compute over.
32#[derive(Debug, Clone, Serialize, Deserialize)]
33pub struct StoredSegment {
34    pub trace_id: String,
35    pub id: String,
36    pub name: String,
37    pub start_time: f64,
38    pub end_time: Option<f64>,
39    pub parent_id: Option<String>,
40    pub origin: Option<String>,
41    pub http_status: Option<i64>,
42    pub http_method: Option<String>,
43    pub http_url: Option<String>,
44    pub fault: bool,
45    pub error: bool,
46    pub throttle: bool,
47    pub downstream: Vec<Downstream>,
48    /// The raw segment document as ingested, echoed by `BatchGetTraces`.
49    pub document: String,
50}
51
52impl StoredSegment {
53    /// Duration in seconds (`end_time - start_time`), when the segment is
54    /// complete (has an `end_time`). An in-progress segment has no duration.
55    pub fn duration(&self) -> Option<f64> {
56        self.end_time.map(|e| (e - self.start_time).max(0.0))
57    }
58}
59
60/// Error parsing a single segment document, mapped to an
61/// `UnprocessedTraceSegment` entry by the handler.
62#[derive(Debug)]
63pub struct ParseError {
64    /// The segment `id`, when it could be read from the malformed document.
65    pub id: Option<String>,
66    pub code: &'static str,
67    pub message: String,
68}
69
70fn bool_field(v: &Value, key: &str) -> bool {
71    v.get(key).and_then(Value::as_bool).unwrap_or(false)
72}
73
74fn f64_field(v: &Value, key: &str) -> Option<f64> {
75    v.get(key).and_then(Value::as_f64)
76}
77
78fn str_field(v: &Value, key: &str) -> Option<String> {
79    v.get(key)
80        .and_then(Value::as_str)
81        .map(std::string::ToString::to_string)
82}
83
84/// Recursively collect `namespace: "remote"` subsegments as downstream calls.
85fn collect_downstream(seg: &Value, out: &mut Vec<Downstream>) {
86    let Some(subs) = seg.get("subsegments").and_then(Value::as_array) else {
87        return;
88    };
89    for sub in subs {
90        let namespace = sub.get("namespace").and_then(Value::as_str);
91        if namespace == Some("remote") {
92            if let Some(name) = str_field(sub, "name") {
93                out.push(Downstream {
94                    name,
95                    start_time: f64_field(sub, "start_time").unwrap_or(0.0),
96                    end_time: f64_field(sub, "end_time"),
97                    fault: bool_field(sub, "fault"),
98                    error: bool_field(sub, "error"),
99                    throttle: bool_field(sub, "throttle"),
100                    kind: str_field(sub, "type"),
101                });
102            }
103        }
104        // Recurse: downstream calls can be nested inside local subsegments.
105        collect_downstream(sub, out);
106    }
107}
108
109/// Parse a single segment document string into a [`StoredSegment`].
110///
111/// A segment document must be a JSON object carrying at minimum a `trace_id`,
112/// an `id`, and a `name` (per the X-Ray segment document schema); anything else
113/// is rejected as a `ParseError` that surfaces as an `UnprocessedTraceSegment`.
114pub fn parse_segment(document: &str) -> Result<StoredSegment, ParseError> {
115    let v: Value = serde_json::from_str(document).map_err(|e| ParseError {
116        id: None,
117        code: "MalformedJson",
118        message: format!("Segment document is not valid JSON: {e}"),
119    })?;
120    if !v.is_object() {
121        return Err(ParseError {
122            id: None,
123            code: "MalformedJson",
124            message: "Segment document must be a JSON object.".to_string(),
125        });
126    }
127    let id = str_field(&v, "id");
128    let trace_id = str_field(&v, "trace_id").ok_or_else(|| ParseError {
129        id: id.clone(),
130        code: "MissingTraceId",
131        message: "Segment document is missing required field 'trace_id'.".to_string(),
132    })?;
133    let seg_id = id.clone().ok_or_else(|| ParseError {
134        id: None,
135        code: "MissingId",
136        message: "Segment document is missing required field 'id'.".to_string(),
137    })?;
138    let name = str_field(&v, "name").ok_or_else(|| ParseError {
139        id: id.clone(),
140        code: "MissingName",
141        message: "Segment document is missing required field 'name'.".to_string(),
142    })?;
143    let start_time = f64_field(&v, "start_time").ok_or_else(|| ParseError {
144        id: id.clone(),
145        code: "MissingStartTime",
146        message: "Segment document is missing required field 'start_time'.".to_string(),
147    })?;
148
149    let http = v.get("http");
150    let http_status = http
151        .and_then(|h| h.get("response"))
152        .and_then(|r| r.get("status"))
153        .and_then(Value::as_i64);
154    let http_method = http
155        .and_then(|h| h.get("request"))
156        .and_then(|r| r.get("method"))
157        .and_then(Value::as_str)
158        .map(std::string::ToString::to_string);
159    let http_url = http
160        .and_then(|h| h.get("request"))
161        .and_then(|r| r.get("url"))
162        .and_then(Value::as_str)
163        .map(std::string::ToString::to_string);
164
165    let mut downstream = Vec::new();
166    collect_downstream(&v, &mut downstream);
167
168    Ok(StoredSegment {
169        trace_id,
170        id: seg_id,
171        name,
172        start_time,
173        end_time: f64_field(&v, "end_time"),
174        parent_id: str_field(&v, "parent_id"),
175        origin: str_field(&v, "origin"),
176        http_status,
177        http_method,
178        http_url,
179        fault: bool_field(&v, "fault"),
180        error: bool_field(&v, "error"),
181        throttle: bool_field(&v, "throttle"),
182        downstream,
183        document: document.to_string(),
184    })
185}
186
187#[cfg(test)]
188mod tests {
189    use super::*;
190
191    #[test]
192    fn parses_minimal_segment() {
193        let doc = r#"{"trace_id":"1-58406520-a006649127e371903a2de979","id":"6b55dcb43fd8c8c1","name":"api","start_time":1.0,"end_time":2.0}"#;
194        let seg = parse_segment(doc).unwrap();
195        assert_eq!(seg.trace_id, "1-58406520-a006649127e371903a2de979");
196        assert_eq!(seg.name, "api");
197        assert_eq!(seg.duration(), Some(1.0));
198        assert_eq!(seg.document, doc);
199    }
200
201    #[test]
202    fn reads_http_and_fault_flags() {
203        let doc = r#"{"trace_id":"1-x","id":"aaaa","name":"web","start_time":1.0,"end_time":2.5,
204            "fault":true,"http":{"request":{"method":"GET","url":"http://x/"},"response":{"status":500}}}"#;
205        let seg = parse_segment(doc).unwrap();
206        assert!(seg.fault);
207        assert_eq!(seg.http_status, Some(500));
208        assert_eq!(seg.http_method.as_deref(), Some("GET"));
209    }
210
211    #[test]
212    fn collects_remote_downstream_calls() {
213        let doc = r#"{"trace_id":"1-x","id":"aaaa","name":"web","start_time":1.0,"end_time":3.0,
214            "subsegments":[
215              {"id":"bbbb","name":"local","start_time":1.1,"end_time":1.2,
216               "subsegments":[{"id":"cccc","name":"dynamo","namespace":"remote","start_time":1.3,"end_time":1.9,"error":true,"type":"AWS::DynamoDB::Table"}]},
217              {"id":"dddd","name":"backend","namespace":"remote","start_time":2.0,"end_time":2.8}
218            ]}"#;
219        let seg = parse_segment(doc).unwrap();
220        assert_eq!(seg.downstream.len(), 2);
221        let names: Vec<&str> = seg.downstream.iter().map(|d| d.name.as_str()).collect();
222        assert!(names.contains(&"dynamo"));
223        assert!(names.contains(&"backend"));
224        let dynamo = seg.downstream.iter().find(|d| d.name == "dynamo").unwrap();
225        assert!(dynamo.error);
226        assert_eq!(dynamo.kind.as_deref(), Some("AWS::DynamoDB::Table"));
227    }
228
229    #[test]
230    fn rejects_missing_trace_id() {
231        let err = parse_segment(r#"{"id":"aaaa","name":"x","start_time":1.0}"#).unwrap_err();
232        assert_eq!(err.code, "MissingTraceId");
233    }
234
235    #[test]
236    fn rejects_malformed_json() {
237        let err = parse_segment("not json").unwrap_err();
238        assert_eq!(err.code, "MalformedJson");
239    }
240}