snowflake.go 8.9 KB

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