82 lines
1.7 KiB
Go
82 lines
1.7 KiB
Go
|
package snowflake
|
||
|
|
|
||
|
|
import (
|
||
|
|
"fmt"
|
||
|
|
"sync"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"github.com/bwmarrin/snowflake"
|
||
|
|
)
|
||
|
|
|
||
|
|
var mu sync.Mutex
|
||
|
|
|
||
|
|
// Snowflake provides a way to NewNode for Generate UID.
|
||
|
|
type Snowflake struct {
|
||
|
|
machineNodeID int64
|
||
|
|
|
||
|
|
startTime time.Time
|
||
|
|
}
|
||
|
|
|
||
|
|
// Option is the options type to configure Snowflake.
|
||
|
|
type Option func(*Snowflake)
|
||
|
|
|
||
|
|
// New returns a new Snowflake instance with the provided options.
|
||
|
|
func New(opts ...Option) *Snowflake {
|
||
|
|
s := &Snowflake{
|
||
|
|
// default machine 1
|
||
|
|
machineNodeID: 1,
|
||
|
|
}
|
||
|
|
|
||
|
|
for _, opt := range opts {
|
||
|
|
opt(s)
|
||
|
|
}
|
||
|
|
|
||
|
|
return s
|
||
|
|
}
|
||
|
|
|
||
|
|
// WithMachineNodeID adds machineID total 10bit = 1024 machine number.
|
||
|
|
func WithMachineNodeID(machineNodeID int64) Option {
|
||
|
|
return func(snowflake *Snowflake) {
|
||
|
|
snowflake.machineNodeID = machineNodeID
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// WithStartTime adds snowflake start timestamp in milliseconds.
|
||
|
|
func WithStartTime(startTime time.Time) Option {
|
||
|
|
return func(snowflake *Snowflake) {
|
||
|
|
snowflake.startTime = startTime
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// GetNowDate return nowTodayDate e.g. 2023-07-20 00:00:00 +0000 UTC.
|
||
|
|
func GetNowDate() (time.Time, error) {
|
||
|
|
startTime := time.Now().UTC().Format(time.DateOnly)
|
||
|
|
|
||
|
|
st, err := time.Parse(time.DateOnly, startTime)
|
||
|
|
if err != nil {
|
||
|
|
return time.Time{}, fmt.Errorf("time.Parse failed :%w", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
return st, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// NewNode return snowflake node use Generate UID.
|
||
|
|
func (s *Snowflake) NewNode() (*snowflake.Node, error) {
|
||
|
|
mu.Lock()
|
||
|
|
defer mu.Unlock()
|
||
|
|
|
||
|
|
snowflake.Epoch = s.startTime.UnixMilli()
|
||
|
|
|
||
|
|
node, err := snowflake.NewNode(s.machineNodeID)
|
||
|
|
if err != nil {
|
||
|
|
return nil, fmt.Errorf("snowflake.NewNode, failed :%w",
|
||
|
|
&NewNodeError{
|
||
|
|
machineNodeID: s.machineNodeID,
|
||
|
|
startTime: s.startTime,
|
||
|
|
Err: err,
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
|
|
return node, nil
|
||
|
|
}
|