use std::{
sync::atomic::{
AtomicU64,
Ordering::{Acquire, Release, SeqCst},
},
time::SystemTime,
};
pub const EPOCH: u64 = 1637806706000;
#[cfg(not(test))]
const SEQUENCE_BITS: u8 = 12;
#[cfg(test)]
const SEQUENCE_BITS: u8 = 1;
const TIMESTAMP_SHIFT: u8 = SEQUENCE_BITS;
const SEQUENCE_MASK: u64 = (1 << SEQUENCE_BITS) - 1;
pub type NID = u64;
pub type SIDGEN = AtomicU64;
pub type SID = u64;
#[derive(Default)]
pub struct Seq(SIDGEN);
impl Seq {
pub fn new() -> Seq {
Seq(SIDGEN::new(0))
}
pub fn next_id(&self) -> SID {
self.0.fetch_add(1, SeqCst)
}
pub fn reset(&self) {
self.0.store(0, Release);
}
}
impl From<SID> for Seq {
fn from(value: SID) -> Self {
Seq(SIDGEN::new(value))
}
}
#[derive(Default)]
pub struct TimestampSeq {
sequence: AtomicU64,
last_cycle_timestamp: AtomicU64,
}
fn get_timestamp() -> u64 {
SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.expect("Clock moved backwards!")
.as_millis() as u64
}
impl TimestampSeq {
pub fn new() -> TimestampSeq {
TimestampSeq::default()
}
fn wait_next_millis(&self) {
let last_timestamp = self.last_cycle_timestamp.load(Acquire);
while get_timestamp() <= last_timestamp {}
}
pub fn next_id(&self) -> u64 {
let sequence = self.sequence.fetch_add(1, SeqCst) & SEQUENCE_MASK;
let mut new_timestamp = get_timestamp();
if sequence == 0 {
let last_timestamp = self.last_cycle_timestamp.load(Acquire);
if last_timestamp == new_timestamp {
self.wait_next_millis();
new_timestamp = get_timestamp();
}
self.last_cycle_timestamp.fetch_max(new_timestamp, Release);
}
(new_timestamp - EPOCH) << TIMESTAMP_SHIFT | sequence
}
}
#[cfg(any(test, feature = "for-test"))]
pub fn into_parts(timestamp: u64) -> (u64, u64) {
(timestamp >> TIMESTAMP_SHIFT, timestamp & SEQUENCE_MASK)
}
#[cfg(test)]
mod tests {
use super::*;
fn wait_for_next_millis() {
std::thread::sleep(std::time::Duration::from_millis(1));
}
#[test]
fn test_next_generates_unique_ids() {
let timestamp_seq = TimestampSeq::new();
let id1 = timestamp_seq.next_id();
let id2 = timestamp_seq.next_id();
let id3 = timestamp_seq.next_id();
assert_ne!(id1, id2);
assert_ne!(id2, id3);
assert_ne!(id1, id3);
}
#[test]
fn test_next_increases_sequence() {
let timestamp_seq = TimestampSeq::new();
let id1 = timestamp_seq.next_id();
let id2 = timestamp_seq.next_id();
assert!(id2 > id1);
}
#[test]
fn test_next_does_not_repeat_ids() {
let timestamp_seq = TimestampSeq::new();
let id1 = timestamp_seq.next_id();
let id2 = timestamp_seq.next_id();
let id3 = timestamp_seq.next_id();
let id4 = timestamp_seq.next_id();
assert_ne!(id1, id2);
assert_ne!(id2, id3);
assert_ne!(id3, id4);
assert_ne!(id1, id4);
println!("{id1} {id2} {id3} {id4}");
println!(
"{:?} {:?} {:?} {:?}",
into_parts(id1),
into_parts(id2),
into_parts(id3),
into_parts(id4)
);
}
#[test]
fn test_next_wait_for_next_millis() {
let timestamp_seq = TimestampSeq::new();
let id1 = timestamp_seq.next_id();
wait_for_next_millis();
let id2 = timestamp_seq.next_id();
let timestamp1 = id1 >> TIMESTAMP_SHIFT;
let timestamp2 = id2 >> TIMESTAMP_SHIFT;
assert!(timestamp2 > timestamp1);
}
}