Files
bl/common/utils/snowflake.go

150 lines
3.9 KiB
Go
Raw Normal View History

package utils
import (
"errors"
"strconv"
"sync"
"time"
)
// Snowflake is a uint32-sized ID with 32 node support and 8.5 year lifespan
type Snowflake uint32
// ID is an alias for Snowflake
type ID = Snowflake
func (s Snowflake) String() string {
return strconv.FormatUint(uint64(s), 16)
}
// ParseSnowflake parses a string into a snowflake, if possible
func ParseSnowflake(s string) (Snowflake, error) {
i, err := strconv.ParseUint(s, 16, 32)
if err != nil {
return 0, err
}
return Snowflake(i), nil
}
// DefaultGenerator is a Generator with the epoch set to Jan 1, 2025 UTC
// Generator holds info needed for generating snowflakes with 32 node support
type Generator struct {
Epoch time.Time
nodeID uint8 // 0-31 (5位支持32个节点)
sequence uint8 // 0-255 (8位支持每秒256个ID)
lastSec uint32 // 上次生成ID的秒级时间戳
mu sync.Mutex
}
// NewGen returns a new generator with specified epoch and node ID
func NewGen(epoch time.Time, nodeID uint8) *Generator {
if nodeID >= 32 { // 5位最多支持0-31
panic(errors.New("node ID must be between 0 and 31"))
}
return &Generator{
Epoch: epoch,
nodeID: nodeID,
}
}
// Get generates a new unique snowflake ID and returns it as uint32
func (g *Generator) Get() uint32 {
g.mu.Lock()
defer g.mu.Unlock()
// 获取当前时间戳(自纪元以来的秒数)
now := time.Now()
sec := uint32(now.Sub(g.Epoch).Seconds())
// 检查时间戳是否超出28位限制 (2^28 - 1 = 268435455秒 ≈ 8.5年)
if sec >= (1 << 28) {
panic(errors.New("timestamp overflow: please update epoch"))
}
// 处理时间回拨
if sec < g.lastSec {
// 计算需要等待的时间
sleepDuration := time.Duration(g.lastSec-sec) * time.Second
// 添加1毫秒缓冲确保时间确实已过
sleepDuration += time.Millisecond
// 等待时间追上
time.Sleep(sleepDuration)
// 重新获取时间戳
now = time.Now()
sec = uint32(now.Sub(g.Epoch).Seconds())
// 再次检查,防止极端情况
if sec < g.lastSec {
panic(errors.New("clock moved backwards beyond recovery"))
}
}
// 处理同一秒内的序列号
if sec == g.lastSec {
g.sequence++
// 序列号溢出8位最大255
if g.sequence > 255 { // 0-255是有效范围
// 等待到下一秒
sleepDuration := time.Second - time.Duration(now.Nanosecond())*time.Nanosecond + time.Microsecond
time.Sleep(sleepDuration)
now = time.Now()
sec = uint32(now.Sub(g.Epoch).Seconds())
g.sequence = 0
}
} else {
// 新的一秒,重置序列号
g.sequence = 0
}
// 更新上次时间戳
g.lastSec = sec
// 组装ID (32位):
// [28位秒级时间戳][5位节点ID][8位序列号]
var id uint32
id |= (sec << 13) // 左移13位5+8放置秒级时间戳
id |= (uint32(g.nodeID) << 8) // 左移8位放置节点ID
id |= uint32(g.sequence) // 放置序列号
return id
}
// Parse extracts timestamp, node ID and sequence from a snowflake
func (g *Generator) Parse(id uint32) (t time.Time, nodeID uint8, sequence uint8) {
// 提取各部分信息
sec := id >> 13
nodeID = uint8((id >> 8) & 0x1F) // 5位掩码0x1F = 31
sequence = uint8(id & 0xFF) // 8位掩码0xFF = 255
// 计算实际时间
totalNanos := g.Epoch.UnixNano() + int64(sec)*int64(time.Second)
t = time.Unix(0, totalNanos).UTC()
return
}
// GetNodeID returns the current node ID
func (g *Generator) GetNodeID() uint8 {
return g.nodeID
}
// GetLastTimestamp returns the last timestamp used (for testing/debugging)
func (g *Generator) GetLastTimestamp() uint32 {
g.mu.Lock()
defer g.mu.Unlock()
return g.lastSec
}
// MaxAge returns the maximum age (in seconds) that this generator can handle before overflow
func (g *Generator) MaxAge() uint32 {
return (1 << 28) - 1 // 268435455秒约8.5年
}
// MaxSequencePerSecond returns the maximum number of sequences per second
func (g *Generator) MaxSequencePerSecond() uint8 {
return 255 // 2^8 - 1
}