backend/internal/pkg/utils/snowflake.go

package utils

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

// 雪花算法常量定义
const (
	// 64位ID结构：1位符号位 + 41位时间戳 + 10位机器ID + 12位序列号

	// 机器ID占用的位数
	nodeIDBits uint64 = 10
	// 序列号占用的位数
	sequenceBits uint64 = 12

	// 机器ID的最大值 (0-1023)
	maxNodeID int64 = -1 ^ (-1 << nodeIDBits)
	// 序列号的最大值 (0-4095)
	maxSequence int64 = -1 ^ (-1 << sequenceBits)

	// 机器ID左移位数
	nodeIDShift uint64 = sequenceBits
	// 时间戳左移位数
	timestampShift uint64 = sequenceBits + nodeIDBits

	// 自定义纪元时间：2024-01-01 00:00:00 UTC
	// 使用自定义纪元可以延长ID有效期约24年（从2024年开始）
	customEpoch int64 = 1704067200000 // 2024-01-01 00:00:00 UTC 的毫秒时间戳
)

// 错误定义
var (
	// ErrInvalidNodeID 机器ID无效
	ErrInvalidNodeID = errors.New("node ID must be between 0 and 1023")
	// ErrClockBackwards 时钟回拨
	ErrClockBackwards = errors.New("clock moved backwards, refusing to generate ID")
)

// IDInfo 解析后的ID信息
type IDInfo struct {
	Timestamp int64 // 生成ID时的时间戳（毫秒）
	NodeID    int64 // 机器ID
	Sequence  int64 // 序列号
}

// Snowflake 雪花算法ID生成器
type Snowflake struct {
	mu            sync.Mutex // 互斥锁，保证线程安全
	nodeID        int64      // 机器ID (0-1023)
	sequence      int64      // 当前序列号 (0-4095)
	lastTimestamp int64      // 上次生成ID的时间戳
}

// 全局雪花算法实例
var (
	globalSnowflake     *Snowflake
	globalSnowflakeOnce sync.Once
	globalSnowflakeErr  error
)

// InitSnowflake 初始化全局雪花算法实例
// nodeID: 机器ID，范围0-1023，可以通过环境变量 NODE_ID 配置
func InitSnowflake(nodeID int64) error {
	globalSnowflake, globalSnowflakeErr = NewSnowflake(nodeID)
	return globalSnowflakeErr
}

// GetSnowflake 获取全局雪花算法实例
// 如果未初始化，会自动使用默认配置初始化
func GetSnowflake() *Snowflake {
	globalSnowflakeOnce.Do(func() {
		if globalSnowflake == nil {
			globalSnowflake, globalSnowflakeErr = NewSnowflake(-1)
		}
	})
	return globalSnowflake
}

// NewSnowflake 创建雪花算法ID生成器实例
// nodeID: 机器ID，范围0-1023，可以通过环境变量 NODE_ID 配置
// 如果nodeID为-1，则尝试从环境变量 NODE_ID 读取
func NewSnowflake(nodeID int64) (*Snowflake, error) {
	// 如果传入-1，尝试从环境变量读取
	if nodeID == -1 {
		nodeIDStr := os.Getenv("NODE_ID")
		if nodeIDStr != "" {
			// 解析环境变量
			parsedID, err := parseInt(nodeIDStr)
			if err != nil {
				return nil, ErrInvalidNodeID
			}
			nodeID = parsedID
		} else {
			// 默认使用0
			nodeID = 0
		}
	}

	// 验证机器ID范围
	if nodeID < 0 || nodeID > maxNodeID {
		return nil, ErrInvalidNodeID
	}

	return &Snowflake{
		nodeID:        nodeID,
		sequence:      0,
		lastTimestamp: 0,
	}, nil
}

// parseInt 辅助函数：解析整数
func parseInt(s string) (int64, error) {
	var result int64
	var negative bool

	if len(s) == 0 {
		return 0, errors.New("empty string")
	}

	i := 0
	if s[0] == '-' {
		negative = true
		i = 1
	}

	for ; i < len(s); i++ {
		if s[i] < '0' || s[i] > '9' {
			return 0, errors.New("invalid character")
		}
		result = result*10 + int64(s[i]-'0')
	}

	if negative {
		result = -result
	}

	return result, nil
}

// GenerateID 生成唯一的雪花算法ID
// 返回值：生成的ID，以及可能的错误（如时钟回拨）
// 线程安全：使用互斥锁保证并发安全
func (s *Snowflake) GenerateID() (int64, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	// 获取当前时间戳（毫秒）
	now := currentTimestamp()

	// 处理时钟回拨
	if now < s.lastTimestamp {
		return 0, ErrClockBackwards
	}

	// 同一毫秒内
	if now == s.lastTimestamp {
		// 序列号递增
		s.sequence = (s.sequence + 1) & maxSequence
		// 序列号溢出，等待下一毫秒
		if s.sequence == 0 {
			now = s.waitNextMillis(now)
		}
	} else {
		// 不同毫秒，序列号重置为0
		s.sequence = 0
	}

	// 更新上次生成时间
	s.lastTimestamp = now

	// 组装ID
	// ID结构：时间戳部分 | 机器ID部分 | 序列号部分
	id := ((now - customEpoch) << timestampShift) |
		(s.nodeID << nodeIDShift) |
		s.sequence

	return id, nil
}

// waitNextMillis 等待到下一毫秒
// 参数：当前时间戳
// 返回值：下一毫秒的时间戳
func (s *Snowflake) waitNextMillis(timestamp int64) int64 {
	now := currentTimestamp()
	for now <= timestamp {
		now = currentTimestamp()
	}
	return now
}

// ParseID 解析雪花算法ID，提取其中的信息
// id: 要解析的雪花算法ID
// 返回值：包含时间戳、机器ID、序列号的结构体
func ParseID(id int64) *IDInfo {
	// 提取序列号（低12位）
	sequence := id & maxSequence

	// 提取机器ID（中间10位）
	nodeID := (id >> nodeIDShift) & maxNodeID

	// 提取时间戳（高41位）
	timestamp := (id >> timestampShift) + customEpoch

	return &IDInfo{
		Timestamp: timestamp,
		NodeID:    nodeID,
		Sequence:  sequence,
	}
}

// currentTimestamp 获取当前时间戳（毫秒）
func currentTimestamp() int64 {
	return time.Now().UnixNano() / 1e6
}

// GetNodeID 获取当前机器ID
func (s *Snowflake) GetNodeID() int64 {
	return s.nodeID
}

// GetCustomEpoch 获取自定义纪元时间
func GetCustomEpoch() int64 {
	return customEpoch
}

// IDToTime 将雪花算法ID转换为生成时间
// 这是一个便捷方法，等价于 ParseID(id).Timestamp
func IDToTime(id int64) time.Time {
	info := ParseID(id)
	return time.Unix(0, info.Timestamp*1e6) // 毫秒转纳秒
}

// ValidateID 验证ID是否为有效的雪花算法ID
// 检查时间戳是否在合理范围内
func ValidateID(id int64) bool {
	if id <= 0 {
		return false
	}

	info := ParseID(id)

	// 检查时间戳是否在合理范围内
	// 不能早于纪元时间，不能晚于当前时间太多（允许1分钟的时钟偏差）
	now := currentTimestamp()
	if info.Timestamp < customEpoch || info.Timestamp > now+60000 {
		return false
	}

	// 检查机器ID和序列号是否在有效范围内
	if info.NodeID < 0 || info.NodeID > maxNodeID {
		return false
	}
	if info.Sequence < 0 || info.Sequence > maxSequence {
		return false
	}

	return true
}