125 lines
2.9 KiB
Go
125 lines
2.9 KiB
Go
package snowflake
|
||
|
||
import (
|
||
"fmt"
|
||
"sync"
|
||
"time"
|
||
)
|
||
|
||
// Snowflake 雪花算法ID生成器
|
||
type Snowflake struct {
|
||
mutex sync.Mutex
|
||
startTime int64 // 起始时间戳（毫秒）
|
||
machineID int64 // 机器ID
|
||
datacenterID int64 // 数据中心ID
|
||
sequence int64 // 序列号
|
||
}
|
||
|
||
const (
|
||
// 雪花算法各个部分的位数
|
||
workerIDBits = 10 // 机器ID位数
|
||
datacenterIDBits = 10 // 数据中心ID位数
|
||
sequenceBits = 12 // 序列号位数
|
||
|
||
// 最大值
|
||
maxWorkerID = -1 ^ (-1 << workerIDBits) // 1023
|
||
maxDatacenterID = -1 ^ (-1 << datacenterIDBits) // 1023
|
||
maxSequence = -1 ^ (-1 << sequenceBits) // 4095
|
||
|
||
// 位移
|
||
workerIDShift = sequenceBits
|
||
datacenterIDShift = sequenceBits + workerIDBits
|
||
timestampShift = sequenceBits + workerIDBits + datacenterIDBits
|
||
)
|
||
|
||
var (
|
||
// 全局雪花算法实例
|
||
instance *Snowflake
|
||
once sync.Once
|
||
)
|
||
|
||
// NewSnowflake 创建新的雪花算法实例
|
||
func NewSnowflake(machineID, datacenterID int64) (*Snowflake, error) {
|
||
if machineID < 0 || machineID > maxWorkerID {
|
||
return nil, fmt.Errorf("machineID必须在0-%d之间", maxWorkerID)
|
||
}
|
||
if datacenterID < 0 || datacenterID > maxDatacenterID {
|
||
return nil, fmt.Errorf("datacenterID必须在0-%d之间", maxDatacenterID)
|
||
}
|
||
|
||
return &Snowflake{
|
||
startTime: 1640995200000, // 2022-01-01 00:00:00 UTC
|
||
machineID: machineID,
|
||
datacenterID: datacenterID,
|
||
}, nil
|
||
}
|
||
|
||
// InitDefault 初始化默认实例
|
||
func InitDefault(machineID, datacenterID int64) error {
|
||
var err error
|
||
once.Do(func() {
|
||
instance, err = NewSnowflake(machineID, datacenterID)
|
||
})
|
||
return err
|
||
}
|
||
|
||
// GetInstance 获取默认实例
|
||
func GetInstance() *Snowflake {
|
||
if instance == nil {
|
||
// 默认初始化
|
||
instance, _ = NewSnowflake(1, 1)
|
||
}
|
||
return instance
|
||
}
|
||
|
||
// NextID 生成下一个ID（返回正数）
|
||
func (s *Snowflake) NextID() int64 {
|
||
s.mutex.Lock()
|
||
defer s.mutex.Unlock()
|
||
|
||
now := time.Now().UnixNano() / 1e6 // 当前时间戳（毫秒）
|
||
|
||
if now < s.startTime {
|
||
panic("时钟回拨，拒绝生成ID")
|
||
}
|
||
|
||
if now == s.startTime {
|
||
// 同一毫秒内
|
||
s.sequence = (s.sequence + 1) & maxSequence
|
||
if s.sequence == 0 {
|
||
// 序列号溢出，等待下一毫秒
|
||
for now <= s.startTime {
|
||
now = time.Now().UnixNano() / 1e6
|
||
}
|
||
}
|
||
} else {
|
||
// 新的毫秒，重置序列号
|
||
s.sequence = 0
|
||
}
|
||
|
||
s.startTime = now
|
||
|
||
// 生成ID - 确保返回正数
|
||
id := ((now - 1640995200000) << timestampShift) |
|
||
(s.datacenterID << datacenterIDShift) |
|
||
(s.machineID << workerIDShift) |
|
||
s.sequence
|
||
|
||
// 确保ID为正数（移除符号位）
|
||
if id < 0 {
|
||
id = id & 0x7FFFFFFFFFFFFFFF // 移除符号位
|
||
}
|
||
|
||
return id
|
||
}
|
||
|
||
// NextIDString 生成下一个ID（字符串格式）
|
||
func (s *Snowflake) NextIDString() string {
|
||
return fmt.Sprintf("%d", s.NextID())
|
||
}
|
||
|
||
// GenerateID 全局函数，生成唯一ID
|
||
func GenerateID() string {
|
||
return GetInstance().NextIDString()
|
||
}
|