From 5675a2b37c8e2a27a04e565118bfdb3fa61adb21 Mon Sep 17 00:00:00 2001 From: Josh Siviour Date: Tue, 25 Aug 2026 08:05:17 +1000 Subject: [PATCH] refactor: convert the ECS scheduler to the shared kvlease package (#867) * refactor: ECS scheduler to used shared kvlease package * fix: wait for the ECS lease release before Run returns --- spinifex/handlers/ecs/scheduler.go | 140 ++++++------------ .../handlers/ecs/scheduler_leader_test.go | 78 ---------- spinifex/handlers/ecs/service_test.go | 66 ++++++++- 3 files changed, 105 insertions(+), 179 deletions(-) delete mode 100644 spinifex/handlers/ecs/scheduler_leader_test.go diff --git a/spinifex/handlers/ecs/scheduler.go b/spinifex/handlers/ecs/scheduler.go index 9e5778ae3..46315a930 100644 --- a/spinifex/handlers/ecs/scheduler.go +++ b/spinifex/handlers/ecs/scheduler.go @@ -9,6 +9,7 @@ import ( "time" "github.com/mulgadc/spinifex/spinifex/handlers/ecs/bus" + "github.com/mulgadc/spinifex/spinifex/kvlease" "github.com/nats-io/nats.go" "github.com/nats-io/nats.go/jetstream" ) @@ -40,36 +41,66 @@ type Scheduler struct { svc *Service holder string - mu sync.Mutex - leader bool - subs []*nats.Subscription + lease *kvlease.Lease + leaseErr error + + mu sync.Mutex + subs []*nats.Subscription } // NewScheduler constructs a Scheduler. holder identifies this daemon in the lease. func NewScheduler(nc *nats.Conn, svc *Service, holder string) *Scheduler { - return &Scheduler{nc: nc, svc: svc, holder: holder} + sc := &Scheduler{nc: nc, svc: svc, holder: holder} + sc.lease, sc.leaseErr = kvlease.New(kvlease.Config{ + Name: "ecs/scheduler", + Bucket: sc.leaderBucket, + Key: schedulerLeaderKey, + Holder: holder, + TTL: KVBucketECSLeaderTTL, + Renew: leaseRefresh, + Retry: leaseRefresh, + // The leader owns the Layer-2 bus subscriptions. A node that cannot wire them stands + // down rather than hold the lease and process nothing. + OnGained: sc.subscribeBus, + OnLost: sc.unsubscribeBus, + }) + return sc +} + +func (sc *Scheduler) leaderBucket(ctx context.Context) (jetstream.KeyValue, error) { + js, err := jetstream.New(sc.nc) + if err != nil { + return nil, err + } + return InitLeaderBucket(ctx, js) } // Run drives the leadership + reaper loop until ctx is cancelled. It is intended // to run as a daemon-boot goroutine; panics are the caller's recover concern. func (sc *Scheduler) Run(ctx context.Context) { - leaseTicker := time.NewTicker(leaseRefresh) + if sc.leaseErr != nil { + slog.ErrorContext(ctx, "ECS scheduler: lease config invalid", "holder", sc.holder, "err", sc.leaseErr) + return + } reaperTicker := time.NewTicker(reaperInterval) reconcileTicker := time.NewTicker(reconcileInterval) sweepTicker := time.NewTicker(sweepInterval) - defer leaseTicker.Stop() defer reaperTicker.Stop() defer reconcileTicker.Stop() defer sweepTicker.Stop() - sc.evaluateLeadership(ctx) + leaseDone := make(chan struct{}) + go func() { + defer close(leaseDone) + sc.lease.Run(ctx) + }() for { select { case <-ctx.Done(): - sc.relinquish() + // The daemon waits on Run, so the lease delete must complete before + // it returns, otherwise the next leader waits out the full TTL. + <-leaseDone return - case <-leaseTicker.C: - sc.evaluateLeadership(ctx) case <-reaperTicker.C: sc.runIfLeader("instance reap", func() error { return sc.reap(ctx) }) case <-reconcileTicker.C: @@ -84,7 +115,7 @@ func (sc *Scheduler) Run(ctx context.Context) { // A pass errors only when it could not observe the whole fleet, so this log is // the operator's signal that a tick was skipped rather than found nothing to do. func (sc *Scheduler) runIfLeader(pass string, run func() error) { - if !sc.isLeader() { + if !sc.lease.Held() { return } if err := run(); err != nil { @@ -92,78 +123,6 @@ func (sc *Scheduler) runIfLeader(pass string, run func() error) { } } -func (sc *Scheduler) isLeader() bool { - sc.mu.Lock() - defer sc.mu.Unlock() - return sc.leader -} - -// evaluateLeadership acquires or refreshes the lease, wiring up (or tearing down) -// the bus subscriptions as leadership changes. -func (sc *Scheduler) evaluateLeadership(ctx context.Context) { - won := sc.acquireOrRefresh(ctx) - sc.mu.Lock() - was := sc.leader - sc.leader = won - sc.mu.Unlock() - - switch { - case won && !was: - if err := sc.subscribeBus(ctx); err != nil { - slog.Error("ECS scheduler: bus subscribe failed", "holder", sc.holder, "err", err) - } else { - slog.Info("ECS scheduler: elected leader, bus subscriptions active", "holder", sc.holder) - } - case !won && was: - sc.unsubscribeBus() - slog.Info("ECS scheduler: lost leadership, bus subscriptions dropped", "holder", sc.holder) - } -} - -// acquireOrRefresh tries to claim the scheduler lease, refreshing it (resetting -// the TTL) when we already hold it. Returns true when we are the leader. -func (sc *Scheduler) acquireOrRefresh(ctx context.Context) bool { - js, err := jetstream.New(sc.nc) - if err != nil { - return false - } - kv, err := InitLeaderBucket(ctx, js) - if err != nil { - return false - } - if _, err := kv.Create(ctx, schedulerLeaderKey, []byte(sc.holder)); err == nil { - return true - } - entry, err := kv.Get(ctx, schedulerLeaderKey) - if err != nil { - return false - } - if string(entry.Value()) != sc.holder { - return false // someone else holds it - } - return refreshLease(ctx, kv, schedulerLeaderKey, sc.holder, entry.Revision()) -} - -// relinquish drops subscriptions and deletes the lease key on shutdown. -func (sc *Scheduler) relinquish() { - sc.unsubscribeBus() - // Run's ctx is already cancelled by the time it calls relinquish, so the - // release runs on its own context — a captured ctx would fail the delete and - // leave the lease for the TTL to reap. - ctx := context.Background() - js, err := jetstream.New(sc.nc) - if err != nil { - return - } - kv, err := InitLeaderBucket(ctx, js) - if err != nil { - return - } - if entry, gerr := kv.Get(ctx, schedulerLeaderKey); gerr == nil && string(entry.Value()) == sc.holder { - _ = dropLease(ctx, kv, schedulerLeaderKey, entry.Revision()) - } -} - // subscribeBus wires the wildcard Layer-2 subscriptions onto the service KV // writers. All clusters/accounts fan into the leader. func (sc *Scheduler) subscribeBus(_ context.Context) error { @@ -316,18 +275,3 @@ func accountIDFromBucket(bucket string) (string, bool) { } return strings.TrimPrefix(bucket, KVBucketECSAccountPrefix), true } - -// refreshLease renews the lease at the revision the caller observed. A revision -// mismatch means the lease turned over between the read and here, so this node -// is no longer the holder. -func refreshLease(ctx context.Context, kv jetstream.KeyValue, key, holder string, rev uint64) bool { - _, err := kv.Update(ctx, key, []byte(holder), rev) - return err == nil -} - -// dropLease drops the lease only if it is still held at the revision the -// caller observed. A mismatch means another node has taken over, so this node -// must not delete a key it no longer owns. -func dropLease(ctx context.Context, kv jetstream.KeyValue, key string, rev uint64) error { - return kv.Delete(ctx, key, jetstream.LastRevision(rev)) -} diff --git a/spinifex/handlers/ecs/scheduler_leader_test.go b/spinifex/handlers/ecs/scheduler_leader_test.go deleted file mode 100644 index e82714e36..000000000 --- a/spinifex/handlers/ecs/scheduler_leader_test.go +++ /dev/null @@ -1,78 +0,0 @@ -package handlers_ecs - -import ( - "testing" - - "github.com/mulgadc/spinifex/spinifex/testutil" - - "github.com/stretchr/testify/require" -) - -func TestSchedulerLeaderRefresh(t *testing.T) { - t.Run("stale revision is refused", func(t *testing.T) { - _, _, js := testutil.StartTestJetStream(t) - kv, err := InitLeaderBucket(t.Context(), js) - require.NoError(t, err) - - staleRev, err := kv.Create(t.Context(), schedulerLeaderKey, []byte("node-a")) - require.NoError(t, err) - - //Bypass TTL - require.NoError(t, kv.Delete(t.Context(), schedulerLeaderKey)) - _, err = kv.Create(t.Context(), schedulerLeaderKey, []byte("node-b")) - require.NoError(t, err) - - require.False(t, refreshLease(t.Context(), kv, schedulerLeaderKey, "node-a", staleRev), - "node-a refreshed a lease that node-b now holds") - - entry, err := kv.Get(t.Context(), schedulerLeaderKey) - require.NoError(t, err) - require.Equal(t, "node-b", string(entry.Value())) - }) - - t.Run("current revision succeeds", func(t *testing.T) { - _, _, js := testutil.StartTestJetStream(t) - kv, err := InitLeaderBucket(t.Context(), js) - require.NoError(t, err) - - rev, err := kv.Create(t.Context(), schedulerLeaderKey, []byte("node-a")) - require.NoError(t, err) - - require.True(t, refreshLease(t.Context(), kv, schedulerLeaderKey, "node-a", rev)) - }) - - t.Run("release at a stale revision is refused", func(t *testing.T) { - _, _, js := testutil.StartTestJetStream(t) - kv, err := InitLeaderBucket(t.Context(), js) - require.NoError(t, err) - - staleRev, err := kv.Create(t.Context(), schedulerLeaderKey, []byte("node-a")) - require.NoError(t, err) - - //Bypass TTL - require.NoError(t, kv.Delete(t.Context(), schedulerLeaderKey)) - _, err = kv.Create(t.Context(), schedulerLeaderKey, []byte("node-b")) - require.NoError(t, err) - - require.Error(t, dropLease(t.Context(), kv, schedulerLeaderKey, staleRev), - "node-a released a lease that node-b now holds") - - entry, err := kv.Get(t.Context(), schedulerLeaderKey) - require.NoError(t, err) - require.Equal(t, "node-b", string(entry.Value())) - }) - - t.Run("release at the current revision succeeds", func(t *testing.T) { - _, _, js := testutil.StartTestJetStream(t) - kv, err := InitLeaderBucket(t.Context(), js) - require.NoError(t, err) - - rev, err := kv.Create(t.Context(), schedulerLeaderKey, []byte("node-a")) - require.NoError(t, err) - - require.NoError(t, dropLease(t.Context(), kv, schedulerLeaderKey, rev)) - - _, err = kv.Get(t.Context(), schedulerLeaderKey) - require.Error(t, err) - }) -} diff --git a/spinifex/handlers/ecs/service_test.go b/spinifex/handlers/ecs/service_test.go index 7f03e01df..93fc61fe6 100644 --- a/spinifex/handlers/ecs/service_test.go +++ b/spinifex/handlers/ecs/service_test.go @@ -23,6 +23,14 @@ func newTestService(t *testing.T) (*Service, *nats.Conn) { return NewService(nc, testRegion, "internal"), nc } +// subCount reads the subscription set under the lock the scheduler writes it +// with. Test-only, so the production struct keeps no accessor it does not need. +func (sc *Scheduler) subCount() int { + sc.mu.Lock() + defer sc.mu.Unlock() + return len(sc.subs) +} + func TestService_CreateCluster_Idempotent(t *testing.T) { svc, _ := newTestService(t) out, err := svc.CreateCluster(context.Background(), &ecs.CreateClusterInput{ClusterName: aws.String("web")}, testAccountID) @@ -444,9 +452,61 @@ func TestScheduler_AcquireLease_SingleLeader(t *testing.T) { svc := NewService(nc, testRegion, "") a := NewScheduler(nc, svc, "holder-a") b := NewScheduler(nc, svc, "holder-b") - assert.True(t, a.acquireOrRefresh(t.Context())) - assert.False(t, b.acquireOrRefresh(t.Context())) - assert.True(t, a.acquireOrRefresh(t.Context())) // refresh keeps leadership + require.True(t, a.lease.TryAcquire(t.Context())) + assert.False(t, b.lease.TryAcquire(t.Context())) + assert.True(t, a.lease.TryAcquire(t.Context()), "a holder re-attempting keeps leadership") +} + +// Leadership owns the Layer-2 bus subscriptions. This pins the wiring — gaining +// the lease subscribes, losing it unsubscribes — which is the regression this +// conversion could introduce and the one thing kvlease cannot check itself. +func TestScheduler_LeadershipWiresBusSubscriptions(t *testing.T) { + _, nc, _ := testutil.StartTestJetStream(t) + svc := NewService(nc, testRegion, "") + sc := NewScheduler(nc, svc, "holder-a") + require.NoError(t, sc.leaseErr) + + require.Zero(t, sc.subCount(), "subscriptions must not exist before election") + + require.True(t, sc.lease.TryAcquire(t.Context())) + assert.Equal(t, 3, sc.subCount(), "leader must own the register, heartbeat and task-state subscriptions") + + sc.lease.Release(t.Context()) + assert.Zero(t, sc.subCount(), "a released leader must drop its subscriptions") +} + +// A node that loses the election must not hold bus subscriptions: both leader +// and loser writing the same KV records is the split-brain this lease prevents. +func TestScheduler_LoserHoldsNoSubscriptions(t *testing.T) { + _, nc, _ := testutil.StartTestJetStream(t) + svc := NewService(nc, testRegion, "") + a := NewScheduler(nc, svc, "holder-a") + b := NewScheduler(nc, svc, "holder-b") + + require.True(t, a.lease.TryAcquire(t.Context())) + require.False(t, b.lease.TryAcquire(t.Context())) + assert.Zero(t, b.subCount()) + assert.Equal(t, 3, a.subCount()) +} + +func TestScheduler_RunReleasesLeaseOnShutdown(t *testing.T) { + _, nc, js := testutil.StartTestJetStream(t) + svc := NewService(nc, testRegion, "") + sc := NewScheduler(nc, svc, "holder-a") + require.NoError(t, sc.leaseErr) + + ctx, cancel := context.WithCancel(t.Context()) + done := make(chan struct{}) + go func() { defer close(done); sc.Run(ctx) }() + require.Eventually(t, sc.lease.Held, 2*time.Second, 20*time.Millisecond) + + cancel() + <-done + + kv, err := InitLeaderBucket(t.Context(), js) + require.NoError(t, err) + _, err = kv.Get(t.Context(), schedulerLeaderKey) + require.Error(t, err, "Run returned with the lease key still present") } // --- GPU placement dimension (Epic C2) ---