use core::{cell::Cell, cmp::Ordering};
#[cfg(feature = "tracing")]
use tracing::instrument;
use crate::{
generator::{Poll, Result, SnowflakeGenerator},
id::SnowflakeId,
time::TimeSource,
};
pub struct BasicSnowflakeGenerator<ID, T>
where
ID: SnowflakeId,
T: TimeSource<ID::Ty>,
{
state: Cell<ID>,
time: T,
}
impl<ID, T> BasicSnowflakeGenerator<ID, T>
where
ID: SnowflakeId,
T: TimeSource<ID::Ty>,
{
pub fn new(machine_id: ID::Ty, time: T) -> Self {
Self::from_components(ID::ZERO, machine_id, ID::ZERO, time)
}
pub fn from_components(
timestamp: ID::Ty,
machine_id: ID::Ty,
sequence: ID::Ty,
time: T,
) -> Self {
let id = ID::from_components(timestamp, machine_id, sequence);
Self {
state: Cell::new(id),
time,
}
}
#[cfg_attr(feature = "tracing", instrument(level = "trace", skip(self, f)))]
pub fn next_id(&self, mut f: impl FnMut(ID::Ty)) -> ID {
loop {
match self.poll_id() {
Poll::Ready { id } => break id,
Poll::Pending { yield_for } => f(yield_for),
}
}
}
#[cfg_attr(feature = "tracing", instrument(level = "trace", skip(self)))]
pub fn poll_id(&self) -> Poll<ID> {
let now = self.time.current_millis();
let state = self.state.get();
let current_ts = state.timestamp();
match now.cmp(¤t_ts) {
Ordering::Equal => {
if state.has_sequence_room() {
let updated = state.increment_sequence();
self.state.set(updated);
Poll::Ready { id: updated }
} else {
Poll::Pending { yield_for: ID::ONE }
}
}
Ordering::Greater => {
let updated = state.rollover_to_timestamp(now);
self.state.set(updated);
Poll::Ready { id: updated }
}
Ordering::Less => Self::cold_clock_behind(now, current_ts),
}
}
#[cold]
#[inline(never)]
fn cold_clock_behind(now: ID::Ty, current_ts: ID::Ty) -> Poll<ID> {
let yield_for = current_ts - now;
debug_assert!(yield_for >= ID::ZERO);
Poll::Pending { yield_for }
}
}
impl<ID, T> SnowflakeGenerator<ID, T> for BasicSnowflakeGenerator<ID, T>
where
ID: SnowflakeId,
T: TimeSource<ID::Ty>,
{
type Err = core::convert::Infallible;
fn new(machine_id: ID::Ty, time: T) -> Self {
Self::new(machine_id, time)
}
fn next_id(&self, f: impl FnMut(ID::Ty)) -> ID {
self.next_id(f)
}
fn try_next_id(&self, f: impl FnMut(ID::Ty)) -> Result<ID, Self::Err> {
Ok(self.next_id(f))
}
fn poll_id(&self) -> Poll<ID> {
self.poll_id()
}
fn try_poll_id(&self) -> Result<Poll<ID>, Self::Err> {
Ok(self.poll_id())
}
}