80 lines
1.9 KiB
Go
80 lines
1.9 KiB
Go
package idgen
|
||
|
||
import (
|
||
"sync"
|
||
"time"
|
||
)
|
||
|
||
const (
|
||
// int64EpochMs 使用固定业务纪元，缩短时间戳占用位数。
|
||
int64EpochMs int64 = 1704067200000
|
||
// maxInt64NodeID 预留 10 bit 给节点，支持 1024 个发号节点。
|
||
maxInt64NodeID int64 = 1023
|
||
// int64SequenceMask 预留 12 bit 给同毫秒序列。
|
||
int64SequenceMask int64 = 4095
|
||
)
|
||
|
||
// Int64Generator 生成 Snowflake-like int64 ID。
|
||
// 格式为 timestamp_ms | node_id | sequence，适合用户主键这种不可变长 ID。
|
||
type Int64Generator struct {
|
||
mu sync.Mutex
|
||
nodeID int64
|
||
lastMs int64
|
||
sequence int64
|
||
now func() time.Time
|
||
sleepUntil func(time.Time)
|
||
}
|
||
|
||
// NewInt64Generator 创建用户主键发号器。
|
||
// nodeID 必须由部署配置保证不同发号节点不冲突。
|
||
func NewInt64Generator(nodeID int64) *Int64Generator {
|
||
if nodeID < 0 {
|
||
nodeID = 0
|
||
}
|
||
if nodeID > maxInt64NodeID {
|
||
nodeID = nodeID & maxInt64NodeID
|
||
}
|
||
|
||
return &Int64Generator{
|
||
nodeID: nodeID,
|
||
now: time.Now,
|
||
sleepUntil: defaultSleepUntil,
|
||
}
|
||
}
|
||
|
||
// NewInt64 返回一个全局唯一倾向的 int64 ID。
|
||
func (g *Int64Generator) NewInt64() int64 {
|
||
g.mu.Lock()
|
||
defer g.mu.Unlock()
|
||
|
||
nowMs := g.now().UnixMilli()
|
||
if nowMs < g.lastMs {
|
||
// 时钟回拨时等待到上次发号时间，避免生成倒退 ID。
|
||
g.sleepUntil(time.UnixMilli(g.lastMs))
|
||
nowMs = g.now().UnixMilli()
|
||
}
|
||
|
||
if nowMs == g.lastMs {
|
||
g.sequence = (g.sequence + 1) & int64SequenceMask
|
||
if g.sequence == 0 {
|
||
// 同毫秒序列耗尽时等到下一毫秒再继续。
|
||
next := time.UnixMilli(g.lastMs + 1)
|
||
g.sleepUntil(next)
|
||
nowMs = g.now().UnixMilli()
|
||
}
|
||
} else {
|
||
g.sequence = 0
|
||
}
|
||
|
||
g.lastMs = nowMs
|
||
|
||
return ((nowMs - int64EpochMs) << 22) | (g.nodeID << 12) | g.sequence
|
||
}
|
||
|
||
// defaultSleepUntil 封装真实等待，测试可通过结构体替换。
|
||
func defaultSleepUntil(target time.Time) {
|
||
if delay := time.Until(target); delay > 0 {
|
||
time.Sleep(delay)
|
||
}
|
||
}
|