go实现雪花算法

发表于

如何在分布式系统中，实现一个有递增性且全局唯一的ID呢？

实现方案之： go实现雪花算法

雪花算法在上一节已经写过了，在代码里也有注释。

直接上代码 –>

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
package snowflake

import (
	"errors"
	"log"
	"sync"
	"time"
)

// 雪花算法
/**
  0 | 00000000 00000000 00000000 00000000 00000000 0 / 00000000 00 / 00000000 0000
  1bit不用 |             41bit时间戳                  / 10bit工作机器ID / 12bit序列号

  1bit不用，二进制中最高位为1代表负数
  41bit代表时间戳，(2^41−1)/(1000∗60∗60∗24∗365)=69 年
  10bit工作机器ID允许分布式最大节点数为1024个（10个bit按场景划分，总计组合1024个）
  12bit序列号表示每毫秒生成的ID序号 2^12-1=4095，即0，1，...4095共4096个序列号

  实际情况下，可以根据具体场景来划分10bit工作机器ID和12bit序列号，
  比如说，把工作机器ID缩减到4位，把bit序列号扩大到18位来满足机器不足，单机可生成更多但序列号ID
*/

const (
	timestampOffset uint8 = 22 // 时间戳偏移量
	workerOffset    uint8 = 12 // 工作机器偏移量

	workerBits   int64 = 10 // 工作机器ID所占bit数
	sequenceBits int64 = 12 // 序列号所占bit数

	maxWorkerId   int64 = -1 ^ (-1 << uint64(workerBits))   // 最大工作机器ID
	maxSequenceId int64 = -1 ^ (-1 << uint64(sequenceBits)) // 最大序列号ID

	epoch int64 = 1599029955445
)

type Worker struct {
	mu         sync.Mutex
	timestamp  int64 // 上一次生成ID的时间戳
	workerId   int64 // 工作机器ID
	sequenceId int64 // 每毫秒的序列号，从0开始，最多4096个
}

func NewSnowflakeWorker(workerId int64) (*Worker, error) {
	if workerId < 0 || workerId > maxWorkerId {
		return nil, errors.New("WorkerId big than maxWorkerId")
	}

	return &Worker{
		timestamp:  0,
		workerId:   workerId,
		sequenceId: 0,
	}, nil
}

func (w *Worker) GenerateId() int64 {
	w.mu.Lock()
	defer w.mu.Unlock()

	timestamp := w.getCurrentTime()
	if timestamp < w.timestamp {
		log.Fatal("can not generate id")
	}

	if timestamp == w.timestamp {
		w.sequenceId++
		if w.sequenceId > maxSequenceId {
			for timestamp <= w.timestamp {
				timestamp = w.getCurrentTime()
			}
			w.sequenceId = 0
			w.timestamp = timestamp
		}
	} else {
		w.sequenceId = 0
		w.timestamp = timestamp
	}

	id := (timestamp-epoch)<<timestampOffset | (w.workerId << workerOffset) | w.sequenceId
	return id
}

func (w *Worker) getCurrentTime() int64 {
	return time.Now().UnixNano() / 1e6
}

如何使用呢？写了一个测试代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
func TestNewSnowflakeWorker(t *testing.T) {
	// 指定workerId，保证不重复即可，实际要结合生产环境来看，最大支持1024个组合
	worker, err := NewSnowflakeWorker(12)
	if err != nil {
		t.Fatal(err)
	}

	ch := make(chan int64)
	defer close(ch)
	count := 10000
	for i := 0; i < count; i++ {
		go func() {
			id := worker.GenerateId()
			ch <- id
		}()
	}

	m := make(map[int64]int)
	for i := 0; i < count; i++ {
		id := <-ch
		_, ok := m[id]
		if ok {
			t.Error("id is not unique")
			return
		}

		m[id] = 1
	}

	fmt.Println("done...")
}

通过执行测试用例，生成的ID不会重复