1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136

use std::fmt;
use std::hash::{Hash, Hasher};
use std::thread::ThreadId;
use std::cell::Cell;

use tracing::Subscriber;
use tracing::span::{Attributes, Id};
use tracing_subscriber::{layer::Context, Layer};
use tracing_subscriber::registry::{LookupSpan, SpanRef};

pub struct RequestIdLayer;

#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)]
pub struct RequestId(u128);

#[derive(Default)]
pub struct IdentHasher(u128);

impl RequestId {
    fn new() -> Self {
        thread_local! {
            pub static COUNTER: Cell<u64> = Cell::new(0);
            pub static THREAD_ID: Cell<Option<ThreadId>> = Cell::new(None);
        }

        let thread_id = THREAD_ID.get().unwrap_or_else(|| {
            let id = std::thread::current().id();
            THREAD_ID.set(Some(id));
            id
        });

        let local_id = COUNTER.get();
        COUNTER.set(local_id.wrapping_add(1));

        let mut hasher = IdentHasher::default();
        thread_id.hash(&mut hasher);
        local_id.hash(&mut hasher);
        RequestId(hasher.0)
    }

    pub fn of<R: for<'a> LookupSpan<'a>>(span: &SpanRef<'_, R>) -> Option<Self> {
        span.extensions().get::<Self>().copied()
    }

    pub fn current() -> Option<Self> {
        RequestIdLayer::current()
    }

    fn short(&self) -> u32 {
        let mut x = ((self.0 & (0xFFFFFFFF << 48)) >> 48) as u32;
        x = (x ^ (x >> 16)).wrapping_mul(0x21f0aaad);
        x = (x ^ (x >> 15)).wrapping_mul(0x735a2d97);
        x = x ^ (x >> 15);
        x
    }

    pub fn layer() -> RequestIdLayer {
        RequestIdLayer
    }
}

impl RequestIdLayer {
    thread_local! {
        static CURRENT_REQUEST_ID: Cell<Option<RequestId>> = Cell::new(None);
    }

    pub fn current() -> Option<RequestId> {
        Self::CURRENT_REQUEST_ID.get()
    }
}

impl<S: Subscriber + for<'a> LookupSpan<'a>> Layer<S> for RequestIdLayer {
    fn on_new_span(&self, _: &Attributes<'_>, id: &Id, ctxt: Context<'_, S>) {
        let span = ctxt.span(id).expect("new_span: span does not exist");
        if span.name() == "request" {
            span.extensions_mut().replace(RequestId::new());
        }
    }

    fn on_enter(&self, id: &Id, ctxt: Context<'_, S>) {
        let span = ctxt.span(id).expect("enter: span does not exist");
        if span.name() == "request" {
            Self::CURRENT_REQUEST_ID.set(RequestId::of(&span));
        }
    }

    fn on_exit(&self, id: &Id, ctxt: Context<'_, S>) {
        let span = ctxt.span(id).expect("enter: span does not exist");
        if span.name() == "request" {
            Self::CURRENT_REQUEST_ID.set(None);
        }
    }
}

impl fmt::Display for RequestId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt(&self.short(), f)
    }
}

impl fmt::LowerHex for RequestId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::LowerHex::fmt(&self.short(), f)
    }
}

impl fmt::UpperHex for RequestId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::UpperHex::fmt(&self.short(), f)
    }
}

impl Hasher for IdentHasher {
    fn finish(&self) -> u64 {
        self.0 as u64
    }

    fn write(&mut self, bytes: &[u8]) {
        for &byte in bytes {
            self.0 = (self.0 << 8) | (byte as u128);
        }
    }

    fn write_u64(&mut self, i: u64) {
        // https://github.com/skeeto/hash-prospector
        fn shuffle(mut x: u64) -> u64 {
            x = x.wrapping_add(1);
            x = (x ^ (x >> 30)).wrapping_mul(0xbf58476d1ce4e5b9);
            x = (x ^ (x >> 27)).wrapping_mul(0x94d049bb133111eb);
            x = x ^ (x >> 31);
            x
        }

        self.0 = (self.0 << 64) | shuffle(i) as u128;
    }
}