snowflake/
lib.rs

1//! Rust version of the `Twitter snowflake algorithm` .
2//!
3
4use std::hint::spin_loop;
5use std::time::{SystemTime, UNIX_EPOCH};
6
7/// The `SnowflakeIdGenerator` type is snowflake algorithm wrapper.
8#[derive(Copy, Clone, Debug)]
9pub struct SnowflakeIdGenerator {
10    /// epoch used by the snowflake algorithm.
11    epoch: SystemTime,
12
13    /// last_time_millis, last time generate id is used times millis.
14    last_time_millis: i64,
15
16    /// machine_id, is use to supplement id machine or sectionalization attribute.
17    pub machine_id: i32,
18
19    /// node_id, is use to supplement id machine-node attribute.
20    pub node_id: i32,
21
22    /// auto-increment record.
23    idx: u16,
24}
25
26/// The `SnowflakeIdBucket` type is snowflake-id-bucket it easy to get id also have a id buffer.
27#[derive(Clone, Debug)]
28pub struct SnowflakeIdBucket {
29    /// Hidden the `SnowflakeIdGenerator` in bucket .
30    snowflake_id_generator: SnowflakeIdGenerator,
31
32    /// The bucket buffer;
33    bucket: Vec<i64>,
34}
35
36impl SnowflakeIdGenerator {
37    /// Constructs a new `SnowflakeIdGenerator` using the UNIX epoch.
38    /// Please make sure that machine_id and node_id is small than 32(2^5);
39    ///
40    /// # Examples
41    ///
42    /// ```
43    /// use snowflake::SnowflakeIdGenerator;
44    ///
45    /// let id_generator = SnowflakeIdGenerator::new(1, 1);
46    /// ```
47    pub fn new(machine_id: i32, node_id: i32) -> SnowflakeIdGenerator {
48        Self::with_epoch(machine_id, node_id, UNIX_EPOCH)
49    }
50
51    /// Constructs a new `SnowflakeIdGenerator` using the specified epoch.
52    /// Please make sure that machine_id and node_id is small than 32(2^5);
53    ///
54    /// # Examples
55    ///
56    /// ```
57    /// use std::time::{Duration, UNIX_EPOCH};
58    /// use snowflake::SnowflakeIdGenerator;
59    ///
60    /// // 1 January 2015 00:00:00
61    /// let discord_epoch = UNIX_EPOCH + Duration::from_millis(1420070400000);
62    /// let id_generator = SnowflakeIdGenerator::with_epoch(1, 1, discord_epoch);
63    /// ```
64    pub fn with_epoch(machine_id: i32, node_id: i32, epoch: SystemTime) -> SnowflakeIdGenerator {
65        //TODO:limit the maximum of input args machine_id and node_id
66        let last_time_millis = get_time_millis(epoch);
67
68        SnowflakeIdGenerator {
69            epoch,
70            last_time_millis,
71            machine_id,
72            node_id,
73            idx: 0,
74        }
75    }
76
77    /// The real_time_generate keep id generate time is eq call method time.
78    ///
79    /// # Examples
80    ///
81    /// ```
82    /// use snowflake::SnowflakeIdGenerator;
83    ///
84    /// let mut id_generator = SnowflakeIdGenerator::new(1, 1);
85    /// id_generator.real_time_generate();
86    /// ```
87    pub fn real_time_generate(&mut self) -> i64 {
88        self.idx = (self.idx + 1) % 4096;
89
90        let mut now_millis = get_time_millis(self.epoch);
91
92        // supplement code for 'clock is moving backwards situation'.
93
94        // If the milliseconds of the current clock are equal to
95        // the number of milliseconds of the most recently generated id,
96        // then check if enough 4096 are generated,
97        // if enough then busy wait until the next millisecond.
98        if now_millis == self.last_time_millis {
99            if self.idx == 0 {
100                now_millis = biding_time_conditions(self.last_time_millis, self.epoch);
101                self.last_time_millis = now_millis;
102            }
103        } else {
104            self.last_time_millis = now_millis;
105            self.idx = 0;
106        }
107
108        // last_time_millis is 64 bits,left shift 22 bit,store 42 bits , machine_id left shift 17 bits,
109        // node_id left shift 12 bits ,idx complementing bits.
110        self.last_time_millis << 22
111            | ((self.machine_id << 17) as i64)
112            | ((self.node_id << 12) as i64)
113            | (self.idx as i64)
114    }
115
116    /// The basic guarantee time punctuality.
117    ///
118    /// Basic guarantee time punctuality.
119    /// sometimes one millis can't use up 4096 ID, the property of the ID isn't real-time.
120    /// But setting time after every 4096 calls.
121    /// # Examples
122    ///
123    /// ```
124    /// use snowflake::SnowflakeIdGenerator;
125    ///
126    /// let mut id_generator = SnowflakeIdGenerator::new(1, 1);
127    /// id_generator.generate();
128    /// ```
129    pub fn generate(&mut self) -> i64 {
130        self.idx = (self.idx + 1) % 4096;
131
132        // Maintenance `last_time_millis` for every 4096 ids generated.
133        if self.idx == 0 {
134            let mut now_millis = get_time_millis(self.epoch);
135
136            if now_millis == self.last_time_millis {
137                now_millis = biding_time_conditions(self.last_time_millis, self.epoch);
138            }
139
140            self.last_time_millis = now_millis;
141        }
142
143        //last_time_millis is 64 bits,left shift 22 bit,store 42 bits , machine_id left shift 17 bits,
144        //node_id left shift 12 bits ,idx complementing bits.
145        self.last_time_millis << 22
146            | ((self.machine_id << 17) as i64)
147            | ((self.node_id << 12) as i64)
148            | (self.idx as i64)
149    }
150
151    /// The lazy generate.
152    ///
153    /// Lazy generate.
154    /// Just start time record last_time_millis it consume every millis ID.
155    /// Maybe faster than standing time.
156    /// # Examples
157    ///
158    /// ```
159    /// use snowflake::SnowflakeIdGenerator;
160    ///
161    /// let mut id_generator = SnowflakeIdGenerator::new(1, 1);
162    /// id_generator.lazy_generate();
163    /// ```
164    pub fn lazy_generate(&mut self) -> i64 {
165        self.idx = (self.idx + 1) % 4096;
166
167        if self.idx == 0 {
168            self.last_time_millis += 1;
169        }
170
171        self.last_time_millis << 22
172            | ((self.machine_id << 17) as i64)
173            | ((self.node_id << 12) as i64)
174            | (self.idx as i64)
175    }
176}
177
178impl SnowflakeIdBucket {
179    /// Constructs a new `SnowflakeIdBucket` using the UNIX epoch.
180    /// Please make sure that machine_id and node_id is small than 32(2^5);
181    ///
182    /// # Examples
183    ///
184    /// ```
185    /// use snowflake::SnowflakeIdBucket;
186    ///
187    /// let id_generator_bucket = SnowflakeIdBucket::new(1, 1);
188    /// ```
189    pub fn new(machine_id: i32, node_id: i32) -> Self {
190        Self::with_epoch(machine_id, node_id, UNIX_EPOCH)
191    }
192
193    /// Constructs a new `SnowflakeIdBucket` using the specified epoch.
194    /// Please make sure that machine_id and node_id is small than 32(2^5);
195    ///
196    /// # Examples
197    ///
198    /// ```
199    /// use std::time::{Duration, UNIX_EPOCH};
200    /// use snowflake::SnowflakeIdBucket;
201    ///
202    /// // 1 January 2015 00:00:00
203    /// let discord_epoch = UNIX_EPOCH + Duration::from_millis(1420070400000);
204    /// let id_generator_bucket = SnowflakeIdBucket::with_epoch(1, 1, discord_epoch);
205    /// ```
206    pub fn with_epoch(machine_id: i32, node_id: i32, epoch: SystemTime) -> Self {
207        let snowflake_id_generator = SnowflakeIdGenerator::with_epoch(machine_id, node_id, epoch);
208        let bucket = Vec::new();
209
210        SnowflakeIdBucket {
211            snowflake_id_generator,
212            bucket,
213        }
214    }
215
216    /// Generate id.
217    ///
218    /// # Examples
219    ///
220    /// ```
221    /// use snowflake::SnowflakeIdBucket;
222    ///
223    /// let mut id_generator_bucket = SnowflakeIdBucket::new(1, 1);
224    /// let id = id_generator_bucket.get_id();
225    ///
226    /// ```
227    pub fn get_id(&mut self) -> i64 {
228        // 247 ns/iter
229        // after self.bucket.push(self.snowflake_id_generator.generate());
230
231        // 7 ns/iter
232        // after self.bucket.push(self.snowflake_id_generator.lazy_generate());
233
234        //500 ns/iter
235        // after self.bucket.push(self.snowflake_id_generator.real_time_generate());
236        if self.bucket.is_empty() {
237            self.generate_ids();
238        }
239        self.bucket.pop().unwrap()
240    }
241
242    fn generate_ids(&mut self) {
243        // 30,350 -- 50,000 ns/iter
244        //self.bucket.push(self.snowflake_id_generator.lazy_generate());
245
246        // 1,107,103 -- 1,035,018 ns/iter
247        //self.bucket.push(self.snowflake_id_generator.generate());
248
249        // 2,201,325 -- 2,082,187 ns/iter
250        //self.bucket.push(self.snowflake_id_generator.real_time_generate());
251
252        for _ in 0..4091 {
253            self.bucket
254                .push(self.snowflake_id_generator.lazy_generate());
255        }
256    }
257}
258
259#[inline(always)]
260/// Get the latest milliseconds of the clock.
261pub fn get_time_millis(epoch: SystemTime) -> i64 {
262    SystemTime::now()
263        .duration_since(epoch)
264        .expect("Time went mackward")
265        .as_millis() as i64
266}
267
268#[inline(always)]
269// Constantly refreshing the latest milliseconds by busy waiting.
270fn biding_time_conditions(last_time_millis: i64, epoch: SystemTime) -> i64 {
271    let mut latest_time_millis: i64;
272    loop {
273        latest_time_millis = get_time_millis(epoch);
274        if latest_time_millis > last_time_millis {
275            return latest_time_millis;
276        }
277        spin_loop();
278    }
279}