ark-member/internal/lib/snackflow/snowflake.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
}