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            if let Some(span_ref) = reg.span(span_id) {
23                if let Some(r) = span_ref.extensions().get::<RequestId>() {
24                    out = Some(r.0.clone());
25                }
26            }
27        }
28    });
29    out
30}
31
32#[cfg(test)]
33mod tests {
34    use super::*;
35    use tracing::info_span;
36    use tracing_subscriber::{fmt, prelude::*};
37
38    #[test]
39    fn set_and_get_request_id_within_span() {
40        tracing::subscriber::with_default(
41            tracing_subscriber::registry().with(fmt::layer()),
42            || {
43                let span = info_span!("test_span");
44                let _guard = span.enter();
45
46                set_request_id("abc-123");
47                assert_eq!(get_request_id().as_deref(), Some("abc-123"));
48            },
49        );
50    }
51
52    #[test]
53    fn overwrite_request_id_replaces_value() {
54        tracing::subscriber::with_default(
55            tracing_subscriber::registry().with(fmt::layer()),
56            || {
57                let span = info_span!("test_span");
58                let _guard = span.enter();
59
60                set_request_id("first");
61                assert_eq!(get_request_id().as_deref(), Some("first"));
62
63                set_request_id("second");
64                assert_eq!(get_request_id().as_deref(), Some("second"));
65            },
66        );
67    }
68}