use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SnowflakeError {
ClockBacktrack {
waited_ts: u64,
last_ts: u64,
},
TimestampOverflow {
timestamp: u64,
},
}
impl std::fmt::Display for SnowflakeError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::ClockBacktrack { waited_ts, last_ts } => {
write!(
f,
"Clock backtrack: waited timestamp {waited_ts} still behind last used {last_ts}"
)
}
Self::TimestampOverflow { timestamp } => {
write!(f, "Timestamp overflow: {timestamp} exceeds 41-bit capacity")
}
}
}
}
impl std::error::Error for SnowflakeError {}
impl crate::i18n::error_ext::LocalizedMsg for SnowflakeError {
fn message_key(&self) -> &'static str {
match self {
Self::ClockBacktrack { .. } => "snowflake-clock-backtrack",
Self::TimestampOverflow { .. } => "snowflake-timestamp-overflow",
}
}
fn message_args(&self) -> Vec<(&str, String)> {
match self {
Self::ClockBacktrack { waited_ts, last_ts } => {
vec![("waited_ts", waited_ts.to_string()), ("last_ts", last_ts.to_string())]
}
Self::TimestampOverflow { timestamp } => vec![("timestamp", timestamp.to_string())],
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct IdComponents {
pub timestamp_ms: u64,
pub machine_id: u32,
pub sequence: u32,
}
pub trait DistributedIdGenerator: Send + Sync {
fn next_id(&self) -> Result<u64, SnowflakeError>;
fn parse_id(&self, id: u64) -> IdComponents;
}
pub struct SnowflakeIdGenerator {
machine_id: u32,
epoch: u64,
ts_seq: AtomicU64,
}
impl SnowflakeIdGenerator {
pub fn new(machine_id: u32, epoch: u64) -> Result<Self, String> {
if machine_id > 1023 {
return Err(format!("machine_id must be 0-1023, got {machine_id}"));
}
Ok(Self {
machine_id,
epoch,
ts_seq: AtomicU64::new(0),
})
}
fn current_timestamp(&self) -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64
- self.epoch
}
fn wait_next_millis(&self, last_ts: u64) -> Option<u64> {
let mut ts = self.current_timestamp();
let mut spins = 0;
while ts <= last_ts {
std::hint::spin_loop();
ts = self.current_timestamp();
spins += 1;
if spins > 100_000 {
return None;
}
}
Some(ts)
}
}
const MAX_TIMESTAMP: u64 = (1 << 41) - 1;
impl DistributedIdGenerator for SnowflakeIdGenerator {
fn next_id(&self) -> Result<u64, SnowflakeError> {
loop {
let timestamp = self.current_timestamp();
if timestamp > MAX_TIMESTAMP {
return Err(SnowflakeError::TimestampOverflow { timestamp });
}
let current = self.ts_seq.load(Ordering::SeqCst);
let last_ts = current >> 12;
if timestamp < last_ts {
match self.wait_next_millis(last_ts) {
Some(waited) if waited > last_ts => continue,
_ => {
return Err(SnowflakeError::ClockBacktrack {
waited_ts: self.current_timestamp(),
last_ts,
});
}
}
}
let new_ts_seq = if timestamp > last_ts {
timestamp << 12
} else {
let seq = (current & 0xFFF) + 1;
if seq > 0xFFF {
match self.wait_next_millis(last_ts) {
Some(_) => continue,
None => {
return Err(SnowflakeError::ClockBacktrack {
waited_ts: self.current_timestamp(),
last_ts,
});
}
}
}
(last_ts << 12) | seq
};
match self
.ts_seq
.compare_exchange(current, new_ts_seq, Ordering::SeqCst, Ordering::SeqCst)
{
Ok(_) => {
let seq = new_ts_seq & 0xFFF;
return Ok((timestamp << 22) | ((self.machine_id as u64) << 12) | seq);
}
Err(_) => {
continue;
}
}
}
}
fn parse_id(&self, id: u64) -> IdComponents {
let timestamp_ms = id >> 22;
let machine_id = ((id >> 12) & 0x3FF) as u32;
let sequence = (id & 0xFFF) as u32;
IdComponents {
timestamp_ms,
machine_id,
sequence,
}
}
}