go-雪花算法(snow_flake)

import (
	"errors"
	"sync"
	"time"
)

// Snowflake 雪花算法生成器
type Snowflake struct {
	mu        sync.Mutex
	timestamp int64 // 毫秒时间戳
	workerID  int64 // 机器ID（0-1023）
	sequence  int64 // 序列号（0-4095）
}

const (
	workerIDBits   = 10                          // 机器ID位数
	sequenceBits   = 12                          // 序列号位数
	maxWorkerID    = -1 ^ (-1 << workerIDBits)   // 最大机器ID（1023）
	sequenceMask   = -1 ^ (-1 << sequenceBits)   // 序列号掩码（4095）
	workerIDShift  = sequenceBits                // 机器ID左移位数
	timestampShift = sequenceBits + workerIDBits // 时间戳左移位数
	epoch          = 1740000000000               // 起始时间戳（2026-02-19，可自定义）
)

// NewSnowflake 创建雪花算法实例（传入机器ID，确保分布式环境唯一）
func NewSnowflake(workerID int64) (*Snowflake, error) {
	if workerID < 0 || workerID > maxWorkerID {
		return nil, errors.New("workerID超出范围（0-1023）")
	}
	return &Snowflake{
		workerID: workerID,
	}, nil
}

// GenUID 生成唯一UID
func (s *Snowflake) GenUID() (uint64, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	// 获取当前毫秒时间戳
	now := time.Now().UnixMilli()

	// 如果时间回拨，报错（避免重复）
	if now < s.timestamp {
		return 0, errors.New("时钟回拨，无法生成UID")
	}

	// 同一毫秒内，序列号递增
	if now == s.timestamp {
		s.sequence = (s.sequence + 1) & sequenceMask
		// 同一毫秒序列号用尽，等待下一毫秒
		if s.sequence == 0 {
			for now <= s.timestamp {
				now = time.Now().UnixMilli()
			}
		}
	} else {
		// 新毫秒，序列号重置为0
		s.sequence = 0
	}

	s.timestamp = now

	// 拼接64位UID
	uid := ((now - epoch) << timestampShift) | (s.workerID << workerIDShift) | s.sequence
	return uint64(uid), nil
}

// 全局雪花算法实例（项目启动时初始化，机器ID可从配置文件读取）
var snowflakeInstance *Snowflake

// InitSnowflake 初始化雪花算法（网关/服务端启动时调用）
func InitSnowflake(workerID int64) error {
	instance, err := NewSnowflake(workerID)
	if err != nil {
		return err
	}
	snowflakeInstance = instance
	return nil
}

// GenUID 全局生成UID函数（简化调用）
func GenUID() uint64 {
	uid, err := snowflakeInstance.GenUID()
	if err != nil {
		panic("生成UID失败：" + err.Error()) // 生产环境可改为返回错误
	}
	return uid
}