clamber_core/snowflake/
mod.rs

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

clamber_core/snowflake/
mod.rs
