[feat] package snowflake: 新增包，分布式唯一性ID生成器

.gitignore

@@ -11,7 +11,6 @@ profile*.pdf
 /pre-commit.sh
 /TODO.md
 
-/pkg/snowflake
 /pkg/tag
 /pkg/nazatime
 /demo/samefile

README.md

@@ -28,12 +28,12 @@ Go语言基础库
 
 ```
 pkg/                    ...... 源码包
+    |-- bininfo/        ...... 将编译时源码的 git 版本信息（当前 commit log 的 sha 值和 commit message），编译时间，Go 版本，平台打入程序中
     |-- nazalog/        ...... 日志库
     |-- slicebytepool/  ...... []byte 内存池
     |-- assert/         ...... 提供了单元测试时的断言功能，减少一些模板代码
     |-- nazastring/     ...... string 和 []byte 相关的操作
     |-- consistenthash/ ...... 一致性哈希
-    |-- bininfo/        ...... 将编译时源码的 git 版本信息（当前 commit log 的 sha 值和 commit message），编译时间，Go 版本，平台打入程序中
     |-- bele/           ...... 提供了大小端的转换操作
     |-- nazaatomic/     ...... 原子操作
     |-- nazajson/       ...... json 操作
@@ -41,6 +41,7 @@ pkg/                    ...... 源码包
     |-- connection/     ...... 对 net.Conn 接口的二次封装
     |-- filebatch/      ...... 文件批处理操作
     |-- nazamd5/        ...... md5 操作
+    |-- snowflake/      ...... 分布式唯一性 ID 生成器
     |-- bitrate/        ...... 计算带宽
     |-- ratelimit/      ...... 限流器，令牌桶
     |-- ic/             ...... 将整型切片压缩成二进制字节切片

pkg/snowflake/snowflake.go

@@ -0,0 +1,160 @@
+// 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
+		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
+}

pkg/snowflake/snowflake_test.go

@@ -0,0 +1,164 @@
+// 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_test
+
+import (
+	"sync"
+	"testing"
+
+	"github.com/q191201771/naza/pkg/assert"
+	"github.com/q191201771/naza/pkg/nazalog"
+	"github.com/q191201771/naza/pkg/snowflake"
+)
+
+func TestNew(t *testing.T) {
+	n, err := snowflake.New(0, 0)
+	assert.Equal(t, nil, err)
+	id, err := n.Gen()
+	assert.Equal(t, nil, err)
+	nazalog.Debug(id)
+	id, err = n.Gen()
+	assert.Equal(t, nil, err)
+	nazalog.Debug(id)
+}
+
+func TestNegative(t *testing.T) {
+	n, err := snowflake.New(0, 0)
+	assert.Equal(t, nil, err)
+	id, err := n.Gen(1288834974657 + (1 << 41))
+	assert.Equal(t, true, id < 0)
+}
+
+func TestAlwaysPositive(t *testing.T) {
+	n, err := snowflake.New(0, 0, func(option *snowflake.Option) {
+		option.AlwaysPositive = true
+	})
+	assert.Equal(t, nil, err)
+	id, err := n.Gen(1288834974657 + (1 << 41))
+	assert.Equal(t, true, id == 0)
+	id, err = n.Gen(1288834974657 + (1 << 41) + 0x1234)
+	assert.Equal(t, true, id >= 0)
+}
+
+func TestErrInitial(t *testing.T) {
+	var (
+		n   *snowflake.Node
+		err error
+		id  int64
+	)
+
+	n, err = snowflake.New(0, 0, func(option *snowflake.Option) {
+		option.SequenceBits = 64
+	})
+	assert.Equal(t, snowflake.ErrInitial, err)
+
+	n, err = snowflake.New(0, 0, func(option *snowflake.Option) {
+		option.WorkerIDBits = 64
+	})
+	assert.Equal(t, snowflake.ErrInitial, err)
+
+	n, err = snowflake.New(0, 0, func(option *snowflake.Option) {
+		option.DataCenterIDBits = 64
+	})
+	assert.Equal(t, snowflake.ErrInitial, err)
+
+	n, err = snowflake.New(0, 0, func(option *snowflake.Option) {
+		option.DataCenterIDBits = 31
+		option.WorkerIDBits = 31
+		option.SequenceBits = 31
+	})
+	assert.Equal(t, snowflake.ErrInitial, err)
+
+	n, err = snowflake.New(100, 0, func(option *snowflake.Option) {
+		option.DataCenterIDBits = 1
+	})
+	assert.Equal(t, snowflake.ErrInitial, err)
+
+	n, err = snowflake.New(0, 100, func(option *snowflake.Option) {
+		option.WorkerIDBits = 1
+	})
+	assert.Equal(t, snowflake.ErrInitial, err)
+
+	if n != nil {
+		id, err = n.Gen()
+		assert.Equal(t, nil, err)
+		nazalog.Debug(id)
+	}
+}
+
+func TestErrGen(t *testing.T) {
+	var (
+		n   *snowflake.Node
+		err error
+		id  int64
+	)
+
+	n, err = snowflake.New(0, 0)
+	assert.Equal(t, nil, err)
+	assert.IsNotNil(t, n)
+
+	id, err = n.Gen(2)
+	assert.Equal(t, nil, err)
+	id, err = n.Gen(1)
+	assert.Equal(t, snowflake.ErrGen, err)
+	nazalog.Debug(id)
+}
+
+func TestMT(t *testing.T) {
+	var (
+		n   *snowflake.Node
+		err error
+	)
+
+	n, err = snowflake.New(0, 0, func(option *snowflake.Option) {
+		option.SequenceBits = 1
+	})
+	assert.Equal(t, nil, err)
+	assert.IsNotNil(t, n)
+
+	ii := 16
+	jj := 16
+	var mu sync.Mutex
+	m := make(map[int64]struct{})
+	var wg sync.WaitGroup
+	wg.Add(ii * jj)
+
+	for i := 0; i < ii; i++ {
+		go func() {
+			for j := 0; j < jj; j++ {
+				id, err := n.Gen()
+				assert.Equal(t, nil, err)
+				mu.Lock()
+				m[id] = struct{}{}
+				mu.Unlock()
+				wg.Done()
+			}
+		}()
+	}
+	wg.Wait()
+	assert.Equal(t, ii*jj, len(m))
+}
+
+func BenchmarkNode_Gen(b *testing.B) {
+	var (
+		n   *snowflake.Node
+		err error
+	)
+
+	n, err = snowflake.New(0, 0)
+	assert.Equal(b, nil, err)
+	assert.IsNotNil(b, n)
+
+	var dummy int64
+	for i := 0; i < b.N; i++ {
+		id, _ := n.Gen()
+		dummy += id
+	}
+	nazalog.Debug(dummy)
+}