69 lines
1.5 KiB
Go
69 lines
1.5 KiB
Go
package id
|
|
|
|
import (
|
|
"fmt"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/bwmarrin/snowflake"
|
|
|
|
"spark-server/internal/pkg/errs"
|
|
)
|
|
|
|
// Generator issues Snowflake MsgIds. Worker ID is fixed per machine. Clock rewind is rejected.
|
|
type Generator struct {
|
|
mu sync.Mutex
|
|
node *snowflake.Node
|
|
lastMs int64
|
|
workerID int64
|
|
}
|
|
|
|
func New(workerID int64) (*Generator, error) {
|
|
node, err := snowflake.NewNode(workerID)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("snowflake node: %w", err)
|
|
}
|
|
return &Generator{node: node, workerID: workerID}, nil
|
|
}
|
|
|
|
func (g *Generator) NextMsgID() (int64, error) {
|
|
return g.NextID()
|
|
}
|
|
|
|
// NextID returns a time-ordered Snowflake ID (user id, conversation id, msg id, mailbox event id).
|
|
func (g *Generator) NextID() (int64, error) {
|
|
g.mu.Lock()
|
|
defer g.mu.Unlock()
|
|
now := time.Now().UnixMilli()
|
|
if now < g.lastMs {
|
|
return 0, errs.ErrClockRewind
|
|
}
|
|
id := g.node.Generate().Int64()
|
|
g.lastMs = now
|
|
return id, nil
|
|
}
|
|
|
|
func (g *Generator) WorkerID() int64 { return g.workerID }
|
|
|
|
const DefaultMsgIDSkew = 5 * time.Minute
|
|
|
|
// IssuedAtMS is the Snowflake unix-ms timestamp encoded in id.
|
|
func IssuedAtMS(id int64) int64 {
|
|
if id <= 0 {
|
|
return 0
|
|
}
|
|
return snowflake.ID(id).Time()
|
|
}
|
|
|
|
// WithinSkew reports whether id's timestamp is within skew of now (SEND client MsgId ±5m).
|
|
func WithinSkew(id int64, now time.Time, skew time.Duration) bool {
|
|
if id <= 0 || skew <= 0 {
|
|
return false
|
|
}
|
|
delta := now.UnixMilli() - IssuedAtMS(id)
|
|
if delta < 0 {
|
|
delta = -delta
|
|
}
|
|
return delta <= skew.Milliseconds()
|
|
}
|