1use std::sync::atomic::{AtomicU64, Ordering};
10use std::time::{SystemTime, UNIX_EPOCH};
11
12pub const DEFAULT_EPOCH_MS: u64 = 1_767_225_600_000;
14
15#[derive(Debug)]
28pub struct Snowflake {
29 worker_id: u16,
30 epoch_ms: u64,
31 sequence: AtomicU64,
32}
33
34impl Snowflake {
35 pub fn new(worker_id: u16) -> Self {
38 Self::with_epoch(worker_id, DEFAULT_EPOCH_MS)
39 }
40
41 pub fn with_epoch(worker_id: u16, epoch_ms: u64) -> Self {
43 Self {
44 worker_id: worker_id & 0x3ff,
45 epoch_ms,
46 sequence: AtomicU64::new(0),
47 }
48 }
49
50 pub fn next_id(&self) -> u64 {
55 let now = SystemTime::now()
56 .duration_since(UNIX_EPOCH)
57 .map(|d| d.as_millis() as u64)
58 .unwrap_or(0);
59 let ts = now.saturating_sub(self.epoch_ms) & 0x1ffffffffff; let seq = self.sequence.fetch_add(1, Ordering::Relaxed) & 0xfff; let worker = self.worker_id as u64 & 0x3ff; (ts << 22) | (worker << 12) | seq
64 }
65}
66
67#[cfg(test)]
68mod tests {
69 use super::*;
70
71 #[test]
72 fn next_id_produces_value() {
73 let gen = Snowflake::new(1);
74 let _ = gen.next_id();
75 }
76
77 #[test]
78 fn worker_id_clamped() {
79 let gen = Snowflake::new(0xffff);
80 let id = gen.next_id();
81 let worker = (id >> 12) & 0x3ff;
82 assert_eq!(worker, 0x3ff);
83 }
84
85 #[test]
86 fn unique_ids() {
87 let gen = Snowflake::new(1);
88 let a = gen.next_id();
89 let b = gen.next_id();
90 assert_ne!(a, b);
91 }
92
93 #[test]
94 fn monotonic_within_ms() {
95 let gen = Snowflake::new(1);
96 let a = gen.next_id();
97 let b = gen.next_id();
98 assert!(b >= a);
99 }
100}