qlib_rs/data/
snowflake.rs

1use std::sync::atomic::{AtomicU64, Ordering};
2use std::time::{SystemTime, UNIX_EPOCH};
3use std::thread::yield_now;
4use std::env;
5
6const NODE_ID_BITS: u64 = 10;
7const SEQUENCE_BITS: u64 = 12;
8
9const MAX_SEQUENCE: u64 = (1 << SEQUENCE_BITS) - 1;
10
11const NODE_ID_SHIFT: u64 = SEQUENCE_BITS;
12const TIMESTAMP_SHIFT: u64 = SEQUENCE_BITS + NODE_ID_BITS;
13
14#[derive(Default)]
15pub struct Snowflake {
16    node_id: u64,
17    state: AtomicU64, // Packed state: (timestamp << 12) | sequence
18}
19
20impl Snowflake {
21    pub fn new() -> Self {
22        match env::var("Q_NODE_ID") {
23            Ok(value) => Self::from(value.parse::<u64>().unwrap_or(0)),
24            Err(_) => Self::from(0),
25        }
26    }
27
28    pub fn generate(&self) -> u64 {
29        loop {
30            let now = current_timestamp();
31
32            let last = self.state.load(Ordering::Relaxed);
33            let last_ts = last >> SEQUENCE_BITS;
34            let last_seq = last & MAX_SEQUENCE;
35
36            if now > last_ts {
37                let next = (now << SEQUENCE_BITS) | 0;
38                if self
39                    .state
40                    .compare_exchange(last, next, Ordering::SeqCst, Ordering::Relaxed)
41                    .is_ok()
42                {
43                    return (now << TIMESTAMP_SHIFT) | (self.node_id << NODE_ID_SHIFT);
44                }
45            } else if now == last_ts {
46                let seq = (last_seq + 1) & MAX_SEQUENCE;
47                if seq == 0 {
48                    // Sequence rollover; wait for next millisecond
49                    while current_timestamp() <= now {
50                        yield_now(); // avoid busy wait
51                    }
52                    continue;
53                }
54
55                let next = (last_ts << SEQUENCE_BITS) | seq;
56                if self
57                    .state
58                    .compare_exchange(last, next, Ordering::SeqCst, Ordering::Relaxed)
59                    .is_ok()
60                {
61                    return (now << TIMESTAMP_SHIFT)
62                        | (self.node_id << NODE_ID_SHIFT)
63                        | seq;
64                }
65            } else {
66                panic!("Clock moved backwards!");
67            }
68        }
69    }
70}
71
72impl From<u64> for Snowflake {
73    fn from(node_id: u64) -> Self {
74        Self {
75            node_id,
76            state: AtomicU64::new(0),
77        }
78    }
79}
80
81fn current_timestamp() -> u64 {
82    SystemTime::now()
83        .duration_since(UNIX_EPOCH)
84        .expect("System time before UNIX_EPOCH")
85        .as_millis() as u64
86}
qlib_rs/data/snowflake.rs

qlib_rs/data/
snowflake.rs
