snow.go 1.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485
  1. package utils
  2. import (
  3. "fmt"
  4. "github.com/google/uuid"
  5. "strconv"
  6. "strings"
  7. "sync"
  8. "time"
  9. )
  10. type Snowflake struct {
  11. epoch int64
  12. timestamp int64
  13. dataCenter int64
  14. workerId int64
  15. sequence int64
  16. mutex sync.Mutex
  17. maxSequence int64
  18. }
  19. const (
  20. epochShift = 22
  21. dataCenterShift = 12
  22. workerIdShift = 10
  23. sequenceShift = 0
  24. sequenceMask = -1 ^ (-1 << sequenceShift)
  25. maxWorkerId = -1 ^ (-1 << workerIdShift)
  26. maxDataCenter = -1 ^ (-1 << dataCenterShift)
  27. )
  28. func NewSnowflake(epoch int64, dataCenter, workerId int64) (*Snowflake, error) {
  29. if dataCenter > maxDataCenter || workerId > maxWorkerId {
  30. return nil, fmt.Errorf("dataCenter or workerId exceeds limit")
  31. }
  32. return &Snowflake{
  33. epoch: epoch,
  34. dataCenter: dataCenter,
  35. workerId: workerId,
  36. maxSequence: -1 ^ (-1 << sequenceShift),
  37. }, nil
  38. }
  39. func (s *Snowflake) NextID() int64 {
  40. s.mutex.Lock()
  41. defer s.mutex.Unlock()
  42. currentTimestamp := time.Now().UnixNano() / 1e6
  43. if s.timestamp == currentTimestamp {
  44. s.sequence = (s.sequence + 1) & sequenceMask
  45. if s.sequence == 0 {
  46. currentTimestamp = s.waitNextTime(currentTimestamp)
  47. }
  48. } else {
  49. s.sequence = 0
  50. }
  51. s.timestamp = currentTimestamp
  52. id := ((currentTimestamp - s.epoch) << epochShift) |
  53. (s.dataCenter << dataCenterShift) |
  54. (s.workerId << workerIdShift) |
  55. s.sequence
  56. return id
  57. }
  58. func (s *Snowflake) waitNextTime(currentTimestamp int64) int64 {
  59. for currentTimestamp <= s.timestamp {
  60. currentTimestamp = time.Now().UnixNano() / 1e6
  61. }
  62. return currentTimestamp
  63. }
  64. var snowflake, _ = NewSnowflake(time.Now().UnixNano()/1e6, 1, 1)
  65. func NextId() int64 {
  66. return snowflake.NextID()
  67. }
  68. func NextIdStr() string {
  69. return strconv.FormatInt(NextId(), 10)
  70. }
  71. func GenerateUUID() string {
  72. return strings.ReplaceAll(uuid.NewString(), "-", "")
  73. }