1use 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
27pub const DEFAULT_SAMPLING_RULE: &str = "Default";
29
30#[derive(Debug, Clone, Serialize, Deserialize)]
32pub struct XrayData {
33 #[serde(default)]
35 pub groups: BTreeMap<String, Value>,
36 #[serde(default)]
39 pub sampling_rules: BTreeMap<String, Value>,
40 #[serde(default)]
43 pub encryption_config: Option<Value>,
44 #[serde(default)]
47 pub resource_policies: BTreeMap<String, Value>,
48 #[serde(default = "default_destination")]
50 pub trace_segment_destination: String,
51 #[serde(default)]
53 pub traces: BTreeMap<String, Vec<StoredSegment>>,
54 #[serde(default)]
57 pub retrievals: BTreeMap<String, Value>,
58 #[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 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
115pub 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}