diff --git a/.gitignore b/.gitignore index c8e5d9c97..7fe86d71f 100644 --- a/.gitignore +++ b/.gitignore @@ -34,6 +34,7 @@ Thumbs.db # Environment .env .venv/ +.cache/ # VScode .vscode/ diff --git a/spinifex/handlers/ecs/scheduler.go b/spinifex/handlers/ecs/scheduler.go index f4cce2eab..9e5778ae3 100644 --- a/spinifex/handlers/ecs/scheduler.go +++ b/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)) +} diff --git a/spinifex/handlers/ecs/scheduler_leader_test.go b/spinifex/handlers/ecs/scheduler_leader_test.go new file mode 100644 index 000000000..e82714e36 --- /dev/null +++ b/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) + }) +}