From 578566b801249eda1fe5e344529a770f1e8a7653 Mon Sep 17 00:00:00 2001 From: Josh Siviour Date: Tue, 25 Aug 2026 08:58:02 +1000 Subject: [PATCH] refactor: convert the bedrock reaper to the shared kvlease package (#869) --- spinifex/handlers/bedrock/reaper.go | 126 ++++++----------------- spinifex/handlers/bedrock/reaper_test.go | 41 ++++++-- 2 files changed, 63 insertions(+), 104 deletions(-) diff --git a/spinifex/handlers/bedrock/reaper.go b/spinifex/handlers/bedrock/reaper.go index bf7684a93..a97c7f308 100644 --- a/spinifex/handlers/bedrock/reaper.go +++ b/spinifex/handlers/bedrock/reaper.go @@ -5,10 +5,10 @@ import ( "errors" "fmt" "log/slog" - "sync" "time" gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock" + "github.com/mulgadc/spinifex/spinifex/kvlease" "github.com/mulgadc/spinifex/spinifex/otelsetup" "github.com/mulgadc/spinifex/spinifex/utils" "github.com/nats-io/nats.go/jetstream" @@ -19,23 +19,9 @@ import ( const ( defaultReaperInterval = 30 * time.Second defaultLeaseRefresh = 20 * time.Second - - // How long an endpoint must have been idle before its GPU is taken back. - // Cold start was measured at 4m12s for the smallest model this platform - // serves, so a TTL of a few multiples of that is what keeps a quiet-but-used - // model from becoming a cold-start generator. Scale-to-zero is a cost - // optimisation; a latency-sensitive workload buys warmth by pinning. - defaultIdleTTL = 15 * time.Minute - - // How many consecutive failed scrapes mean the serving process is wedged - // rather than briefly busy. Long enough that a GC pause or a dropped packet - // is not an outage, short enough that a dead VM does not hold a device for a - // whole idleTTL. - maxScrapeFailures = 5 - - // Per-scrape budget. Generous next to the sweep interval, because a slow - // answer is still an answer and a timeout here counts as a failure. - defaultScrapeTimeout = 5 * time.Second + defaultIdleTTL = 15 * time.Minute // How long an endpoint must have been idle before its GPU is taken back. + maxScrapeFailures = 5 // How many consecutive failed scrapes mean the serving process is wedged rather than briefly busy. + defaultScrapeTimeout = 5 * time.Second // Per-scrape budget. ) // ReaperDeps are the reaper's overridable knobs. Every zero value takes the @@ -50,23 +36,27 @@ type ReaperDeps struct { // Reaper is the leader-elected idle-reclaim loop. One node holds the lease and // does the scraping and reaping; every node keeps serving the API, so a // leaderless gap delays a reclaim rather than failing a request. -// -// It reads idleness from the serving process's own Prometheus /metrics rather -// than from gateway-reported activity: that keeps the daemon the single writer -// of endpoint state, survives a gateway node dying mid-call, and works -// unchanged with any number of gateways. type Reaper struct { - svc *Service - holder string - deps ReaperDeps - - mu sync.Mutex - leader bool + svc *Service + holder string + deps ReaperDeps + lease *kvlease.Lease + leaseErr error } // NewReaper builds the reaper for svc; holder identifies this daemon in the lease. func NewReaper(svc *Service, holder string, deps ReaperDeps) *Reaper { - return &Reaper{svc: svc, holder: holder, deps: deps} + r := &Reaper{svc: svc, holder: holder, deps: deps} + r.lease, r.leaseErr = kvlease.New(kvlease.Config{ + Name: "bedrock/reaper", + Bucket: r.leaderBucket, + Key: leaderKey, + Holder: holder, + TTL: KVBucketLeaderTTL, + Renew: r.leaseRefresh(), + Retry: r.leaseRefresh(), + }) + return r } func (r *Reaper) interval() time.Duration { @@ -100,19 +90,24 @@ func (r *Reaper) scrapeTimeout() time.Duration { // Run drives the leadership and sweep loop until ctx is cancelled. Intended as // a daemon-boot goroutine; panics are the caller's recover concern. func (r *Reaper) Run(ctx context.Context) { - leaseTicker := time.NewTicker(r.leaseRefresh()) + if r.leaseErr != nil { + slog.ErrorContext(ctx, "bedrock reaper: lease config invalid", "holder", r.holder, "err", r.leaseErr) + return + } sweepTicker := time.NewTicker(r.interval()) - defer leaseTicker.Stop() defer sweepTicker.Stop() - r.evaluateLeadership(ctx) + leaseDone := make(chan struct{}) + go func() { + defer close(leaseDone) + r.lease.Run(ctx) + }() + for { select { case <-ctx.Done(): - r.relinquish() + <-leaseDone return - case <-leaseTicker.C: - r.evaluateLeadership(ctx) case <-sweepTicker.C: if !r.IsLeader() { continue @@ -126,64 +121,7 @@ func (r *Reaper) Run(ctx context.Context) { // IsLeader reports whether this node currently holds the reaper lease. func (r *Reaper) IsLeader() bool { - r.mu.Lock() - defer r.mu.Unlock() - return r.leader -} - -func (r *Reaper) evaluateLeadership(ctx context.Context) { - won := r.acquireOrRefresh(ctx) - r.mu.Lock() - was := r.leader - r.leader = won - r.mu.Unlock() - - switch { - case won && !was: - slog.Info("bedrock reaper: elected leader", "holder", r.holder) - case !won && was: - slog.Info("bedrock reaper: lost leadership", "holder", r.holder) - } -} - -// acquireOrRefresh claims the lease, or refreshes it (resetting the TTL) when -// this node already holds it. -func (r *Reaper) acquireOrRefresh(ctx context.Context) bool { - kv, err := r.leaderBucket(ctx) - if err != nil { - return false - } - if _, err := kv.Create(ctx, leaderKey, []byte(r.holder)); err == nil { - return true - } - entry, err := kv.Get(ctx, leaderKey) - if err != nil { - return false - } - if string(entry.Value()) != r.holder { - return false - } - if _, err := kv.Put(ctx, leaderKey, []byte(r.holder)); err != nil { - return false - } - return true -} - -// relinquish releases the lease on shutdown so the next leader is elected -// immediately rather than after the TTL. -func (r *Reaper) relinquish() { - // Run's ctx is already cancelled by the time this is called, so the release - // runs on its own — a cancelled ctx would fail the delete. - ctx := context.Background() - kv, err := r.leaderBucket(ctx) - if err != nil { - return - } - if entry, gerr := kv.Get(ctx, leaderKey); gerr == nil && string(entry.Value()) == r.holder { - if err := kv.Delete(ctx, leaderKey); err != nil { - slog.Debug("bedrock reaper: release lease failed", "holder", r.holder, "err", err) - } - } + return r.lease.Held() } func (r *Reaper) leaderBucket(ctx context.Context) (jetstream.KeyValue, error) { diff --git a/spinifex/handlers/bedrock/reaper_test.go b/spinifex/handlers/bedrock/reaper_test.go index 936dd185e..4cb1cb5b6 100644 --- a/spinifex/handlers/bedrock/reaper_test.go +++ b/spinifex/handlers/bedrock/reaper_test.go @@ -395,40 +395,41 @@ func TestReaper_LeaseAdmitsOneLeader(t *testing.T) { f := newReaperFixture(t, ReaperDeps{}) other := NewReaper(f.svc, "node-b", ReaperDeps{}) - f.reaper.evaluateLeadership(t.Context()) - other.evaluateLeadership(t.Context()) + require.True(t, f.reaper.lease.TryAcquire(t.Context())) + require.False(t, other.lease.TryAcquire(t.Context()), + "two nodes must never sweep the same endpoints at once") assert.True(t, f.reaper.IsLeader()) - assert.False(t, other.IsLeader(), "two nodes must never sweep the same endpoints at once") + assert.False(t, other.IsLeader()) // A refresh by the holder keeps it, and does not hand it away. - f.reaper.evaluateLeadership(t.Context()) + require.True(t, f.reaper.lease.TryAcquire(t.Context())) assert.True(t, f.reaper.IsLeader()) } // Releasing on shutdown is what makes a rolling restart cost seconds rather // than a full lease TTL of no reclaim. -func TestReaper_RelinquishFreesTheLeaseImmediately(t *testing.T) { +func TestReaper_ReleaseFreesTheLeaseImmediately(t *testing.T) { f := newReaperFixture(t, ReaperDeps{}) other := NewReaper(f.svc, "node-b", ReaperDeps{}) - f.reaper.evaluateLeadership(t.Context()) + f.reaper.lease.TryAcquire(t.Context()) require.True(t, f.reaper.IsLeader()) - f.reaper.relinquish() - other.evaluateLeadership(t.Context()) + f.reaper.lease.Release(t.Context()) + other.lease.TryAcquire(t.Context()) assert.True(t, other.IsLeader()) } // A node that lost the election does no work at all, so a leaderless gap // delays a reclaim rather than duplicating one. func TestReaper_NonLeaderSweepsNothing(t *testing.T) { - f := newReaperFixture(t, ReaperDeps{Interval: 5 * time.Millisecond, LeaseRefresh: time.Hour, IdleTTL: time.Nanosecond}) + f := newReaperFixture(t, ReaperDeps{Interval: 5 * time.Millisecond, IdleTTL: time.Nanosecond}) f.ready(t) f.age(t, time.Hour, time.Hour) holder := NewReaper(f.svc, "node-b", ReaperDeps{}) - holder.evaluateLeadership(t.Context()) + holder.lease.TryAcquire(t.Context()) require.True(t, holder.IsLeader()) ctx, cancel := context.WithTimeout(t.Context(), 60*time.Millisecond) @@ -450,3 +451,23 @@ func TestEndpointRecord_LastActive(t *testing.T) { assert.Equal(t, ready, EndpointRecord{ReadyAt: ready}.LastActive(), "an endpoint quiet since launch is idle since launch, not since the zero time") } + +// Shutdown must delete the lease key before Run returns. The daemon waits on Run +// via its shutdown group, so a key left behind parks the next election for the +// full TTL rather than freeing it immediately. +func TestReaper_RunReleasesLeaseOnShutdown(t *testing.T) { + f := newReaperFixture(t, ReaperDeps{Interval: time.Hour}) + require.NoError(t, f.reaper.leaseErr) + + ctx, cancel := context.WithCancel(t.Context()) + done := make(chan struct{}) + go func() { defer close(done); f.reaper.Run(ctx) }() + require.Eventually(t, f.reaper.IsLeader, 2*time.Second, 20*time.Millisecond) + + cancel() + <-done + + other := NewReaper(f.svc, "node-b", ReaperDeps{}) + assert.True(t, other.lease.TryAcquire(t.Context()), + "Run returned with the lease key still present") +}