rocket/trace/subscriber/
request_id.rs

Help
1use std::fmt;
2use std::hash::{Hash, Hasher};
3use std::thread::ThreadId;
4use std::cell::Cell;
5
6use tracing::Subscriber;
7use tracing::span::{Attributes, Id};
8use tracing_subscriber::{layer::Context, Layer};
9use tracing_subscriber::registry::{LookupSpan, SpanRef};
10
11pub struct RequestIdLayer;
12
13#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)]
14pub struct RequestId(u128);
15
16#[derive(Default)]
17pub struct IdentHasher(u128);
18
19impl RequestId {
20    fn new() -> Self {
21        thread_local! {
22            pub static COUNTER: Cell<u64> = Cell::new(0);
23            pub static THREAD_ID: Cell<Option<ThreadId>> = Cell::new(None);
24        }
25
26        let thread_id = THREAD_ID.get().unwrap_or_else(|| {
27            let id = std::thread::current().id();
28            THREAD_ID.set(Some(id));
29            id
30        });
31
32        let local_id = COUNTER.get();
33        COUNTER.set(local_id.wrapping_add(1));
34
35        let mut hasher = IdentHasher::default();
36        thread_id.hash(&mut hasher);
37        local_id.hash(&mut hasher);
38        RequestId(hasher.0)
39    }
40
41    pub fn of<R: for<'a> LookupSpan<'a>>(span: &SpanRef<'_, R>) -> Option<Self> {
42        span.extensions().get::<Self>().copied()
43    }
44
45    pub fn current() -> Option<Self> {
46        RequestIdLayer::current()
47    }
48
49    fn short(&self) -> u32 {
50        let mut x = ((self.0 & (0xFFFFFFFF << 48)) >> 48) as u32;
51        x = (x ^ (x >> 16)).wrapping_mul(0x21f0aaad);
52        x = (x ^ (x >> 15)).wrapping_mul(0x735a2d97);
53        x = x ^ (x >> 15);
54        x
55    }
56
57    pub fn layer() -> RequestIdLayer {
58        RequestIdLayer
59    }
60}
61
62impl RequestIdLayer {
63    thread_local! {
64        static CURRENT_REQUEST_ID: Cell<Option<RequestId>> = Cell::new(None);
65    }
66
67    pub fn current() -> Option<RequestId> {
68        Self::CURRENT_REQUEST_ID.get()
69    }
70}
71
72impl<S: Subscriber + for<'a> LookupSpan<'a>> Layer<S> for RequestIdLayer {
73    fn on_new_span(&self, _: &Attributes<'_>, id: &Id, ctxt: Context<'_, S>) {
74        let span = ctxt.span(id).expect("new_span: span does not exist");
75        if span.name() == "request" {
76            span.extensions_mut().replace(RequestId::new());
77        }
78    }
79
80    fn on_enter(&self, id: &Id, ctxt: Context<'_, S>) {
81        let span = ctxt.span(id).expect("enter: span does not exist");
82        if span.name() == "request" {
83            Self::CURRENT_REQUEST_ID.set(RequestId::of(&span));
84        }
85    }
86
87    fn on_exit(&self, id: &Id, ctxt: Context<'_, S>) {
88        let span = ctxt.span(id).expect("enter: span does not exist");
89        if span.name() == "request" {
90            Self::CURRENT_REQUEST_ID.set(None);
91        }
92    }
93}
94
95impl fmt::Display for RequestId {
96    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
97        fmt::Display::fmt(&self.short(), f)
98    }
99}
100
101impl fmt::LowerHex for RequestId {
102    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
103        fmt::LowerHex::fmt(&self.short(), f)
104    }
105}
106
107impl fmt::UpperHex for RequestId {
108    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
109        fmt::UpperHex::fmt(&self.short(), f)
110    }
111}
112
113impl Hasher for IdentHasher {
114    fn finish(&self) -> u64 {
115        self.0 as u64
116    }
117
118    fn write(&mut self, bytes: &[u8]) {
119        for &byte in bytes {
120            self.0 = (self.0 << 8) | (byte as u128);
121        }
122    }
123
124    fn write_u64(&mut self, i: u64) {
125        // https://github.com/skeeto/hash-prospector
126        fn shuffle(mut x: u64) -> u64 {
127            x = x.wrapping_add(1);
128            x = (x ^ (x >> 30)).wrapping_mul(0xbf58476d1ce4e5b9);
129            x = (x ^ (x >> 27)).wrapping_mul(0x94d049bb133111eb);
130            x = x ^ (x >> 31);
131            x
132        }
133
134        self.0 = (self.0 << 64) | shuffle(i) as u128;
135    }
136}
rocket/trace/subscriber/request_id.rs

rocket/trace/subscriber/
request_id.rs
