use std::sync::{
Arc,
Mutex,
};
use std::thread;
use std::time::{
Duration,
SystemTime,
UNIX_EPOCH,
};
use super::constants::{
DEFAULT_MAX_SKEW_MILLIS,
DEFAULT_QUBIT_EPOCH_MILLIS,
};
use super::time_slice::TimeSlice;
use super::{
IdMode,
QubitSnowflakeBuilder,
TimestampPrecision,
};
use crate::{
IdError,
IdGenerator,
};
pub struct QubitSnowflakeGenerator {
builder: QubitSnowflakeBuilder,
epoch: SystemTime,
max_skew_millis: u64,
clock: Arc<dyn Fn() -> SystemTime + Send + Sync>,
state: Mutex<TimeSlice>,
}
impl QubitSnowflakeGenerator {
pub fn new(host: u64) -> Result<Self, IdError> {
Self::with_options(
IdMode::Sequential,
TimestampPrecision::Second,
host,
UNIX_EPOCH + Duration::from_millis(DEFAULT_QUBIT_EPOCH_MILLIS),
)
}
pub fn with_options(
mode: IdMode,
precision: TimestampPrecision,
host: u64,
epoch: SystemTime,
) -> Result<Self, IdError> {
Self::with_clock(
mode,
precision,
host,
epoch,
DEFAULT_MAX_SKEW_MILLIS,
SystemTime::now,
)
}
pub fn with_clock<F>(
mode: IdMode,
precision: TimestampPrecision,
host: u64,
epoch: SystemTime,
max_skew_millis: u64,
clock: F,
) -> Result<Self, IdError>
where
F: Fn() -> SystemTime + Send + Sync + 'static,
{
Ok(Self {
builder: QubitSnowflakeBuilder::new(mode, precision, host)?,
epoch,
max_skew_millis,
clock: Arc::new(clock),
state: Mutex::new(TimeSlice::new(0)),
})
}
pub const fn builder(&self) -> &QubitSnowflakeBuilder {
&self.builder
}
pub const fn epoch(&self) -> SystemTime {
self.epoch
}
pub fn generate_at(&self, time: SystemTime, sequence: u64) -> Result<u64, IdError> {
let timestamp = self.timestamp_for(time)?;
self.builder.build(timestamp, sequence)
}
fn timestamp_for(&self, time: SystemTime) -> Result<u64, IdError> {
let elapsed = time
.duration_since(self.epoch)
.map_err(|_| IdError::TimeBeforeEpoch)?;
let timestamp = elapsed.as_millis() / u128::from(self.builder.precision().divisor_millis());
if timestamp > u128::from(self.builder.max_timestamp()) {
return Err(IdError::TimestampOverflow {
timestamp: u64::try_from(timestamp).unwrap_or(u64::MAX),
max: self.builder.max_timestamp(),
});
}
Ok(timestamp as u64)
}
fn current_timestamp(&self) -> Result<u64, IdError> {
self.timestamp_for((self.clock)())
}
fn wait_for_next_timestamp(&self, last_timestamp: u64) -> Result<u64, IdError> {
let mut timestamp = self.current_timestamp()?;
while timestamp <= last_timestamp {
thread::sleep(Duration::from_millis(
self.builder.precision().wait_duration_millis(),
));
timestamp = self.current_timestamp()?;
}
Ok(timestamp)
}
}
impl IdGenerator<u64> for QubitSnowflakeGenerator {
type Error = IdError;
fn next_id(&self) -> Result<u64, Self::Error> {
loop {
let mut state = self
.state
.lock()
.expect("generator state mutex should not be poisoned");
let mut timestamp = self.current_timestamp()?;
if state.timestamp > timestamp {
let skew = state.timestamp - timestamp;
let skew_millis = skew * self.builder.precision().divisor_millis();
if skew_millis > self.max_skew_millis {
return Err(IdError::ClockMovedBackwards {
last_timestamp: state.timestamp,
current_timestamp: timestamp,
skew_millis,
max_skew_millis: self.max_skew_millis,
});
}
drop(state);
thread::sleep(Duration::from_millis(skew_millis));
continue;
}
let sequence = if timestamp == state.timestamp {
let next_sequence = (state.sequence + 1) & self.builder.max_sequence();
if next_sequence == 0 {
drop(state);
timestamp = self.wait_for_next_timestamp(timestamp)?;
let mut state = self
.state
.lock()
.expect("generator state mutex should not be poisoned");
state.timestamp = timestamp;
state.sequence = 0;
return self.builder.build(timestamp, 0);
}
next_sequence
} else {
0
};
state.timestamp = timestamp;
state.sequence = sequence;
return self.builder.build(timestamp, sequence);
}
}
}