shard.go

package lockd

import (
	"sync"
	"sync/atomic"
	"time"

	"github.com/imharjot/lockd/clock"
	"github.com/imharjot/lockd/lease"
)

// shard owns a disjoint subset of keys. It holds the live actors, the timing
// wheel that expires their leases, and a monotonic token counter. The only
// lock on the hot path is mu, guarding the actor map; everything else is
// either atomic or confined to a single goroutine.
type shard struct {
	id  int
	clk clock.Clock
	eng *Engine

	mu     sync.Mutex
	actors map[string]*keyActor

	// wheel mutations are serialized by wmu, which is held very briefly.
	wmu   sync.Mutex
	wheel *timingWheel

	tokens atomic.Uint64

	tickStop chan struct{}
	tickDone chan struct{}
	manual   bool

	actorBuf int
}

func newShard(id int, eng *Engine, clk clock.Clock, tick time.Duration, slots int, manual bool) *shard {
	sh := &shard{
		id:       id,
		clk:      clk,
		eng:      eng,
		actors:   make(map[string]*keyActor),
		wheel:    newTimingWheel(tick, slots, clk.Now()),
		tickStop: make(chan struct{}),
		tickDone: make(chan struct{}),
		manual:   manual,
		actorBuf: 64,
	}
	if manual {
		close(sh.tickDone)
	} else {
		go sh.tickLoop(tick)
	}
	return sh
}

// nextToken returns a fencing token unique and monotonically increasing within
// the shard. Tokens never repeat for the lifetime of the process, which is the
// property the fencing argument relies on.
func (sh *shard) nextToken() lease.Token {
	return lease.Token(sh.tokens.Add(1))
}

// seedToken makes sure future tokens exceed t, used during WAL replay so a
// recovered process never reissues a token already handed out before a crash.
func (sh *shard) seedToken(t lease.Token) {
	for {
		cur := sh.tokens.Load()
		if uint64(t) <= cur {
			return
		}
		if sh.tokens.CompareAndSwap(cur, uint64(t)) {
			return
		}
	}
}

// actorFor returns the live actor for key, creating one on first use.
func (sh *shard) actorFor(key string) *keyActor {
	sh.mu.Lock()
	a, ok := sh.actors[key]
	if !ok {
		a = newKeyActor(key, sh, sh.clk, sh.actorBuf)
		sh.actors[key] = a
	}
	sh.mu.Unlock()
	return a
}

// lookupActor returns the actor only if it already exists.
func (sh *shard) lookupActor(key string) (*keyActor, bool) {
	sh.mu.Lock()
	a, ok := sh.actors[key]
	sh.mu.Unlock()
	return a, ok
}

// schedule registers an expiry for key in the wheel.
func (sh *shard) schedule(key string, deadline time.Time) {
	sh.wmu.Lock()
	sh.wheel.add(key, deadline)
	sh.wmu.Unlock()
}

// reschedule updates an existing expiry; identical to schedule but named for
// intent at call sites.
func (sh *shard) reschedule(key string, deadline time.Time) {
	sh.schedule(key, deadline)
}

// cancelSchedule drops a pending expiry for key.
func (sh *shard) cancelSchedule(key string) {
	sh.wmu.Lock()
	sh.wheel.remove(key)
	sh.wmu.Unlock()
}

// tickLoop advances the wheel on the clock's cadence and fans expiries out to
// the owning actors. Each due key is sent an opExpire command, keeping all
// state changes on the actor goroutine.
func (sh *shard) tickLoop(tick time.Duration) {
	defer close(sh.tickDone)
	for {
		select {
		case <-sh.tickStop:
			return
		case <-sh.clk.After(tick):
			sh.runTick()
		}
	}
}

func (sh *shard) runTick() {
	now := sh.clk.Now()
	sh.wmu.Lock()
	expired := sh.wheel.advance(now)
	sh.wmu.Unlock()
	for _, key := range expired {
		if a, ok := sh.lookupActor(key); ok {
			cmd := newCommand(opExpire)
			if a.submit(cmd) {
				<-cmd.reply
			}
		}
	}
}

// tickOnce drives a single wheel advance synchronously. Tests on the fake
// clock call this after advancing time instead of waiting for the loop.
func (sh *shard) tickOnce() { sh.runTick() }

func (sh *shard) stop() {
	if !sh.manual {
		close(sh.tickStop)
	}
	<-sh.tickDone
	sh.mu.Lock()
	actors := make([]*keyActor, 0, len(sh.actors))
	for _, a := range sh.actors {
		actors = append(actors, a)
	}
	sh.actors = make(map[string]*keyActor)
	sh.mu.Unlock()
	for _, a := range actors {
		close(a.quit)
		<-a.done
	}
}

// journal hooks forward durable events to the engine's WAL. They are no-ops
// when no WAL is configured.

func (sh *shard) journalGrant(l lease.Lease) { sh.eng.journalGrant(l) }
func (sh *shard) journalRenew(l lease.Lease) { sh.eng.journalRenew(l) }
func (sh *shard) journalRelease(key, holder string, tok lease.Token) {
	sh.eng.journalRelease(key, holder, tok)
}
func (sh *shard) journalExpire(key string) { sh.eng.journalExpire(key) }