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

     1  package snow
     2  
     3  import (
     4  	"time"
     5  
     6  	"github.com/bingoohuang/gg/pkg/goip"
     7  )
     8  
     9  // Option for the snowflake
    10  type Option struct {
    11  	// Epoch is set to the twitter snowflake epoch of Nov 04 2010 01:42:54 UTC in milliseconds
    12  	// You may customize this to set a different epoch for your application.
    13  	Epoch int64 // 1288834974657
    14  
    15  	// NodeBits holds the number of bits to use for Node
    16  	// Remember, you have a total 22 bits to share between Node/Step
    17  	NodeBits int8 // 10
    18  
    19  	// StepBits holds the number of bits to use for Step
    20  	// Remember, you have a total 22 bits to share between Node/Step
    21  	StepBits int8 // 12
    22  
    23  	// NodeID for the snowflake.
    24  	NodeID int64
    25  
    26  	// TimestampUnit for the time goes unit, default is 1ms.
    27  	TimestampUnit time.Duration
    28  }
    29  
    30  // Apply applies the option functions to the option.
    31  // nolint gomnd
    32  func (o *Option) Apply(fns ...OptionFn) {
    33  	for _, fn := range fns {
    34  		fn(o)
    35  	}
    36  
    37  	if o.Epoch == 0 {
    38  		o.Epoch = 1288834974657 // 2010-11-04T09:42:54+08:00
    39  	}
    40  
    41  	if o.NodeBits < 0 {
    42  		o.NodeBits = 10
    43  	}
    44  
    45  	if o.StepBits < 0 {
    46  		o.StepBits = 12
    47  	}
    48  
    49  	if o.NodeID < 0 {
    50  		o.NodeID = defaultIPNodeID()
    51  		var nodeMax int64 = -1 ^ (-1 << o.NodeBits)
    52  		o.NodeID &= nodeMax
    53  	}
    54  }
    55  
    56  // OptionFn defines the function prototype to apply options.
    57  type OptionFn func(*Option)
    58  
    59  // WithNodeID set the customized nodeID.
    60  func WithNodeID(nodeID int64) OptionFn { return func(o *Option) { o.NodeID = nodeID } }
    61  
    62  // WithNodeIDLocalIP set the customized nodeID  with the last 8 bits of local IP v4 and first 2 bits of p.
    63  func WithNodeIDLocalIP(p int64, ip string) OptionFn {
    64  	if ip == "" {
    65  		ip, _ = goip.MainIP()
    66  	}
    67  
    68  	return func(o *Option) { o.NodeID = (p << 8) | ipNodeID(ip) }
    69  }
    70  
    71  // WithEpoch set the customized epoch.
    72  func WithEpoch(epoch int64) OptionFn { return func(o *Option) { o.Epoch = epoch } }
    73  
    74  // WithNodeBits set the customized NodeBits n.
    75  func WithNodeBits(n int8) OptionFn { return func(o *Option) { o.NodeBits = n } }
    76  
    77  // WithStepBits set the customized StepBits n.
    78  func WithStepBits(n int8) OptionFn { return func(o *Option) { o.StepBits = n } }
    79  
    80  // WithTimestampUnit set the customized TimestampUnit n.
    81  func WithTimestampUnit(n time.Duration) OptionFn { return func(o *Option) { o.TimestampUnit = n } }