78 lines
1.6 KiB
Go
78 lines
1.6 KiB
Go
package idutil
|
||
|
||
import (
|
||
"fmt"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/sony/sonyflake"
|
||
)
|
||
|
||
// SnowflakeGenerator 雪花算法 ID 生成器（单例）
|
||
type SnowflakeGenerator struct {
|
||
flake *sonyflake.Sonyflake
|
||
mu sync.Mutex
|
||
}
|
||
|
||
// globalGen 全局生成器实例
|
||
var (
|
||
globalGen *SnowflakeGenerator
|
||
once sync.Once
|
||
)
|
||
|
||
// GetGenerator 获取全局雪花算法生成器（线程安全）
|
||
func GetGenerator() (*SnowflakeGenerator, error) {
|
||
var err error
|
||
once.Do(func() {
|
||
globalGen, err = NewSnowflakeGenerator()
|
||
})
|
||
return globalGen, err
|
||
}
|
||
|
||
// NewSnowflakeGenerator 创建新的雪花算法生成器
|
||
func NewSnowflakeGenerator() (*SnowflakeGenerator, error) {
|
||
// Sonyflake 配置
|
||
st := sonyflake.Settings{
|
||
StartTime: time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC), // 自定义起始时间
|
||
}
|
||
|
||
flake := sonyflake.NewSonyflake(st)
|
||
if flake == nil {
|
||
return nil, fmt.Errorf("snowflake generator creation failed")
|
||
}
|
||
|
||
return &SnowflakeGenerator{
|
||
flake: flake,
|
||
}, nil
|
||
}
|
||
|
||
// NextID 生成下一个 ID（uint64）
|
||
func (g *SnowflakeGenerator) NextID() (uint64, error) {
|
||
g.mu.Lock()
|
||
defer g.mu.Unlock()
|
||
|
||
id, err := g.flake.NextID()
|
||
if err != nil {
|
||
return 0, fmt.Errorf("failed to generate snowflake id: %w", err)
|
||
}
|
||
|
||
return id, nil
|
||
}
|
||
|
||
// NextString 生成字符串格式 ID（便于 JSON 传输）
|
||
func (g *SnowflakeGenerator) NextString() (string, error) {
|
||
id, err := g.NextID()
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
|
||
return fmt.Sprintf("%d", id), nil
|
||
}
|
||
|
||
// ParseID 从字符串解析 ID
|
||
func ParseID(s string) (uint64, error) {
|
||
var id uint64
|
||
_, err := fmt.Sscanf(s, "%d", &id)
|
||
return id, err
|
||
}
|