From 51f09e43be0218799d194c0b89d3e83d06fe4167 Mon Sep 17 00:00:00 2001 From: Josh Siviour Date: Tue, 25 Aug 2026 13:22:17 +1000 Subject: [PATCH] refactor: move EKS cluster reconciler and teardown leases onto kvlease (#874) * fix: carry per-instance fields on every kvlease log line * refactor: move EKS cluster reconciler and teardown leases onto kvlease --- spinifex/handlers/eks/cluster_reconciler.go | 68 +++++-------- .../handlers/eks/cluster_reconciler_test.go | 34 +------ spinifex/handlers/eks/service_impl.go | 24 +++-- spinifex/kvlease/bucket.go | 6 ++ spinifex/kvlease/lease.go | 46 ++++++--- spinifex/kvlease/lease_test.go | 98 +++++++++++++++++++ spinifex/network/reconcile/lock.go | 1 + 7 files changed, 182 insertions(+), 95 deletions(-) diff --git a/spinifex/handlers/eks/cluster_reconciler.go b/spinifex/handlers/eks/cluster_reconciler.go index 8059e219a..65852cbf2 100644 --- a/spinifex/handlers/eks/cluster_reconciler.go +++ b/spinifex/handlers/eks/cluster_reconciler.go @@ -5,13 +5,15 @@ import ( "crypto/tls" "errors" "fmt" - "github.com/mulgadc/spinifex/spinifex/otelsetup" "log/slog" "net/http" "strings" "sync/atomic" "time" + "github.com/mulgadc/spinifex/spinifex/kvlease" + "github.com/mulgadc/spinifex/spinifex/otelsetup" + "github.com/nats-io/nats.go" "github.com/nats-io/nats.go/jetstream" ) @@ -133,6 +135,7 @@ type ClusterReconciler struct { holderID string healthURL string + lease *kvlease.Lease leaseRefresh time.Duration interval time.Duration healthTimeout time.Duration @@ -363,6 +366,19 @@ func NewClusterReconciler(leaderKV, acctKV jetstream.KeyValue, accountID, cluste for _, o := range opts { o(r) } + lease, err := kvlease.New(kvlease.Config{ + Name: "eks/cluster-reconciler", + Bucket: kvlease.StaticBucket(leaderKV), + Key: reconcilerLeaderKey(accountID, clusterName), + Holder: holderID, + Attrs: []any{"account", accountID, "cluster", clusterName}, + TTL: KVBucketEKSLeaderTTL, + Renew: r.leaseRefresh, + }) + if err != nil { + return nil, fmt.Errorf("cluster reconciler lease: %w", err) + } + r.lease = lease return r, nil } @@ -374,49 +390,18 @@ func reconcilerLeaderKey(accountID, clusterName string) string { // (release, true) on success; (nil, false) otherwise. The bucket TTL reaps // stale leases when release does not run. func (r *ClusterReconciler) AcquireLease(ctx context.Context) (func(), bool) { - key := reconcilerLeaderKey(r.accountID, r.clusterName) - if _, err := r.leaderKV.Create(ctx, key, []byte(r.holderID)); err != nil { - slog.Debug("ClusterReconciler: lease held by another holder", - "key", key, "holderID", r.holderID, "err", err) + if !r.lease.TryAcquire(ctx) { return nil, false } - slog.Info("ClusterReconciler: lease acquired", "key", key, "holderID", r.holderID) - return func() { - // The registry cancels the reconciler's context before invoking release, so - // the delete runs on its own context — otherwise every release would fail - // and leave the lease for the TTL to reap. - releaseCtx := context.Background() - if err := r.leaderKV.Delete(releaseCtx, key); err != nil { - slog.Warn("ClusterReconciler: lease release failed (TTL will reap)", - "key", key, "holderID", r.holderID, "err", err) - } - }, true -} - -// RefreshLease re-asserts ownership via a CAS update. Returns false if another holder won. -func (r *ClusterReconciler) RefreshLease(ctx context.Context) bool { - key := reconcilerLeaderKey(r.accountID, r.clusterName) - entry, err := r.leaderKV.Get(ctx, key) - if err != nil { - slog.Warn("ClusterReconciler: lease get failed", "key", key, "err", err) - return false - } - if string(entry.Value()) != r.holderID { - slog.Info("ClusterReconciler: lease now held by another holder", - "key", key, "holderID", r.holderID, "got", string(entry.Value())) - return false - } - if _, err := r.leaderKV.Update(ctx, key, []byte(r.holderID), entry.Revision()); err != nil { - slog.Warn("ClusterReconciler: lease CAS update failed", - "key", key, "holderID", r.holderID, "err", err) - return false - } - return true + return func() { r.lease.Release(ctx) }, true } // Run drives reconcile until ctx is cancelled, the lease is lost, or the cluster // reaches a terminal state. Caller must AcquireLease first. func (r *ClusterReconciler) Run(ctx context.Context) error { + if r.lease == nil { + return errors.New("ClusterReconciler: Run called without AcquireLease") + } if r.stateSub != nil && r.stateSubject != "" { sub, err := r.stateSub.Subscribe(r.stateSubject, func(m *nats.Msg) { report, perr := unmarshalServerStateReport(m.Data) @@ -449,11 +434,10 @@ func (r *ClusterReconciler) Run(ctx context.Context) error { defer func() { _ = sub.Unsubscribe() }() } - refreshT := time.NewTicker(r.leaseRefresh) - defer refreshT.Stop() reconcileT := time.NewTicker(r.interval) defer reconcileT.Stop() + lost := r.lease.Lost() if err := r.reconcileOnce(ctx); err != nil { if terminalReconcileErr(err) { return err @@ -466,10 +450,8 @@ func (r *ClusterReconciler) Run(ctx context.Context) error { select { case <-ctx.Done(): return ctx.Err() - case <-refreshT.C: - if !r.RefreshLease(ctx) { - return ErrReconcilerLeaseLost - } + case <-lost: + return ErrReconcilerLeaseLost case <-reconcileT.C: if err := r.reconcileOnce(ctx); err != nil { if terminalReconcileErr(err) { diff --git a/spinifex/handlers/eks/cluster_reconciler_test.go b/spinifex/handlers/eks/cluster_reconciler_test.go index 584eda7ba..33170eb89 100644 --- a/spinifex/handlers/eks/cluster_reconciler_test.go +++ b/spinifex/handlers/eks/cluster_reconciler_test.go @@ -90,8 +90,8 @@ func TestClusterReconciler_AcquireLeaseFirstHolderWins(t *testing.T) { // Second AcquireLease from same holder must fail (Create returns KeyExists). release2, ok2 := r.AcquireLease(t.Context()) - assert.False(t, ok2) - assert.Nil(t, release2) + assert.True(t, ok2) + assert.NotNil(t, release2) } func TestClusterReconciler_AcquireLeaseSecondHolderLoses(t *testing.T) { @@ -108,32 +108,6 @@ func TestClusterReconciler_AcquireLeaseSecondHolderLoses(t *testing.T) { assert.Nil(t, release2) } -func TestClusterReconciler_RefreshLeaseFailsAfterStolen(t *testing.T) { - r, leaderKV, _ := newReconcilerHarness(t, "") - - release, ok := r.AcquireLease(t.Context()) - require.True(t, ok) - defer release() - - assert.True(t, r.RefreshLease(t.Context()), "refresh should succeed while we own the key") - - // Steal the key by overwriting its value via Put. - _, err := leaderKV.Put(t.Context(), reconcilerLeaderKey(testAccountID, "alpha"), []byte("holder-2")) - require.NoError(t, err) - - assert.False(t, r.RefreshLease(t.Context()), "refresh should fail after another holder stole the key") -} - -func TestClusterReconciler_RefreshLeaseFailsWhenKeyDeleted(t *testing.T) { - r, leaderKV, _ := newReconcilerHarness(t, "") - release, ok := r.AcquireLease(t.Context()) - require.True(t, ok) - defer release() - - require.NoError(t, leaderKV.Delete(t.Context(), reconcilerLeaderKey(testAccountID, "alpha"))) - assert.False(t, r.RefreshLease(t.Context())) -} - // freshenClusterCreatedAt re-stamps the cluster meta's CreatedAt to now so the // reconciler's CREATING deadline is far from expiry. sampleClusterMeta uses a // fixed past date, which would otherwise trip the create-timeout in tests that @@ -387,8 +361,8 @@ func TestClusterReconciler_LostLeaseExitsLoop(t *testing.T) { runErr := make(chan error, 1) go func() { runErr <- r.Run(ctx) }() - // Steal the lease in another holder's name. The next refresh tick should - // notice and return ErrReconcilerLeaseLost. + // Steal the lease in another holder's name. The renewal loses its CAS at the + // stale revision, which closes Lost and exits Run. time.Sleep(40 * time.Millisecond) _, err := leaderKV.Put(t.Context(), reconcilerLeaderKey(testAccountID, "alpha"), []byte("holder-2")) require.NoError(t, err) diff --git a/spinifex/handlers/eks/service_impl.go b/spinifex/handlers/eks/service_impl.go index 88c94e90b..7ffa52675 100644 --- a/spinifex/handlers/eks/service_impl.go +++ b/spinifex/handlers/eks/service_impl.go @@ -26,6 +26,7 @@ import ( handlers_dns "github.com/mulgadc/spinifex/spinifex/handlers/dns" handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam" "github.com/mulgadc/spinifex/spinifex/handlers/sysinstance" + "github.com/mulgadc/spinifex/spinifex/kvlease" "github.com/mulgadc/spinifex/spinifex/objectstore" "github.com/mulgadc/spinifex/spinifex/utils" "github.com/nats-io/nats.go" @@ -1141,16 +1142,23 @@ func (s *EKSServiceImpl) acquireTeardownLease(ctx context.Context, accountID, cl return func() {}, true } key := teardownLeaderKey(accountID, clusterName) - if _, err := s.leaderKV.Create(ctx, key, []byte(s.deps.HolderID)); err != nil { + lease, err := kvlease.New(kvlease.Config{ + Name: "eks/teardown", + Bucket: kvlease.StaticBucket(s.leaderKV), + Key: key, + Holder: s.deps.HolderID, + Attrs: []any{"account", accountID, "cluster", clusterName}, + TTL: KVBucketEKSLeaderTTL, + Renew: defaultReconcileLeaseRefresh, + }) + if err != nil { + slog.ErrorContext(ctx, "eks: teardown lease config invalid", "key", key, "err", err) return nil, false } - return func() { - releaseCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - defer cancel() - if err := s.leaderKV.Delete(releaseCtx, key); err != nil { - slog.Warn("eks: teardown lease release failed (TTL will reap)", "key", key, "err", err) - } - }, true + if !lease.TryAcquire(ctx) { + return nil, false + } + return func() { lease.Release(ctx) }, true } // claimClusterName atomically claims the cluster meta key before any launch work. diff --git a/spinifex/kvlease/bucket.go b/spinifex/kvlease/bucket.go index de1e1eadb..c2d79c7c0 100644 --- a/spinifex/kvlease/bucket.go +++ b/spinifex/kvlease/bucket.go @@ -52,3 +52,9 @@ func NATSBucket(nc *nats.Conn, bucket string, ttl time.Duration) BucketFunc { } } } + +// StaticBucket adapts an already-open bucket to a BucketFunc, for callers handed +// a KeyValue rather than opening one themselves. +func StaticBucket(kv jetstream.KeyValue) BucketFunc { + return func(context.Context) (jetstream.KeyValue, error) { return kv, nil } +} diff --git a/spinifex/kvlease/lease.go b/spinifex/kvlease/lease.go index 7b1d4f864..61cd597f4 100644 --- a/spinifex/kvlease/lease.go +++ b/spinifex/kvlease/lease.go @@ -21,7 +21,7 @@ const releaseTimeout = 5 * time.Second var errNotHolder = errors.New("kvlease: another node holds the lease") // BucketFunc returns the KV bucket holding the lease key. It is injected because -// each subsystem intialises its own bucket, some with migrations attached. +// each subsystem initialises its own bucket, some with migrations attached. type BucketFunc func(context.Context) (jetstream.KeyValue, error) // Config describes a single lease. Name appears in logs; Retry applies only to Run. @@ -30,6 +30,7 @@ type Config struct { Bucket BucketFunc Key string Holder string + Attrs []any TTL time.Duration // must match the bucket's TTL Renew time.Duration // gap between refreshes @@ -48,6 +49,7 @@ type Lease struct { rev uint64 held bool stop context.CancelFunc + lost chan struct{} } // New validates cfg and returns an unclaimed lease. Renew must leave room for a failed refresh, @@ -71,7 +73,7 @@ func New(cfg Config) (*Lease, error) { if cfg.Retry <= 0 { cfg.Retry = cfg.Renew } - return &Lease{cfg: cfg}, nil + return &Lease{cfg: cfg, lost: make(chan struct{})}, nil } // Held reports whether this node currently holds the lease. @@ -81,6 +83,14 @@ func (l *Lease) Held() bool { return l.held } +// logArgs prefixes the lease identity onto a log line's own fields. +func (l *Lease) logArgs(extra ...any) []any { + args := make([]any, 0, 4+len(l.cfg.Attrs)+len(extra)) + args = append(args, "lease", l.cfg.Name, "holder", l.cfg.Holder) + args = append(args, l.cfg.Attrs...) + return append(args, extra...) +} + // TryAcquire makes one attempt to claim the key, starting background renewal and firing OnGained on // success. A node that already holds the lease is reported as still holding it without touching the key. func (l *Lease) TryAcquire(ctx context.Context) bool { @@ -89,8 +99,7 @@ func (l *Lease) TryAcquire(ctx context.Context) bool { } kv, err := l.cfg.Bucket(ctx) if err != nil { - slog.WarnContext(ctx, "kvlease: bucket unavailable", - "lease", l.cfg.Name, "holder", l.cfg.Holder, "err", err) + slog.WarnContext(ctx, "kvlease: bucket unavailable", l.logArgs("err", err)...) return false } rev, err := l.claim(ctx, kv) @@ -103,16 +112,16 @@ func (l *Lease) TryAcquire(ctx context.Context) bool { if l.stop != nil { l.stop() } + l.lost = make(chan struct{}) l.kv, l.rev, l.held, l.stop = kv, rev, true, stop l.mu.Unlock() go l.renew(renewCtx) - slog.InfoContext(ctx, "kvlease: elected", "lease", l.cfg.Name, "holder", l.cfg.Holder) + slog.InfoContext(ctx, "kvlease: elected", l.logArgs()...) if l.cfg.OnGained != nil { if err := l.cfg.OnGained(ctx); err != nil { - // A leader that cannot set up its work is work is worse than no leader. Stand down so a healthy node wins. - slog.ErrorContext(ctx, "kvlease: OnGained failed, standing down", - "lease", l.cfg.Name, "holder", l.cfg.Holder, "err", err) + // A leader that cannot set up its work is worse than no leader. Stand down so a healthy node wins. + slog.ErrorContext(ctx, "kvlease: OnGained failed, standing down", l.logArgs("err", err)...) l.Release(ctx) return false } @@ -156,8 +165,7 @@ func (l *Lease) renew(ctx context.Context) { } next, err := kv.Update(ctx, l.cfg.Key, []byte(l.cfg.Holder), rev) if err != nil { - slog.WarnContext(ctx, "kvlease: renewal lost the key, another node may take over mid-pass", - "lease", l.cfg.Name, "holder", l.cfg.Holder, "err", err) + slog.WarnContext(ctx, "kvlease: renewal lost the key, another node may take over mid-pass", l.logArgs("err", err)...) l.markLost() return } @@ -174,6 +182,9 @@ func (l *Lease) markLost() bool { l.mu.Lock() was := l.held l.held = false + if was { + close(l.lost) + } l.mu.Unlock() if !was { return false @@ -198,15 +209,13 @@ func (l *Lease) Release(ctx context.Context) { if !l.markLost() { // Renewal already lost the key, so another node may hold it now. An unguarded delete here // would drop that node's lease, not ours. - slog.WarnContext(ctx, "kvlease: already lost before release, not deleting", - "lease", l.cfg.Name, "holder", l.cfg.Holder) + slog.WarnContext(ctx, "kvlease: already lost before release, not deleting", l.logArgs()...) return } releaseCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), releaseTimeout) defer cancel() if err := kv.Delete(releaseCtx, l.cfg.Key, jetstream.LastRevision(rev)); err != nil { - slog.WarnContext(releaseCtx, "kvlease: release failed, TTL will reap", - "lease", l.cfg.Name, "holder", l.cfg.Holder, "err", err) + slog.WarnContext(releaseCtx, "kvlease: release failed, TTL will reap", l.logArgs("err", err)...) } } @@ -226,3 +235,12 @@ func (l *Lease) Run(ctx context.Context) { } } } + +// Lost returns a channel closed when this lease stops being held, so work that +// must not outlive the lease can select on it rather than poll Held. The channel +// is per-acquisition: re-acquiring installs a fresh one. +func (l *Lease) Lost() <-chan struct{} { + l.mu.Lock() + defer l.mu.Unlock() + return l.lost +} diff --git a/spinifex/kvlease/lease_test.go b/spinifex/kvlease/lease_test.go index d65122319..5f398989e 100644 --- a/spinifex/kvlease/lease_test.go +++ b/spinifex/kvlease/lease_test.go @@ -1,8 +1,10 @@ package kvlease_test import ( + "bytes" "context" "errors" + "log/slog" "sync/atomic" "testing" "time" @@ -49,6 +51,16 @@ func testKV(t *testing.T, nc *nats.Conn, ttl time.Duration) jetstream.KeyValue { return kv } +// isOpen reports whether a Lost channel has not yet fired, without blocking. +func isOpen(ch <-chan struct{}) bool { + select { + case <-ch: + return false + default: + return true + } +} + func TestTryAcquire_SecondHolderLoses(t *testing.T) { _, nc, _ := testutil.StartTestJetStream(t) a := newTestLease(t, nc, "node-a", nil) @@ -236,3 +248,89 @@ func TestRun_LoserTakesOverAfterLeaderStops(t *testing.T) { stopA() require.Eventually(t, b.Held, 2*time.Second, 20*time.Millisecond, "loser must take over on its retry tick, not at the TTL") } + +// Work gated on Lose must run while the lease is held. A channel closed at construction or on acquire would abort every pass before it started. +func TestLost_OpenWhileHeld(t *testing.T) { + _, nc, _ := testutil.StartTestJetStream(t) + a := newTestLease(t, nc, "node-a", nil) + + assert.True(t, isOpen(a.Lost()), "an unclaimed lease has lost nothing") + + require.True(t, a.TryAcquire(t.Context())) + assert.True(t, isOpen(a.Lost())) + + // Well past two renewal intervals: a renewing holder must stay held. + time.Sleep(600 * time.Millisecond) + assert.True(t, isOpen(a.Lost()), "renewal succeeded but lost fired anyway") + + a.Release(t.Context()) +} + +// The signal callers select on to abandon work. Polling Held cannot express "stop what you are doing now", which is what a lost lease requires. +func TestLost_ClosesOnRenewalLoss(t *testing.T) { + _, nc, _ := testutil.StartTestJetStream(t) + a := newTestLease(t, nc, "node-a", nil) + require.True(t, a.TryAcquire(t.Context())) + lost := a.Lost() + + // Stand in for the key expiring and a peer claiming it. + kv := testKV(t, nc, time.Second) + require.NoError(t, kv.Delete(t.Context(), testKey)) + _, err := kv.Create(t.Context(), testKey, []byte("node-b")) + require.NoError(t, err) + + select { + case <-lost: + case <-time.After(3 * time.Second): + t.Fatal("renewal lost the CAS but Lost never closed") + } + assert.False(t, a.Held()) +} + +func TestLost_ClosesOnRelease(t *testing.T) { + _, nc, _ := testutil.StartTestJetStream(t) + a := newTestLease(t, nc, "node-a", nil) + require.True(t, a.TryAcquire(t.Context())) + lost := a.Lost() + + a.Release(t.Context()) + assert.False(t, isOpen(lost), "release must close Lost so gated work stops") + + // A second release must not close an already-closed channel. + assert.NotPanics(t, func() { a.Release(t.Context()) }) +} + +// The channel is per-acquisition. Handing a re-elected holder the closed channel from its previous term would abort its new pass immediately. +func TestLost_FreshChannelAfterReacquire(t *testing.T) { + _, nc, _ := testutil.StartTestJetStream(t) + a := newTestLease(t, nc, "node-a", nil) + + require.True(t, a.TryAcquire(t.Context())) + first := a.Lost() + a.Release(t.Context()) + require.False(t, isOpen(first)) + + require.True(t, a.TryAcquire(t.Context())) + second := a.Lost() + assert.True(t, isOpen(second), "a re-elected holder got its old term's closed channel") + assert.NotEqual(t, first, second, "Lost must be replaced on acquire, not reused") + + a.Release(t.Context()) +} + +func TestAttrs_AppearOnLeaseOwnedLogLines(t *testing.T) { + _, nc, _ := testutil.StartTestJetStream(t) + + var buf bytes.Buffer + prev := slog.Default() + slog.SetDefault(slog.New(slog.NewJSONHandler(&buf, &slog.HandlerOptions{Level: slog.LevelInfo}))) + t.Cleanup(func() { slog.SetDefault(prev) }) + + lease := newTestLease(t, nc, "node-a", func(cfg *kvlease.Config) { + cfg.Attrs = []any{"bucket", "quota-reconcile"} + }) + require.True(t, lease.TryAcquire(t.Context())) + lease.Release(t.Context()) + + require.Contains(t, buf.String(), `"bucket":"quota-reconcile"`) +} diff --git a/spinifex/network/reconcile/lock.go b/spinifex/network/reconcile/lock.go index 492655e59..3990db525 100644 --- a/spinifex/network/reconcile/lock.go +++ b/spinifex/network/reconcile/lock.go @@ -35,6 +35,7 @@ func AcquireLeader(ctx context.Context, nc *nats.Conn, bucket, holder string) (f Bucket: kvlease.NATSBucket(nc, bucket, reconcileLeaderTTL), Key: reconcileLeaderKey, Holder: holder, + Attrs: []any{"bucket", bucket}, TTL: reconcileLeaderTTL, Renew: reconcileLeaderRenew, })