#![allow(dead_code)]
use std::hint::spin_loop;
use std::ops::Deref;
use std::sync::atomic::{AtomicI16, AtomicI64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
use serde::de::Visitor;
use serde::{de, Serialize, Serializer};
#[derive(Debug)]
pub struct NumericalUniqueIdGenerator {
epoch: SystemTime,
pub datacenter_id: i32,
pub machine_id: i32,
timestamp: AtomicI64,
sequence_num: AtomicI16,
}
#[derive(Debug)]
pub struct NumericalUniqueIdBucket {
snowflake_id_generator: NumericalUniqueIdGenerator,
bucket: Vec<i64>,
}
impl NumericalUniqueIdGenerator {
pub fn new(datacenter_id: i32, machine_id: i32) -> NumericalUniqueIdGenerator {
Self::with_epoch(datacenter_id, machine_id, UNIX_EPOCH)
}
pub fn with_epoch(datacenter_id: i32, machine_id: i32, epoch: SystemTime) -> NumericalUniqueIdGenerator {
let timestamp = current_time_in_milli(epoch);
NumericalUniqueIdGenerator {
epoch,
timestamp: AtomicI64::new(timestamp),
datacenter_id,
machine_id,
sequence_num: AtomicI16::new(0),
}
}
fn get_snowflake(&self) -> i64 {
self.timestamp.load(Ordering::Relaxed) << 22 | ((self.datacenter_id << 17) as i64) | ((self.machine_id << 12) as i64) | (self.sequence_num.load(Ordering::Relaxed) as i64)
}
pub fn generate(&self) -> i64 {
self.sequence_num.store((self.sequence_num.load(Ordering::Relaxed) + 1) % 4096, Ordering::Relaxed);
let mut now_millis = current_time_in_milli(self.epoch);
if self.timestamp.load(Ordering::Relaxed) == now_millis {
if self.sequence_num.load(Ordering::Relaxed) == 0 {
now_millis = race_next_milli(self.timestamp.load(Ordering::Relaxed), self.epoch);
self.timestamp.store(now_millis, Ordering::Relaxed);
}
} else {
self.timestamp.store(now_millis, Ordering::Relaxed);
self.sequence_num.store(0, Ordering::Relaxed);
}
self.get_snowflake()
}
}
impl NumericalUniqueIdBucket {
pub fn new(datacenter_id: i32, machine_id: i32) -> Self {
Self::with_epoch(datacenter_id, machine_id, UNIX_EPOCH)
}
pub fn with_epoch(datacenter_id: i32, machine_id: i32, epoch: SystemTime) -> Self {
let snowflake_id_generator = NumericalUniqueIdGenerator::with_epoch(datacenter_id, machine_id, epoch);
let bucket = Vec::new();
NumericalUniqueIdBucket { snowflake_id_generator, bucket }
}
pub fn get_id(&mut self) -> i64 {
if self.bucket.is_empty() {
self.fill_bucket();
}
self.bucket.pop().unwrap()
}
fn fill_bucket(&mut self) {
for _ in 0..4091 {
self.bucket.push(self.snowflake_id_generator.generate());
}
}
}
#[inline(always)]
pub fn current_time_in_milli(epoch: SystemTime) -> i64 {
SystemTime::now().duration_since(epoch).expect("System Time Error!").as_millis() as i64
}
#[inline(always)]
fn race_next_milli(timestamp: i64, epoch: SystemTime) -> i64 {
let mut latest_time_millis: i64;
loop {
latest_time_millis = current_time_in_milli(epoch);
if latest_time_millis > timestamp {
return latest_time_millis;
}
spin_loop();
}
}
pub fn id_generator() -> &'static NumericalUniqueIdGenerator {
use std::sync::OnceLock;
static SNOWFLAKE: OnceLock<NumericalUniqueIdGenerator> = OnceLock::new();
SNOWFLAKE.get_or_init(|| {
NumericalUniqueIdGenerator::new(
std::env::var("DATACENTER_ID").expect("DATACENTER_ID MUST BE SET").parse::<i32>().expect("Parsing Failed!"),
std::env::var("MACHINE_ID").expect("MACHINE_ID MUST BE SET").parse::<i32>().expect("Parsing Failed!"),
)
})
}
#[derive(Clone, Hash, PartialEq, Debug, Eq, Ord, PartialOrd, Copy, Default)]
pub struct SnowFlake(pub i64);
impl SnowFlake {
pub fn generate() -> Self {
id_generator().generate().into()
}
}
impl Deref for SnowFlake {
type Target = i64;
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl From<i64> for SnowFlake {
fn from(value: i64) -> Self {
Self(value)
}
}
impl From<SnowFlake> for String {
fn from(value: SnowFlake) -> Self {
value.0.to_string()
}
}
impl From<SnowFlake> for i64 {
fn from(value: SnowFlake) -> Self {
value.0
}
}
impl std::fmt::Display for SnowFlake {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.0)
}
}
impl<'de> serde::Deserialize<'de> for SnowFlake {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
struct SnowflakeVisitor;
impl<'de> Visitor<'de> for SnowflakeVisitor {
type Value = SnowFlake;
fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
f.write_str("Snowflake as a number or string")
}
fn visit_i64<E>(self, id: i64) -> Result<Self::Value, E>
where
E: de::Error,
{
Ok(SnowFlake(id))
}
fn visit_u64<E>(self, id: u64) -> Result<Self::Value, E>
where
E: de::Error,
{
if id < i64::MAX as u64 {
Ok(SnowFlake(id.try_into().unwrap()))
} else {
Err(E::custom(format!("Snowflake out of range: {}", id)))
}
}
fn visit_str<E>(self, id: &str) -> Result<Self::Value, E>
where
E: de::Error,
{
match id.parse::<u64>() {
Ok(val) => self.visit_u64(val),
Err(_) => Err(E::custom("Failed to parse snowflake")),
}
}
}
deserializer.deserialize_any(SnowflakeVisitor)
}
}
impl Serialize for SnowFlake {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
let s = self.0.to_string();
serializer.serialize_str(&s)
}
}
#[test]
fn test_generate() {
let id_generator = NumericalUniqueIdGenerator::new(1, 2);
let mut ids = Vec::with_capacity(10000);
for _ in 0..99 {
for _ in 0..10000 {
ids.push(id_generator.generate());
}
ids.sort();
ids.dedup();
assert_eq!(10000, ids.len());
ids.clear();
}
}
#[test]
fn test_generate_not_sequential_value_when_sleep() {
let id_generator = NumericalUniqueIdGenerator::new(1, 2);
let first = id_generator.generate();
std::thread::sleep(std::time::Duration::from_millis(1));
let second = id_generator.generate();
assert!(first < second);
assert_ne!(first + 1, second);
}
#[test]
fn test_singleton_generate() {
let id_generator = id_generator();
let mut ids = Vec::with_capacity(1000000);
for _ in 0..99 {
for _ in 0..1000000 {
ids.push(id_generator.generate());
}
assert_eq!(1000000, ids.len());
assert!(ids.first().unwrap() < ids.last().unwrap());
assert!(ids.get(999998).unwrap() < ids.get(999999).unwrap());
ids.clear();
}
}