package utils import ( "errors" "os" "sync" "time" ) // 雪花算法常量定义 const ( // 64位ID结构:1位符号位 + 41位时间戳 + 10位机器ID + 12位序列号 // 机器ID占用的位数 nodeIDBits uint64 = 10 // 序列号占用的位数 sequenceBits uint64 = 12 // 机器ID的最大值 (0-1023) maxNodeID int64 = -1 ^ (-1 << nodeIDBits) // 序列号的最大值 (0-4095) maxSequence int64 = -1 ^ (-1 << sequenceBits) // 机器ID左移位数 nodeIDShift uint64 = sequenceBits // 时间戳左移位数 timestampShift uint64 = sequenceBits + nodeIDBits // 自定义纪元时间:2024-01-01 00:00:00 UTC // 使用自定义纪元可以延长ID有效期约24年(从2024年开始) customEpoch int64 = 1704067200000 // 2024-01-01 00:00:00 UTC 的毫秒时间戳 ) // 错误定义 var ( // ErrInvalidNodeID 机器ID无效 ErrInvalidNodeID = errors.New("node ID must be between 0 and 1023") // ErrClockBackwards 时钟回拨 ErrClockBackwards = errors.New("clock moved backwards, refusing to generate ID") ) // IDInfo 解析后的ID信息 type IDInfo struct { Timestamp int64 // 生成ID时的时间戳(毫秒) NodeID int64 // 机器ID Sequence int64 // 序列号 } // Snowflake 雪花算法ID生成器 type Snowflake struct { mu sync.Mutex // 互斥锁,保证线程安全 nodeID int64 // 机器ID (0-1023) sequence int64 // 当前序列号 (0-4095) lastTimestamp int64 // 上次生成ID的时间戳 } // 全局雪花算法实例 var ( globalSnowflake *Snowflake globalSnowflakeOnce sync.Once globalSnowflakeErr error ) // InitSnowflake 初始化全局雪花算法实例 // nodeID: 机器ID,范围0-1023,可以通过环境变量 NODE_ID 配置 func InitSnowflake(nodeID int64) error { globalSnowflake, globalSnowflakeErr = NewSnowflake(nodeID) return globalSnowflakeErr } // GetSnowflake 获取全局雪花算法实例 // 如果未初始化,会自动使用默认配置初始化 func GetSnowflake() *Snowflake { globalSnowflakeOnce.Do(func() { if globalSnowflake == nil { globalSnowflake, globalSnowflakeErr = NewSnowflake(-1) } }) return globalSnowflake } // NewSnowflake 创建雪花算法ID生成器实例 // nodeID: 机器ID,范围0-1023,可以通过环境变量 NODE_ID 配置 // 如果nodeID为-1,则尝试从环境变量 NODE_ID 读取 func NewSnowflake(nodeID int64) (*Snowflake, error) { // 如果传入-1,尝试从环境变量读取 if nodeID == -1 { nodeIDStr := os.Getenv("NODE_ID") if nodeIDStr != "" { // 解析环境变量 parsedID, err := parseInt(nodeIDStr) if err != nil { return nil, ErrInvalidNodeID } nodeID = parsedID } else { // 默认使用0 nodeID = 0 } } // 验证机器ID范围 if nodeID < 0 || nodeID > maxNodeID { return nil, ErrInvalidNodeID } return &Snowflake{ nodeID: nodeID, sequence: 0, lastTimestamp: 0, }, nil } // parseInt 辅助函数:解析整数 func parseInt(s string) (int64, error) { var result int64 var negative bool if len(s) == 0 { return 0, errors.New("empty string") } i := 0 if s[0] == '-' { negative = true i = 1 } for ; i < len(s); i++ { if s[i] < '0' || s[i] > '9' { return 0, errors.New("invalid character") } result = result*10 + int64(s[i]-'0') } if negative { result = -result } return result, nil } // GenerateID 生成唯一的雪花算法ID // 返回值:生成的ID,以及可能的错误(如时钟回拨) // 线程安全:使用互斥锁保证并发安全 func (s *Snowflake) GenerateID() (int64, error) { s.mu.Lock() defer s.mu.Unlock() // 获取当前时间戳(毫秒) now := currentTimestamp() // 处理时钟回拨 if now < s.lastTimestamp { return 0, ErrClockBackwards } // 同一毫秒内 if now == s.lastTimestamp { // 序列号递增 s.sequence = (s.sequence + 1) & maxSequence // 序列号溢出,等待下一毫秒 if s.sequence == 0 { now = s.waitNextMillis(now) } } else { // 不同毫秒,序列号重置为0 s.sequence = 0 } // 更新上次生成时间 s.lastTimestamp = now // 组装ID // ID结构:时间戳部分 | 机器ID部分 | 序列号部分 id := ((now - customEpoch) << timestampShift) | (s.nodeID << nodeIDShift) | s.sequence return id, nil } // waitNextMillis 等待到下一毫秒 // 参数:当前时间戳 // 返回值:下一毫秒的时间戳 func (s *Snowflake) waitNextMillis(timestamp int64) int64 { now := currentTimestamp() for now <= timestamp { now = currentTimestamp() } return now } // ParseID 解析雪花算法ID,提取其中的信息 // id: 要解析的雪花算法ID // 返回值:包含时间戳、机器ID、序列号的结构体 func ParseID(id int64) *IDInfo { // 提取序列号(低12位) sequence := id & maxSequence // 提取机器ID(中间10位) nodeID := (id >> nodeIDShift) & maxNodeID // 提取时间戳(高41位) timestamp := (id >> timestampShift) + customEpoch return &IDInfo{ Timestamp: timestamp, NodeID: nodeID, Sequence: sequence, } } // currentTimestamp 获取当前时间戳(毫秒) func currentTimestamp() int64 { return time.Now().UnixNano() / 1e6 } // GetNodeID 获取当前机器ID func (s *Snowflake) GetNodeID() int64 { return s.nodeID } // GetCustomEpoch 获取自定义纪元时间 func GetCustomEpoch() int64 { return customEpoch } // IDToTime 将雪花算法ID转换为生成时间 // 这是一个便捷方法,等价于 ParseID(id).Timestamp func IDToTime(id int64) time.Time { info := ParseID(id) return time.Unix(0, info.Timestamp*1e6) // 毫秒转纳秒 } // ValidateID 验证ID是否为有效的雪花算法ID // 检查时间戳是否在合理范围内 func ValidateID(id int64) bool { if id <= 0 { return false } info := ParseID(id) // 检查时间戳是否在合理范围内 // 不能早于纪元时间,不能晚于当前时间太多(允许1分钟的时钟偏差) now := currentTimestamp() if info.Timestamp < customEpoch || info.Timestamp > now+60000 { return false } // 检查机器ID和序列号是否在有效范围内 if info.NodeID < 0 || info.NodeID > maxNodeID { return false } if info.Sequence < 0 || info.Sequence > maxSequence { return false } return true }