go-kv/wal/sequence.go

package wal

import (
	"sync/atomic"

	"github.com/dailz/go-kv"
)

// SequenceManager manages monotonic sequence number allocation for the WAL.
// Invariant: durableSequence <= publishedSequence <= nextSequence
type SequenceManager struct {
	nextSequence      atomic.Uint64
	publishedSequence atomic.Uint64
	durableSequence   atomic.Uint64
	exhausted         atomic.Bool
}

// NewSequenceManager creates a SequenceManager initialised from a recovered
// sequence number. All three watermarks start at recoveredSequence.
func NewSequenceManager(recoveredSequence uint64) *SequenceManager {
	sm := &SequenceManager{}
	sm.nextSequence.Store(recoveredSequence)
	sm.publishedSequence.Store(recoveredSequence)
	sm.durableSequence.Store(recoveredSequence)
	return sm
}

// AllocateBatch atomically reserves [base, base+count-1] sequence numbers.
// Returns ErrSequenceExhausted if count is 0 or the allocation would overflow uint64.
func (sm *SequenceManager) AllocateBatch(count uint32) (baseSequence uint64, err error) {
	if count == 0 {
		return 0, go_kv.ErrSequenceExhausted
	}
	for {
		if sm.exhausted.Load() {
			return 0, go_kv.ErrSequenceExhausted
		}
		base := sm.nextSequence.Load()
		last := base + uint64(count) - 1
		if last < base {
			return 0, go_kv.ErrSequenceExhausted
		}
		newNext := last + 1
		if !sm.nextSequence.CompareAndSwap(base, newNext) {
			continue
		}
		if newNext == 0 {
			sm.exhausted.Store(true)
		}
		return base, nil
	}
}

// Publish advances the publishedSequence watermark to seq (only forward).
func (sm *SequenceManager) Publish(seq uint64) {
	for {
		current := sm.publishedSequence.Load()
		if seq <= current {
			return
		}
		if sm.publishedSequence.CompareAndSwap(current, seq) {
			return
		}
	}
}

// MarkDurable advances the durableSequence watermark to seq (only forward).
func (sm *SequenceManager) MarkDurable(seq uint64) {
	for {
		current := sm.durableSequence.Load()
		if seq <= current {
			return
		}
		if sm.durableSequence.CompareAndSwap(current, seq) {
			return
		}
	}
}

// Published returns the current publishedSequence watermark.
func (sm *SequenceManager) Published() uint64 { return sm.publishedSequence.Load() }

// Durable returns the current durableSequence watermark.
func (sm *SequenceManager) Durable() uint64 { return sm.durableSequence.Load() }

// NextSequence returns the next sequence number to be allocated.
func (sm *SequenceManager) NextSequence() uint64 { return sm.nextSequence.Load() }