laic-backend/tool/snowflake.go

package tool

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

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

	timestampBits = 28 // 时间戳位数（约17年）
	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 1023"))
	}
	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()

	now := time.Now().UnixMilli()
	now = (now - epoch) & (-1 ^ (-1 << timestampBits))
	if now <= 0 {
		now = 1
	}

	if now < s.lastTime {
		waitTime := s.lastTime - now
		if waitTime > 100 {
			now = s.lastTime + 1
			if now > (1<<timestampBits)-1 {
				return -1, errors.New("timestamp overflow due to clock moved backwards")
			}
		} else {
			time.Sleep(time.Duration(waitTime) * time.Millisecond)
			now = (time.Now().UnixMilli() - epoch) & (-1 ^ (-1 << timestampBits))
			if now <= 0 {
				now = 1
			}
			if now < s.lastTime {
				now = s.lastTime + 1
				if now > (1<<timestampBits)-1 {
					return -1, errors.New("timestamp overflow due to clock moved backwards")
				}
			}
		}
	}

	if s.lastTime == now {
		s.sequence = (s.sequence + 1) & maxSequence
		if s.sequence == 0 {
			for now <= s.lastTime {
				now = (time.Now().UnixMilli() - epoch) & (-1 ^ (-1 << timestampBits))
				if now <= 0 {
					now = s.lastTime + 1
					break
				}
			}
		}
	} else {
		s.sequence = 0
	}

	s.lastTime = now

	id := (now << timestampShift) |
		(s.workerID << workerIDShift) |
		s.sequence

	if id == 0 {
		s.sequence = 1
		id = (now << timestampShift) |
			(s.workerID << workerIDShift) |
			s.sequence
	}

	return id, nil
}

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