150 lines
4.0 KiB
Go
150 lines
4.0 KiB
Go
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
|
||
var DefaultGenerator = NewGen(time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC), 0)
|
||
|
||
// 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
|
||
}
|