snowflake.go 7.7 KB

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