intelligent_monitoring_backend/Code/backend/internal/utils/snowflake.go

60 lines
1.3 KiB
Go

package utils
import (
"fmt"
"sync"
"time"
)
const (
epoch = int64(1609459200000) // 2021-01-01 00:00:00 的时间戳(毫秒)
workerBits = 5
sequenceBits = 12
maxWorkerID = -1 ^ (-1 << workerBits)
maxSequenceID = -1 ^ (-1 << sequenceBits)
workerShift = sequenceBits
timestampShift = sequenceBits + workerBits
)
type Snowflake struct {
mu sync.Mutex
timestamp int64
workerID int
sequenceID int
}
func NewSnowflake(workerID int) *Snowflake {
if workerID < 0 || workerID > maxWorkerID {
panic(fmt.Sprintf("worker ID must be between 0 and %d", maxWorkerID))
}
return &Snowflake{
timestamp: 0,
workerID: workerID,
sequenceID: 0,
}
}
func (s *Snowflake) NextID() int64 {
s.mu.Lock()
defer s.mu.Unlock()
currentTime := time.Now().UnixNano() / 1000000 // 转换为毫秒
if s.timestamp == currentTime {
s.sequenceID = (s.sequenceID + 1) & maxSequenceID
if s.sequenceID == 0 {
// 当前毫秒内的序列号用尽,等待下一毫秒
for currentTime <= s.timestamp {
currentTime = time.Now().UnixNano() / 1000000
}
}
} else {
s.sequenceID = 0
}
s.timestamp = currentTime
id := ((currentTime - epoch) << timestampShift) |
(int64(s.workerID) << workerShift) |
int64(s.sequenceID)
return id
}