pika_id/
snowflake.rs

1use crate::utils::now_timestamp;
2
3#[derive(Clone, Debug, Copy)]
4pub struct Snowflake {
5    epoch: u64,
6    node_id: u32,
7    seq: u32,
8    last_sequence_exhaustion: u64,
9}
10
11#[derive(Clone, Debug, Copy)]
12pub struct DecodedSnowflake {
13    pub id: u64,
14    pub timestamp: u64,
15    pub node_id: u32,
16    pub seq: u64,
17    pub epoch: u64,
18}
19
20impl Snowflake {
21    pub fn new_with_nodeid(epoch: u64, node_id: u32) -> Self {
22        Self {
23            epoch,
24            node_id,
25            seq: 0,
26            last_sequence_exhaustion: 0,
27        }
28    }
29
30    #[inline]
31    pub fn gen(self) -> String {
32        self.gen_with_ts(now_timestamp())
33    }
34
35    pub fn gen_with_ts(mut self, timestamp: u64) -> String {
36        if self.seq >= 4095 && timestamp == self.last_sequence_exhaustion {
37            while now_timestamp() - timestamp < 1 {
38                continue;
39            }
40        }
41
42        self.seq = if self.seq >= 4095 { 0 } else { self.seq + 1 };
43
44        if self.seq == 4095 {
45            self.last_sequence_exhaustion = timestamp;
46        }
47
48        let sf = ((timestamp - self.epoch) << 22)
49            | (u64::from(self.node_id) << 12)
50            | u64::from(self.seq);
51
52        sf.to_string()
53    }
54
55    pub fn decode(self, sf: &str) -> DecodedSnowflake {
56        let sf = sf.parse::<u64>().unwrap();
57        let timestamp = (sf >> 22) + self.epoch;
58        let node_id = (sf >> 12) & 0b11_1111_1111;
59        let seq = sf & 0b1111_1111_1111;
60
61        DecodedSnowflake {
62            id: sf,
63            timestamp,
64            node_id: node_id as u32,
65            seq,
66            epoch: self.epoch,
67        }
68    }
69}
70
71mod test {
72    #[test]
73    fn generate_snowflake() {
74        // if the node_id >= 1024 it will go to 0?
75        let sf = super::Snowflake::new_with_nodeid(650_153_600_000, 1023);
76        let snowflake = sf.gen();
77
78        let deconstruct = sf.decode(&snowflake);
79
80        assert_eq!(deconstruct.epoch, 650_153_600_000);
81        assert_eq!(deconstruct.node_id, 1023);
82    }
83}
pika_id/snowflake.rs

pika_id/
snowflake.rs
