From 602247af8018bd2811ecb1ca58adcae76ad35a65 Mon Sep 17 00:00:00 2001 From: Josh Siviour Date: Tue, 25 Aug 2026 15:37:07 +1000 Subject: [PATCH] refactor: promote the daemon CAS primitives into kvutil (#878) * refactor: promote the daemon CAS primitives into kvutil * fix: do not warn about a lost lease when the node never held one --- spinifex/daemon/jetstream.go | 175 +++--------------- spinifex/kvlease/lease.go | 6 +- spinifex/kvlease/lease_test.go | 51 +++++ spinifex/kvutil/cas.go | 161 ++++++++++++++++ .../{daemon => kvutil}/cas_conflict_test.go | 56 +++--- 5 files changed, 261 insertions(+), 188 deletions(-) create mode 100644 spinifex/kvutil/cas.go rename spinifex/{daemon => kvutil}/cas_conflict_test.go (70%) diff --git a/spinifex/daemon/jetstream.go b/spinifex/daemon/jetstream.go index b87bfdd04..1bce43039 100644 --- a/spinifex/daemon/jetstream.go +++ b/spinifex/daemon/jetstream.go @@ -10,6 +10,7 @@ import ( "sync" "time" + "github.com/mulgadc/spinifex/spinifex/kvutil" "github.com/mulgadc/spinifex/spinifex/migrate" "github.com/mulgadc/spinifex/spinifex/otelsetup" "github.com/mulgadc/spinifex/spinifex/utils" @@ -258,145 +259,6 @@ func (m *JetStreamManager) recoverTerminatedKVBucket(ctx context.Context) (jetst }, &m.terminatedKV, TerminatedInstanceBucketVersion) } -// maxCASRetries bounds the optimistic-concurrency retry loop in casUpdate. -// A conflict this many times in a row means sustained write contention on -// the key, not a single overlapping update — callers should surface an error. -const maxCASRetries = 5 - -// casUpdate performs an optimistic-concurrency read-modify-write against kv: -// Get the current entry (or start from a zero value when absent and -// createIfAbsent is set), apply mutate to the decoded value, then commit via -// Update/Create using the observed revision. A concurrent writer that changes -// the key between Get and Update/Create is detected via a revision conflict -// and retried, bounded by maxCASRetries. -func casUpdate[T any](ctx context.Context, kv jetstream.KeyValue, key string, mutate func(*T), createIfAbsent bool) (*T, error) { - var lastErr error - for range maxCASRetries { - var value T - var revision uint64 - - entry, err := kv.Get(ctx, key) - switch { - case err == nil: - if uerr := json.Unmarshal(entry.Value(), &value); uerr != nil { - return nil, uerr - } - revision = entry.Revision() - case errors.Is(err, jetstream.ErrKeyNotFound) && createIfAbsent: - revision = 0 - default: - return nil, err - } - - mutate(&value) - - data, merr := json.Marshal(&value) - if merr != nil { - return nil, merr - } - - if revision == 0 { - _, err = kv.Create(ctx, key, data) - } else { - _, err = kv.Update(ctx, key, data, revision) - } - if err == nil { - return &value, nil - } - // Create reports a lost race as ErrKeyExists, Update as - // ErrKeyRevisionMismatch — and only the latter holds on a replicated - // bucket, where the conflict carries a different API error code. - if errors.Is(err, jetstream.ErrKeyExists) || errors.Is(err, jetstream.ErrKeyRevisionMismatch) { - lastErr = err - continue - } - return nil, err - } - return nil, fmt.Errorf("CAS update exhausted %d retries for key %s: %w", maxCASRetries, key, lastErr) -} - -// casPut writes data to key under kv using optimistic concurrency: it reads -// the current revision (or treats the key as absent, using Create, when -// createIfAbsent is set) then commits via Update/Create, retrying on a -// revision conflict up to maxCASRetries. This is the "replace wholesale" -// counterpart to casUpdate, for values (like vm.VM, which embeds a -// sync.Mutex) that cannot be safely decoded into and copied out of a shared -// struct value. -func casPut(ctx context.Context, kv jetstream.KeyValue, key string, data []byte, createIfAbsent bool) error { - var lastErr error - for range maxCASRetries { - var revision uint64 - - entry, err := kv.Get(ctx, key) - switch { - case err == nil: - revision = entry.Revision() - case errors.Is(err, jetstream.ErrKeyNotFound) && createIfAbsent: - revision = 0 - default: - return err - } - - if revision == 0 { - _, err = kv.Create(ctx, key, data) - } else { - _, err = kv.Update(ctx, key, data, revision) - } - if err == nil { - return nil - } - // Create reports a lost race as ErrKeyExists, Update as - // ErrKeyRevisionMismatch — and only the latter holds on a replicated - // bucket, where the conflict carries a different API error code. - if errors.Is(err, jetstream.ErrKeyExists) || errors.Is(err, jetstream.ErrKeyRevisionMismatch) { - lastErr = err - continue - } - return err - } - return fmt.Errorf("CAS put exhausted %d retries for key %s: %w", maxCASRetries, key, lastErr) -} - -// casClaim atomically removes key from kv, decoding its current value into T -// first — the delete-side counterpart to casPut/casUpdate. It is the -// primitive an exclusive "claim" is built on: at most one caller can ever -// observe a successful delete for a given revision, so at most one caller -// gets back a non-nil value. A concurrent writer that changes the key -// between Get and Delete (another claim, or an unrelated update) is -// detected via a revision conflict and retried, bounded by maxCASRetries, -// so a losing racer only fails outright when the key is truly gone -// (notFound) or retries are exhausted. -func casClaim[T any](ctx context.Context, kv jetstream.KeyValue, key string) (value *T, notFound bool, err error) { - var lastErr error - for range maxCASRetries { - entry, gerr := kv.Get(ctx, key) - if gerr != nil { - if errors.Is(gerr, jetstream.ErrKeyNotFound) { - return nil, true, nil - } - return nil, false, gerr - } - - var v T - if uerr := json.Unmarshal(entry.Value(), &v); uerr != nil { - return nil, false, uerr - } - - if derr := kv.Delete(ctx, key, jetstream.LastRevision(entry.Revision())); derr != nil { - // A revision-guarded Delete reports a lost race only as - // ErrKeyRevisionMismatch; ErrKeyExists never matches on R>1. - if errors.Is(derr, jetstream.ErrKeyRevisionMismatch) { - lastErr = derr - continue - } - return nil, false, derr - } - - return &v, false, nil - } - return nil, false, fmt.Errorf("CAS claim exhausted %d retries for key %s: %w", maxCASRetries, key, lastErr) -} - // Heartbeat represents a daemon's periodic health status published to cluster KV. // // AvailableVCPU / AvailableMem are observability-only (host - allocated, @@ -466,9 +328,9 @@ func (m *JetStreamManager) WriteClusterShutdown(state *ClusterShutdownState) err if m.clusterKV == nil { return errors.New("cluster state KV not initialized") } - _, err := casUpdate(context.Background(), m.clusterKV, "cluster.shutdown", func(s *ClusterShutdownState) { - *s = *state - }, true) + _, err := kvutil.Update(context.Background(), m.clusterKV, "cluster.shutdown", + kvutil.CASConfig{CreateIfAbsent: true}, + func(s *ClusterShutdownState) (bool, error) { *s = *state; return true, nil }) return err } @@ -480,7 +342,8 @@ func (m *JetStreamManager) UpdateClusterShutdown(mutate func(*ClusterShutdownSta if m.clusterKV == nil { return nil, errors.New("cluster state KV not initialized") } - return casUpdate(context.Background(), m.clusterKV, "cluster.shutdown", mutate, false) + return kvutil.Update(context.Background(), m.clusterKV, "cluster.shutdown", kvutil.CASConfig{}, + func(s *ClusterShutdownState) (bool, error) { mutate(s); return true, nil }) } // UpdateMgmtIPAM atomically applies mutate to the mgmt-ipam record for @@ -496,7 +359,9 @@ func (m *JetStreamManager) UpdateMgmtIPAM(subnet string, mutate func(*MgmtIPReco if m.clusterKV == nil { return nil, errors.New("cluster state KV not initialized") } - return casUpdate(context.Background(), m.clusterKV, mgmtIPAMKeyPrefix+subnet, mutate, createIfAbsent) + return kvutil.Update(context.Background(), m.clusterKV, mgmtIPAMKeyPrefix+subnet, + kvutil.CASConfig{CreateIfAbsent: createIfAbsent}, + func(r *MgmtIPRecord) (bool, error) { mutate(r); return true, nil }) } // ReadClusterShutdown reads the cluster shutdown state from KV. @@ -817,7 +682,7 @@ func (m *JetStreamManager) WriteStoppedInstance(instanceID string, instance *vm. } key := StoppedInstancePrefix + instanceID - err = casPut(context.Background(), m.kv, key, jsonData, true) + err = kvutil.Put(context.Background(), m.kv, key, kvutil.CASConfig{CreateIfAbsent: true}, jsonData) if err != nil { if isStreamUnavailable(err) { slog.Warn("KV stream unavailable, attempting recovery", "operation", "WriteStoppedInstance", "key", key, "err", err) @@ -825,7 +690,7 @@ func (m *JetStreamManager) WriteStoppedInstance(instanceID string, instance *vm. if recoverErr != nil { return err } - if retryErr := casPut(context.Background(), kv, key, jsonData, true); retryErr != nil { + if retryErr := kvutil.Put(context.Background(), kv, key, kvutil.CASConfig{CreateIfAbsent: true}, jsonData); retryErr != nil { return retryErr } slog.Debug("Wrote stopped instance to JetStream KV (after recovery)", "key", key, "instanceId", instanceID) @@ -923,7 +788,7 @@ func (m *JetStreamManager) ClaimStoppedInstance(instanceID string) (*vm.VM, erro } key := StoppedInstancePrefix + instanceID - instance, notFound, err := casClaim[vm.VM](context.Background(), m.kv, key) + instance, notFound, err := kvutil.Claim[vm.VM](context.Background(), m.kv, key, kvutil.CASConfig{}) if err != nil { if isStreamUnavailable(err) { slog.Warn("KV stream unavailable, attempting recovery", "operation", "ClaimStoppedInstance", "key", key, "err", err) @@ -931,7 +796,7 @@ func (m *JetStreamManager) ClaimStoppedInstance(instanceID string) (*vm.VM, erro if recoverErr != nil { return nil, err } - instance, notFound, err = casClaim[vm.VM](context.Background(), kv, key) + instance, notFound, err = kvutil.Claim[vm.VM](context.Background(), kv, key, kvutil.CASConfig{}) if err != nil { return nil, err } @@ -961,7 +826,8 @@ func (m *JetStreamManager) UpdateStoppedInstance(instanceID string, mutate func( } key := StoppedInstancePrefix + instanceID - updated, err := casUpdate(context.Background(), m.kv, key, mutate, false) + apply := func(v *vm.VM) (bool, error) { mutate(v); return true, nil } + updated, err := kvutil.Update(context.Background(), m.kv, key, kvutil.CASConfig{}, apply) if err != nil { if isStreamUnavailable(err) { slog.Warn("KV stream unavailable, attempting recovery", "operation", "UpdateStoppedInstance", "key", key, "err", err) @@ -969,7 +835,7 @@ func (m *JetStreamManager) UpdateStoppedInstance(instanceID string, mutate func( if recoverErr != nil { return nil, err } - return casUpdate(context.Background(), kv, key, mutate, false) + return kvutil.Update(context.Background(), kv, key, kvutil.CASConfig{}, apply) } return nil, err } @@ -1055,7 +921,7 @@ func (m *JetStreamManager) WriteTerminatedInstance(instanceID string, instance * } key := TerminatedInstancePrefix + instanceID - err = casPut(context.Background(), m.terminatedKV, key, jsonData, true) + err = kvutil.Put(context.Background(), m.terminatedKV, key, kvutil.CASConfig{CreateIfAbsent: true}, jsonData) if err != nil { if isStreamUnavailable(err) { slog.Warn("KV stream unavailable, attempting recovery", "operation", "WriteTerminatedInstance", "key", key, "err", err) @@ -1063,7 +929,7 @@ func (m *JetStreamManager) WriteTerminatedInstance(instanceID string, instance * if recoverErr != nil { return err } - if retryErr := casPut(context.Background(), kv, key, jsonData, true); retryErr != nil { + if retryErr := kvutil.Put(context.Background(), kv, key, kvutil.CASConfig{CreateIfAbsent: true}, jsonData); retryErr != nil { return retryErr } slog.Debug("Wrote terminated instance to JetStream KV (after recovery)", "key", key, "instanceId", instanceID) @@ -1088,7 +954,8 @@ func (m *JetStreamManager) UpdateTerminatedInstance(instanceID string, mutate fu } key := TerminatedInstancePrefix + instanceID - updated, err := casUpdate(context.Background(), m.terminatedKV, key, mutate, false) + apply := func(v *vm.VM) (bool, error) { mutate(v); return true, nil } + updated, err := kvutil.Update(context.Background(), m.terminatedKV, key, kvutil.CASConfig{}, apply) if err != nil { if isStreamUnavailable(err) { slog.Warn("KV stream unavailable, attempting recovery", "operation", "UpdateTerminatedInstance", "key", key, "err", err) @@ -1096,7 +963,7 @@ func (m *JetStreamManager) UpdateTerminatedInstance(instanceID string, mutate fu if recoverErr != nil { return nil, err } - return casUpdate(context.Background(), kv, key, mutate, false) + return kvutil.Update(context.Background(), kv, key, kvutil.CASConfig{}, apply) } return nil, err } diff --git a/spinifex/kvlease/lease.go b/spinifex/kvlease/lease.go index 61cd597f4..7412be9f5 100644 --- a/spinifex/kvlease/lease.go +++ b/spinifex/kvlease/lease.go @@ -201,11 +201,15 @@ func (l *Lease) markLost() bool { func (l *Lease) Release(ctx context.Context) { l.mu.Lock() kv, rev, stop := l.kv, l.rev, l.stop - l.stop = nil + l.kv, l.stop = nil, nil l.mu.Unlock() if stop != nil { stop() } + if kv == nil { + // Never acquired, or already released. + return + } 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. diff --git a/spinifex/kvlease/lease_test.go b/spinifex/kvlease/lease_test.go index 5f398989e..1e73a1141 100644 --- a/spinifex/kvlease/lease_test.go +++ b/spinifex/kvlease/lease_test.go @@ -334,3 +334,54 @@ func TestAttrs_AppearOnLeaseOwnedLogLines(t *testing.T) { require.Contains(t, buf.String(), `"bucket":"quota-reconcile"`) } + +// A node that never won the lease still unwinds its Run loop through Release. +// That is not a lost key, so it must not raise the warning that says one. +func TestRelease_WithoutAcquireIsSilent(t *testing.T) { + _, nc, _ := testutil.StartTestJetStream(t) + + holder := newTestLease(t, nc, "node-a", nil) + require.True(t, holder.TryAcquire(t.Context())) + t.Cleanup(func() { holder.Release(t.Context()) }) + + 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) }) + + loser := newTestLease(t, nc, "node-b", nil) + require.False(t, loser.TryAcquire(t.Context())) + loser.Release(t.Context()) + + assert.NotContains(t, buf.String(), "already lost before release") +} + +// The warning still has to fire where it means something: renewal lost the key, +// so a delete here would drop the successor's lease rather than. +func TestRelease_AfterRenewalLossWarns(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())) + + 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") + } + + 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) }) + + a.Release(t.Context()) + assert.Contains(t, buf.String(), "already lost before release") +} diff --git a/spinifex/kvutil/cas.go b/spinifex/kvutil/cas.go new file mode 100644 index 000000000..8464f0b41 --- /dev/null +++ b/spinifex/kvutil/cas.go @@ -0,0 +1,161 @@ +package kvutil + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "math/rand/v2" + "time" + + "github.com/nats-io/nats.go/jetstream" +) + +// casAttempts bounds an optimistic-concurrency loop. +const casAttempts = 5 + +// casBackoffBase scales the pause between attempts. +const casBackoffBase = 2 * time.Millisecond + +// CASConfig tunes an optimistic-concurrency loop. The zero value retreis casAttempts times and fails when the key is absent. +type CASConfig struct { + Attempts int // 0 selects casAttempts + CreateIfAbsent bool // start from a zero value when the key is missing +} + +// isConflict reports a lost CAS race. Create reports one as ErrKeyExists and Update as ErrKeyRevisionMismatch. +// A revision-guarded Delete only ever reports the latter. +func isConflict(err error) bool { + return errors.Is(err, jetstream.ErrKeyExists) || errors.Is(err, jetstream.ErrKeyRevisionMismatch) +} + +// retryCAS runs attempt until is succeeds, returns a non-conflict error, or spends its buudget. It pauses between attempts, +// jittered so contending writers do not re-collide on the same revision. +func retryCAS(ctx context.Context, cfg CASConfig, op, key string, attempt func() error) error { + attempts := cfg.Attempts + if attempts <= 0 { + attempts = casAttempts + } + var lastErr error + for i := range attempts { + err := attempt() + if err == nil { + return nil + } + if !isConflict(err) { + return err + } + lastErr = err + backoff := casBackoffBase * time.Duration(i+1) + jitter := time.Duration(rand.Int64N(int64(backoff))) //nolint:gosec // jitter, not cryptographic + select { + case <-time.After(backoff/2 + jitter): + case <-ctx.Done(): + return ctx.Err() + } + } + return fmt.Errorf("CAS %s exhausted %d attempts for key %s: %w", op, attempts, key, lastErr) +} + +// Update is an optimistic-concurrency read-modify-write: +// decode the current entry, apply mutate, commit at the observed revision. +func Update[T any](ctx context.Context, kv jetstream.KeyValue, key string, cfg CASConfig, mutate func(*T) (bool, error)) (*T, error) { + var result *T + err := retryCAS(ctx, cfg, "update", key, func() error { + var value T + var revision uint64 + + entry, err := kv.Get(ctx, key) + switch { + case err == nil: + if uerr := json.Unmarshal(entry.Value(), &value); uerr != nil { + return uerr + } + revision = entry.Revision() + case errors.Is(err, jetstream.ErrKeyNotFound) && cfg.CreateIfAbsent: + revision = 0 + default: + return err + } + + changed, merr := mutate(&value) + if merr != nil { + return merr + } + if !changed { + result = &value + return nil + } + + data, merr := json.Marshal(&value) + if merr != nil { + return merr + } + if revision == 0 { + _, err = kv.Create(ctx, key, data) + } else { + _, err = kv.Update(ctx, key, data, revision) + } + if err != nil { + return err + } + result = &value + return nil + }) + if err != nil { + return nil, err + } + return result, nil +} + +// Put is the replace-wholesale counterpart to Update, for values that cannot safely be decoded into and copied out of a shared struct. +func Put(ctx context.Context, kv jetstream.KeyValue, key string, cfg CASConfig, data []byte) error { + return retryCAS(ctx, cfg, "put", key, func() error { + var revision uint64 + entry, err := kv.Get(ctx, key) + switch { + case err == nil: + revision = entry.Revision() + case errors.Is(err, jetstream.ErrKeyNotFound) && cfg.CreateIfAbsent: + revision = 0 + default: + return err + } + if revision == 0 { + _, err = kv.Create(ctx, key, data) + } else { + _, err = kv.Update(ctx, key, data, revision) + } + return err + }) +} + +// Claim atomically removes key, decoding its value first. At most one caller can observe a successful delete at a given revision, so +// at most one gets a non-nil value back - the primitive an exclusive claim is built on. +func Claim[T any](ctx context.Context, kv jetstream.KeyValue, key string, cfg CASConfig) (value *T, notFound bool, err error) { + var result *T + var absent bool + rerr := retryCAS(ctx, cfg, "claim", key, func() error { + entry, gerr := kv.Get(ctx, key) + if gerr != nil { + if errors.Is(gerr, jetstream.ErrKeyNotFound) { + absent = true + return nil + } + return gerr + } + var v T + if uerr := json.Unmarshal(entry.Value(), &v); uerr != nil { + return uerr + } + if derr := kv.Delete(ctx, key, jetstream.LastRevision(entry.Revision())); derr != nil { + return derr + } + result = &v + return nil + }) + if rerr != nil { + return nil, false, rerr + } + return result, absent, nil +} diff --git a/spinifex/daemon/cas_conflict_test.go b/spinifex/kvutil/cas_conflict_test.go similarity index 70% rename from spinifex/daemon/cas_conflict_test.go rename to spinifex/kvutil/cas_conflict_test.go index 7e89f607a..837f80a5e 100644 --- a/spinifex/daemon/cas_conflict_test.go +++ b/spinifex/kvutil/cas_conflict_test.go @@ -1,25 +1,19 @@ -//test:in-package — drives the unexported CAS primitives (casUpdate, casPut and -//the generic casClaim[T]) against maxCASRetries. A generic function cannot be -//re-exported through export_test.go without a wrapper per instantiation. - -package daemon +package kvutil_test import ( "context" "fmt" "testing" - "github.com/nats-io/nats.go" + "github.com/mulgadc/spinifex/spinifex/kvutil" + "github.com/mulgadc/spinifex/spinifex/testutil" "github.com/nats-io/nats.go/jetstream" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) -// wrongLastSequence builds the error nats.go returns for a lost CAS race on a -// bucket whose stream reports the conflict under code. Single-replica streams -// report 10071, replicated ones 10164. Update and revision-guarded Delete wrap -// both with ErrKeyRevisionMismatch; only 10071 also satisfies ErrKeyExists, so -// a predicate written against ErrKeyExists alone is blind on a real cluster. +const casTestAttempts = 5 + func wrongLastSequence(code jetstream.ErrorCode) error { apiErr := &jetstream.APIError{ ErrorCode: code, @@ -29,9 +23,6 @@ func wrongLastSequence(code jetstream.ErrorCode) error { return fmt.Errorf("%w: %w", apiErr, jetstream.ErrKeyRevisionMismatch) } -// casConflictKV forces the next failuresLeft writes of each kind to fail with a -// wrong-last-sequence response, so the CAS primitives can be driven against a -// replicated bucket's error shape without standing up a real cluster. type casConflictKV struct { jetstream.KeyValue @@ -56,8 +47,7 @@ func (k *casConflictKV) Delete(ctx context.Context, key string, opts ...jetstrea return k.KeyValue.Delete(ctx, key, opts...) } -// conflictCodes names the two wrong-last-sequence codes a bucket can report, so -// every CAS test covers a replicated bucket as well as a single-replica one. +// conflictCodes names the two wrong-last-sequence codes a bucket can report, so every CAS test covers a replicated bucket as well as single-replica one. var conflictCodes = map[string]jetstream.ErrorCode{ "single replica": jetstream.JSErrCodeStreamWrongLastSequence, "replicated": jetstream.JSErrCodeStreamWrongLastSequenceConstant, @@ -65,16 +55,11 @@ var conflictCodes = map[string]jetstream.ErrorCode{ func casTestBucket(t *testing.T, name string) jetstream.KeyValue { t.Helper() - nc, err := nats.Connect(sharedJSNATSURL) - require.NoError(t, err) - t.Cleanup(nc.Close) - - js, err := jetstream.New(nc) - require.NoError(t, err) + _, nc, _ := testutil.StartTestJetStream(t) + js := testutil.NewJetStream(t, nc) kv, err := js.CreateKeyValue(t.Context(), jetstream.KeyValueConfig{Bucket: name, History: 1}) require.NoError(t, err) - t.Cleanup(func() { _ = js.DeleteKeyValue(context.Background(), name) }) return kv } @@ -82,6 +67,11 @@ type casRecord struct { Counter int `json:"counter"` } +func bumpCounter(r *casRecord) (bool, error) { + r.Counter++ + return true, nil +} + // A lost Update race must be retried on both replica counts. Matching only // jetstream.ErrKeyExists passes the single-replica case and surfaces a raw // "wrong last sequence" error on the replicated one. @@ -92,8 +82,8 @@ func TestCASUpdate_RetriesRevisionConflict(t *testing.T) { _, err := kv.Put(t.Context(), "key", []byte(`{"counter":1}`)) require.NoError(t, err) - stub := &casConflictKV{KeyValue: kv, conflictCode: code, updateFailures: maxCASRetries - 1} - got, err := casUpdate(t.Context(), stub, "key", func(r *casRecord) { r.Counter++ }, false) + stub := &casConflictKV{KeyValue: kv, conflictCode: code, updateFailures: casTestAttempts - 1} + got, err := kvutil.Update(t.Context(), stub, "key", kvutil.CASConfig{Attempts: casTestAttempts}, bumpCounter) require.NoError(t, err) assert.Equal(t, 2, got.Counter) assert.Zero(t, stub.updateFailures, "all injected conflicts must have been consumed") @@ -108,8 +98,8 @@ func TestCASUpdate_ExhaustsRetriesOnSustainedConflict(t *testing.T) { _, err := kv.Put(t.Context(), "key", []byte(`{"counter":1}`)) require.NoError(t, err) - stub := &casConflictKV{KeyValue: kv, conflictCode: code, updateFailures: maxCASRetries + 1} - _, err = casUpdate(t.Context(), stub, "key", func(r *casRecord) { r.Counter++ }, false) + stub := &casConflictKV{KeyValue: kv, conflictCode: code, updateFailures: casTestAttempts + 1} + _, err = kvutil.Update(t.Context(), stub, "key", kvutil.CASConfig{Attempts: casTestAttempts}, bumpCounter) require.Error(t, err) assert.Contains(t, err.Error(), "CAS update exhausted") assert.ErrorIs(t, err, jetstream.ErrKeyRevisionMismatch) @@ -124,8 +114,8 @@ func TestCASPut_RetriesRevisionConflict(t *testing.T) { _, err := kv.Put(t.Context(), "key", []byte(`{"counter":1}`)) require.NoError(t, err) - stub := &casConflictKV{KeyValue: kv, conflictCode: code, updateFailures: maxCASRetries - 1} - require.NoError(t, casPut(t.Context(), stub, "key", []byte(`{"counter":9}`), false)) + stub := &casConflictKV{KeyValue: kv, conflictCode: code, updateFailures: casTestAttempts - 1} + require.NoError(t, kvutil.Put(t.Context(), stub, "key", kvutil.CASConfig{Attempts: casTestAttempts}, []byte(`{"counter":9}`))) assert.Zero(t, stub.updateFailures, "all injected conflicts must have been consumed") entry, err := kv.Get(t.Context(), "key") @@ -135,8 +125,8 @@ func TestCASPut_RetriesRevisionConflict(t *testing.T) { } } -// casClaim's write is a revision-guarded Delete, which reports a lost race only -// as ErrKeyRevisionMismatch — ErrKeyExists never matches it on any replica count +// Claim's write is a revision-guarded Delete, which reports a lost race only +// as ErrKeyRevisionMismatch — ErrKeyExists never matches it on a // once the conflict carries the replicated code. func TestCASClaim_RetriesRevisionConflict(t *testing.T) { for name, code := range conflictCodes { @@ -145,8 +135,8 @@ func TestCASClaim_RetriesRevisionConflict(t *testing.T) { _, err := kv.Put(t.Context(), "key", []byte(`{"counter":7}`)) require.NoError(t, err) - stub := &casConflictKV{KeyValue: kv, conflictCode: code, deleteFailures: maxCASRetries - 1} - got, notFound, err := casClaim[casRecord](t.Context(), stub, "key") + stub := &casConflictKV{KeyValue: kv, conflictCode: code, deleteFailures: casTestAttempts - 1} + got, notFound, err := kvutil.Claim[casRecord](t.Context(), stub, "key", kvutil.CASConfig{Attempts: casTestAttempts}) require.NoError(t, err) require.False(t, notFound) require.NotNil(t, got)