Skip to content
Permalink
Branch: master
Find file Copy path
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
161 lines (137 sloc) 4.1 KB
// Copyright 2019, Chef. All rights reserved.
// https://github.com/q191201771/naza
//
// Use of this source code is governed by a MIT-style license
// that can be found in the License file.
//
// Author: Chef (191201771@qq.com)
package snowflake
import (
"errors"
"sync"
"time"
)
var (
ErrInitial = errors.New("lal.snowflake: initial error")
ErrGen = errors.New("lal.snowflake: gen error")
)
type Option struct {
DataCenterIDBits int // 数据中心编号字段在所生成 ID 所占的位数,取值范围见 validate 函数
WorkerIDBits int // 节点编号
SequenceBits int // 递增序列
Twepoch int64 // 基准时间点
AlwaysPositive bool // 是否只生成正数 ID,如果是,则时间戳所占位数会减少1位
}
var defaultOption = Option{
DataCenterIDBits: 5,
WorkerIDBits: 5,
SequenceBits: 12,
Twepoch: int64(1288834974657), // 对应现实时间: 2010/11/4 9:42:54.657
AlwaysPositive: false,
}
type Node struct {
dataCenterID int64
workerID int64
option Option
seqMask uint32
workerIDShift uint32
dataCenterIDShift uint32
timestampShift uint32
mu sync.Mutex
lastTs int64
seq uint32
}
type ModOption func(option *Option)
// dataCenterID 和 workerID 的取值范围取决于 DataCenterIDBits 和 WorkerIDBits
// 假设 DataCenterIDBits 为 5,则 dataCenterID 取值范围为 [0, 32]
func New(dataCenterID int, workerID int, modOptions ...ModOption) (*Node, error) {
option := defaultOption
for _, fn := range modOptions {
fn(&option)
}
if err := validate(dataCenterID, workerID, option); err != nil {
return nil, err
}
return &Node{
dataCenterID: int64(dataCenterID),
workerID: int64(workerID),
option: option,
seqMask: uint32(bitsToMax(option.SequenceBits)),
workerIDShift: uint32(option.SequenceBits),
dataCenterIDShift: uint32(option.SequenceBits + option.WorkerIDBits),
timestampShift: uint32(option.SequenceBits + option.WorkerIDBits + option.DataCenterIDBits),
lastTs: -1,
}, nil
}
func (n *Node) Gen(nowUnixMSec ...int64) (int64, error) {
n.mu.Lock()
defer n.mu.Unlock()
// 当前 Unix 时间戳可由外部传入
var now int64
if len(nowUnixMSec) == 0 {
now = time.Now().UnixNano() / 1e6
} else {
now = nowUnixMSec[0]
}
// 时间戳回退,返回错误
if now < n.lastTs {
return -1, ErrGen
}
// 时间戳相同时,使用递增序号解决冲突
if now == n.lastTs {
n.seq = (n.seq + 1) & n.seqMask
// 递增序号翻转为 0,表示该时间戳下的序号已经全部用完,阻塞等待系统时间增长
if n.seq == 0 {
for now <= n.lastTs {
now = time.Now().UnixNano() / 1e6
}
}
} else {
n.seq = 0
}
n.lastTs = now
// 如果保证只返回正数,则生成的 ID 的最高位,也即时间戳的最高位保持为 0
ts := now - n.option.Twepoch
if n.option.AlwaysPositive {
ts = clearBit(ts, 63-n.timestampShift)
}
ts <<= n.timestampShift
// 用所有字段组合生成 ID 返回
return ts | (n.dataCenterID << n.dataCenterIDShift) | (n.workerID << n.workerIDShift) | int64(n.seq), nil
}
func validate(dataCenterID int, workerID int, option Option) error {
if option.DataCenterIDBits < 0 || option.DataCenterIDBits > 31 {
return ErrInitial
}
if option.WorkerIDBits < 0 || option.WorkerIDBits > 31 {
return ErrInitial
}
if option.SequenceBits < 0 || option.SequenceBits > 31 {
return ErrInitial
}
if option.DataCenterIDBits+option.WorkerIDBits+option.SequenceBits >= 64 {
return ErrInitial
}
if option.DataCenterIDBits > 0 {
if dataCenterID > bitsToMax(option.DataCenterIDBits) {
return ErrInitial
}
}
if option.WorkerIDBits > 0 {
if workerID > bitsToMax(option.WorkerIDBits) {
return ErrInitial
}
}
return nil
}
// 位的数量对应的最大值,该函数也可以叫做 bitsToMask
func bitsToMax(bits int) int {
// -1 表示所有位都为 1
return int(int32(-1) ^ (int32(-1) << uint32(bits)))
}
// 将 <num> 的第 <index> 设置为 0
func clearBit(num int64, index uint32) int64 {
bit := int64(1 << index)
mask := int64(-1) ^ bit
return num & mask
}
You can’t perform that action at this time.