refactor: extract the KV leader lease into one package (#865)

2026-10-15 09:01:31 +00:00 · 2026-08-24 14:38:54 +10:00
parent 6fc7bdfab8
commit 4bbc96f1e0
4 changed files with 534 additions and 123 deletions
@@ -0,0 +1,54 @@
+package kvlease
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"time"
+
+	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
+)
+
+// Bounded wait for JetStream quorum on cold multi-node start. Vars (not consts) so tests can shrink them.
+var (
+	bucketRetryFor  = 60 * time.Second
+	bucketRetryStep = 1 * time.Second
+)
+
+// NATSBucket attaches to the named bucket, creating it only when genuinely absent, and
+// retrying while JetStream is still forming quorum. Callers with their own bucket initialiser
+// should pass that instead.
+func NATSBucket(nc *nats.Conn, bucket string, ttl time.Duration) BucketFunc {
+	return func(ctx context.Context) (jetstream.KeyValue, error) {
+		js, err := jetstream.New(nc)
+		if err != nil {
+			return nil, fmt.Errorf("kvlease: JetStream unavailable: %w", err)
+		}
+		deadline := time.Now().Add(bucketRetryFor)
+		for {
+			// CreateKeyValue reports "stream name already in use" when the bucket exists, so
+			// attach first and create only when genuinely absent.
+			kv, err := js.KeyValue(ctx, bucket)
+			if errors.Is(err, jetstream.ErrBucketNotFound) {
+				kv, err = js.CreateKeyValue(ctx, jetstream.KeyValueConfig{
+					Bucket:  bucket,
+					History: 1,
+					TTL:     ttl,
+				})
+			}
+			if err == nil {
+				return kv, nil
+			}
+			if time.Now().After(deadline) {
+				return nil, fmt.Errorf("kvlease: KV bucket %q unreachable after %s: %w", bucket, bucketRetryFor, err)
+			}
+			// A shutdown mid-wait must not sit out the remaining retry window.
+			select {
+			case <-ctx.Done():
+				return nil, ctx.Err()
+			case <-time.After(bucketRetryStep):
+			}
+		}
+	}
+}
@@ -0,0 +1,228 @@
+// Package kvlease elects a single holder of a JetStream KV key. The key's TTL
+// bounds crash recovery; background renewal keeps it alive while the holder
+// works, so a pass that outlives the TTL does not lose the lease underneath it.
+package kvlease
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"log/slog"
+	"sync"
+	"time"
+
+	"github.com/nats-io/nats.go/jetstream"
+)
+
+// releaseTimeout bounds the delete on the way out, which runs on a detached
+// context and so cannot inherit the caller's deadline.
+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.
+type BucketFunc func(context.Context) (jetstream.KeyValue, error)
+
+// Config describes a single lease. Name appears in logs; Retry applies only to Run.
+type Config struct {
+	Name   string
+	Bucket BucketFunc
+	Key    string
+	Holder string
+
+	TTL   time.Duration // must match the bucket's TTL
+	Renew time.Duration // gap between refreshes
+	Retry time.Duration // session mode only: non-leader re-attempt interval
+
+	// Edges fire per transition, never under the lock
+	OnGained func(context.Context) error
+	OnLost   func()
+}
+
+// Lease is a CAS-guarded claim on a single key. The zero value is unusable; call New.
+type Lease struct {
+	cfg  Config
+	mu   sync.Mutex
+	kv   jetstream.KeyValue
+	rev  uint64
+	held bool
+	stop context.CancelFunc
+}
+
+// New validates cfg and returns an unclaimed lease. Renew must leave room for a failed refresh,
+// so it may be at most half the TTL.
+func New(cfg Config) (*Lease, error) {
+	switch {
+	case cfg.Bucket == nil:
+		return nil, errors.New("kvlease: Bucket is required")
+	case cfg.Key == "":
+		return nil, errors.New("kvlease: Key is required")
+	case cfg.Holder == "":
+		return nil, errors.New("kvlease: Holder is required")
+	case cfg.TTL <= 0:
+		return nil, errors.New("kvlease: TTL must be positive")
+	case cfg.Renew <= 0 || cfg.Renew*2 > cfg.TTL:
+		return nil, fmt.Errorf("kvlease: Renew must be in (0, TTL/2]; got renew=%s ttl=%s", cfg.Renew, cfg.TTL)
+	}
+	if cfg.Name == "" {
+		cfg.Name = cfg.Key
+	}
+	if cfg.Retry <= 0 {
+		cfg.Retry = cfg.Renew
+	}
+	return &Lease{cfg: cfg}, nil
+}
+
+// Held reports whether this node currently holds the lease.
+func (l *Lease) Held() bool {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	return l.held
+}
+
+// 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 {
+	if l.Held() {
+		return true
+	}
+	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)
+		return false
+	}
+	rev, err := l.claim(ctx, kv)
+	if err != nil {
+		return false
+	}
+
+	renewCtx, stop := context.WithCancel(ctx)
+	l.mu.Lock()
+	if l.stop != nil {
+		l.stop()
+	}
+	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)
+	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)
+			l.Release(ctx)
+			return false
+		}
+	}
+	return true
+}
+
+// claim creates the key or adopts one this holder left behind. A restarted
+// process finds its own key still present and must refresh it rather than
+// sit out the TTL with no leader elected.
+func (l *Lease) claim(ctx context.Context, kv jetstream.KeyValue) (uint64, error) {
+	rev, err := kv.Create(ctx, l.cfg.Key, []byte(l.cfg.Holder))
+	if err == nil {
+		return rev, nil
+	}
+	entry, gerr := kv.Get(ctx, l.cfg.Key)
+	if gerr != nil {
+		return 0, gerr
+	}
+	if string(entry.Value()) != l.cfg.Holder {
+		return 0, errNotHolder
+	}
+	return kv.Update(ctx, l.cfg.Key, []byte(l.cfg.Holder), entry.Revision())
+}
+
+// renew refreshes the key at the revision last observed, until ctx is cancelled or a refresh
+// loses the CAS - which means the key expired and another node claimed it.
+func (l *Lease) renew(ctx context.Context) {
+	ticker := time.NewTicker(l.cfg.Renew)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			l.mu.Lock()
+			kv, rev, held := l.kv, l.rev, l.held
+			l.mu.Unlock()
+			if !held {
+				return
+			}
+			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)
+				l.markLost()
+				return
+			}
+			l.mu.Lock()
+			l.rev = next
+			l.mu.Unlock()
+		}
+	}
+}
+
+// markLost clears held state and fires OnLost. It reports whether this call was the transition,
+// so renewal and release never double-fire the edge.
+func (l *Lease) markLost() bool {
+	l.mu.Lock()
+	was := l.held
+	l.held = false
+	l.mu.Unlock()
+	if !was {
+		return false
+	}
+	if l.cfg.OnLost != nil {
+		l.cfg.OnLost()
+	}
+	return true
+}
+
+// Release deletes the key if this node still owns it at the revision it last observed.
+// It runs on a detached context: shutdown cancels ctx first, and skipping the delete would park
+// the lease for the full TTL.
+func (l *Lease) Release(ctx context.Context) {
+	l.mu.Lock()
+	kv, rev, stop := l.kv, l.rev, l.stop
+	l.stop = nil
+	l.mu.Unlock()
+	if stop != nil {
+		stop()
+	}
+	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)
+		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)
+	}
+}
+
+// Run holds the lease across many passes, re-attempting on a ticker while this node is not the
+// holder. A holder does not re-attempt: renewal keeps the key.
+func (l *Lease) Run(ctx context.Context) {
+	ticker := time.NewTicker(l.cfg.Retry)
+	defer ticker.Stop()
+	l.TryAcquire(ctx)
+	for {
+		select {
+		case <-ctx.Done():
+			l.Release(ctx)
+			return
+		case <-ticker.C:
+			l.TryAcquire(ctx)
+		}
+	}
+}
@@ -0,0 +1,238 @@
+package kvlease_test
+
+import (
+	"context"
+	"errors"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/mulgadc/spinifex/spinifex/kvlease"
+	"github.com/mulgadc/spinifex/spinifex/testutil"
+	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+const (
+	testBucket = "kvlease-test"
+	testKey    = "leader"
+)
+
+// newTestLease builds a lease over a short-lived bucket so a test can outlast
+// the TTL in a second rather than a minute. Bucket is built after mut so a test
+// that lengthens the TTL gets a bucket that matches it.
+func newTestLease(t *testing.T, nc *nats.Conn, holder string, mut func(*kvlease.Config)) *kvlease.Lease {
+	t.Helper()
+	cfg := kvlease.Config{
+		Name:   "test",
+		Key:    testKey,
+		Holder: holder,
+		TTL:    time.Second,
+		Renew:  250 * time.Millisecond,
+		Retry:  100 * time.Millisecond,
+	}
+	if mut != nil {
+		mut(&cfg)
+	}
+	cfg.Bucket = kvlease.NATSBucket(nc, testBucket, cfg.TTL)
+	l, err := kvlease.New(cfg)
+	require.NoError(t, err)
+	return l
+}
+
+func testKV(t *testing.T, nc *nats.Conn, ttl time.Duration) jetstream.KeyValue {
+	t.Helper()
+	kv, err := kvlease.NATSBucket(nc, testBucket, ttl)(t.Context())
+	require.NoError(t, err)
+	return kv
+}
+
+func TestTryAcquire_SecondHolderLoses(t *testing.T) {
+	_, nc, _ := testutil.StartTestJetStream(t)
+	a := newTestLease(t, nc, "node-a", nil)
+	b := newTestLease(t, nc, "node-b", nil)
+
+	require.True(t, a.TryAcquire(t.Context()), "first acquire must win and create the bucket")
+	assert.False(t, b.TryAcquire(t.Context()), "second acquire must lose while the lease is held")
+
+	a.Release(t.Context())
+	assert.True(t, b.TryAcquire(t.Context()), "acquire after release must win")
+	b.Release(t.Context())
+}
+
+// A restarted process finds its own key still present. Refusing it would leave the subsystem leaderless for the rest of
+// the TTL for no reason.
+func TestTryAcquire_AdoptsOwnKeyAfterRestart(t *testing.T) {
+	_, nc, _ := testutil.StartTestJetStream(t)
+	kv := testKV(t, nc, time.Second)
+
+	rev, err := kv.Create(t.Context(), testKey, []byte("node-a"))
+	require.NoError(t, err)
+
+	restarted := newTestLease(t, nc, "node-a", nil)
+	require.True(t, restarted.TryAcquire(t.Context()), "a holder must adopt the key it left behind")
+
+	entry, err := kv.Get(t.Context(), testKey)
+	require.NoError(t, err)
+	assert.Greater(t, entry.Revision(), rev, "adopting must refresh the key, resetting its TTL")
+	restarted.Release(t.Context())
+}
+
+// A pass can outlive the TTL on its own, so the key must be renewed while the pass runs. Without this a peer elects
+// itself mid-pass and both act at once.
+func TestRenew_KeySurvivesBeyondTTL(t *testing.T) {
+	_, nc, _ := testutil.StartTestJetStream(t)
+	a := newTestLease(t, nc, "node-a", nil)
+	require.True(t, a.TryAcquire(t.Context()))
+	defer a.Release(t.Context())
+
+	time.Sleep(2500 * time.Millisecond)
+
+	b := newTestLease(t, nc, "node-b", nil)
+	assert.False(t, b.TryAcquire(t.Context()), "lease expired mid-pass: a peer can now act concurrently")
+	assert.True(t, a.Held())
+}
+
+// If the key was lost anyway, the stale holder's release must not delete the successor's key -
+// that hand the lease to a third node silently.
+func TestRelease_DoesNotDeleteSuccessorKey(t *testing.T) {
+	_, nc, _ := testutil.StartTestJetStream(t)
+	a := newTestLease(t, nc, "node-a", func(c *kvlease.Config) {
+		c.TTL, c.Renew = time.Minute, 20*time.Second // no renewal during the test
+	})
+	require.True(t, a.TryAcquire(t.Context()))
+
+	kv := testKV(t, nc, time.Minute)
+	_, err := kv.Put(t.Context(), testKey, []byte("node-b"))
+	require.NoError(t, err)
+
+	a.Release(t.Context())
+
+	entry, err := kv.Get(t.Context(), testKey)
+	require.NoError(t, err, "stale release deleted the successor's key")
+	assert.Equal(t, "node-b", string(entry.Value()))
+}
+
+// Shutdown cancels the acquiring context first. A release bound to it would no-op and park the lease for the full TTL.
+func TestRelease_SurvivesCancelledContext(t *testing.T) {
+	_, nc, _ := testutil.StartTestJetStream(t)
+	ctx, cancel := context.WithCancel(t.Context())
+	a := newTestLease(t, nc, "node-a", nil)
+	require.True(t, a.TryAcquire(ctx))
+
+	cancel()
+	a.Release(ctx)
+
+	b := newTestLease(t, nc, "node-b", nil)
+	assert.True(t, b.TryAcquire(t.Context()), "lease must be free immediately, not at the TTL")
+	b.Release(t.Context())
+}
+
+func TestEdges_FireOncePerTransition(t *testing.T) {
+	_, nc, _ := testutil.StartTestJetStream(t)
+	var gained, lost atomic.Int32
+	a := newTestLease(t, nc, "node-a", func(c *kvlease.Config) {
+		c.OnGained = func(context.Context) error { gained.Add(1); return nil }
+		c.OnLost = func() { lost.Add(1) }
+	})
+
+	require.True(t, a.TryAcquire(t.Context()))
+	require.True(t, a.TryAcquire(t.Context()), "a holder re-attempting must stay held")
+	assert.EqualValues(t, 1, gained.Load(), "OnGained fired twice for one election")
+	assert.EqualValues(t, 0, lost.Load())
+
+	a.Release(t.Context())
+	assert.EqualValues(t, 1, lost.Load())
+	a.Release(t.Context())
+	assert.EqualValues(t, 1, lost.Load(), "a second release must not re-fire OnLost")
+}
+
+func TestRenew_LossFiresOnLost(t *testing.T) {
+	_, nc, _ := testutil.StartTestJetStream(t)
+	lost := make(chan struct{}, 4)
+	a := newTestLease(t, nc, "node-a", func(c *kvlease.Config) {
+		c.OnLost = func() { lost <- struct{}{} }
+	})
+	require.True(t, a.TryAcquire(t.Context()))
+
+	// 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 OnLost never fired")
+	}
+	assert.False(t, a.Held())
+}
+
+// A leader that cannot set up its own work is worse than no leader.
+func TestTryAcquire_OnGainedFailureStandsDown(t *testing.T) {
+	_, nc, _ := testutil.StartTestJetStream(t)
+	var lost atomic.Int32
+	a := newTestLease(t, nc, "node-a", func(c *kvlease.Config) {
+		c.OnGained = func(context.Context) error { return errors.New("subscribe failed") }
+		c.OnLost = func() { lost.Add(1) }
+	})
+
+	assert.False(t, a.TryAcquire(t.Context()), "a holder that cannot set up its work must not report leadership")
+	assert.False(t, a.Held())
+	assert.EqualValues(t, 1, lost.Load(), "OnLost must fire so the caller tears down partial statea")
+
+	b := newTestLease(t, nc, "node-b", nil)
+	assert.True(t, b.TryAcquire(t.Context()), "standing down must free the key immediately")
+	b.Release(t.Context())
+}
+
+func TestNew_RejectUnworkableConfig(t *testing.T) {
+	base := func() kvlease.Config {
+		return kvlease.Config{
+			Bucket: func(context.Context) (jetstream.KeyValue, error) { return nil, nil },
+			Key:    testKey,
+			Holder: "node-a",
+			TTL:    time.Minute,
+			Renew:  20 * time.Second,
+		}
+	}
+	_, err := kvlease.New(base())
+	require.NoError(t, err)
+
+	cases := map[string]func(*kvlease.Config){
+		"no bucket":           func(c *kvlease.Config) { c.Bucket = nil },
+		"no key":              func(c *kvlease.Config) { c.Key = "" },
+		"no holder":           func(c *kvlease.Config) { c.Holder = "" },
+		"no ttl":              func(c *kvlease.Config) { c.TTL = 0 },
+		"renew equals ttl":    func(c *kvlease.Config) { c.Renew = c.TTL },
+		"renew over half ttl": func(c *kvlease.Config) { c.Renew = 31 * time.Second },
+	}
+	for name, mut := range cases {
+		t.Run(name, func(t *testing.T) {
+			cfg := base()
+			mut(&cfg)
+			_, err := kvlease.New(cfg)
+			assert.Error(t, err)
+		})
+	}
+}
+
+// Session mode: a loser must take over on its retry tick rather than wait out the TTL, which is what makes a rolling restart a sub-second gap.
+func TestRun_LoserTakesOverAfterLeaderStops(t *testing.T) {
+	_, nc, _ := testutil.StartTestJetStream(t)
+	aCtx, stopA := context.WithCancel(t.Context())
+	a := newTestLease(t, nc, "node-a", nil)
+	go a.Run(aCtx)
+	require.Eventually(t, a.Held, 2*time.Second, 20*time.Millisecond)
+
+	b := newTestLease(t, nc, "node-b", nil)
+	go b.Run(t.Context())
+	time.Sleep(300 * time.Millisecond)
+	require.False(t, b.Held(), "two holders at once")
+
+	stopA()
+	require.Eventually(t, b.Held, 2*time.Second, 20*time.Millisecond, "loser must take over on its retry tick, not at the TTL")
+}
@@ -2,24 +2,17 @@ package reconcile

 import (
 	"context"
-	"errors"
-	"github.com/mulgadc/spinifex/spinifex/otelsetup"
 	"log/slog"
-	"sync"
 	"time"

+	"github.com/mulgadc/spinifex/spinifex/kvlease"
 	"github.com/nats-io/nats.go"
-	"github.com/nats-io/nats.go/jetstream"
 )

 // Single CAS-elected leader key; TTL bounds crash-recovery.
 const (
 	KVBucketVPCDReconcile = "spinifex-vpcd-reconcile"
 	reconcileLeaderKey    = "leader"
-
-	// reconcileReleaseTimeout bounds the lock delete on the way out, which runs
-	// on a detached context and so cannot inherit the caller's deadline.
-	reconcileReleaseTimeout = 5 * time.Second
 )

 // Leader-key lifetime. Vars (not consts) so tests can shrink them.
