rust_distributed_id/
rid_generator.rs

1use std::sync::{Arc, Mutex};
2use std::time::{SystemTime, UNIX_EPOCH};
3use rbatis::rbatis::Rbatis;
4
5use crate::worker::worker_assigner;
6use crate::bits_allocator;
7use crate::config::rid_config;
8
9#[derive(Debug)]
10pub struct UidGenerator {
11    bits_allocator:  bits_allocator::BitsAllocator,
12    config:          rid_config::UidConfig,
13    worker_id:       i64,
14    last_second:     Arc<Mutex<i64>>,
15    sequence :       i64,
16}
17
18impl UidGenerator {
19    //New create the default uid generator instance
20    pub  async fn new(config: &rid_config::UidConfig, rb: &'static Rbatis) -> Self {
21        let new_config = config.clone();
22        let id_assigner = worker_assigner::Assigner::new(new_config.port.clone(), rb);
23        let allocator = bits_allocator::BitsAllocator::new(
24            config.time_bits,
25            config.worker_bits,
26            config.seq_bits
27        );
28        let mut worker_id = id_assigner.assign_worker_id().await;
29
30        if worker_id > allocator.max_worker_id {
31            worker_id = worker_id % allocator.max_worker_id
32        }
33
34        return UidGenerator {
35            bits_allocator: allocator,
36            config: new_config,
37            worker_id,
38            last_second: Arc::new(Mutex::new(0)),
39            sequence: 0,
40        }
41    }
42
43    // get_uid generate the unique id
44    pub fn get_uid(& mut self) -> i64 {
45        let c = &self.config;
46        return self.next_id(c.epoch_seconds, c.max_backward_seconds,c.enable_backward);
47    }
48
49    // parse_uid parse the generated unique id then get the meta information
50    // +------+----------------------+----------------+-----------+
51    // | sign |     delta seconds    | worker node id | sequence  |
52    // +------+----------------------+----------------+-----------+
53    //   1bit          30bits              7bits         13bits
54    pub fn parse_uid(&self, uid: i64) -> String {
55        let total_bits = bits_allocator::TOTAL_BITS;
56        let sign_bits = self.bits_allocator.sign_bits;
57        let timestamp_bits = self.bits_allocator.timestamp_bits;
58        let worker_id_bits = self.bits_allocator.worker_id_bits;
59        let sequence_bits = self.bits_allocator.sequence_bits;
60        let allocate_total_bits = self.bits_allocator.allocate_total_bits;
61
62        let sequence = (uid << (total_bits - sequence_bits)) >> (total_bits - sequence_bits);
63        let worker_id = (uid << (timestamp_bits + sign_bits + total_bits - allocate_total_bits))
64            >> (total_bits - worker_id_bits);
65        let delta_seconds = uid >> (worker_id_bits + sequence_bits);
66
67        // format as string
68        return format!(
69            r#"{{\"uid\":\"{}\",\"timestamp\":\"{}\",\"worker_id\":\"{}\",\"sequence\":\"{}\"}}"#,
70            uid, self.config.epoch_seconds + delta_seconds, worker_id, sequence
71        );
72    }
73
74    fn next_id(
75        & mut self,
76        epoch_seconds: i64,
77        max_backward_seconds: i64,
78        enable_backward: bool
79    ) -> i64 {
80        let mut last_second = self.last_second.lock().unwrap();
81        let mut current_second = self.get_current_second(epoch_seconds);
82
83        if current_second < *last_second {
84            if !enable_backward {
85                panic!("Clock moved backwards. But backward not enabled");
86            }
87
88            let refused_seconds = *last_second - current_second;
89            if refused_seconds <= max_backward_seconds {
90                while current_second < *last_second {
91                    current_second = self.get_current_second(epoch_seconds)
92                }
93            } else {
94                panic!("Clock moved backwards. Refused seconds bigger than max backward seconds")
95            }
96        } else if current_second == *last_second { // At the same second, increase sequence
97            self.sequence = (self.sequence + 1) & self.bits_allocator.max_sequence;
98            // Exceed the max sequence, we wait the next second to generate uid
99            if self.sequence == 0 {
100                current_second = self.get_next_second(*last_second, epoch_seconds);
101            }
102        } else {
103            // At the different second, sequence restart from zero
104            self.sequence = 0i64;
105        }
106
107        *last_second = current_second;
108
109        // Allocate bits for uid
110        self.bits_allocator.allocate(
111            current_second - epoch_seconds,
112            self.worker_id,
113            self.sequence
114        )
115    }
116
117    #[inline(always)]
118    fn get_current_second(&self, epoch_seconds: i64) -> i64 {
119        let current_seconds = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs() as i64;
120        if current_seconds - epoch_seconds > self.bits_allocator.max_delta_seconds {
121            panic!("Timestamp bits is exhausted. Refusing uid generation.")
122        }
123
124        return current_seconds
125    }
126
127    #[inline(always)]
128    fn get_next_second(&self, last_timestamp: i64, epoch_seconds: i64) ->i64 {
129        let mut timestamp = self.get_current_second(epoch_seconds);
130
131        while timestamp <= last_timestamp {
132            timestamp = self.get_current_second(epoch_seconds)
133        }
134
135        return timestamp
136    }
137
138}
139
140
141
142
143