Skip to main content

fakecloud_xray/
state.rs

1//! Account-partitioned, serializable state for AWS X-Ray (`xray`).
2//!
3//! Control-plane resources (groups, sampling rules, the encryption config,
4//! resource policies, the indexing rule, the trace-segment destination) are
5//! stored as their already-output-valid wire JSON objects so reads echo exactly
6//! what writes persisted. The data plane stores ingested [`StoredSegment`]s
7//! keyed by trace id. Tags are keyed by resource ARN in a single map covering
8//! every taggable resource (groups + sampling rules).
9//!
10//! Every map key is a plain `String` (name, ARN, trace id), so the snapshot
11//! never depends on the tuple-key serde adapter that has silently broken
12//! snapshot serialization on other services.
13
14use std::collections::BTreeMap;
15use std::sync::Arc;
16
17use parking_lot::RwLock;
18use serde::{Deserialize, Serialize};
19use serde_json::{json, Value};
20
21use fakecloud_core::multi_account::{AccountState, MultiAccountState};
22
23use crate::segment::StoredSegment;
24
25pub const XRAY_SNAPSHOT_SCHEMA_VERSION: u32 = 1;
26
27/// The name of the built-in sampling rule X-Ray always provides.
28pub const DEFAULT_SAMPLING_RULE: &str = "Default";
29
30/// Per-account AWS X-Ray state.
31#[derive(Debug, Clone, Serialize, Deserialize)]
32pub struct XrayData {
33    /// Groups keyed by `GroupName`, stored as their `Group` wire object.
34    #[serde(default)]
35    pub groups: BTreeMap<String, Value>,
36    /// Sampling rules keyed by `RuleName`, stored as their `SamplingRuleRecord`
37    /// wire object.
38    #[serde(default)]
39    pub sampling_rules: BTreeMap<String, Value>,
40    /// The account `EncryptionConfig` wire object, once `PutEncryptionConfig`
41    /// has set it; `None` means the default (`Type: NONE`).
42    #[serde(default)]
43    pub encryption_config: Option<Value>,
44    /// Resource policies keyed by `PolicyName`, stored as their `ResourcePolicy`
45    /// wire object.
46    #[serde(default)]
47    pub resource_policies: BTreeMap<String, Value>,
48    /// The trace-segment destination (`XRay` or `CloudWatchLogs`).
49    #[serde(default = "default_destination")]
50    pub trace_segment_destination: String,
51    /// Ingested trace segments keyed by trace id.
52    #[serde(default)]
53    pub traces: BTreeMap<String, Vec<StoredSegment>>,
54    /// Active trace retrievals (Transaction Search) keyed by retrieval token,
55    /// each `{ "TraceIds": [...], "StartTime": .., "EndTime": .. }`.
56    #[serde(default)]
57    pub retrievals: BTreeMap<String, Value>,
58    /// Tags keyed by resource ARN (groups + sampling rules).
59    #[serde(default)]
60    pub tags: BTreeMap<String, BTreeMap<String, String>>,
61}
62
63fn default_destination() -> String {
64    "XRay".to_string()
65}
66
67impl Default for XrayData {
68    fn default() -> Self {
69        Self {
70            groups: BTreeMap::new(),
71            sampling_rules: BTreeMap::new(),
72            encryption_config: None,
73            resource_policies: BTreeMap::new(),
74            trace_segment_destination: default_destination(),
75            traces: BTreeMap::new(),
76            retrievals: BTreeMap::new(),
77            tags: BTreeMap::new(),
78        }
79    }
80}
81
82impl XrayData {
83    /// Seed the built-in, undeletable `Default` sampling rule that every X-Ray
84    /// account carries (matched last, 1 req/s reservoir + 5% of the rest).
85    fn seed_default_rule(&mut self, region: &str, account: &str) {
86        if self.sampling_rules.contains_key(DEFAULT_SAMPLING_RULE) {
87            return;
88        }
89        let arn = format!("arn:aws:xray:{region}:{account}:sampling-rule/{DEFAULT_SAMPLING_RULE}");
90        let now = now_epoch();
91        let record = json!({
92            "SamplingRule": {
93                "RuleName": DEFAULT_SAMPLING_RULE,
94                "RuleARN": arn,
95                "ResourceARN": "*",
96                "Priority": 10000,
97                "FixedRate": 0.05,
98                "ReservoirSize": 1,
99                "ServiceName": "*",
100                "ServiceType": "*",
101                "Host": "*",
102                "HTTPMethod": "*",
103                "URLPath": "*",
104                "Version": 1,
105                "Attributes": {},
106            },
107            "CreatedAt": now,
108            "ModifiedAt": now,
109        });
110        self.sampling_rules
111            .insert(DEFAULT_SAMPLING_RULE.to_string(), record);
112    }
113}
114
115/// Current time as restJson1 epoch-seconds (a floating-point number). X-Ray's
116/// timestamp members carry no `@timestampFormat`, so restJson1's default
117/// epoch-seconds applies.
118pub fn now_epoch() -> f64 {
119    chrono::Utc::now().timestamp_millis() as f64 / 1000.0
120}
121
122impl AccountState for XrayData {
123    fn new_for_account(account_id: &str, region: &str, _endpoint: &str) -> Self {
124        let mut data = Self::default();
125        data.seed_default_rule(region, account_id);
126        data
127    }
128}
129
130pub type SharedXrayState = Arc<RwLock<MultiAccountState<XrayData>>>;
131
132#[derive(Debug, Serialize, Deserialize)]
133pub struct XraySnapshot {
134    pub schema_version: u32,
135    pub accounts: MultiAccountState<XrayData>,
136}
137
138#[cfg(test)]
139mod tests {
140    use super::*;
141
142    #[test]
143    fn new_account_has_default_sampling_rule() {
144        let data = XrayData::new_for_account("000000000000", "us-east-1", "");
145        assert!(data.sampling_rules.contains_key(DEFAULT_SAMPLING_RULE));
146        let rule = &data.sampling_rules[DEFAULT_SAMPLING_RULE]["SamplingRule"];
147        assert_eq!(rule["ResourceARN"], json!("*"));
148        assert_eq!(rule["Priority"], json!(10000));
149    }
150
151    #[test]
152    fn default_destination_is_xray() {
153        let data = XrayData::default();
154        assert_eq!(data.trace_segment_destination, "XRay");
155    }
156}