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

@@ -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) {

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)

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.

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 }
+}

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
+}

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"`)
+}

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,
 	})