|
| 1 | +use std::sync::atomic::{AtomicU32, Ordering}; |
| 2 | +use std::time::{SystemTime, UNIX_EPOCH}; |
| 3 | + |
| 4 | +use once_cell::sync::OnceCell; |
| 5 | +use rand::RngCore; |
| 6 | + |
| 7 | +use crate::id::{Id, RAW_LEN}; |
| 8 | +use crate::machine_id; |
| 9 | +use crate::pid; |
| 10 | + |
| 11 | +#[derive(Debug)] |
| 12 | +pub struct Generator { |
| 13 | + counter: AtomicU32, |
| 14 | + machine_id: [u8; 3], |
| 15 | + pid: [u8; 2], |
| 16 | +} |
| 17 | + |
| 18 | +pub fn get() -> &'static Generator { |
| 19 | + static INSTANCE: OnceCell<Generator> = OnceCell::new(); |
| 20 | + |
| 21 | + INSTANCE.get_or_init(|| Generator { |
| 22 | + counter: AtomicU32::new(init_random()), |
| 23 | + machine_id: machine_id::get(), |
| 24 | + pid: pid::get().to_be_bytes(), |
| 25 | + }) |
| 26 | +} |
| 27 | + |
| 28 | +impl Generator { |
| 29 | + pub fn new_id(&self) -> Id { |
| 30 | + self.with_time(&SystemTime::now()) |
| 31 | + } |
| 32 | + |
| 33 | + fn with_time(&self, time: &SystemTime) -> Id { |
| 34 | + // Panic if the time is before the epoch. |
| 35 | + let unix_ts = time |
| 36 | + .duration_since(UNIX_EPOCH) |
| 37 | + .expect("Clock may have gone backwards"); |
| 38 | + self.generate(unix_ts.as_secs() as u32) |
| 39 | + } |
| 40 | + |
| 41 | + fn generate(&self, unix_ts: u32) -> Id { |
| 42 | + let counter = self.counter.fetch_add(1, Ordering::SeqCst); |
| 43 | + |
| 44 | + let mut raw = [0u8; RAW_LEN]; |
| 45 | + // 4 bytes of Timestamp (big endian) |
| 46 | + raw[0..=3].copy_from_slice(&unix_ts.to_be_bytes()); |
| 47 | + // 3 bytes of Machine ID |
| 48 | + raw[4..=6].copy_from_slice(&self.machine_id); |
| 49 | + // 2 bytes of PID |
| 50 | + raw[7..=8].copy_from_slice(&self.pid); |
| 51 | + // 3 bytes of increment counter (big endian) |
| 52 | + raw[9..].copy_from_slice(&counter.to_be_bytes()[1..]); |
| 53 | + |
| 54 | + Id(raw) |
| 55 | + } |
| 56 | +} |
| 57 | + |
| 58 | +// https://github.com/rs/xid/blob/efa678f304ab65d6d57eedcb086798381ae22206/id.go#L136 |
| 59 | +fn init_random() -> u32 { |
| 60 | + let mut bs = [0u8; 3]; |
| 61 | + rand::thread_rng().fill_bytes(&mut bs); |
| 62 | + u32::from_be_bytes([0, bs[0], bs[1], bs[2]]) |
| 63 | +} |
0 commit comments