nginx-ui/server/utils/snow.go

package utils

import (
	"fmt"
	"strconv"
	"sync"
	"time"
)

type Snowflake struct {
	epoch       int64
	timestamp   int64
	dataCenter  int64
	workerId    int64
	sequence    int64
	mutex       sync.Mutex
	maxSequence int64
}

const (
	epochShift      = 22
	dataCenterShift = 12
	workerIdShift   = 10
	sequenceShift   = 0
	sequenceMask    = -1 ^ (-1 << sequenceShift)
	maxWorkerId     = -1 ^ (-1 << workerIdShift)
	maxDataCenter   = -1 ^ (-1 << dataCenterShift)
)

func NewSnowflake(epoch int64, dataCenter, workerId int64) (*Snowflake, error) {
	if dataCenter > maxDataCenter || workerId > maxWorkerId {
		return nil, fmt.Errorf("dataCenter or workerId exceeds limit")
	}
	return &Snowflake{
		epoch:       epoch,
		dataCenter:  dataCenter,
		workerId:    workerId,
		maxSequence: -1 ^ (-1 << sequenceShift),
	}, nil
}

func (s *Snowflake) NextID() int64 {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	currentTimestamp := time.Now().UnixNano() / 1e6
	if s.timestamp == currentTimestamp {
		s.sequence = (s.sequence + 1) & sequenceMask
		if s.sequence == 0 {
			currentTimestamp = s.waitNextTime(currentTimestamp)
		}
	} else {
		s.sequence = 0
	}
	s.timestamp = currentTimestamp

	id := ((currentTimestamp - s.epoch) << epochShift) |
		(s.dataCenter << dataCenterShift) |
		(s.workerId << workerIdShift) |
		s.sequence
	return id
}

func (s *Snowflake) waitNextTime(currentTimestamp int64) int64 {
	for currentTimestamp <= s.timestamp {
		currentTimestamp = time.Now().UnixNano() / 1e6
	}
	return currentTimestamp
}

var snowflake, _ = NewSnowflake(time.Now().UnixNano()/1e6, 1, 1)

func NextId() int64 {
	return snowflake.NextID()
}

func NextIdStr() string {
	return strconv.FormatInt(NextId(), 10)
}