rust_mcp_extra/id_generator/
snow_flake_id_generator.rs

1//! Medium size ,Globally unique , Time-sortable , Compact (64 bits),
2//! Use case: Distributed systems needing high-throughput, unique IDs without collisions.
3//! [ timestamp (41 bits) | machine id (10 bits) | sequence (12 bits) ]
4
5use once_cell::sync::Lazy;
6use rust_mcp_sdk::id_generator::IdGenerator;
7use std::sync::atomic::{AtomicU64, Ordering};
8use std::time::{SystemTime, UNIX_EPOCH};
9
10/// Epoch (customizable to reduce total bits needed)
11static SHORTER_EPOCH: Lazy<u64> = Lazy::new(|| {
12    SystemTime::now()
13        .duration_since(UNIX_EPOCH)
14        .expect("invalid system time!")
15        .as_millis() as u64
16});
17
18/// A Snowflake ID generator implementation producing 64-bit unique IDs.
19///
20/// Snowflake IDs are composed of:
21/// - A timestamp in milliseconds since a custom epoch (usually a fixed past time),
22/// - A machine ID (or worker ID) to differentiate between nodes,
23/// - A sequence number that increments within the same millisecond to avoid collisions.
24///
25/// Format (64 bits total):
26/// - 41 bits: timestamp (ms since SHORTER_EPOCH)
27/// - 10 bits: machine ID (0-1023)
28/// - 12 bits: sequence number (per ms)
29///
30/// This generator ensures:
31/// - Uniqueness across multiple machines (given unique machine IDs),
32/// - Monotonic increasing IDs when generated in the same process,
33/// - Thread safety with internal locking.
34pub struct SnowflakeIdGenerator {
35    machine_id: u16, // 10 bits max
36    last_timestamp: AtomicU64,
37    sequence: AtomicU64,
38}
39
40impl SnowflakeIdGenerator {
41    pub fn new(machine_id: u16) -> Self {
42        assert!(
43            machine_id < 1024,
44            "Machine ID must be less than 1024 (10 bits)"
45        );
46        SnowflakeIdGenerator {
47            machine_id,
48            last_timestamp: AtomicU64::new(0),
49            sequence: AtomicU64::new(0),
50        }
51    }
52
53    fn current_timestamp(&self) -> u64 {
54        let now = SystemTime::now()
55            .duration_since(UNIX_EPOCH)
56            .expect("invalid system time!")
57            .as_millis() as u64;
58
59        now.saturating_sub(*SHORTER_EPOCH)
60    }
61
62    fn next_id(&self) -> u64 {
63        let mut timestamp = self.current_timestamp();
64
65        let last_ts = self.last_timestamp.load(Ordering::Relaxed);
66
67        let sequence = if timestamp == last_ts {
68            // same millisecond - increment sequence
69            let seq = self.sequence.fetch_add(1, Ordering::Relaxed) & 0xFFF; // 12 bits
70            if seq == 0 {
71                // Sequence overflow - wait for next ms
72                while timestamp <= last_ts {
73                    timestamp = self.current_timestamp();
74                }
75                self.sequence.store(0, Ordering::Relaxed);
76                self.last_timestamp.store(timestamp, Ordering::Relaxed);
77                0
78            } else {
79                seq
80            }
81        } else {
82            // new timestamp
83            self.sequence.store(0, Ordering::Relaxed);
84            self.last_timestamp.store(timestamp, Ordering::Relaxed);
85            0
86        };
87
88        // Compose ID: [timestamp][machine_id][sequence]
89        ((timestamp & 0x1FFFFFFFFFF) << 22)  // 41 bits
90            | ((self.machine_id as u64 & 0x3FF) << 12) // 10 bits
91            | (sequence & 0xFFF) // 12 bits
92    }
93}
94
95impl<T> IdGenerator<T> for SnowflakeIdGenerator
96where
97    T: From<String>,
98{
99    fn generate(&self) -> T {
100        let id = self.next_id();
101        T::from(id.to_string()) // We could optionally encode it to base64 or base62
102    }
103}
104
105#[cfg(test)]
106mod tests {
107    use super::*;
108
109    #[test]
110    fn generates_id() {
111        let generator = SnowflakeIdGenerator::new(1);
112        let id: String = generator.generate();
113        assert!(!id.is_empty(), "Generated ID should not be empty");
114    }
115
116    #[test]
117    fn generates_unique_ids() {
118        let generator = SnowflakeIdGenerator::new(1);
119        let mut ids = std::collections::HashSet::new();
120        for _ in 0..1000 {
121            let id: String = generator.generate();
122            assert!(ids.insert(id), "Duplicate ID generated");
123        }
124    }
125
126    #[test]
127    fn ids_are_monotonic_increasing() {
128        let generator = SnowflakeIdGenerator::new(1);
129        let mut prev_id = 0u64;
130
131        for _ in 0..1000 {
132            let id: String = generator.generate();
133            let current_id: u64 = id.parse().expect("ID should be a valid u64");
134            assert!(
135                current_id > prev_id,
136                "ID not strictly increasing: {current_id} <= {prev_id}"
137            );
138            prev_id = current_id;
139        }
140    }
141
142    #[test]
143    fn handles_sequence_rollover() {
144        // Try to simulate a sequence rollover by generating many IDs quickly
145        // just ensuring it doesn't panic
146        let generator = SnowflakeIdGenerator::new(1);
147        for _ in 0..2000 {
148            let _id: String = generator.generate();
149        }
150    }
151
152    #[test]
153    fn respects_machine_id_limit() {
154        // Valid machine ID
155        let _ = SnowflakeIdGenerator::new(1023);
156    }
157
158    #[test]
159    #[should_panic(expected = "Machine ID must be less than 1024")]
160    fn rejects_invalid_machine_id() {
161        // Invalid machine ID (greater than 1023)
162        let _ = SnowflakeIdGenerator::new(1024);
163    }
164}
rust_mcp_extra/id_generator/snow_flake_id_generator.rs

rust_mcp_extra/id_generator/
snow_flake_id_generator.rs