@@ -32,130 +25,28 @@ var (
 	reconcileLeaderRenew = 20 * time.Second
 )

-// Bounded wait for JetStream quorum on cold multi-node start. Vars (not
-// consts) so tests can shrink them.
-var (
-	leaderRetryFor  = 60 * time.Second
-	leaderRetryStep = 1 * time.Second
-)
-
 // AcquireLeader elects one leader on the named lock bucket. Independent
 // reconcile loops pass distinct buckets so they never share a single mutex: the
 // gateway quota reconcile must not block vpcd's network reconcile, and vice
 // versa.
 func AcquireLeader(ctx context.Context, nc *nats.Conn, bucket, holder string) (func(), bool) {
-	js, err := jetstream.New(nc)
+	lease, err := kvlease.New(kvlease.Config{
+		Name:   "reconcile/lock",
+		Bucket: kvlease.NATSBucket(nc, bucket, reconcileLeaderTTL),
+		Key:    reconcileLeaderKey,
+		Holder: holder,
+		TTL:    reconcileLeaderTTL,
+		Renew:  reconcileLeaderRenew,
+	})
 	if err != nil {
-		slog.Error("reconcile/lock: JetStream unavailable, skipping reconcile",
+		slog.ErrorContext(ctx, "reconcile/lock: lease config invalid",
 			"holder", holder, "bucket", bucket, "err", err)
 		return nil, false
 	}
-
-	var kv jetstream.KeyValue
-	deadline := time.Now().Add(leaderRetryFor)
-	for {
-		// Get-or-create: CreateKeyValue returns "stream name already in use" if
-		// the bucket exists; attach first, create only when genuinely absent.
-		kv, err = js.KeyValue(ctx, bucket)
-		if errors.Is(err, jetstream.ErrBucketNotFound) {
-			kv, err = js.CreateKeyValue(ctx, jetstream.KeyValueConfig{
-				Bucket:  bucket,
-				History: 1,
-				TTL:     reconcileLeaderTTL,
-			})
-		}
-		if err == nil {
-			break
-		}
-		if time.Now().After(deadline) {
-			slog.Error("reconcile/lock: JetStream KV unreachable after retry, skipping reconcile",
-				"holder", holder, "bucket", bucket, "waited_ms", otelsetup.Millis(leaderRetryFor), "err", err)
-			return nil, false
-		}
-		slog.Debug("reconcile/lock: JetStream KV not ready, retrying", "holder", holder, "bucket", bucket, "err", err)
-		// A shutdown mid-wait must not sit out the remaining retry window.
-		select {
-		case <-ctx.Done():
-			slog.Info("reconcile/lock: cancelled while waiting for JetStream KV", "holder", holder, "bucket", bucket)
-			return nil, false
-		case <-time.After(leaderRetryStep):
-		}
-	}
-
-	rev, err := kv.Create(ctx, reconcileLeaderKey, []byte(holder))
-	if err != nil {
-		slog.Info("reconcile/lock: another holder is leader, skipping reconcile", "holder", holder, "bucket", bucket, "err", err)
+	if !lease.TryAcquire(ctx) {
+		slog.InfoContext(ctx, "reconcile/lock: not leader, skipping reconcile",
+			"holder", holder, "bucket", bucket)
 		return nil, false
 	}
-
-	slog.Info("reconcile/lock: elected", "holder", holder, "bucket", bucket)
-	lease := &leaderLease{kv: kv, rev: rev}
-	renewCtx, stopRenew := context.WithCancel(ctx)
-	go lease.renew(renewCtx, holder, bucket)
-
-	return func() {
-		stopRenew()
-		rev, held := lease.revision()
-		if !held {
-			// Renewal lost the key, so another node may already hold it. An
-			// unguarded delete here would drop that node's lock, not ours.
-			slog.Warn("reconcile/lock: lock already lost before release; not deleting", "holder", holder, "bucket", bucket)
-			return
-		}
-		// Release outlives ctx: shutdown is the common reason to release, and
-		// skipping the delete would park the lock for the full TTL and stall
-		// every other node's reconcile.
-		releaseCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), reconcileReleaseTimeout)
-		defer cancel()
-		if err := kv.Delete(releaseCtx, reconcileLeaderKey, jetstream.LastRevision(rev)); err != nil {
-			slog.Warn("reconcile/lock: failed to release lock (TTL will reap)", "holder", holder, "bucket", bucket, "err", err)
-		}
-	}, true
-}
-
-// leaderLease tracks the revision of a held leader key. Renewal and the final
-// delete are both CAS-guarded on it, so a pass that outlives the TTL can never
-// delete the key of the node that took over from it.
-type leaderLease struct {
-	mu   sync.Mutex
-	kv   jetstream.KeyValue
-	rev  uint64
-	lost bool
-}
-
-// revision returns the current revision and whether the lease is still held.
-func (l *leaderLease) revision() (uint64, bool) {
-	l.mu.Lock()
-	defer l.mu.Unlock()
-	return l.rev, !l.lost
-}
-
-// renew refreshes the leader key until ctx is cancelled or a refresh loses the
-// CAS, which means the key expired and another node claimed it.
-func (l *leaderLease) renew(ctx context.Context, holder, bucket string) {
-	ticker := time.NewTicker(reconcileLeaderRenew)
-	defer ticker.Stop()
-	for {
-		select {
-		case <-ctx.Done():
-			return
-		case <-ticker.C:
-			rev, held := l.revision()
-			if !held {
-				return
-			}
-			next, err := l.kv.Update(ctx, reconcileLeaderKey, []byte(holder), rev)
-			if err != nil {
-				l.mu.Lock()
-				l.lost = true
-				l.mu.Unlock()
-				slog.Warn("reconcile/lock: leader key renewal failed; another node may take over mid-pass",
-					"holder", holder, "bucket", bucket, "err", err)
-				return
-			}
-			l.mu.Lock()
-			l.rev = next
-			l.mu.Unlock()
-		}
-	}
+	return func() { lease.Release(ctx) }, true
 }