use crate::error::{ClamberError, Result};
use serde::{Deserialize, Serialize};
use std::sync::Mutex;
use twitter_snowflake::Snowflake;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SnowflakeConfig {
pub worker_id: u64,
pub epoch: Option<u64>,
}
impl Default for SnowflakeConfig {
fn default() -> Self {
Self {
worker_id: 1,
epoch: None, }
}
}
impl SnowflakeConfig {
pub fn new(worker_id: u64) -> Result<Self> {
Self::validate_worker_id(worker_id)?;
Ok(Self {
worker_id,
epoch: None,
})
}
pub fn with_epoch(worker_id: u64, epoch: u64) -> Result<Self> {
Self::validate_worker_id(worker_id)?;
Ok(Self {
worker_id,
epoch: Some(epoch),
})
}
pub fn worker_id(mut self, worker_id: u64) -> Result<Self> {
Self::validate_worker_id(worker_id)?;
self.worker_id = worker_id;
Ok(self)
}
pub fn epoch(mut self, epoch: u64) -> Self {
self.epoch = Some(epoch);
self
}
fn validate_worker_id(worker_id: u64) -> Result<()> {
if worker_id > 1023 {
return Err(ClamberError::SnowflakeConfigError {
details: format!("工作者ID必须在0-1023范围内,当前值: {}", worker_id),
});
}
Ok(())
}
}
pub struct SnowflakeManager {
generator: Mutex<Snowflake>,
config: SnowflakeConfig,
}
impl SnowflakeManager {
pub fn new(config: SnowflakeConfig) -> Result<Self> {
let generator = if let Some(epoch) = config.epoch {
Snowflake::builder()
.with_worker_id(config.worker_id)
.with_epoch(epoch)
.build()
} else {
Snowflake::new(config.worker_id)
};
let generator = generator.map_err(|e| ClamberError::SnowflakeInitError {
details: format!("初始化Snowflake生成器失败: {:?}", e),
})?;
Ok(Self {
generator: Mutex::new(generator),
config,
})
}
pub fn default() -> Result<Self> {
Self::new(SnowflakeConfig::default())
}
pub fn generate_id(&self) -> Result<u64> {
let mut generator =
self.generator
.lock()
.map_err(|e| ClamberError::SnowflakeGenerateError {
details: format!("获取生成器锁失败: {}", e),
})?;
generator
.generate()
.map_err(|e| ClamberError::SnowflakeGenerateError {
details: format!("生成ID失败: {:?}", e),
})
}
pub fn generate_ids(&self, count: usize) -> Result<Vec<u64>> {
let mut ids = Vec::with_capacity(count);
for _ in 0..count {
ids.push(self.generate_id()?);
}
Ok(ids)
}
pub fn get_config(&self) -> &SnowflakeConfig {
&self.config
}
pub fn worker_id(&self) -> u64 {
self.config.worker_id
}
pub fn parse_id(&self, id: u64) -> SnowflakeIdInfo {
let timestamp = (id >> 22) & 0x1FFFFFFFFFF; let worker_id = (id >> 12) & 0x3FF; let sequence = id & 0xFFF;
SnowflakeIdInfo {
id,
timestamp,
worker_id,
sequence: sequence as u16,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SnowflakeIdInfo {
pub id: u64,
pub timestamp: u64,
pub worker_id: u64,
pub sequence: u16,
}
impl SnowflakeIdInfo {
pub fn generation_time(&self, epoch: Option<u64>) -> u64 {
let epoch = epoch.unwrap_or(1288834974657); self.timestamp + epoch
}
pub fn generation_time_string(&self, epoch: Option<u64>) -> String {
let timestamp_ms = self.generation_time(epoch);
let datetime = chrono::DateTime::from_timestamp_millis(timestamp_ms as i64);
match datetime {
Some(dt) => dt.format("%Y-%m-%d %H:%M:%S%.3f").to_string(),
None => "Invalid timestamp".to_string(),
}
}
pub fn to_string(&self) -> String {
self.id.to_string()
}
pub fn from_string(id_str: &str) -> Result<u64> {
id_str
.parse::<u64>()
.map_err(|e| ClamberError::SnowflakeConfigError {
details: format!("无法解析ID字符串: {}", e),
})
}
}
use once_cell::sync::Lazy;
static DEFAULT_MANAGER: Lazy<Result<SnowflakeManager>> = Lazy::new(|| SnowflakeManager::default());
fn get_default_manager() -> Result<&'static SnowflakeManager> {
match &*DEFAULT_MANAGER {
Ok(manager) => Ok(manager),
Err(e) => Err(ClamberError::SnowflakeInitError {
details: format!("无法初始化默认Snowflake管理器: {}", e),
}),
}
}
pub fn generate_id() -> Result<u64> {
get_default_manager()?.generate_id()
}
pub fn generate_ids(count: usize) -> Result<Vec<u64>> {
get_default_manager()?.generate_ids(count)
}
pub fn parse_id(id: u64) -> Result<SnowflakeIdInfo> {
Ok(get_default_manager()?.parse_id(id))
}
pub fn generate_string_id() -> Result<String> {
Ok(generate_id()?.to_string())
}
pub fn parse_string_id(id_str: &str) -> Result<SnowflakeIdInfo> {
let id = SnowflakeIdInfo::from_string(id_str)?;
parse_id(id)
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::HashSet;
#[test]
fn test_snowflake_config_creation() {
let config = SnowflakeConfig::new(1).unwrap();
assert_eq!(config.worker_id, 1);
assert!(config.epoch.is_none());
let config_with_epoch = SnowflakeConfig::with_epoch(3, 1609459200000).unwrap();
assert_eq!(config_with_epoch.worker_id, 3);
assert_eq!(config_with_epoch.epoch, Some(1609459200000));
assert!(SnowflakeConfig::new(1024).is_err());
}
#[test]
fn test_snowflake_manager_creation() {
let config = SnowflakeConfig::new(1).unwrap();
let manager = SnowflakeManager::new(config).unwrap();
assert_eq!(manager.worker_id(), 1);
}
#[test]
fn test_id_generation() {
let config = SnowflakeConfig::new(1).unwrap();
let manager = SnowflakeManager::new(config).unwrap();
let id = manager.generate_id().unwrap();
assert!(id > 0);
let ids = manager.generate_ids(100).unwrap();
assert_eq!(ids.len(), 100);
let unique_ids: HashSet<u64> = ids.into_iter().collect();
assert_eq!(unique_ids.len(), 100); }
#[test]
fn test_id_parsing() {
let config = SnowflakeConfig::new(1).unwrap();
let manager = SnowflakeManager::new(config).unwrap();
let id = manager.generate_id().unwrap();
let info = manager.parse_id(id);
assert_eq!(info.id, id);
assert!(info.timestamp > 0);
}
#[test]
fn test_convenience_functions() {
let id = generate_id().unwrap();
assert!(id > 0);
let ids = generate_ids(10).unwrap();
assert_eq!(ids.len(), 10);
let info = parse_id(id).unwrap();
assert_eq!(info.id, id);
}
#[test]
fn test_string_functions() {
let id = generate_string_id().unwrap();
assert!(!id.is_empty());
let info = parse_string_id(&id).unwrap();
assert_eq!(info.to_string(), id);
}
#[test]
fn test_id_info_time_functions() {
let id = generate_id().unwrap();
let info = parse_id(id).unwrap();
let gen_time = info.generation_time(None);
assert!(gen_time > 0);
let time_str = info.generation_time_string(None);
assert!(!time_str.is_empty());
assert!(time_str.contains("-")); }
}