use core::time; use std::sync::{Arc, Mutex}; use chrono::NaiveDateTime; use once_cell::sync::{Lazy, OnceCell}; use thiserror::Error; const HAIL_PERIOD_START: Lazy = Lazy::new(|| { crate::time::date(2022, 2, 22) .map(|d| d.and_hms_opt(22, 22, 22)) .flatten() .map(|dt| crate::time::attach_asia_shanghai(dt)) .map(|dt| dt.timestamp()) .unwrap_or_else(|| NaiveDateTime::MIN.timestamp()) }); type TimestampValidator = fn(i64) -> bool; type TimestampGenerator = fn() -> i64; static INSTANCE: OnceCell = OnceCell::new(); #[derive(Debug, Error)] pub enum HailSerialCodeAlgorithmError { #[error("Algorithm is already initialized")] AlgorithmAlreadyInitialized, } /// 冰雹序列ID算法。 /// 缩减了时间戳的位数,相比雪花算法可以额外支持近40年。 pub struct HailSerialCodeAlgorithm { validator: Option, generator: Option, host_id: i64, last_timestamp: Arc>, counter: Arc>, } impl HailSerialCodeAlgorithm { /// 获取一个算法实例用于获取序列ID。 pub fn get() -> &'static Self { INSTANCE.get().unwrap() } /// 初始化整个序列ID算法。 /// ! 注意,如果选择使用内置的主机独立时间戳生成器和验证器,那么将不能保证多主机状态下的序列ID一致性。可能会存在个别主机时间回拨现象。 /// /// - `host_id`:主机ID,取值范围为0~65535。 /// - `timestamp_generator`:时间戳生成器,用于生成时间戳。如果不提供,则使用算法内置的主机独立时间戳生成器。 /// - `timestamp_validatoe`:时间戳验证器,用于验证时间戳是否有效。如果不提供,则使用算法内置的主机独立时间戳验证器。 pub fn initialize_algorithm( host_id: i64, timestamp_generator: Option, timestamp_validatoe: Option, ) -> Result<(), HailSerialCodeAlgorithmError> { let algorithm = HailSerialCodeAlgorithm { validator: timestamp_validatoe, generator: timestamp_generator, host_id, last_timestamp: Arc::new(Mutex::new(0)), counter: Arc::new(Mutex::new(0)), }; INSTANCE .set(algorithm) .map_err(|_| HailSerialCodeAlgorithmError::AlgorithmAlreadyInitialized) } /// 生成一个自计时起点以来的时间戳。 fn generate_timestamp(&self) -> i64 { let current_time = crate::time::now_asia_shanghai().timestamp(); current_time - *HAIL_PERIOD_START } /// 生成一个64位长整型序列ID。 pub fn generate_serial(&self) -> i64 { let last_timestamp = self.last_timestamp.clone(); let mut last_timestamp = last_timestamp.lock().unwrap(); let counter = self.counter.clone(); let mut counter = counter.lock().unwrap(); loop { let timestamp = if let Some(generator) = self.generator { generator() } else { self.generate_timestamp() }; if let Some(validator) = self.validator { if !validator(timestamp) { std::thread::sleep(time::Duration::from_secs(1)); continue; } } else if timestamp < *last_timestamp { std::thread::sleep(time::Duration::from_secs(1)); continue; } if *last_timestamp < timestamp { // 对齐时间戳并重置序列计数器 *last_timestamp = timestamp; *counter = 0; } *counter += 1; return (timestamp << 20) | ((self.host_id & 0xFFFF) << 16) | (*counter & 0xFFFF_FFFF); } } /// 生成一个17位长前补零的序列ID字符串。 pub fn generate_string_serial(&self) -> String { let serial = self.generate_serial(); format!("{:017}", serial) } /// 生成一个带字符串前缀17位长前补零的序列ID字符串。 pub fn generate_prefixed_string_serial(&self, prefix: &str) -> String { let serial = self.generate_serial(); format!("{}{:017}", prefix, serial) } }