-
Notifications
You must be signed in to change notification settings - Fork 0
/
snowflake.go
364 lines (288 loc) · 8.7 KB
/
snowflake.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
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
package snowflake
import (
"encoding/base64"
"encoding/binary"
"errors"
"fmt"
"strconv"
"sync"
"time"
)
var (
// Epoch is set to the twitter snowflake epoch of Nov 04 2010 01:42:54 UTC in milliseconds
// You may customize this to set a different epoch for your application.
Epoch int64 = 1288834974657
// NodeBits holds the number of bits to use for Node
// Remember, you have a total 22 bits to share between Node/Step
NodeBits uint8 = 10
// StepBits holds the number of bits to use for Step
// Remember, you have a total 22 bits to share between Node/Step
StepBits uint8 = 12
// DEPRECATED: the below four variables will be removed in a future release.
mu sync.Mutex
nodeMax int64 = -1 ^ (-1 << NodeBits)
nodeMask = nodeMax << StepBits
stepMask int64 = -1 ^ (-1 << StepBits)
timeShift = NodeBits + StepBits
nodeShift = StepBits
)
const encodeBase32Map = "ybndrfg8ejkmcpqxot1uwisza345h769"
var decodeBase32Map [256]byte
const encodeBase58Map = "123456789abcdefghijkmnopqrstuvwxyzABCDEFGHJKLMNPQRSTUVWXYZ"
var decodeBase58Map [256]byte
// A JSONSyntaxError is returned from UnmarshalJSON if an invalid ID is provided.
type JSONSyntaxError struct{ original []byte }
func (j JSONSyntaxError) Error() string {
return fmt.Sprintf("invalid snowflake ID %q", string(j.original))
}
// ErrInvalidBase58 is returned by ParseBase58 when given an invalid []byte
var ErrInvalidBase58 = errors.New("invalid base58")
// ErrInvalidBase32 is returned by ParseBase32 when given an invalid []byte
var ErrInvalidBase32 = errors.New("invalid base32")
// Create maps for decoding Base58/Base32.
// This speeds up the process tremendously.
func init() {
for i := 0; i < len(encodeBase58Map); i++ {
decodeBase58Map[i] = 0xFF
}
for i := 0; i < len(encodeBase58Map); i++ {
decodeBase58Map[encodeBase58Map[i]] = byte(i)
}
for i := 0; i < len(encodeBase32Map); i++ {
decodeBase32Map[i] = 0xFF
}
for i := 0; i < len(encodeBase32Map); i++ {
decodeBase32Map[encodeBase32Map[i]] = byte(i)
}
}
// A Node struct holds the basic information needed for a snowflake generator
// node
type Node struct {
mu sync.Mutex
epoch time.Time
time int64
node int64
step int64
nodeMax int64
nodeMask int64
stepMask int64
timeShift uint8
nodeShift uint8
}
// An ID is a custom type used for a snowflake ID. This is used so we can
// attach methods onto the ID.
type ID int64
// NewNode returns a new snowflake node that can be used to generate snowflake
// IDs
func NewNode(node int64) (*Node, error) {
// re-calc in case custom NodeBits or StepBits were set
// DEPRECATED: the below block will be removed in a future release.
mu.Lock()
nodeMax = -1 ^ (-1 << NodeBits)
nodeMask = nodeMax << StepBits
stepMask = -1 ^ (-1 << StepBits)
timeShift = NodeBits + StepBits
nodeShift = StepBits
mu.Unlock()
n := Node{}
n.node = node
n.nodeMax = -1 ^ (-1 << NodeBits)
n.nodeMask = n.nodeMax << StepBits
n.stepMask = -1 ^ (-1 << StepBits)
n.timeShift = NodeBits + StepBits
n.nodeShift = StepBits
if n.node < 0 || n.node > n.nodeMax {
return nil, errors.New("Node number must be between 0 and " + strconv.FormatInt(n.nodeMax, 10))
}
var curTime = time.Now()
// add time.Duration to curTime to make sure we use the monotonic clock if available
n.epoch = curTime.Add(time.Unix(Epoch/1000, (Epoch%1000)*1000000).Sub(curTime))
return &n, nil
}
// Generate creates and returns a unique snowflake ID
// To help guarantee uniqueness
// - Make sure your system is keeping accurate system time
// - Make sure you never have multiple nodes running with the same node ID
func (n *Node) Generate() ID {
n.mu.Lock()
now := time.Since(n.epoch).Nanoseconds() / 1000000
if now == n.time {
n.step = (n.step + 1) & n.stepMask
if n.step == 0 {
for now <= n.time {
now = time.Since(n.epoch).Nanoseconds() / 1000000
}
}
} else {
n.step = 0
}
n.time = now
r := ID((now)<<n.timeShift |
(n.node << n.nodeShift) |
(n.step),
)
n.mu.Unlock()
return r
}
// Int64 returns an int64 of the snowflake ID
func (f ID) Int64() int64 {
return int64(f)
}
// ParseInt64 converts an int64 into a snowflake ID
func ParseInt64(id int64) ID {
return ID(id)
}
// String returns a string of the snowflake ID
func (f ID) String() string {
return strconv.FormatInt(int64(f), 10)
}
// ParseString converts a string into a snowflake ID
func ParseString(id string) (ID, error) {
i, err := strconv.ParseInt(id, 10, 64)
return ID(i), err
}
// Base2 returns a string base2 of the snowflake ID
func (f ID) Base2() string {
return strconv.FormatInt(int64(f), 2)
}
// ParseBase2 converts a Base2 string into a snowflake ID
func ParseBase2(id string) (ID, error) {
i, err := strconv.ParseInt(id, 2, 64)
return ID(i), err
}
// Base32 uses the z-base-32 character set but encodes and decodes similar
// to base58, allowing it to create an even smaller result string.
// NOTE: There are many different base32 implementations so becareful when
// doing any interoperation.
func (f ID) Base32() string {
if f < 32 {
return string(encodeBase32Map[f])
}
b := make([]byte, 0, 12)
for f >= 32 {
b = append(b, encodeBase32Map[f%32])
f /= 32
}
b = append(b, encodeBase32Map[f])
for x, y := 0, len(b)-1; x < y; x, y = x+1, y-1 {
b[x], b[y] = b[y], b[x]
}
return string(b)
}
// ParseBase32 parses a base32 []byte into a snowflake ID
// NOTE: There are many different base32 implementations so becareful when
// doing any interoperation.
func ParseBase32(b []byte) (ID, error) {
var id int64
for i := range b {
if decodeBase32Map[b[i]] == 0xFF {
return -1, ErrInvalidBase32
}
id = id*32 + int64(decodeBase32Map[b[i]])
}
return ID(id), nil
}
// Base36 returns a base36 string of the snowflake ID
func (f ID) Base36() string {
return strconv.FormatInt(int64(f), 36)
}
// ParseBase36 converts a Base36 string into a snowflake ID
func ParseBase36(id string) (ID, error) {
i, err := strconv.ParseInt(id, 36, 64)
return ID(i), err
}
// Base58 returns a base58 string of the snowflake ID
func (f ID) Base58() string {
if f < 58 {
return string(encodeBase58Map[f])
}
b := make([]byte, 0, 11)
for f >= 58 {
b = append(b, encodeBase58Map[f%58])
f /= 58
}
b = append(b, encodeBase58Map[f])
for x, y := 0, len(b)-1; x < y; x, y = x+1, y-1 {
b[x], b[y] = b[y], b[x]
}
return string(b)
}
// ParseBase58 parses a base58 []byte into a snowflake ID
func ParseBase58(b []byte) (ID, error) {
var id int64
for i := range b {
if decodeBase58Map[b[i]] == 0xFF {
return -1, ErrInvalidBase58
}
id = id*58 + int64(decodeBase58Map[b[i]])
}
return ID(id), nil
}
// Base64 returns a base64 string of the snowflake ID
func (f ID) Base64() string {
return base64.StdEncoding.EncodeToString(f.Bytes())
}
// ParseBase64 converts a base64 string into a snowflake ID
func ParseBase64(id string) (ID, error) {
b, err := base64.StdEncoding.DecodeString(id)
if err != nil {
return -1, err
}
return ParseBytes(b)
}
// Bytes returns a byte slice of the snowflake ID
func (f ID) Bytes() []byte {
return []byte(f.String())
}
// ParseBytes converts a byte slice into a snowflake ID
func ParseBytes(id []byte) (ID, error) {
i, err := strconv.ParseInt(string(id), 10, 64)
return ID(i), err
}
// IntBytes returns an array of bytes of the snowflake ID, encoded as a
// big endian integer.
func (f ID) IntBytes() [8]byte {
var b [8]byte
binary.BigEndian.PutUint64(b[:], uint64(f))
return b
}
// ParseIntBytes converts an array of bytes encoded as big endian integer as
// a snowflake ID
func ParseIntBytes(id [8]byte) ID {
return ID(int64(binary.BigEndian.Uint64(id[:])))
}
// Time returns an int64 unix timestamp in milliseconds of the snowflake ID time
// DEPRECATED: the below function will be removed in a future release.
func (f ID) Time() int64 {
return (int64(f) >> timeShift) + Epoch
}
// Node returns an int64 of the snowflake ID node number
// DEPRECATED: the below function will be removed in a future release.
func (f ID) Node() int64 {
return int64(f) & nodeMask >> nodeShift
}
// Step returns an int64 of the snowflake step (or sequence) number
// DEPRECATED: the below function will be removed in a future release.
func (f ID) Step() int64 {
return int64(f) & stepMask
}
// MarshalJSON returns a json byte array string of the snowflake ID.
func (f ID) MarshalJSON() ([]byte, error) {
buff := make([]byte, 0, 22)
buff = append(buff, '"')
buff = strconv.AppendInt(buff, int64(f), 10)
buff = append(buff, '"')
return buff, nil
}
// UnmarshalJSON converts a json byte array of a snowflake ID into an ID type.
func (f *ID) UnmarshalJSON(b []byte) error {
if len(b) < 3 || b[0] != '"' || b[len(b)-1] != '"' {
return JSONSyntaxError{b}
}
i, err := strconv.ParseInt(string(b[1:len(b)-1]), 10, 64)
if err != nil {
return err
}
*f = ID(i)
return nil
}