github.com/godaddy-x/freego@v1.0.156/utils/snowflake/id_worker.go (about)

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