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

use std::{
    hash::Hasher,
    sync::{atomic::AtomicU64, OnceLock},
};

use hyper::{http::HeaderValue, Request, Response};

use crate::{helper_layers::function::Inner, SgBody};

/// Generate a `x-request-id` header for the request and response.
pub trait XRequestIdAlgo {
    fn generate() -> HeaderValue;
}
/// The header name for `x-request-id`.
pub const X_REQUEST_ID_HEADER_NAME: &str = "x-request-id";
/// Add a `x-request-id` header to the request and then response.
///
/// If the request already has a `x-request-id` header, it will be used.
pub async fn x_request_id<A: XRequestIdAlgo>(mut request: Request<SgBody>, inner: Inner) -> Response<SgBody> {
    let id = if let Some(id) = request.headers().get(X_REQUEST_ID_HEADER_NAME) {
        id.clone()
    } else {
        let id = A::generate();
        request.headers_mut().insert(X_REQUEST_ID_HEADER_NAME, id.clone());
        id
    };
    let mut resp = inner.call(request).await;
    resp.headers_mut().insert(X_REQUEST_ID_HEADER_NAME, id);
    resp
}
/// # Reference
/// - discord: https://discord.com/developers/docs/reference#snowflakes
/// - instagram: https://instagram-engineering.com/sharding-ids-at-instagram-1cf5a71e5a5c
/// # Bits
/// - 42: timestamp
/// - 10: machine id
/// - 12: increment id
pub struct Snowflake;

impl XRequestIdAlgo for Snowflake {
    fn generate() -> HeaderValue {
        static INC: AtomicU64 = AtomicU64::new(0);
        let ts_id = unsafe { std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap_unchecked().as_millis() as u64 } << 22;
        let mach_id = machine_id() << 12;
        let inc = INC.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        let id = ts_id | mach_id | inc;
        unsafe { HeaderValue::from_str(&format!("{:016x}", id)).unwrap_unchecked() }
    }
}

fn machine_id() -> u64 {
    static MACHINE_ID: OnceLock<u64> = OnceLock::new();
    *MACHINE_ID.get_or_init(|| {
        let mid = std::env::var("MACHINE_ID");
        let mut hasher = std::hash::DefaultHasher::new();
        if let Ok(mid) = mid {
            if let Ok(mid) = mid.parse::<u64>() {
                mid
            } else {
                hasher.write(mid.as_bytes());
                hasher.finish()
            }
        } else {
            #[cfg(target_os = "linux")]
            {
                // let's try to read system mid
                let mid = std::fs::read_to_string("/var/lib/dbus/machine-id").unwrap_or(rand::random::<u64>().to_string());
                hasher.write(mid.as_bytes());
                hasher.finish()
            }
            #[cfg(not(target_os = "linux"))]
            {
                // let's generate random one
                rand::random::<u64>()
            }
        }
    })
}