144 lines
3.4 KiB
Go
144 lines
3.4 KiB
Go
package common
|
||
|
||
import (
|
||
"errors"
|
||
"strconv"
|
||
"sync"
|
||
"time"
|
||
)
|
||
|
||
/**
|
||
分布式/多实例运行（重要）
|
||
如果你有多个 Pod 或服务器，必须在程序启动时给它们分配不同的 ID，否则还是会冲突
|
||
|
||
func main() {
|
||
// 获取当前机器的编号，比如从配置文件或环境变量读取
|
||
// 假设这是第 2 号机器
|
||
myMachineID := int64(2)
|
||
|
||
// 初始化
|
||
err := common.InitGenerator(myMachineID)
|
||
if err != nil {
|
||
panic(err)
|
||
}
|
||
|
||
// 之后随处调用
|
||
println(common.GenerateStringID())
|
||
}
|
||
*/
|
||
|
||
// 定义常量 (标准雪花算法配置)
|
||
const (
|
||
epoch = int64(1704067200000) // 起始时间戳 2024-01-01
|
||
workerBits = uint(10) // 机器ID位数
|
||
sequenceBits = uint(12) // 序列号位数
|
||
maxWorker = -1 ^ (-1 << workerBits)
|
||
maxSequence = -1 ^ (-1 << sequenceBits)
|
||
workerShift = sequenceBits
|
||
timestampShift = sequenceBits + workerBits
|
||
)
|
||
|
||
type IDGenerator struct {
|
||
mu sync.Mutex
|
||
lastTime int64
|
||
workerID int64
|
||
sequence int64
|
||
}
|
||
|
||
var (
|
||
defaultGenerator *IDGenerator
|
||
once sync.Once
|
||
)
|
||
|
||
// InitGenerator 初始化单例生成器
|
||
// 修复：校验逻辑移到 once.Do 外部，防止校验失败消耗掉 once 的执行机会
|
||
func InitGenerator(workerID int64) error {
|
||
// 1. 先校验，如果失败直接返回，不要触碰 once
|
||
if workerID < 0 || workerID > int64(maxWorker) {
|
||
return errors.New("worker ID excess of limit (0-1023)")
|
||
}
|
||
|
||
// 2. 执行初始化
|
||
once.Do(func() {
|
||
defaultGenerator = &IDGenerator{
|
||
workerID: workerID,
|
||
lastTime: 0,
|
||
sequence: 0,
|
||
}
|
||
})
|
||
return nil
|
||
}
|
||
|
||
// getInstance 获取单例
|
||
func getInstance() *IDGenerator {
|
||
// 双重检查，虽然 once.Do 是线程安全的，但如果 InitGenerator 没被调用过，
|
||
// 我们需要确保这里能兜底初始化
|
||
if defaultGenerator == nil {
|
||
once.Do(func() {
|
||
defaultGenerator = &IDGenerator{
|
||
workerID: 1, // 默认机器ID，防止未初始化导致 panic
|
||
lastTime: 0,
|
||
sequence: 0,
|
||
}
|
||
})
|
||
}
|
||
|
||
// 【关键修复】如果经过上面的逻辑 defaultGenerator 还是 nil
|
||
// (这种情况极少见，除非 InitGenerator 曾经被错误调用且没有赋值)
|
||
// 强制创建一个临时的或抛出 panic，避免空指针崩溃
|
||
if defaultGenerator == nil {
|
||
// 最后的兜底，防止崩溃
|
||
return &IDGenerator{workerID: 1}
|
||
}
|
||
|
||
return defaultGenerator
|
||
}
|
||
|
||
// GenerateLongID 全局辅助函数
|
||
func GenerateLongID() int64 {
|
||
return getInstance().NextID()
|
||
}
|
||
|
||
// GenerateStringID 全局辅助函数
|
||
func GenerateStringID() string {
|
||
return strconv.FormatInt(getInstance().NextID(), 10)
|
||
}
|
||
|
||
// NextID 生成下一个 ID
|
||
func (g *IDGenerator) NextID() int64 {
|
||
// 防御性编程：防止 g 为 nil
|
||
if g == nil {
|
||
// 如果实例是 nil，尝试获取默认实例
|
||
if defaultGenerator != nil {
|
||
g = defaultGenerator
|
||
} else {
|
||
// 极端情况，创建一个临时对象（虽然锁不住全局，但能防崩）
|
||
g = &IDGenerator{workerID: 1}
|
||
}
|
||
}
|
||
|
||
g.mu.Lock()
|
||
defer g.mu.Unlock()
|
||
|
||
now := time.Now().UnixMilli()
|
||
|
||
if now < g.lastTime {
|
||
now = g.lastTime
|
||
}
|
||
|
||
if now == g.lastTime {
|
||
g.sequence = (g.sequence + 1) & int64(maxSequence)
|
||
if g.sequence == 0 {
|
||
for now <= g.lastTime {
|
||
now = time.Now().UnixMilli()
|
||
}
|
||
}
|
||
} else {
|
||
g.sequence = 0
|
||
}
|
||
|
||
g.lastTime = now
|
||
|
||
return ((now - epoch) << timestampShift) | (g.workerID << workerShift) | g.sequence
|
||
}
|