rustrails_support/
current_attributes.rs

1use serde_json::Value;
2use std::cell::RefCell;
3use std::collections::HashMap;
4
5thread_local! {
6    static CURRENT_ATTRIBUTES: RefCell<HashMap<String, Value>> = RefCell::new(HashMap::new());
7}
8
9/// A thread-local request-scoped attribute store.
10#[derive(Debug, Default, Clone, Copy)]
11pub struct CurrentAttributes;
12
13impl CurrentAttributes {
14    /// Creates a new handle to the thread-local attribute store.
15    #[must_use]
16    pub fn new() -> Self {
17        Self
18    }
19
20    /// Stores `value` for `key`, returning the previous value when present.
21    pub fn set(&self, key: impl Into<String>, value: Value) -> Option<Value> {
22        CURRENT_ATTRIBUTES.with(|attributes| attributes.borrow_mut().insert(key.into(), value))
23    }
24
25    /// Returns the cloned value for `key` when present.
26    #[must_use]
27    pub fn get(&self, key: &str) -> Option<Value> {
28        CURRENT_ATTRIBUTES.with(|attributes| attributes.borrow().get(key).cloned())
29    }
30
31    /// Clears all attributes for the current thread.
32    pub fn reset(&self) {
33        CURRENT_ATTRIBUTES.with(|attributes| attributes.borrow_mut().clear());
34    }
35
36    /// Temporarily applies `attrs` for the duration of `f`, then restores the prior state.
37    pub fn with_attributes<F, R>(&self, attrs: HashMap<String, Value>, f: F) -> R
38    where
39        F: FnOnce() -> R,
40    {
41        let previous = CURRENT_ATTRIBUTES.with(|attributes| {
42            let mut attributes = attributes.borrow_mut();
43            let previous = attributes.clone();
44            attributes.extend(attrs);
45            previous
46        });
47        let _guard = ResetGuard { previous };
48        f()
49    }
50
51    fn snapshot(&self) -> HashMap<String, Value> {
52        CURRENT_ATTRIBUTES.with(|attributes| attributes.borrow().clone())
53    }
54}
55
56struct ResetGuard {
57    previous: HashMap<String, Value>,
58}
59
60impl Drop for ResetGuard {
61    fn drop(&mut self) {
62        CURRENT_ATTRIBUTES.with(|attributes| {
63            *attributes.borrow_mut() = std::mem::take(&mut self.previous);
64        });
65    }
66}
67
68#[cfg(test)]
69mod tests {
70    use super::CurrentAttributes;
71    use serde_json::json;
72    use std::collections::HashMap;
73    use std::thread;
74
75    fn run_isolated<R>(test: impl FnOnce() -> R + Send + 'static) -> R
76    where
77        R: Send + 'static,
78    {
79        match thread::spawn(test).join() {
80            Ok(result) => result,
81            Err(payload) => std::panic::resume_unwind(payload),
82        }
83    }
84
85    #[test]
86    fn current_attributes_starts_empty() {
87        run_isolated(|| {
88            let current = CurrentAttributes::new();
89            assert!(current.snapshot().is_empty());
90            assert_eq!(current.get("missing"), None);
91        });
92    }
93
94    #[test]
95    fn set_and_get_round_trip_values() {
96        run_isolated(|| {
97            let current = CurrentAttributes::new();
98            current.set("request_id", json!("abc123"));
99
100            assert_eq!(current.get("request_id"), Some(json!("abc123")));
101        });
102    }
103
104    #[test]
105    fn set_returns_previous_value() {
106        run_isolated(|| {
107            let current = CurrentAttributes::new();
108            assert_eq!(current.set("count", json!(1)), None);
109            assert_eq!(current.set("count", json!(2)), Some(json!(1)));
110        });
111    }
112
113    #[test]
114    fn get_returns_none_for_missing_keys() {
115        run_isolated(|| {
116            let current = CurrentAttributes::new();
117            current.set("present", json!(true));
118
119            assert_eq!(current.get("missing"), None);
120        });
121    }
122
123    #[test]
124    fn reset_clears_all_attributes() {
125        run_isolated(|| {
126            let current = CurrentAttributes::new();
127            current.set("request_id", json!("abc123"));
128            current.set("user_id", json!(42));
129
130            current.reset();
131
132            assert!(current.snapshot().is_empty());
133        });
134    }
135
136    #[test]
137    fn with_attributes_adds_values_for_the_scope() {
138        run_isolated(|| {
139            let current = CurrentAttributes::new();
140            let attrs = HashMap::from([
141                (String::from("request_id"), json!("abc123")),
142                (String::from("user_id"), json!(42)),
143            ]);
144
145            let values = current.with_attributes(attrs, || current.snapshot());
146
147            assert_eq!(values.get("request_id"), Some(&json!("abc123")));
148            assert_eq!(values.get("user_id"), Some(&json!(42)));
149            assert!(current.snapshot().is_empty());
150        });
151    }
152
153    #[test]
154    fn with_attributes_restores_existing_values_after_scope() {
155        run_isolated(|| {
156            let current = CurrentAttributes::new();
157            current.set("request_id", json!("outer"));
158
159            current.with_attributes(
160                HashMap::from([(String::from("request_id"), json!("inner"))]),
161                || {
162                    assert_eq!(current.get("request_id"), Some(json!("inner")));
163                },
164            );
165
166            assert_eq!(current.get("request_id"), Some(json!("outer")));
167        });
168    }
169
170    #[test]
171    fn with_attributes_keeps_unrelated_values_visible_inside_scope() {
172        run_isolated(|| {
173            let current = CurrentAttributes::new();
174            current.set("tenant_id", json!("tenant-1"));
175
176            current.with_attributes(
177                HashMap::from([(String::from("request_id"), json!("abc123"))]),
178                || {
179                    assert_eq!(current.get("tenant_id"), Some(json!("tenant-1")));
180                    assert_eq!(current.get("request_id"), Some(json!("abc123")));
181                },
182            );
183        });
184    }
185
186    #[test]
187    fn with_attributes_can_be_nested() {
188        run_isolated(|| {
189            let current = CurrentAttributes::new();
190            current.set("request_id", json!("outer"));
191
192            current.with_attributes(
193                HashMap::from([(String::from("request_id"), json!("middle"))]),
194                || {
195                    assert_eq!(current.get("request_id"), Some(json!("middle")));
196                    current.with_attributes(
197                        HashMap::from([(String::from("request_id"), json!("inner"))]),
198                        || {
199                            assert_eq!(current.get("request_id"), Some(json!("inner")));
200                        },
201                    );
202                    assert_eq!(current.get("request_id"), Some(json!("middle")));
203                },
204            );
205
206            assert_eq!(current.get("request_id"), Some(json!("outer")));
207        });
208    }
209
210    #[test]
211    fn with_attributes_restores_state_after_panic() {
212        run_isolated(|| {
213            let current = CurrentAttributes::new();
214            current.set("request_id", json!("outer"));
215
216            let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
217                current.with_attributes(
218                    HashMap::from([(String::from("request_id"), json!("inner"))]),
219                    || panic!("boom"),
220                );
221            }));
222
223            assert!(result.is_err());
224            assert_eq!(current.get("request_id"), Some(json!("outer")));
225        });
226    }
227
228    #[test]
229    fn current_attributes_are_isolated_per_thread() {
230        let current = CurrentAttributes::new();
231        current.set("request_id", json!("main"));
232
233        let child_value = run_isolated(|| {
234            let current = CurrentAttributes::new();
235            assert_eq!(current.get("request_id"), None);
236            current.set("request_id", json!("child"));
237            current.get("request_id")
238        });
239
240        assert_eq!(child_value, Some(json!("child")));
241        assert_eq!(current.get("request_id"), Some(json!("main")));
242    }
243
244    #[test]
245    fn reset_only_affects_the_current_thread() {
246        let current = CurrentAttributes::new();
247        current.set("request_id", json!("main"));
248
249        run_isolated(|| {
250            let current = CurrentAttributes::new();
251            current.set("request_id", json!("child"));
252            current.reset();
253            assert_eq!(current.get("request_id"), None);
254        });
255
256        assert_eq!(current.get("request_id"), Some(json!("main")));
257    }
258
259    #[test]
260    fn snapshot_returns_a_clone_of_the_store() {
261        run_isolated(|| {
262            let current = CurrentAttributes::new();
263            current.set("request_id", json!("abc123"));
264            let mut snapshot = current.snapshot();
265            snapshot.insert(String::from("request_id"), json!("changed"));
266
267            assert_eq!(current.get("request_id"), Some(json!("abc123")));
268        });
269    }
270}
rustrails_support/current_attributes.rs

rustrails_support/
current_attributes.rs
