bl/common/utils/snowflake.go

package utils

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

// Snowflake is a uint32-sized ID with 32 node support and 8.5 year lifespan
type Snowflake uint32

// ID is an alias for Snowflake
type ID = Snowflake

func (s Snowflake) String() string {
	return strconv.FormatUint(uint64(s), 16)
}

// ParseSnowflake parses a string into a snowflake, if possible
func ParseSnowflake(s string) (Snowflake, error) {
	i, err := strconv.ParseUint(s, 16, 32)
	if err != nil {
		return 0, err
	}
	return Snowflake(i), nil
}

// DefaultGenerator is a Generator with the epoch set to Jan 1, 2025 UTC


// Generator holds info needed for generating snowflakes with 32 node support
type Generator struct {
	Epoch    time.Time
	nodeID   uint8  // 0-31 (5位，支持32个节点)
	sequence uint8  // 0-255 (8位，支持每秒256个ID)
	lastSec  uint32 // 上次生成ID的秒级时间戳
	mu       sync.Mutex
}

// NewGen returns a new generator with specified epoch and node ID
func NewGen(epoch time.Time, nodeID uint8) *Generator {
	if nodeID >= 32 { // 5位最多支持0-31
		panic(errors.New("node ID must be between 0 and 31"))
	}
	return &Generator{
		Epoch:  epoch,
		nodeID: nodeID,
	}
}

// Get generates a new unique snowflake ID and returns it as uint32
func (g *Generator) Get() uint32 {
	g.mu.Lock()
	defer g.mu.Unlock()

	// 获取当前时间戳（自纪元以来的秒数）
	now := time.Now()
	sec := uint32(now.Sub(g.Epoch).Seconds())

	// 检查时间戳是否超出28位限制 (2^28 - 1 = 268435455秒 ≈ 8.5年)
	if sec >= (1 << 28) {
		panic(errors.New("timestamp overflow: please update epoch"))
	}

	// 处理时间回拨
	if sec < g.lastSec {
		// 计算需要等待的时间
		sleepDuration := time.Duration(g.lastSec-sec) * time.Second
		// 添加1毫秒缓冲，确保时间确实已过
		sleepDuration += time.Millisecond
		// 等待时间追上
		time.Sleep(sleepDuration)

		// 重新获取时间戳
		now = time.Now()
		sec = uint32(now.Sub(g.Epoch).Seconds())

		// 再次检查，防止极端情况
		if sec < g.lastSec {
			panic(errors.New("clock moved backwards beyond recovery"))
		}
	}

	// 处理同一秒内的序列号
	if sec == g.lastSec {
		g.sequence++
		// 序列号溢出（8位最大255）
		if g.sequence > 255 { // 0-255是有效范围
			// 等待到下一秒
			sleepDuration := time.Second - time.Duration(now.Nanosecond())*time.Nanosecond + time.Microsecond
			time.Sleep(sleepDuration)

			now = time.Now()
			sec = uint32(now.Sub(g.Epoch).Seconds())
			g.sequence = 0
		}
	} else {
		// 新的一秒，重置序列号
		g.sequence = 0
	}

	// 更新上次时间戳
	g.lastSec = sec

	// 组装ID (32位):
	// [28位秒级时间戳][5位节点ID][8位序列号]
	var id uint32
	id |= (sec << 13)             // 左移13位（5+8）放置秒级时间戳
	id |= (uint32(g.nodeID) << 8) // 左移8位放置节点ID
	id |= uint32(g.sequence)      // 放置序列号

	return id
}

// Parse extracts timestamp, node ID and sequence from a snowflake
func (g *Generator) Parse(id uint32) (t time.Time, nodeID uint8, sequence uint8) {
	// 提取各部分信息
	sec := id >> 13
	nodeID = uint8((id >> 8) & 0x1F) // 5位掩码（0x1F = 31）
	sequence = uint8(id & 0xFF)      // 8位掩码（0xFF = 255）

	// 计算实际时间
	totalNanos := g.Epoch.UnixNano() + int64(sec)*int64(time.Second)
	t = time.Unix(0, totalNanos).UTC()
	return
}

// GetNodeID returns the current node ID
func (g *Generator) GetNodeID() uint8 {
	return g.nodeID
}

// GetLastTimestamp returns the last timestamp used (for testing/debugging)
func (g *Generator) GetLastTimestamp() uint32 {
	g.mu.Lock()
	defer g.mu.Unlock()
	return g.lastSec
}

// MaxAge returns the maximum age (in seconds) that this generator can handle before overflow
func (g *Generator) MaxAge() uint32 {
	return (1 << 28) - 1 // 268435455秒，约8.5年
}

// MaxSequencePerSecond returns the maximum number of sequences per second
func (g *Generator) MaxSequencePerSecond() uint8 {
	return 255 // 2^8 - 1
}