moduforge_model/
id_generator.rs

1use std::sync::{Mutex, OnceLock};
2use std::time::{SystemTime, UNIX_EPOCH};
3
4pub struct IdGenerator {
5    data_center_id_shift: u64,
6    worker_id_shift: u64,
7    timestamp_left_shift: u64,
8    sequence: u64,
9    last_timestamp: i64,
10    options: Options,
11}
12
13struct Options {
14    start_time: i64,
15    data_center_id_bits: u64,
16    worker_id_bits: u64,
17    sequence_bits: u64,
18    worker_id: u64,
19    data_center_id: u64,
20}
21
22impl IdGenerator {
23    fn new() -> Self {
24        let options = Options {
25            start_time: 0,
26            data_center_id_bits: 5,
27            worker_id_bits: 5,
28            sequence_bits: 12,
29            worker_id: 0,
30            data_center_id: 0,
31        };
32
33        let data_center_id_shift =
34            options.worker_id_bits + options.sequence_bits;
35        let worker_id_shift = options.sequence_bits;
36        let timestamp_left_shift = options.worker_id_bits
37            + options.sequence_bits
38            + options.data_center_id_bits;
39        IdGenerator {
40            data_center_id_shift,
41            worker_id_shift,
42            timestamp_left_shift,
43            sequence: 0,
44            last_timestamp: -1,
45            options,
46        }
47    }
48
49    pub fn get_instance() -> &'static Mutex<IdGenerator> {
50        static INSTANCE: OnceLock<Mutex<IdGenerator>> = OnceLock::new();
51
52        INSTANCE.get_or_init(|| Mutex::new(IdGenerator::new()))
53    }
54    pub fn get_id() -> String {
55        let id = {
56            let mut id_generator = IdGenerator::get_instance().lock().unwrap();
57            id_generator.get_next_id()
58        };
59        id
60    }
61
62    pub fn get_next_id(&mut self) -> String {
63        let timestamp = self.get_timestamp();
64
65        if timestamp < self.last_timestamp {
66            panic!("Clock moved backwards");
67        }
68
69        if timestamp == self.last_timestamp {
70            self.sequence = (self.sequence + 1) & self.max_sequence();
71            if self.sequence == 0 {
72                self.last_timestamp = self.next_millis(self.last_timestamp);
73            }
74        } else {
75            self.sequence = 0;
76        }
77
78        self.last_timestamp = timestamp;
79
80        self.generate_id(timestamp)
81    }
82
83    fn next_millis(
84        &self,
85        last_timestamp: i64,
86    ) -> i64 {
87        let mut timestamp = self.get_timestamp();
88
89        while timestamp <= last_timestamp {
90            timestamp = self.get_timestamp();
91        }
92        timestamp
93    }
94
95    fn generate_id(
96        &self,
97        timestamp: i64,
98    ) -> String {
99        let mut id = (timestamp as u128) << self.timestamp_left_shift;
100        id |=
101            (self.options.data_center_id as u128) << self.data_center_id_shift;
102        id |= (self.options.worker_id as u128) << self.worker_id_shift;
103        id |= self.sequence as u128;
104        id.to_string()
105    }
106
107    fn get_timestamp(&self) -> i64 {
108        let now = SystemTime::now()
109            .duration_since(UNIX_EPOCH)
110            .expect("Time went backwards");
111        now.as_millis() as i64 - self.options.start_time
112    }
113
114    fn max_sequence(&self) -> u64 {
115        !((-1_i64 << self.options.sequence_bits) as u64)
116    }
117}