laic-backend/tool/snowflake.go

package tool

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

const (
	// 起始时间戳 (2023-01-01 00:00:00 UTC)
	epoch int64 = 1672531200000

	timestampBits = 41 // 时间戳位数（约69年）
	workerIDBits  = 5  // 工作机器ID所占位数
	sequenceBits  = 12 // 序列号所占位数

	maxWorkerID = -1 ^ (-1 << workerIDBits)
	maxSequence = -1 ^ (-1 << sequenceBits)

	workerIDShift  = sequenceBits
	timestampShift = sequenceBits + workerIDBits
)

type Snowflake struct {
	mu        sync.Mutex
	timestamp int64
	workerID  int64
	sequence  int64
	lastTime  int64
}

var DefaultSnowflake *Snowflake

func init() {
	DefaultSnowflake = NewSnowflake(1)
}

func NewSnowflake(workerID int64) *Snowflake {
	if workerID < 0 || workerID > maxWorkerID {
		panic(errors.New("worker ID must be between 0 and 31"))
	}
	return &Snowflake{
		timestamp: 0,
		workerID:  workerID,
		sequence:  0,
		lastTime:  -1,
	}
}

// NextID 生成下一个ID
func (s *Snowflake) NextID() (int64, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	maxTimestamp := int64(1<<timestampBits) - 1
	now := time.Now().UnixMilli() - epoch
	if now <= 0 {
		return -1, errors.New("timestamp is before snowflake epoch")
	}
	if now > maxTimestamp {
		return -1, errors.New("timestamp overflow")
	}

	if now < s.lastTime {
		waitTime := s.lastTime - now
		if waitTime > 100 {
			return -1, errors.New("clock moved backwards")
		}
		time.Sleep(time.Duration(waitTime) * time.Millisecond)
		now = time.Now().UnixMilli() - epoch
		if now < s.lastTime {
			return -1, errors.New("clock moved backwards")
		}
		if now > maxTimestamp {
			return -1, errors.New("timestamp overflow")
		}
	}

	if now == s.lastTime {
		s.sequence = (s.sequence + 1) & maxSequence
		if s.sequence == 0 {
			for now <= s.lastTime {
				time.Sleep(time.Millisecond)
				now = time.Now().UnixMilli() - epoch
				if now > maxTimestamp {
					return -1, errors.New("timestamp overflow")
				}
			}
		}
	} else {
		s.sequence = 0
	}

	s.lastTime = now
	return (now << timestampShift) | (s.workerID << workerIDShift) | s.sequence, nil
}

// NextID 全局函数，使用默认实例生成ID
func NextID() (int64, error) {
	return DefaultSnowflake.NextID()
}