Fix/ecs scheduler leader cas (#861)

* fix(ecs): guard scheduler leader refresh and release with a revision CAS

* Ignore .cache

* update comment

.gitignore

@@ -34,6 +34,7 @@ Thumbs.db
 # Environment
 .env
 .venv/
+.cache/
 
 # VScode
 .vscode/

spinifex/handlers/ecs/scheduler.go

@@ -141,11 +141,7 @@ func (sc *Scheduler) acquireOrRefresh(ctx context.Context) bool {
 	if string(entry.Value()) != sc.holder {
 		return false // someone else holds it
 	}
-	// We hold it — refresh to reset the TTL.
-	if _, err := kv.Put(ctx, schedulerLeaderKey, []byte(sc.holder)); err != nil {
-		return false
-	}
-	return true
+	return refreshLease(ctx, kv, schedulerLeaderKey, sc.holder, entry.Revision())
 }
 
 // relinquish drops subscriptions and deletes the lease key on shutdown.
@@ -164,7 +160,7 @@ func (sc *Scheduler) relinquish() {
 		return
 	}
 	if entry, gerr := kv.Get(ctx, schedulerLeaderKey); gerr == nil && string(entry.Value()) == sc.holder {
-		_ = kv.Delete(ctx, schedulerLeaderKey)
+		_ = dropLease(ctx, kv, schedulerLeaderKey, entry.Revision())
 	}
 }
 
@@ -320,3 +316,18 @@ 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))
+}

spinifex/handlers/ecs/scheduler_leader_test.go

@@ -0,0 +1,78 @@
+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)
+	})
+}