ntool/nrandom/snowflake/snowflake_offset.go

package snowflake

import (
    "sync"
    "time"
)

// SnowWorkerOffset 雪花漂移算法
type SnowWorkerOffset struct {
    baseTime          int64  // 基础时间（ms单位）
    workerId          uint16 // 机器码
    workerIdBitLength byte   // 机器码位长
    seqBitLength      byte   // 序列数位长
    maxSeqNumber      uint32 // 最大序列数（含）
    minSeqNumber      uint32 // 最小序列数（含）
    topOverCostCount  uint32 // 最大漂移次数（含）

    timestampShift   byte
    currentSeqNumber uint32

    lastTimeTick           int64
    turnBackTimeTick       int64
    turnBackIndex          byte
    isOverCost             bool
    overCostCountInOneTerm uint32

    sync.Mutex
}

func NewSnowWorkerOffset(options *Options) *SnowWorkerOffset {
    if options == nil {
        options = defaultOptions
    }
    baseTime := options.BaseTime
    if baseTime <= 0 {
        baseTime = defaultOptions.BaseTime
    }

    workerIdBitLength := options.WorkerIdBitLength
    if workerIdBitLength == 0 {
        workerIdBitLength = defaultOptions.WorkerIdBitLength
    }

    seqBitLength := options.SeqBitLength
    if seqBitLength == 0 {
        seqBitLength = defaultOptions.SeqBitLength
    }

    // 序列数位长+机器码位长不超过22
    if workerIdBitLength+seqBitLength > 22 {
        workerIdBitLength = defaultOptions.WorkerIdBitLength
        if workerIdBitLength+seqBitLength > 22 {
            seqBitLength = defaultOptions.SeqBitLength
        }
    }

    minSeqNumber := options.MinSeqNumber
    if minSeqNumber < 5 {
        minSeqNumber = defaultOptions.MinSeqNumber
    }

    maxSeqNumber := options.MaxSeqNumber
    if maxSeqNumber <= 0 {
        maxSeqNumber = (1 << seqBitLength) - 1
    }

    topOverCostCount := options.TopOverCostCount
    if topOverCostCount == 0 {
        topOverCostCount = defaultOptions.TopOverCostCount
    }

    return &SnowWorkerOffset{
        baseTime:          baseTime,
        workerId:          options.WorkerId,
        workerIdBitLength: workerIdBitLength,
        seqBitLength:      seqBitLength,
        maxSeqNumber:      maxSeqNumber,
        minSeqNumber:      minSeqNumber,
        topOverCostCount:  topOverCostCount,
        timestampShift:    workerIdBitLength + seqBitLength,
        currentSeqNumber:  minSeqNumber,

        lastTimeTick:           0,
        turnBackTimeTick:       0,
        turnBackIndex:          0,
        isOverCost:             false,
        overCostCountInOneTerm: 0,
    }
}

func (s *SnowWorkerOffset) NextID() int64 {
    s.Lock()
    defer s.Unlock()
    if s.isOverCost {
        return s.nextOverCost()
    } else {
        return s.nextNormal()
    }
}

func (s *SnowWorkerOffset) nextOverCost() int64 {
    currentTimeTick := s.currentTimeTick()
    if currentTimeTick > s.lastTimeTick {
        s.lastTimeTick = currentTimeTick
        s.currentSeqNumber = s.minSeqNumber
        s.isOverCost = false
        s.overCostCountInOneTerm = 0
        return s.calcId(s.lastTimeTick)
    }
    if s.overCostCountInOneTerm >= s.topOverCostCount {
        s.lastTimeTick = s.nextTimeTick()
        s.currentSeqNumber = s.minSeqNumber
        s.isOverCost = false
        s.overCostCountInOneTerm = 0
        return s.calcId(s.lastTimeTick)
    }
    if s.currentSeqNumber > s.maxSeqNumber {
        s.lastTimeTick++
        s.currentSeqNumber = s.minSeqNumber
        s.isOverCost = true
        s.overCostCountInOneTerm++
        return s.calcId(s.lastTimeTick)
    }
    return s.calcId(s.lastTimeTick)
}

func (s *SnowWorkerOffset) nextNormal() int64 {
    currentTimeTick := s.currentTimeTick()
    if currentTimeTick < s.lastTimeTick {
        if s.turnBackTimeTick < 1 {
            s.turnBackTimeTick = s.lastTimeTick - 1
            s.turnBackIndex++
            // 每毫秒序列数的前5位是预留位，0用于手工新值，1-4是时间回拨次序
            // 支持4次回拨次序（避免回拨重叠导致ID重复），可无限次回拨（次序循环使用）。
            if s.turnBackIndex > 4 {
                s.turnBackIndex = 1
            }
        }
        return s.calcTurnBackId(s.turnBackTimeTick)
    }
    // 时间追平时，turnBackTimeTick 清零
    if s.turnBackTimeTick > 0 {
        s.turnBackTimeTick = 0
    }
    if currentTimeTick > s.lastTimeTick {
        s.lastTimeTick = currentTimeTick
        s.currentSeqNumber = s.minSeqNumber
        return s.calcId(s.lastTimeTick)
    }
    if s.currentSeqNumber > s.maxSeqNumber {
        s.lastTimeTick++
        s.currentSeqNumber = s.minSeqNumber
        s.isOverCost = true
        s.overCostCountInOneTerm = 1
        return s.calcId(s.lastTimeTick)
    }
    return s.calcId(s.lastTimeTick)
}

func (s *SnowWorkerOffset) calcId(timeTick int64) int64 {
    id := timeTick<<s.timestampShift + int64(s.workerId<<s.seqBitLength) + int64(s.currentSeqNumber)
    s.currentSeqNumber++
    return id
}

func (s *SnowWorkerOffset) calcTurnBackId(timeTick int64) int64 {
    id := timeTick<<s.timestampShift + int64(s.workerId<<s.seqBitLength) + int64(s.turnBackIndex)
    s.turnBackIndex++
    return id
}

func (s *SnowWorkerOffset) currentTimeTick() int64 {
    return (time.Now().UnixNano() / 1e6) - s.baseTime
}

func (s *SnowWorkerOffset) nextTimeTick() int64 {
    tmp := s.currentTimeTick()
    for tmp <= s.lastTimeTick {
        tmp = s.currentTimeTick()
    }
    return tmp
}