wwweww
95 lines
1.9 KiB
Go
95 lines
1.9 KiB
Go
package generator
|
|
|
|
import (
|
|
"errors"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
const (
|
|
epoch = int64(1609459200000)
|
|
datacenterIdBits = uint(5)
|
|
workerIdBits = uint(5)
|
|
sequenceBits = uint(12)
|
|
|
|
maxDatacenterId = -1 ^ (-1 << datacenterIdBits)
|
|
maxWorkerId = -1 ^ (-1 << workerIdBits)
|
|
maxSequence = -1 ^ (-1 << sequenceBits)
|
|
|
|
workerIdShift = sequenceBits
|
|
datacenterIdShift = sequenceBits + workerIdBits
|
|
timestampLeftShift = sequenceBits + workerIdBits + datacenterIdBits
|
|
)
|
|
|
|
type Snowflake struct {
|
|
mu sync.Mutex
|
|
timestamp int64
|
|
datacenterId int64
|
|
workerId int64
|
|
sequence int64
|
|
}
|
|
|
|
func NewSnowflake(datacenterId, workerId int64) (*Snowflake, error) {
|
|
if datacenterId < 0 || datacenterId > maxDatacenterId {
|
|
return nil, errors.New("datacenter id must be between 0 and 31")
|
|
}
|
|
if workerId < 0 || workerId > maxWorkerId {
|
|
return nil, errors.New("worker id must be between 0 and 31")
|
|
}
|
|
|
|
return &Snowflake{
|
|
timestamp: 0,
|
|
datacenterId: datacenterId,
|
|
workerId: workerId,
|
|
sequence: 0,
|
|
}, nil
|
|
}
|
|
|
|
func (s *Snowflake) NextID() (int64, error) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
now := time.Now().UnixMilli()
|
|
|
|
if now < s.timestamp {
|
|
return 0, errors.New("clock moved backwards")
|
|
}
|
|
|
|
if now == s.timestamp {
|
|
s.sequence = (s.sequence + 1) & maxSequence
|
|
if s.sequence == 0 {
|
|
for now <= s.timestamp {
|
|
now = time.Now().UnixMilli()
|
|
}
|
|
}
|
|
} else {
|
|
s.sequence = 0
|
|
}
|
|
|
|
s.timestamp = now
|
|
|
|
id := ((now - epoch) << timestampLeftShift) |
|
|
(s.datacenterId << datacenterIdShift) |
|
|
(s.workerId << workerIdShift) |
|
|
s.sequence
|
|
|
|
return id, nil
|
|
}
|
|
|
|
func (s *Snowflake) NextIDs(count int) ([]int64, error) {
|
|
if count <= 0 || count > 1000 {
|
|
return nil, errors.New("count must be between 1 and 1000")
|
|
}
|
|
|
|
ids := make([]int64, count)
|
|
for i := 0; i < count; i++ {
|
|
id, err := s.NextID()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
ids[i] = id
|
|
}
|
|
|
|
return ids, nil
|
|
}
|