snowgen/
snowgen.rs

1//! Snowgen is a module that provides a struct for generating unique, distributed IDs.
2
3use std::time::{Duration, Instant};
4
5/// Snowgen struct represents a unique ID generator.
6pub struct Snowgen {
7    pub(crate) node_id: i32,
8    pub(crate) machine_id: Option<i32>,
9    pub(crate) epoch: Instant,
10    #[allow(dead_code)]
11    pub(crate) timestamp_bits: u8,
12    pub(crate) node_id_bits: u8,
13    pub(crate) machine_id_bits: u8,
14    pub(crate) sequence_bits: u8,
15    pub(crate) last_timestamp: i64,
16    pub(crate) sequence: u16,
17}
18
19impl Snowgen {
20    /// Generate the next unique ID.
21    ///
22    /// # Returns
23    ///
24    /// A Result containing a u64 unique ID or an error message.
25    pub fn next_id(&mut self) -> Result<u64, &'static str> {
26        let timestamp = self.current_timestamp();
27        let seq;
28
29        if timestamp == self.last_timestamp {
30            seq = self.sequence + 1;
31
32            if seq >= 1 << self.sequence_bits {
33                std::thread::sleep(Duration::from_micros(1));
34                return self.next_id();
35            }
36        } else {
37            seq = 0;
38            self.last_timestamp = timestamp;
39        }
40
41        self.sequence = seq;
42
43        // Combine the timestamp, node_id, machine_id (if provided), and sequence to generate a unique ID
44        let id = ((timestamp as u64)
45            << (self.node_id_bits + self.machine_id_bits + self.sequence_bits))
46            | ((self.node_id as u64) << (self.machine_id_bits + self.sequence_bits))
47            | ((self.machine_id.unwrap_or(0) as u64) << self.sequence_bits)
48            | seq as u64;
49
50        Ok(id)
51    }
52
53    /// Get the current timestamp based on the epoch.
54    ///
55    /// # Returns
56    ///
57    /// An i64 timestamp.
58    #[inline(always)]
59    fn current_timestamp(&self) -> i64 {
60        self.epoch.elapsed().as_millis() as i64
61    }
62}