rusty_snowflake/
generator.rs

1use super::Snowflake;
2
3#[derive(Debug, Clone, Copy, Eq, PartialEq)]
4/// A snowflake generator that can be used to generate snowflake IDs.
5pub struct SnowflakeGenerator {
6    /// The last snowflake ID generated by the generator.
7    pub last_snowflake: Snowflake,
8}
9
10impl SnowflakeGenerator {
11    /// Create a new snowflake generator with the given worker ID
12    ///
13    /// # Arguments
14    /// * `worker_id` - The worker ID of the snowflake generator
15    ///
16    /// # Returns
17    /// A new `SnowflakeGenerator`
18    ///
19    /// # Example
20    /// ```rust
21    /// use rusty_snowflake::{SnowflakeGenerator, Snowflake};
22    ///
23    /// const WORKER_ID: u64 = 420;
24    ///
25    /// let generator = SnowflakeGenerator::new(WORKER_ID);
26    ///
27    /// assert_eq!(generator.last_snowflake, Snowflake::new(WORKER_ID));
28    /// ```
29    pub fn new(worker_id: u64) -> SnowflakeGenerator {
30        SnowflakeGenerator {
31            last_snowflake: Snowflake::new(worker_id),
32        }
33    }
34
35    /// Start at the given snowflake
36    ///
37    /// # Arguments
38    /// * `snowflake` - The snowflake to start at
39    ///
40    /// # Returns
41    /// A `SnowflakeGenerator`
42    ///
43    /// # Example
44    /// ```rust
45    /// use rusty_snowflake::{Snowflake, SnowflakeGenerator};
46    ///
47    /// const WORKER_ID: u64 = 420;
48    ///
49    /// let generator = SnowflakeGenerator::start_at(&Snowflake::new(WORKER_ID));
50    ///
51    /// assert_eq!(generator.last_snowflake, Snowflake::new(WORKER_ID));
52    /// ```
53    pub fn start_at(snowflake: &Snowflake) -> SnowflakeGenerator {
54        SnowflakeGenerator {
55            last_snowflake: *snowflake,
56        }
57    }
58
59    /// Generates the next snowflake ID and returns a reference to it.
60    ///
61    /// This method generates the next snowflake ID by updating the `last_snowflake` stored in the generator.
62    /// It increments the sequence number and, if necessary, adjusts the timestamp to ensure uniqueness.
63    /// The generated snowflake ID is then returned as a reference.
64    ///
65    /// # Examples
66    ///
67    /// ```rust
68    /// use rusty_snowflake::SnowflakeGenerator;
69    ///
70    /// let mut generator = SnowflakeGenerator::new(420);
71    ///
72    /// // Generate the next snowflake ID
73    /// let snowflake = generator.next();
74    /// println!("Generated snowflake ID: {:?}", snowflake);
75    /// ```
76    ///
77    /// # Returns
78    ///
79    /// A reference to the generated snowflake ID.
80    ///
81    /// # Panics
82    ///
83    /// This method does not panic.
84    pub fn next(&mut self) -> &Snowflake {
85        self.last_snowflake = self.last_snowflake.next();
86        &self.last_snowflake
87    }
88
89    /// Get the current timestamp in seconds since the epoch (1970-01-01 00:00:00 UTC).
90    ///
91    /// # Returns
92    /// The current timestamp in seconds
93    pub fn get_timestamp() -> u64 {
94        std::time::SystemTime::now()
95            .duration_since(std::time::UNIX_EPOCH)
96            .expect("SystemTime before UNIX EPOCH!")
97            .as_secs()
98    }
99
100    /// Wait for the next second and return the timestamp
101    ///
102    /// # Arguments
103    /// * `current_timestamp` - The current timestamp in seconds
104    ///
105    /// # Returns
106    /// The timestamp of the next second
107    pub fn wait_next_timestamp(last_timestamp: u64) -> u64 {
108        let mut timestamp = SnowflakeGenerator::get_timestamp();
109        while timestamp <= last_timestamp {
110            std::thread::sleep(std::time::Duration::from_millis(1));
111            timestamp = SnowflakeGenerator::get_timestamp();
112        }
113        timestamp
114    }
115}
116
117#[cfg(test)]
118mod tests {
119
120    use super::*;
121
122    #[test]
123    fn test_new() {
124        let generator = SnowflakeGenerator::new(420);
125        assert_eq!(generator.last_snowflake, Snowflake::new(420));
126
127        let generator = SnowflakeGenerator::new(69);
128        assert_eq!(generator.last_snowflake, Snowflake::new(69));
129    }
130
131    #[test]
132    fn test_start_at() {
133        let snowflake = Snowflake::new(420);
134        let generator = SnowflakeGenerator::start_at(&snowflake);
135        assert_eq!(generator.last_snowflake, snowflake);
136    }
137
138    #[test]
139    fn test_next() {
140        let mut generator = SnowflakeGenerator::new(420);
141
142        for _ in 0..10 {
143            let snowflake = *generator.next();
144            assert_eq!(snowflake, generator.last_snowflake);
145        }
146    }
147
148    #[test]
149    fn test_eq() {
150        let gen1 = SnowflakeGenerator::new(420);
151        let gen2 = SnowflakeGenerator::new(420);
152        assert_eq!(gen1, gen2);
153
154        let gen1 = Snowflake::new(420);
155        let gen2 = Snowflake::new(69);
156        assert_ne!(gen1, gen2);
157    }
158
159    #[test]
160    fn test_snowflake_wait_next_timestamp() {
161        let timestamp = SnowflakeGenerator::get_timestamp();
162
163        let next = SnowflakeGenerator::wait_next_timestamp(timestamp);
164
165        assert!(
166            next > timestamp,
167            "Snowflake.wait_next_timestamp didn't return a new timestamp"
168        );
169    }
170
171    #[test]
172    fn test_get_timestamp() {
173        let timestamp = SnowflakeGenerator::get_timestamp();
174
175        assert!(
176            timestamp
177                == std::time::SystemTime::now()
178                    .duration_since(std::time::UNIX_EPOCH)
179                    .unwrap()
180                    .as_secs()
181        );
182    }
183}
rusty_snowflake/generator.rs

rusty_snowflake/
generator.rs
