github.com/bingoohuang/gg@v0.0.0-20240325092523-45da7dee9335/pkg/snow/snow.go (about)

     1  // Package snow provides a very simple Twitter snowflake generator and parser.
     2  package snow
     3  
     4  import (
     5  	"fmt"
     6  	"sync"
     7  	"time"
     8  )
     9  
    10  // A Node struct holds the basic information needed for a snowflake generator node
    11  type Node struct {
    12  	mu sync.Mutex
    13  
    14  	option    Option
    15  	nodeID    int64
    16  	stepMask  int64
    17  	timeShift uint8
    18  	nodeShift uint8
    19  
    20  	time int64
    21  	step int64
    22  
    23  	// epoch is snowflake epoch in milliseconds.
    24  	epoch    int64
    25  	nodeMask int64
    26  
    27  	epochTime time.Time
    28  	unit      time.Duration
    29  }
    30  
    31  // GetOption return the option.
    32  func (n *Node) GetOption() Option { return n.option }
    33  
    34  // GetEpoch returns an int64 epoch is snowflake epoch in milliseconds.
    35  func (n *Node) GetEpoch() int64 { return n.epoch }
    36  
    37  // GetTime returns an int64 unix timestamp in milliseconds of the snowflake ID time.
    38  func (n *Node) GetTime() int64 { return n.time }
    39  
    40  // GetNodeID returns an int64 of the snowflake ID node number
    41  func (n *Node) GetNodeID() int64 { return n.nodeID }
    42  
    43  // GetStep returns an int64 of the snowflake step (or sequence) number
    44  func (n *Node) GetStep() int64 { return n.step }
    45  
    46  func (n *Node) applyOption(o Option) error {
    47  	n.option = o
    48  	n.nodeID = o.NodeID
    49  
    50  	var nodeMax int64 = -1 ^ (-1 << o.NodeBits)
    51  
    52  	n.nodeMask = nodeMax << o.StepBits
    53  	n.stepMask = -1 ^ (-1 << o.StepBits)
    54  	n.timeShift = uint8(o.NodeBits + o.StepBits)
    55  	n.nodeShift = uint8(o.StepBits)
    56  
    57  	curTime := time.Now()
    58  	// add time.Duration to curTime to make sure we use the monotonic clock if available
    59  	n.epochTime = curTime.Add(time.Unix(o.Epoch/1e3, (o.Epoch%1000)*1e6).Sub(curTime))
    60  
    61  	n.unit = o.TimestampUnit / time.Millisecond
    62  	if n.unit == 0 {
    63  		n.unit = 1
    64  	}
    65  	n.epoch = int64(n.epochTime.Nanosecond() / 1e6)
    66  
    67  	if o.NodeID >= 0 && o.NodeID <= nodeMax {
    68  		return nil
    69  	}
    70  
    71  	return fmt.Errorf("NodeID %d must be between 0 and %d", o.NodeID, nodeMax)
    72  }
    73  
    74  // NewNode returns a new snowflake node that can be used to generate snowflake IDs.
    75  func NewNode(optionFns ...OptionFn) (*Node, error) {
    76  	option := Option{NodeBits: -1, StepBits: -1, NodeID: -1}
    77  	option.Apply(optionFns...)
    78  
    79  	n := &Node{}
    80  	if err := n.applyOption(option); err != nil {
    81  		return nil, err
    82  	}
    83  
    84  	return n, nil
    85  }
    86  
    87  // Next creates and returns a unique snowflake ID
    88  // To help guarantee uniqueness
    89  // - Make sure your system is keeping accurate system time
    90  // - Make sure you never have multiple nodes running with the same node ID
    91  func (n *Node) Next() ID {
    92  	n.mu.Lock()
    93  	defer n.mu.Unlock()
    94  
    95  	return n.next()
    96  }
    97  
    98  func (n *Node) next() ID {
    99  	now := n.now()
   100  	if now == n.time {
   101  		if n.step = (n.step + 1) & n.stepMask; n.step == 0 {
   102  			for now <= n.time {
   103  				n.sleep()
   104  				now = n.now()
   105  			}
   106  		}
   107  	} else {
   108  		n.step = 0
   109  	}
   110  
   111  	n.time = now
   112  
   113  	return ID(now<<n.timeShift | n.nodeID<<n.nodeShift | n.step)
   114  }
   115  
   116  func (n *Node) sleep() {
   117  	if n.unit > 1 {
   118  		time.Sleep(n.unit * time.Millisecond)
   119  	}
   120  }
   121  
   122  func (n *Node) now() int64 {
   123  	return time.Since(n.epochTime).Nanoseconds() / 1e6 / int64(n.unit)
   124  }