snowflake.go 8.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365
  1. // Package snowflake provides a very simple Twitter snowflake generator and parser.
  2. package snowflake
  3. import (
  4. "encoding/base64"
  5. "encoding/binary"
  6. "errors"
  7. "fmt"
  8. "strconv"
  9. "sync"
  10. "time"
  11. )
  12. var (
  13. // Epoch is set to the twitter snowflake epoch of Nov 04 2010 01:42:54 UTC in milliseconds
  14. // You may customize this to set a different epoch for your application.
  15. Epoch int64 = 1288834974657
  16. // NodeBits holds the number of bits to use for Node
  17. // Remember, you have a total 22 bits to share between Node/Step
  18. NodeBits uint8 = 10
  19. // StepBits holds the number of bits to use for Step
  20. // Remember, you have a total 22 bits to share between Node/Step
  21. StepBits uint8 = 12
  22. // DEPRECATED: the below four variables will be removed in a future release.
  23. mu sync.Mutex
  24. nodeMax int64 = -1 ^ (-1 << NodeBits)
  25. nodeMask = nodeMax << StepBits
  26. stepMask int64 = -1 ^ (-1 << StepBits)
  27. timeShift = NodeBits + StepBits
  28. nodeShift = StepBits
  29. )
  30. const encodeBase32Map = "ybndrfg8ejkmcpqxot1uwisza345h769"
  31. var decodeBase32Map [256]byte
  32. const encodeBase58Map = "123456789abcdefghijkmnopqrstuvwxyzABCDEFGHJKLMNPQRSTUVWXYZ"
  33. var decodeBase58Map [256]byte
  34. // A JSONSyntaxError is returned from UnmarshalJSON if an invalid ID is provided.
  35. type JSONSyntaxError struct{ original []byte }
  36. func (j JSONSyntaxError) Error() string {
  37. return fmt.Sprintf("invalid snowflake ID %q", string(j.original))
  38. }
  39. // ErrInvalidBase58 is returned by ParseBase58 when given an invalid []byte
  40. var ErrInvalidBase58 = errors.New("invalid base58")
  41. // ErrInvalidBase32 is returned by ParseBase32 when given an invalid []byte
  42. var ErrInvalidBase32 = errors.New("invalid base32")
  43. // Create maps for decoding Base58/Base32.
  44. // This speeds up the process tremendously.
  45. func init() {
  46. for i := 0; i < len(encodeBase58Map); i++ {
  47. decodeBase58Map[i] = 0xFF
  48. }
  49. for i := 0; i < len(encodeBase58Map); i++ {
  50. decodeBase58Map[encodeBase58Map[i]] = byte(i)
  51. }
  52. for i := 0; i < len(encodeBase32Map); i++ {
  53. decodeBase32Map[i] = 0xFF
  54. }
  55. for i := 0; i < len(encodeBase32Map); i++ {
  56. decodeBase32Map[encodeBase32Map[i]] = byte(i)
  57. }
  58. }
  59. // A Node struct holds the basic information needed for a snowflake generator
  60. // node
  61. type Node struct {
  62. mu sync.Mutex
  63. epoch time.Time
  64. time int64
  65. node int64
  66. step int64
  67. nodeMax int64
  68. nodeMask int64
  69. stepMask int64
  70. timeShift uint8
  71. nodeShift uint8
  72. }
  73. // An ID is a custom type used for a snowflake ID. This is used so we can
  74. // attach methods onto the ID.
  75. type ID int64
  76. // NewNode returns a new snowflake node that can be used to generate snowflake
  77. // IDs
  78. func NewNode(node int64) (*Node, error) {
  79. // re-calc in case custom NodeBits or StepBits were set
  80. // DEPRECATED: the below block will be removed in a future release.
  81. mu.Lock()
  82. nodeMax = -1 ^ (-1 << NodeBits)
  83. nodeMask = nodeMax << StepBits
  84. stepMask = -1 ^ (-1 << StepBits)
  85. timeShift = NodeBits + StepBits
  86. nodeShift = StepBits
  87. mu.Unlock()
  88. n := Node{}
  89. n.node = node
  90. n.nodeMax = -1 ^ (-1 << NodeBits)
  91. n.nodeMask = n.nodeMax << StepBits
  92. n.stepMask = -1 ^ (-1 << StepBits)
  93. n.timeShift = NodeBits + StepBits
  94. n.nodeShift = StepBits
  95. if n.node < 0 || n.node > n.nodeMax {
  96. return nil, errors.New("Node number must be between 0 and " + strconv.FormatInt(n.nodeMax, 10))
  97. }
  98. var curTime = time.Now()
  99. // add time.Duration to curTime to make sure we use the monotonic clock if available
  100. n.epoch = curTime.Add(time.Unix(Epoch/1000, (Epoch%1000)*1000000).Sub(curTime))
  101. return &n, nil
  102. }
  103. // Generate creates and returns a unique snowflake ID
  104. // To help guarantee uniqueness
  105. // - Make sure your system is keeping accurate system time
  106. // - Make sure you never have multiple nodes running with the same node ID
  107. func (n *Node) Generate() ID {
  108. n.mu.Lock()
  109. now := time.Since(n.epoch).Nanoseconds() / 1000000
  110. if now == n.time {
  111. n.step = (n.step + 1) & n.stepMask
  112. if n.step == 0 {
  113. for now <= n.time {
  114. now = time.Since(n.epoch).Nanoseconds() / 1000000
  115. }
  116. }
  117. } else {
  118. n.step = 0
  119. }
  120. n.time = now
  121. r := ID((now)<<n.timeShift |
  122. (n.node << n.nodeShift) |
  123. (n.step),
  124. )
  125. n.mu.Unlock()
  126. return r
  127. }
  128. // Int64 returns an int64 of the snowflake ID
  129. func (f ID) Int64() int64 {
  130. return int64(f)
  131. }
  132. // ParseInt64 converts an int64 into a snowflake ID
  133. func ParseInt64(id int64) ID {
  134. return ID(id)
  135. }
  136. // String returns a string of the snowflake ID
  137. func (f ID) String() string {
  138. return strconv.FormatInt(int64(f), 10)
  139. }
  140. // ParseString converts a string into a snowflake ID
  141. func ParseString(id string) (ID, error) {
  142. i, err := strconv.ParseInt(id, 10, 64)
  143. return ID(i), err
  144. }
  145. // Base2 returns a string base2 of the snowflake ID
  146. func (f ID) Base2() string {
  147. return strconv.FormatInt(int64(f), 2)
  148. }
  149. // ParseBase2 converts a Base2 string into a snowflake ID
  150. func ParseBase2(id string) (ID, error) {
  151. i, err := strconv.ParseInt(id, 2, 64)
  152. return ID(i), err
  153. }
  154. // Base32 uses the z-base-32 character set but encodes and decodes similar
  155. // to base58, allowing it to create an even smaller result string.
  156. // NOTE: There are many different base32 implementations so becareful when
  157. // doing any interoperation.
  158. func (f ID) Base32() string {
  159. if f < 32 {
  160. return string(encodeBase32Map[f])
  161. }
  162. b := make([]byte, 0, 12)
  163. for f >= 32 {
  164. b = append(b, encodeBase32Map[f%32])
  165. f /= 32
  166. }
  167. b = append(b, encodeBase32Map[f])
  168. for x, y := 0, len(b)-1; x < y; x, y = x+1, y-1 {
  169. b[x], b[y] = b[y], b[x]
  170. }
  171. return string(b)
  172. }
  173. // ParseBase32 parses a base32 []byte into a snowflake ID
  174. // NOTE: There are many different base32 implementations so becareful when
  175. // doing any interoperation.
  176. func ParseBase32(b []byte) (ID, error) {
  177. var id int64
  178. for i := range b {
  179. if decodeBase32Map[b[i]] == 0xFF {
  180. return -1, ErrInvalidBase32
  181. }
  182. id = id*32 + int64(decodeBase32Map[b[i]])
  183. }
  184. return ID(id), nil
  185. }
  186. // Base36 returns a base36 string of the snowflake ID
  187. func (f ID) Base36() string {
  188. return strconv.FormatInt(int64(f), 36)
  189. }
  190. // ParseBase36 converts a Base36 string into a snowflake ID
  191. func ParseBase36(id string) (ID, error) {
  192. i, err := strconv.ParseInt(id, 36, 64)
  193. return ID(i), err
  194. }
  195. // Base58 returns a base58 string of the snowflake ID
  196. func (f ID) Base58() string {
  197. if f < 58 {
  198. return string(encodeBase58Map[f])
  199. }
  200. b := make([]byte, 0, 11)
  201. for f >= 58 {
  202. b = append(b, encodeBase58Map[f%58])
  203. f /= 58
  204. }
  205. b = append(b, encodeBase58Map[f])
  206. for x, y := 0, len(b)-1; x < y; x, y = x+1, y-1 {
  207. b[x], b[y] = b[y], b[x]
  208. }
  209. return string(b)
  210. }
  211. // ParseBase58 parses a base58 []byte into a snowflake ID
  212. func ParseBase58(b []byte) (ID, error) {
  213. var id int64
  214. for i := range b {
  215. if decodeBase58Map[b[i]] == 0xFF {
  216. return -1, ErrInvalidBase58
  217. }
  218. id = id*58 + int64(decodeBase58Map[b[i]])
  219. }
  220. return ID(id), nil
  221. }
  222. // Base64 returns a base64 string of the snowflake ID
  223. func (f ID) Base64() string {
  224. return base64.StdEncoding.EncodeToString(f.Bytes())
  225. }
  226. // ParseBase64 converts a base64 string into a snowflake ID
  227. func ParseBase64(id string) (ID, error) {
  228. b, err := base64.StdEncoding.DecodeString(id)
  229. if err != nil {
  230. return -1, err
  231. }
  232. return ParseBytes(b)
  233. }
  234. // Bytes returns a byte slice of the snowflake ID
  235. func (f ID) Bytes() []byte {
  236. return []byte(f.String())
  237. }
  238. // ParseBytes converts a byte slice into a snowflake ID
  239. func ParseBytes(id []byte) (ID, error) {
  240. i, err := strconv.ParseInt(string(id), 10, 64)
  241. return ID(i), err
  242. }
  243. // IntBytes returns an array of bytes of the snowflake ID, encoded as a
  244. // big endian integer.
  245. func (f ID) IntBytes() [8]byte {
  246. var b [8]byte
  247. binary.BigEndian.PutUint64(b[:], uint64(f))
  248. return b
  249. }
  250. // ParseIntBytes converts an array of bytes encoded as big endian integer as
  251. // a snowflake ID
  252. func ParseIntBytes(id [8]byte) ID {
  253. return ID(int64(binary.BigEndian.Uint64(id[:])))
  254. }
  255. // Time returns an int64 unix timestamp in milliseconds of the snowflake ID time
  256. // DEPRECATED: the below function will be removed in a future release.
  257. func (f ID) Time() int64 {
  258. return (int64(f) >> timeShift) + Epoch
  259. }
  260. // Node returns an int64 of the snowflake ID node number
  261. // DEPRECATED: the below function will be removed in a future release.
  262. func (f ID) Node() int64 {
  263. return int64(f) & nodeMask >> nodeShift
  264. }
  265. // Step returns an int64 of the snowflake step (or sequence) number
  266. // DEPRECATED: the below function will be removed in a future release.
  267. func (f ID) Step() int64 {
  268. return int64(f) & stepMask
  269. }
  270. // MarshalJSON returns a json byte array string of the snowflake ID.
  271. func (f ID) MarshalJSON() ([]byte, error) {
  272. buff := make([]byte, 0, 22)
  273. buff = append(buff, '"')
  274. buff = strconv.AppendInt(buff, int64(f), 10)
  275. buff = append(buff, '"')
  276. return buff, nil
  277. }
  278. // UnmarshalJSON converts a json byte array of a snowflake ID into an ID type.
  279. func (f *ID) UnmarshalJSON(b []byte) error {
  280. if len(b) < 3 || b[0] != '"' || b[len(b)-1] != '"' {
  281. return JSONSyntaxError{b}
  282. }
  283. i, err := strconv.ParseInt(string(b[1:len(b)-1]), 10, 64)
  284. if err != nil {
  285. return err
  286. }
  287. *f = ID(i)
  288. return nil
  289. }