123 lines
3.6 KiB
Go
123 lines
3.6 KiB
Go
package snowflake
|
||
|
||
import (
|
||
"errors"
|
||
"fmt"
|
||
"sync"
|
||
"time"
|
||
)
|
||
|
||
// 定义常量
|
||
const (
|
||
// 位数分配
|
||
sequenceBits = 12 // 序列号占用的位数
|
||
workerIdBits = 5 // 工作机器ID占用的位数
|
||
datacenterIdBits = 5 // 数据中心ID占用的位数
|
||
|
||
// 最大值
|
||
maxSequence = -1 ^ (-1 << sequenceBits) // 4095
|
||
maxWorkerId = -1 ^ (-1 << workerIdBits) // 31
|
||
maxDatacenterId = -1 ^ (-1 << datacenterIdBits) // 31
|
||
|
||
// 位移偏移量
|
||
workerIdShift = sequenceBits // 12
|
||
datacenterIdShift = sequenceBits + workerIdBits // 12 + 5 = 17
|
||
timestampLeftShift = sequenceBits + workerIdBits + datacenterIdBits // 12 + 5 + 5 = 22
|
||
)
|
||
|
||
// 起始时间戳 (纪元)，可以使用程序上线的时间，这里设置为 2020-01-01 00:00:00 UTC
|
||
var epoch int64 = 1577836800000
|
||
|
||
// Snowflake 结构体
|
||
type Snowflake struct {
|
||
mu sync.Mutex // 互斥锁，保证并发安全
|
||
lastTime int64 // 上次生成ID的时间戳
|
||
workerId int64 // 工作机器ID
|
||
datacenterId int64 // 数据中心ID
|
||
sequence int64 // 当前毫秒内的序列号
|
||
}
|
||
|
||
// NewSnowflake 初始化一个 Snowflake 实例
|
||
// workerId: 工作机器ID (0 ~ 31)
|
||
// datacenterId: 数据中心ID (0 ~ 31)
|
||
func NewSnowflake(workerId, datacenterId int64) (*Snowflake, error) {
|
||
if workerId < 0 || workerId > maxWorkerId {
|
||
return nil, errors.New(fmt.Sprintf("worker Id can't be greater than %d or less than 0", maxWorkerId))
|
||
}
|
||
if datacenterId < 0 || datacenterId > maxDatacenterId {
|
||
return nil, errors.New(fmt.Sprintf("datacenter Id can't be greater than %d or less than 0", maxDatacenterId))
|
||
}
|
||
|
||
return &Snowflake{
|
||
lastTime: 0,
|
||
workerId: workerId,
|
||
datacenterId: datacenterId,
|
||
sequence: 0,
|
||
}, nil
|
||
}
|
||
|
||
// NextId 生成下一个 ID
|
||
func (s *Snowflake) NextId() (int64, error) {
|
||
s.mu.Lock()
|
||
defer s.mu.Unlock()
|
||
|
||
// 获取当前时间戳（毫秒）
|
||
now := time.Now().UnixMilli()
|
||
|
||
// 如果当前时间小于上次生成ID的时间，说明时钟回拨，抛出异常
|
||
if now < s.lastTime {
|
||
return 0, errors.New(fmt.Sprintf("Clock moved backwards. Refusing to generate id for %d milliseconds", s.lastTime-now))
|
||
}
|
||
|
||
// 如果是同一毫秒内
|
||
if now == s.lastTime {
|
||
// 序列号自增
|
||
s.sequence = (s.sequence + 1) & maxSequence
|
||
// 如果序列号溢出（超过4095），则等待下一毫秒
|
||
if s.sequence == 0 {
|
||
now = s.waitNextMillis(now)
|
||
}
|
||
} else {
|
||
// 不同毫秒，序列号重置为0
|
||
s.sequence = 0
|
||
}
|
||
|
||
// 更新最后时间戳
|
||
s.lastTime = now
|
||
|
||
// 组装 ID
|
||
// (当前时间 - 起始时间) << 时间戳位移 | 数据中心ID << 数据中心位移 | 工作ID << 工作位移 | 序列号
|
||
id := ((now - epoch) << timestampLeftShift) |
|
||
(s.datacenterId << datacenterIdShift) |
|
||
(s.workerId << workerIdShift) |
|
||
s.sequence
|
||
|
||
return id, nil
|
||
}
|
||
|
||
// waitNextMillis 阻塞等待下一毫秒
|
||
func (s *Snowflake) waitNextMillis(lastTime int64) int64 {
|
||
now := time.Now().UnixMilli()
|
||
for now <= lastTime {
|
||
now = time.Now().UnixMilli()
|
||
}
|
||
return now
|
||
}
|
||
|
||
// ParseId 解析 ID，用于调试或查看 ID 组成部分
|
||
func ParseId(id int64) map[string]interface{} {
|
||
timestamp := (id >> timestampLeftShift) + epoch
|
||
datacenterId := (id >> datacenterIdShift) & maxDatacenterId
|
||
workerId := (id >> workerIdShift) & maxWorkerId
|
||
sequence := id & maxSequence
|
||
|
||
return map[string]interface{}{
|
||
"id": id,
|
||
"timestamp": timestamp,
|
||
"time_str": time.UnixMilli(timestamp).Format("2006-01-02 15:04:05.000"),
|
||
"datacenterId": datacenterId,
|
||
"workerId": workerId,
|
||
"sequence": sequence,
|
||
}
|
||
}
|