nl-im-service/internal/utils/snowflake.go

/**
 * package utils
 * 作用：雪花ID生成器，生成全局唯一的ID
 * 说明：使用Twitter的雪花算法，生成64位整数ID
 */
package utils

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

const (
	// 时间戳占用位数（41位，可以使用69年）
	timestampBits = 41
	// 数据中心ID占用位数（5位，最多32个数据中心）
	datacenterIDBits = 5
	// 机器ID占用位数（5位，每个数据中心最多32台机器）
	machineIDBits = 5
	// 序列号占用位数（12位，每毫秒最多4096个ID）
	sequenceBits = 12

	// 最大值
	maxDatacenterID = -1 ^ (-1 << datacenterIDBits)
	maxMachineID    = -1 ^ (-1 << machineIDBits)
	maxSequence     = -1 ^ (-1 << sequenceBits)

	// 位移
	machineIDShift    = sequenceBits
	datacenterIDShift = sequenceBits + machineIDBits
	timestampShift    = sequenceBits + machineIDBits + datacenterIDBits

	// 起始时间戳（2024-01-01 00:00:00）
	epoch int64 = 1704067200000
)

// Snowflake 雪花ID生成器
type Snowflake struct {
	mutex        sync.Mutex
	datacenterID int64
	machineID    int64
	sequence     int64
	lastStamp    int64
}

var (
	// 全局雪花ID生成器实例
	globalSnowflake *Snowflake
	once            sync.Once
)

/**
 * InitSnowflake
 * 功能：初始化全局雪花ID生成器
 * @param datacenterID 数据中心ID（0-31）
 * @param machineID 机器ID（0-31）
 */
func InitSnowflake(datacenterID, machineID int64) error {
	if datacenterID < 0 || datacenterID > maxDatacenterID {
		return errors.New("datacenter ID must be between 0 and 31")
	}
	if machineID < 0 || machineID > maxMachineID {
		return errors.New("machine ID must be between 0 and 31")
	}

	once.Do(func() {
		globalSnowflake = &Snowflake{
			datacenterID: datacenterID,
			machineID:    machineID,
			sequence:     0,
			lastStamp:    -1,
		}
	})

	return nil
}

/**
 * NextID
 * 功能：生成下一个ID
 * @returns 64位整数ID
 */
func NextID() (int64, error) {
	if globalSnowflake == nil {
		// 默认使用datacenterID=1, machineID=1
		if err := InitSnowflake(1, 1); err != nil {
			return 0, err
		}
	}

	return globalSnowflake.nextID()
}

/**
 * nextID
 * 功能：生成下一个ID（内部方法）
 */
func (s *Snowflake) nextID() (int64, error) {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	now := time.Now().UnixMilli()

	// 如果当前时间小于上次时间，说明时钟回拨
	if now < s.lastStamp {
		return 0, errors.New("clock moved backwards")
	}

	// 如果是同一毫秒内
	if now == s.lastStamp {
		s.sequence = (s.sequence + 1) & maxSequence
		// 序列号溢出，等待下一毫秒
		if s.sequence == 0 {
			now = s.waitNextMillis(s.lastStamp)
		}
	} else {
		// 新的毫秒，序列号重置
		s.sequence = 0
	}

	s.lastStamp = now

	// 生成ID
	id := ((now - epoch) << timestampShift) |
		(s.datacenterID << datacenterIDShift) |
		(s.machineID << machineIDShift) |
		s.sequence

	return id, nil
}

/**
 * waitNextMillis
 * 功能：等待下一毫秒
 */
func (s *Snowflake) waitNextMillis(lastStamp int64) int64 {
	now := time.Now().UnixMilli()
	for now <= lastStamp {
		now = time.Now().UnixMilli()
	}
	return now
}

/**
 * NextIDString
 * 功能：生成下一个ID（字符串格式）
 */
func NextIDString() (string, error) {
	id, err := NextID()
	if err != nil {
		return "", err
	}
	return int64ToString(id), nil
}

/**
 * int64ToString
 * 功能：将int64转换为字符串
 */
func int64ToString(id int64) string {
	if id == 0 {
		return "0"
	}
	negative := id < 0
	if negative {
		id = -id
	}

	var result []byte
	for id > 0 {
		result = append([]byte{byte('0' + id%10)}, result...)
		id /= 10
	}

	if negative {
		result = append([]byte{'-'}, result...)
	}

	return string(result)
}