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

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

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.

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

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

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)