openzeppelin_relayer/observability/
request_id.rs

1use tracing::Span;
2use tracing_subscriber::{registry::LookupSpan, Registry};
3
4#[derive(Clone, Debug)]
5pub struct RequestId(pub String);
6
7pub fn set_request_id(id: impl Into<String>) {
8    let id = id.into();
9    Span::current().with_subscriber(|(span_id, subscriber)| {
10        if let Some(reg) = subscriber.downcast_ref::<Registry>() {
11            if let Some(span_ref) = reg.span(span_id) {
12                span_ref.extensions_mut().replace(RequestId(id));
13            }
14        }
15    });
16}
17
18pub fn get_request_id() -> Option<String> {
19    let mut out = None;
20    Span::current().with_subscriber(|(span_id, subscriber)| {
21        if let Some(reg) = subscriber.downcast_ref::<Registry>() {
22            // Walk up the span tree: current → parent → grandparent → ...
23            // The RequestId is stored on the root span by RequestIdMiddleware,
24            // but child spans (e.g. from #[instrument]) don't inherit extensions.
25            let mut current = reg.span(span_id);
26            while let Some(span_ref) = current {
27                if let Some(r) = span_ref.extensions().get::<RequestId>() {
28                    out = Some(r.0.clone());
29                    return;
30                }
31                current = span_ref.parent();
32            }
33        }
34    });
35    out
36}
37
38#[cfg(test)]
39mod tests {
40    use super::*;
41    use tracing::info_span;
42    use tracing_subscriber::{fmt, prelude::*};
43
44    #[test]
45    fn set_and_get_request_id_within_span() {
46        tracing::subscriber::with_default(
47            tracing_subscriber::registry().with(fmt::layer()),
48            || {
49                let span = info_span!("test_span");
50                let _guard = span.enter();
51
52                set_request_id("abc-123");
53                assert_eq!(get_request_id().as_deref(), Some("abc-123"));
54            },
55        );
56    }
57
58    #[test]
59    fn overwrite_request_id_replaces_value() {
60        tracing::subscriber::with_default(
61            tracing_subscriber::registry().with(fmt::layer()),
62            || {
63                let span = info_span!("test_span");
64                let _guard = span.enter();
65
66                set_request_id("first");
67                assert_eq!(get_request_id().as_deref(), Some("first"));
68
69                set_request_id("second");
70                assert_eq!(get_request_id().as_deref(), Some("second"));
71            },
72        );
73    }
74
75    #[test]
76    fn get_request_id_traverses_parent_spans() {
77        tracing::subscriber::with_default(
78            tracing_subscriber::registry().with(fmt::layer()),
79            || {
80                let parent = info_span!("parent_span");
81                let _parent_guard = parent.enter();
82
83                set_request_id("parent-req-id");
84
85                // Enter a child span (simulates #[instrument] on a downstream function)
86                let child = info_span!("child_span");
87                let _child_guard = child.enter();
88
89                // Should find the RequestId on the parent span
90                assert_eq!(get_request_id().as_deref(), Some("parent-req-id"));
91            },
92        );
93    }
94
95    #[test]
96    fn get_request_id_prefers_closest_span() {
97        tracing::subscriber::with_default(
98            tracing_subscriber::registry().with(fmt::layer()),
99            || {
100                let parent = info_span!("parent_span");
101                let _parent_guard = parent.enter();
102                set_request_id("parent-id");
103
104                let child = info_span!("child_span");
105                let _child_guard = child.enter();
106                set_request_id("child-id");
107
108                // Should find the RequestId on the child span first
109                assert_eq!(get_request_id().as_deref(), Some("child-id"));
110            },
111        );
112    }
113
114    #[test]
115    fn get_request_id_traverses_multiple_levels() {
116        tracing::subscriber::with_default(
117            tracing_subscriber::registry().with(fmt::layer()),
118            || {
119                let root = info_span!("root_span");
120                let _root_guard = root.enter();
121                set_request_id("root-req-id");
122
123                let mid = info_span!("mid_span");
124                let _mid_guard = mid.enter();
125
126                let leaf = info_span!("leaf_span");
127                let _leaf_guard = leaf.enter();
128
129                // Should traverse leaf → mid → root and find it on root
130                assert_eq!(get_request_id().as_deref(), Some("root-req-id"));
131            },
132        );
133    }
134}
openzeppelin_relayer/observability/request_id.rs

openzeppelin_relayer/observability/
request_id.rs
