/
worker.go
84 lines (73 loc) · 2.03 KB
/
worker.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
package snowflake
import (
"fmt"
"sync"
"sync/atomic"
"github.com/chanxuehong/rand"
)
const (
// 0(1bit) + 41bits timestamp + 10bits workerId + 12bits sequence
workerIdBits = 10
maxWorkerId = int64(-1) ^ (int64(-1) << workerIdBits)
sequenceBits = 12
sequenceMask = int64(-1) ^ (int64(-1) << sequenceBits)
workerIdShift = sequenceBits
timestampShift = workerIdBits + sequenceBits
snowflakeIdBits = 63
snowflakeIdMask = int64(-1) ^ (int64(-1) << snowflakeIdBits)
)
var gSequenceStart = rand.Int63() & sequenceMask // Reduce the collision probability
type Worker struct {
workerId int64
mutex sync.Mutex
sequenceStart int64
lastTimestamp int64
lastSequence int64
}
func NewWorker(workerId int64) (*Worker, error) {
if workerId < 0 || workerId > maxWorkerId {
return nil, fmt.Errorf("workerId can't be less than 0 or greater than %d", maxWorkerId)
}
sequenceStart := atomic.LoadInt64(&gSequenceStart)
wk := &Worker{
workerId: workerId,
sequenceStart: sequenceStart,
lastTimestamp: -1,
lastSequence: sequenceStart,
}
return wk, nil
}
func (wk *Worker) NextId() int64 {
var (
timestamp = twepochTimestamp()
workerId = wk.workerId
sequence int64
)
wk.mutex.Lock() // Lock
switch {
case timestamp > wk.lastTimestamp:
sequence = wk.sequenceStart
wk.lastTimestamp = timestamp
wk.lastSequence = sequence
wk.mutex.Unlock() // Unlock
case timestamp == wk.lastTimestamp:
sequence = (wk.lastSequence + 1) & sequenceMask
if sequence == wk.sequenceStart {
timestamp = tillNextMillis(timestamp)
wk.lastTimestamp = timestamp
}
wk.lastSequence = sequence
wk.mutex.Unlock() // Unlock
default: // timestamp < wk.lastTimestamp
sequenceStart := rand.Int63() & sequenceMask
atomic.StoreInt64(&gSequenceStart, sequenceStart)
wk.sequenceStart = sequenceStart
sequence = wk.sequenceStart
wk.lastTimestamp = timestamp
wk.lastSequence = sequence
wk.mutex.Unlock() // Unlock
}
id := timestamp<<timestampShift | workerId<<workerIdShift | sequence
id &= snowflakeIdMask
return id
}