clamber_core/snowflake/
mod.rs

1use crate::error::{ClamberError, Result};
2use serde::{Deserialize, Serialize};
3use std::sync::Mutex;
4use twitter_snowflake::Snowflake;
5
6/// Snowflake配置结构
7#[derive(Debug, Clone, Serialize, Deserialize)]
8pub struct SnowflakeConfig {
9    /// 工作者ID (0-1023)
10    pub worker_id: u64,
11    /// 自定义纪元时间戳（毫秒，可选）
12    pub epoch: Option<u64>,
13}
14
15impl Default for SnowflakeConfig {
16    fn default() -> Self {
17        Self {
18            worker_id: 1,
19            epoch: None, // 使用默认纪元
20        }
21    }
22}
23
24impl SnowflakeConfig {
25    /// 创建新的Snowflake配置
26    pub fn new(worker_id: u64) -> Result<Self> {
27        Self::validate_worker_id(worker_id)?;
28
29        Ok(Self {
30            worker_id,
31            epoch: None,
32        })
33    }
34
35    /// 创建带自定义纪元的配置
36    pub fn with_epoch(worker_id: u64, epoch: u64) -> Result<Self> {
37        Self::validate_worker_id(worker_id)?;
38
39        Ok(Self {
40            worker_id,
41            epoch: Some(epoch),
42        })
43    }
44
45    /// 设置工作者ID
46    pub fn worker_id(mut self, worker_id: u64) -> Result<Self> {
47        Self::validate_worker_id(worker_id)?;
48        self.worker_id = worker_id;
49        Ok(self)
50    }
51
52    /// 设置自定义纪元
53    pub fn epoch(mut self, epoch: u64) -> Self {
54        self.epoch = Some(epoch);
55        self
56    }
57
58    /// 验证工作者ID有效性
59    fn validate_worker_id(worker_id: u64) -> Result<()> {
60        if worker_id > 1023 {
61            return Err(ClamberError::SnowflakeConfigError {
62                details: format!("工作者ID必须在0-1023范围内，当前值: {}", worker_id),
63            });
64        }
65        Ok(())
66    }
67}
68
69/// Snowflake ID生成器封装
70pub struct SnowflakeManager {
71    generator: Mutex<Snowflake>,
72    config: SnowflakeConfig,
73}
74
75impl SnowflakeManager {
76    /// 使用自定义配置创建Snowflake管理器
77    pub fn new(config: SnowflakeConfig) -> Result<Self> {
78        let generator = if let Some(epoch) = config.epoch {
79            Snowflake::builder()
80                .with_worker_id(config.worker_id)
81                .with_epoch(epoch)
82                .build()
83        } else {
84            Snowflake::new(config.worker_id)
85        };
86
87        let generator = generator.map_err(|e| ClamberError::SnowflakeInitError {
88            details: format!("初始化Snowflake生成器失败: {:?}", e),
89        })?;
90
91        Ok(Self {
92            generator: Mutex::new(generator),
93            config,
94        })
95    }
96
97    /// 使用默认配置创建Snowflake管理器
98    pub fn default() -> Result<Self> {
99        Self::new(SnowflakeConfig::default())
100    }
101
102    /// 生成新的Snowflake ID
103    pub fn generate_id(&self) -> Result<u64> {
104        let mut generator =
105            self.generator
106                .lock()
107                .map_err(|e| ClamberError::SnowflakeGenerateError {
108                    details: format!("获取生成器锁失败: {}", e),
109                })?;
110
111        generator
112            .generate()
113            .map_err(|e| ClamberError::SnowflakeGenerateError {
114                details: format!("生成ID失败: {:?}", e),
115            })
116    }
117
118    /// 生成多个ID
119    pub fn generate_ids(&self, count: usize) -> Result<Vec<u64>> {
120        let mut ids = Vec::with_capacity(count);
121        for _ in 0..count {
122            ids.push(self.generate_id()?);
123        }
124        Ok(ids)
125    }
126
127    /// 获取当前配置
128    pub fn get_config(&self) -> &SnowflakeConfig {
129        &self.config
130    }
131
132    /// 获取工作者ID
133    pub fn worker_id(&self) -> u64 {
134        self.config.worker_id
135    }
136
137    /// 解析Snowflake ID的各个组成部分
138    pub fn parse_id(&self, id: u64) -> SnowflakeIdInfo {
139        // Twitter Snowflake ID结构：1位符号位 + 41位时间戳 + 10位工作者ID + 12位序列号
140        let timestamp = (id >> 22) & 0x1FFFFFFFFFF; // 41位时间戳
141        let worker_id = (id >> 12) & 0x3FF; // 10位工作者ID
142        let sequence = id & 0xFFF; // 12位序列号
143
144        SnowflakeIdInfo {
145            id,
146            timestamp,
147            worker_id,
148            sequence: sequence as u16,
149        }
150    }
151}
152
153/// Snowflake ID信息结构
154#[derive(Debug, Clone, Serialize, Deserialize)]
155pub struct SnowflakeIdInfo {
156    /// 原始ID
157    pub id: u64,
158    /// 时间戳部分
159    pub timestamp: u64,
160    /// 工作者ID部分
161    pub worker_id: u64,
162    /// 序列号部分
163    pub sequence: u16,
164}
165
166impl SnowflakeIdInfo {
167    /// 获取生成时间（毫秒时间戳）
168    pub fn generation_time(&self, epoch: Option<u64>) -> u64 {
169        let epoch = epoch.unwrap_or(1288834974657); // Twitter纪元 (2010-11-04T01:42:54.657Z)
170        self.timestamp + epoch
171    }
172
173    /// 获取可读的时间字符串
174    pub fn generation_time_string(&self, epoch: Option<u64>) -> String {
175        let timestamp_ms = self.generation_time(epoch);
176        let datetime = chrono::DateTime::from_timestamp_millis(timestamp_ms as i64);
177        match datetime {
178            Some(dt) => dt.format("%Y-%m-%d %H:%M:%S%.3f").to_string(),
179            None => "Invalid timestamp".to_string(),
180        }
181    }
182
183    /// 转换为字符串格式的ID
184    pub fn to_string(&self) -> String {
185        self.id.to_string()
186    }
187
188    /// 从字符串解析ID
189    pub fn from_string(id_str: &str) -> Result<u64> {
190        id_str
191            .parse::<u64>()
192            .map_err(|e| ClamberError::SnowflakeConfigError {
193                details: format!("无法解析ID字符串: {}", e),
194            })
195    }
196}
197
198// 便利函数：使用默认配置
199use once_cell::sync::Lazy;
200
201static DEFAULT_MANAGER: Lazy<Result<SnowflakeManager>> = Lazy::new(|| SnowflakeManager::default());
202
203/// 获取或创建默认的Snowflake管理器
204fn get_default_manager() -> Result<&'static SnowflakeManager> {
205    match &*DEFAULT_MANAGER {
206        Ok(manager) => Ok(manager),
207        Err(e) => Err(ClamberError::SnowflakeInitError {
208            details: format!("无法初始化默认Snowflake管理器: {}", e),
209        }),
210    }
211}
212
213/// 使用默认配置生成ID
214pub fn generate_id() -> Result<u64> {
215    get_default_manager()?.generate_id()
216}
217
218/// 使用默认配置生成多个ID
219pub fn generate_ids(count: usize) -> Result<Vec<u64>> {
220    get_default_manager()?.generate_ids(count)
221}
222
223/// 解析ID信息
224pub fn parse_id(id: u64) -> Result<SnowflakeIdInfo> {
225    Ok(get_default_manager()?.parse_id(id))
226}
227
228/// 生成字符串格式的ID
229pub fn generate_string_id() -> Result<String> {
230    Ok(generate_id()?.to_string())
231}
232
233/// 从字符串解析并获取ID信息
234pub fn parse_string_id(id_str: &str) -> Result<SnowflakeIdInfo> {
235    let id = SnowflakeIdInfo::from_string(id_str)?;
236    parse_id(id)
237}
238
239#[cfg(test)]
240mod tests {
241    use super::*;
242    use std::collections::HashSet;
243
244    #[test]
245    fn test_snowflake_config_creation() {
246        // 测试正常创建
247        let config = SnowflakeConfig::new(1).unwrap();
248        assert_eq!(config.worker_id, 1);
249        assert!(config.epoch.is_none());
250
251        // 测试带纪元的创建
252        let config_with_epoch = SnowflakeConfig::with_epoch(3, 1609459200000).unwrap();
253        assert_eq!(config_with_epoch.worker_id, 3);
254        assert_eq!(config_with_epoch.epoch, Some(1609459200000));
255
256        // 测试无效工作者ID
257        assert!(SnowflakeConfig::new(1024).is_err());
258    }
259
260    #[test]
261    fn test_snowflake_manager_creation() {
262        let config = SnowflakeConfig::new(1).unwrap();
263        let manager = SnowflakeManager::new(config).unwrap();
264        assert_eq!(manager.worker_id(), 1);
265    }
266
267    #[test]
268    fn test_id_generation() {
269        let config = SnowflakeConfig::new(1).unwrap();
270        let manager = SnowflakeManager::new(config).unwrap();
271
272        // 生成单个ID
273        let id = manager.generate_id().unwrap();
274        assert!(id > 0);
275
276        // 生成多个ID并检查唯一性
277        let ids = manager.generate_ids(100).unwrap();
278        assert_eq!(ids.len(), 100);
279
280        let unique_ids: HashSet<u64> = ids.into_iter().collect();
281        assert_eq!(unique_ids.len(), 100); // 所有ID应该是唯一的
282    }
283
284    #[test]
285    fn test_id_parsing() {
286        let config = SnowflakeConfig::new(1).unwrap();
287        let manager = SnowflakeManager::new(config).unwrap();
288
289        let id = manager.generate_id().unwrap();
290        let info = manager.parse_id(id);
291
292        assert_eq!(info.id, id);
293        assert!(info.timestamp > 0);
294    }
295
296    #[test]
297    fn test_convenience_functions() {
298        // 测试便利函数
299        let id = generate_id().unwrap();
300        assert!(id > 0);
301
302        let ids = generate_ids(10).unwrap();
303        assert_eq!(ids.len(), 10);
304
305        let info = parse_id(id).unwrap();
306        assert_eq!(info.id, id);
307    }
308
309    #[test]
310    fn test_string_functions() {
311        let id = generate_string_id().unwrap();
312        assert!(!id.is_empty());
313
314        let info = parse_string_id(&id).unwrap();
315        assert_eq!(info.to_string(), id);
316    }
317
318    #[test]
319    fn test_id_info_time_functions() {
320        let id = generate_id().unwrap();
321        let info = parse_id(id).unwrap();
322
323        let gen_time = info.generation_time(None);
324        assert!(gen_time > 0);
325
326        let time_str = info.generation_time_string(None);
327        assert!(!time_str.is_empty());
328        assert!(time_str.contains("-")); // 应该包含日期格式
329    }
330}
clamber_core/snowflake/mod.rs

clamber_core/snowflake/
mod.rs
