rust_snowflake/
lib.rs

1use std::time::{SystemTime, UNIX_EPOCH};
2
3const WORKER_ID_BITS: i64 = 5;
4const DATACENTER_ID_BITS: i64 = 5;
5const SEQUENCE_BITS: i64 = 12;
6
7const MAX_WORKER_ID: i64 = -1 ^ (-1 << WORKER_ID_BITS);
8const MAX_DATACENTER_ID: i64 = -1 ^ (-1 << DATACENTER_ID_BITS);
9
10const WORKER_ID_SHIFT: i64 = SEQUENCE_BITS;
11const DATACENTER_ID_SHIFT: i64 = SEQUENCE_BITS + WORKER_ID_BITS;
12const TIMESTAMP_LEFT_SHIFT: i64 = SEQUENCE_BITS + WORKER_ID_BITS + DATACENTER_ID_BITS;
13
14const SEQUENCE_MASK: i64 = -1 ^ (-1 << SEQUENCE_BITS);
15
16#[derive(Debug)]
17pub struct IdWorker {
18    worker_id: i64,
19    datacenter_id: i64,
20    sequence: i64,
21    last_timestamp: i64,
22    start_offset_timestamp: i64,
23}
24
25impl IdWorker {
26    pub fn new(worker_id: i64, datacenter_id: i64, start_offset_timestamp: i64) -> Self {
27        if worker_id < 0 || worker_id > MAX_WORKER_ID {
28            panic!(
29                "IdWorker: worker_id check failed: {}, MAX: {}",
30                worker_id, MAX_WORKER_ID
31            );
32        }
33
34        if datacenter_id < 0 || datacenter_id > MAX_DATACENTER_ID {
35            panic!(
36                "IdWorker: datacenter_id check failed: {}, MAX: {}",
37                datacenter_id, MAX_DATACENTER_ID
38            );
39        }
40
41        IdWorker {
42            worker_id,
43            datacenter_id,
44            sequence: 0,
45            last_timestamp: 0,
46            start_offset_timestamp,
47        }
48    }
49
50    pub fn next_id(&mut self) -> i64 {
51        let mut timestamp = self.gen_time();
52        assert!(timestamp >= self.last_timestamp);
53
54        if timestamp == self.last_timestamp {
55            self.sequence = (self.sequence + 1) & SEQUENCE_MASK;
56
57            // overflow and block until next millisecond
58            if self.sequence == 0 {
59                loop {
60                    timestamp = self.gen_time();
61                    if timestamp > self.last_timestamp {
62                        break;
63                    }
64                }
65            }
66        } else {
67            self.sequence = 0;
68        }
69
70        self.last_timestamp = timestamp;
71
72        (timestamp << TIMESTAMP_LEFT_SHIFT)
73            | (self.datacenter_id << DATACENTER_ID_SHIFT)
74            | (self.worker_id << WORKER_ID_SHIFT)
75            | self.sequence
76    }
77
78    pub fn get_location(id: i64) -> (i64, i64) {
79        let mut c_id = id;
80        c_id >>= SEQUENCE_BITS;
81        c_id = Self::slice(c_id, 64 - WORKER_ID_BITS - DATACENTER_ID_BITS);
82
83        let worker_id = c_id >> 5;
84        let dc_id = Self::slice(c_id, 64 - DATACENTER_ID_BITS);
85
86        (worker_id, dc_id)
87    }
88
89    fn slice(id: i64, offset: i64) -> i64 {
90        let c_id = id << offset - 1;
91        let d_id = if c_id < 0 {
92            c_id ^ (0 - std::i64::MAX - 1)
93        } else {
94            c_id
95        };
96
97        d_id >> offset - 1
98    }
99
100
101    fn gen_time(&self) -> i64 {
102        let duration = SystemTime::now().duration_since(UNIX_EPOCH).unwrap();
103        let time = (duration.as_secs() * 1000 + duration.subsec_nanos() as u64 / 1_000_000) as i64
104            - self.start_offset_timestamp;
105
106        if time < 0 {
107            panic!("IdWorker: can't get correct time. current: {}", time);
108        }
109
110        time as i64
111    }
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117    #[test]
118    fn it_works() {
119        let mut worker = IdWorker::new(31, 31, 1586497735375);
120        let (dc_id, worker_id) = IdWorker::get_location(worker.next_id());
121        
122        assert_eq!(worker_id, 31);
123        assert_eq!(dc_id, 31);
124    }
125}
rust_snowflake/lib.rs

rust_snowflake/
lib.rs
