refactor: route the last hand-rolled KV codecs and CAS loops through kvstore (#1061)

spinifex/daemon/jetstream.go

@@ -182,34 +182,31 @@ func (m *JetStreamManager) InitKVBucket() error {
 	return err
 }
 
-// InitClusterStateBucket initializes the cluster-state KV bucket, creating it if it doesn't exist.
-func (m *JetStreamManager) InitClusterStateBucket() error {
-	ctx := context.Background()
-	kv, err := m.js.KeyValue(ctx, ClusterStateBucket)
-	if err != nil {
-		if errors.Is(err, jetstream.ErrBucketNotFound) {
-			slog.Debug("Creating JetStream KV bucket", "bucket", ClusterStateBucket, "replicas", m.replicas)
-			kv, err = m.js.CreateKeyValue(ctx, jetstream.KeyValueConfig{
-				Bucket:      ClusterStateBucket,
-				Description: "Spinifex cluster state (heartbeats, shutdown markers, service maps)",
-				History:     1,
-				Replicas:    m.replicas,
-				TTL:         1 * time.Hour,
-			})
-			if err != nil {
-				return err
-			}
-		} else {
-			return err
-		}
-	} else {
-		slog.Debug("Connected to existing JetStream KV bucket", "bucket", ClusterStateBucket)
+// clusterStateConfig describes the cluster-state bucket.
+func clusterStateConfig(replicas int) kvstore.Config {
+	return kvstore.Config{
+		Name:        ClusterStateBucket,
+		Description: "Spinifex cluster state (heartbeats, shutdown markers, service maps)",
+		History:     1,
+		Replicas:    replicas,
+		TTL:         1 * time.Hour,
+		OnOpen: func(ctx context.Context, kv jetstream.KeyValue) error {
+			return migrate.DefaultRegistry.RunKV(ctx, ClusterStateBucket, kv, ClusterStateBucketVersion)
+		},
+		Missing: "cluster state KV bucket not initialized",
 	}
+}
 
-	m.clusterKV = kv
-	if err := migrate.DefaultRegistry.RunKV(ctx, ClusterStateBucket, kv, ClusterStateBucketVersion); err != nil {
-		return fmt.Errorf("migrate %s: %w", ClusterStateBucket, err)
+// InitClusterStateBucket initializes the cluster-state KV bucket, creating it if it doesn't exist.
+//
+// The handle is kept rather than the Bucket: this one holds several unrelated
+// record types under one namespace, and its callers still read it as raw keys.
+func (m *JetStreamManager) InitClusterStateBucket() error {
+	kv, err := kvstore.NewBucket(m.js, clusterStateConfig(m.replicas)).KV(context.Background())
+	if err != nil {
+		return err
 	}
+	m.clusterKV = kv
 	return nil
 }
 
@@ -679,7 +676,7 @@ func (m *JetStreamManager) ClaimStoppedInstance(instanceID string) (*vm.VM, erro
 
 	claimed := *record
 	claimed.Spec.DesiredState = vm.DesiredRunning
-	if err := m.records.CompareAndSet(context.Background(), key, &claimed, rev); err != nil {
+	if _, err := m.records.CompareAndSet(context.Background(), key, &claimed, rev); err != nil {
 		if errors.Is(err, kvstore.ErrConflict) {
 			return nil, vm.ErrStoppedInstanceClaimed
 		}

spinifex/gateway/bedrock/guardrail.go

@@ -148,7 +148,7 @@ func (s *GuardrailStore) put(ctx context.Context, rec GuardrailRecord) error {
 // concurrent mutations of the same guardrail never silently clobber one
 // another. A lost race is reported as a retryable ConflictException.
 func (s *GuardrailStore) update(ctx context.Context, key string, rec GuardrailRecord, rev uint64) error {
-	err := s.store.CompareAndSet(ctx, key, &rec, rev)
+	_, err := s.store.CompareAndSet(ctx, key, &rec, rev)
 	if errors.Is(err, kvstore.ErrConflict) {
 		return awserrors.Errorf(awserrors.ErrorConflictException,
 			"bedrock: guardrail %s was modified concurrently; retry the request", rec.GuardrailID)

spinifex/gateway/bedrock/provisioned.go

@@ -128,7 +128,7 @@ func (s *ProvisionedStore) getRevision(ctx context.Context, key string) (Provisi
 // race is reported as a retryable ConflictException rather than silently
 // clobbering the winner's write.
 func (s *ProvisionedStore) update(ctx context.Context, key string, rec ProvisionedModelRecord, rev uint64) error {
-	err := s.store.CompareAndSet(ctx, key, &rec, rev)
+	_, err := s.store.CompareAndSet(ctx, key, &rec, rev)
 	if errors.Is(err, kvstore.ErrConflict) {
 		return awserrors.Errorf(awserrors.ErrorConflictException,
 			"bedrock: provisioned throughput %s was modified concurrently; retry the request", rec.ARN)

spinifex/handlers/acm/store.go

@@ -520,7 +520,7 @@ func (s *Store) AcquireLease(ctx context.Context, certArn, holderID string, ttl
 	}
 	rec.LeaseHolder = holderID
 	rec.LeaseExpiresAt = now.Add(ttl)
-	if err := s.store.CompareAndSet(ctx, key, rec, rev); err != nil {
+	if _, err := s.store.CompareAndSet(ctx, key, rec, rev); err != nil {
 		if errors.Is(err, kvstore.ErrConflict) {
 			// Lost the race to a concurrent acquirer between Get and Update;
 			// the caller skips this tick rather than retrying immediately.
@@ -552,7 +552,7 @@ func (s *Store) ReleaseLease(ctx context.Context, certArn, holderID string) erro
 	}
 	rec.LeaseHolder = ""
 	rec.LeaseExpiresAt = time.Time{}
-	if err := s.store.CompareAndSet(ctx, key, rec, rev); err != nil {
+	if _, err := s.store.CompareAndSet(ctx, key, rec, rev); err != nil {
 		if errors.Is(err, kvstore.ErrConflict) {
 			return nil // record changed concurrently; nothing to clean up
 		}

spinifex/handlers/bedrock/eviction_test.go

@@ -234,5 +234,6 @@ func pinEndpoint(t *testing.T, s *Service, modelID string) {
 	require.NoError(t, err)
 	require.True(t, found)
 	rec.Pinned = true
-	require.NoError(t, s.store.CompareAndSet(t.Context(), key, &rec, rev))
+	_, casErr := s.store.CompareAndSet(t.Context(), key, &rec, rev)
+	require.NoError(t, casErr)
 }

spinifex/handlers/bedrock/reaper.go

@@ -286,7 +286,7 @@ func (r *Reaper) persistObservation(ctx context.Context, prev, next EndpointReco
 		return nil
 	}
 	next.Generation = current.Generation + 1
-	if err := r.svc.store.CompareAndSet(ctx, key, &next, rev); err != nil {
+	if _, err := r.svc.store.CompareAndSet(ctx, key, &next, rev); err != nil {
 		if errors.Is(err, kvstore.ErrConflict) {
 			return nil
 		}

spinifex/handlers/bedrock/reaper_test.go

@@ -130,7 +130,8 @@ func (f *reaperFixture) age(t *testing.T, readyAgo, activeAgo time.Duration) End
 	if activeAgo > 0 {
 		rec.LastActiveAt = now.Add(-activeAgo)
 	}
-	require.NoError(t, f.store.CompareAndSet(t.Context(), key, &rec, rev))
+	_, casErr := f.store.CompareAndSet(t.Context(), key, &rec, rev)
+	require.NoError(t, casErr)
 	return rec
 }
 
@@ -178,7 +179,8 @@ func TestReaper_NeverReapsPinnedEndpoint(t *testing.T) {
 	rec, rev, _, err := f.store.getRevision(t.Context(), key)
 	require.NoError(t, err)
 	rec.Pinned = true
-	require.NoError(t, f.store.CompareAndSet(t.Context(), key, &rec, rev))
+	_, casErr := f.store.CompareAndSet(t.Context(), key, &rec, rev)
+	require.NoError(t, casErr)
 
 	require.NoError(t, f.reaper.sweepOnce(t.Context()))
 
@@ -292,7 +294,8 @@ func TestReaper_PinnedEndpointNotReapedOnScrapeFailure(t *testing.T) {
 	require.NoError(t, err)
 	require.True(t, found)
 	rec.Pinned = true
-	require.NoError(t, f.store.CompareAndSet(t.Context(), key, &rec, rev))
+	_, casErr := f.store.CompareAndSet(t.Context(), key, &rec, rev)
+	require.NoError(t, casErr)
 
 	f.transport.breakScrapes()
 	for range maxScrapeFailures + 2 {

spinifex/handlers/bedrock/service.go

@@ -351,7 +351,7 @@ func (s *Service) runLaunch(ctx context.Context, key string, rec EndpointRecord,
 		slog.ErrorContext(ctx, "bedrock: re-read endpoint before READY write failed", "group", spec.GroupID, "err", gerr)
 		return
 	}
-	if err := s.store.CompareAndSet(ctx, key, &rec, rev); err != nil {
+	if _, err := s.store.CompareAndSet(ctx, key, &rec, rev); err != nil {
 		slog.ErrorContext(ctx, "bedrock: CAS write of READY state failed", "group", spec.GroupID, "err", err)
 	}
 }
@@ -461,7 +461,7 @@ func (s *Service) Delete(ctx context.Context, in *DeleteEndpointInput, _ string)
 		}
 		rec.State = StateDraining
 		rec.Generation++
-		if err := s.store.CompareAndSet(ctx, key, &rec, rev); err != nil {
+		if _, err := s.store.CompareAndSet(ctx, key, &rec, rev); err != nil {
 			return nil, fmt.Errorf("bedrock: mark endpoint %s draining: %w", in.ModelID, err)
 		}
 	}

spinifex/handlers/bedrock/store_test.go

@@ -47,11 +47,12 @@ func TestStore_CreateGetUpdateDelete(t *testing.T) {
 
 	got.State = StateReady
 	got.Generation = 2
-	require.NoError(t, store.CompareAndSet(t.Context(), key, &got, gotRev))
+	_, casErr := store.CompareAndSet(t.Context(), key, &got, gotRev)
+	require.NoError(t, casErr)
 
 	// The CAS update must fail against the now-stale revision — this is what
 	// stops a launch goroutine and a concurrent writer stomping each other.
-	err = store.CompareAndSet(t.Context(), key, &got, gotRev)
+	_, err = store.CompareAndSet(t.Context(), key, &got, gotRev)
 	assert.ErrorIs(t, err, kvstore.ErrConflict)
 
 	require.NoError(t, store.Purge(t.Context(), key))

spinifex/handlers/ec2/snapshot/service_impl.go

@@ -2,7 +2,6 @@ package handlers_ec2_snapshot
 
 import (
 	"context"
-	"encoding/json"
 	"errors"
 	"fmt"
 	"log/slog"
@@ -23,7 +22,7 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/ebsprovider"
 	"github.com/mulgadc/spinifex/spinifex/filterutil"
 	"github.com/mulgadc/spinifex/spinifex/handlers/ec2/volumestate"
-	"github.com/mulgadc/spinifex/spinifex/kvutil"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 	"github.com/mulgadc/spinifex/spinifex/migrate"
 	"github.com/mulgadc/spinifex/spinifex/objectstore"
 	"github.com/mulgadc/spinifex/spinifex/types"
@@ -38,7 +37,13 @@ var _ SnapshotService = (*SnapshotServiceImpl)(nil)
 const (
 	KVBucketVolumeSnapshots        = "spinifex-volume-snapshots"
 	KVBucketVolumeSnapshotsVersion = 1
-	snapshotCleanupTimeout         = 5 * time.Second
+	KVBucketVolumeSnapshotsHistory = 10
+
+	// snapshotRefCASAttempts bounds one index write's retry on a revision
+	// conflict. Every snapshot of a volume shares one record, so simultaneous
+	// creates on one volume contend once each; the default five sits below that.
+	snapshotRefCASAttempts = 25
+	snapshotCleanupTimeout = 5 * time.Second
 )
 
 // SnapshotServiceImpl implements SnapshotService with S3-backed storage.
@@ -46,7 +51,7 @@ type SnapshotServiceImpl struct {
 	config   *config.Config
 	store    objectstore.ObjectStore
 	natsConn *nats.Conn
-	snapKV   jetstream.KeyValue
+	snapRefs *kvstore.Store[[]string]
 	mutex    sync.RWMutex
 	metadata *ebsmetadata.Store
 	provider ebsprovider.EBSProvider
@@ -84,13 +89,13 @@ func NewSnapshotServiceImplWithNATS(ctx context.Context, cfg *config.Config, nat
 		return nil, nil, fmt.Errorf("create JetStream: %w", err)
 	}
 
-	kv, err := kvutil.GetOrCreateBucket(ctx, js, KVBucketVolumeSnapshots, 10)
+	refs := kvstore.New[[]string](js, snapshotRefsConfig())
+	// Opened eagerly: the handle is returned to the caller, which wires it into
+	// the volume service, so it has to exist before this function returns.
+	kv, err := refs.KV(ctx)
 	if err != nil {
 		return nil, nil, fmt.Errorf("create KV bucket %s: %w", KVBucketVolumeSnapshots, err)
 	}
-	if err := migrate.DefaultRegistry.RunKV(ctx, KVBucketVolumeSnapshots, kv, KVBucketVolumeSnapshotsVersion); err != nil {
-		return nil, nil, fmt.Errorf("migrate %s: %w", KVBucketVolumeSnapshots, err)
-	}
 
 	slog.Info("Snapshot service initialized with JetStream KV", "bucket", KVBucketVolumeSnapshots)
 
@@ -98,7 +103,7 @@ func NewSnapshotServiceImplWithNATS(ctx context.Context, cfg *config.Config, nat
 		config:   cfg,
 		store:    store,
 		natsConn: natsConn,
-		snapKV:   kv,
+		snapRefs: refs,
 		metadata: ebsmetadata.NewStore(store, cfg.Predastore.Bucket),
 	}, kv, nil
 }
@@ -113,11 +118,25 @@ func NewSnapshotServiceImplWithStore(cfg *config.Config, store objectstore.Objec
 		metadata: ebsmetadata.NewStore(store, cfg.Predastore.Bucket),
 	}
 	if len(snapshotKV) > 0 {
-		svc.snapKV = snapshotKV[0]
+		svc.snapRefs = kvstore.Over[[]string](nil, snapshotKV[0], snapshotRefsConfig())
 	}
 	return svc
 }
 
+// snapshotRefsConfig describes the volume-to-snapshot index bucket. The value
+// is the list of snapshot IDs taken from one volume, keyed by volume ID.
+func snapshotRefsConfig() kvstore.Config {
+	return kvstore.Config{
+		Name:     KVBucketVolumeSnapshots,
+		History:  KVBucketVolumeSnapshotsHistory,
+		Missing:  "volume snapshot KV bucket not initialized",
+		Attempts: snapshotRefCASAttempts,
+		OnOpen: func(ctx context.Context, kv jetstream.KeyValue) error {
+			return migrate.DefaultRegistry.RunKV(ctx, KVBucketVolumeSnapshots, kv, KVBucketVolumeSnapshotsVersion)
+		},
+	}
+}
+
 // getSnapshotConfig reads the caller's own snapshot document. A snapshot
 // outside the caller's prefix is absent, so ownership needs no second check.
 func (s *SnapshotServiceImpl) getSnapshotConfig(ctx context.Context, accountID, snapshotID string) (ebsmetadata.Snapshot, error) {
@@ -781,58 +800,21 @@ func (s *SnapshotServiceImpl) CopySnapshot(ctx context.Context, input *ec2.CopyS
 // addSnapshotRef adds snapshotID to the volume's snapshot list in KV.
 // Uses CAS (Create/Update with revision) to prevent lost updates under concurrency.
 func (s *SnapshotServiceImpl) addSnapshotRef(ctx context.Context, volumeID, snapshotID string) error {
-	if s.snapKV == nil {
+	if s.snapRefs == nil {
 		slog.Debug("addSnapshotRef: snapshotKV is nil, skipping", "volumeId", volumeID, "snapshotId", snapshotID)
 		return nil
 	}
 
-	const maxRetries = 5
-	for attempt := range maxRetries {
-		entry, err := s.snapKV.Get(ctx, volumeID)
-		var snapshots []string
-
-		if err != nil {
-			if !errors.Is(err, jetstream.ErrKeyNotFound) {
-				return fmt.Errorf("addSnapshotRef: failed to get KV key %s: %w", volumeID, err)
-			}
-			// Key doesn't exist yet — create with just this snapshot
-			data, err := json.Marshal([]string{snapshotID})
-			if err != nil {
-				return fmt.Errorf("addSnapshotRef: failed to marshal snapshot list: %w", err)
-			}
-			if _, err := s.snapKV.Create(ctx, volumeID, data); err != nil {
-				if attempt < maxRetries-1 {
-					continue // concurrent Create/Update — retry
-				}
-				return fmt.Errorf("addSnapshotRef: failed to create KV key %s: %w", volumeID, err)
-			}
-			slog.Info("addSnapshotRef: added snapshot ref", "volumeId", volumeID, "snapshotId", snapshotID)
-			return nil
-		}
-
-		if err := json.Unmarshal(entry.Value(), &snapshots); err != nil {
-			return fmt.Errorf("addSnapshotRef: failed to unmarshal KV value for %s: %w", volumeID, err)
-		}
-
-		snapshots = append(snapshots, snapshotID)
-
-		data, err := json.Marshal(snapshots)
-		if err != nil {
-			return fmt.Errorf("addSnapshotRef: failed to marshal snapshot list: %w", err)
-		}
-
-		if _, err := s.snapKV.Update(ctx, volumeID, data, entry.Revision()); err != nil {
-			if attempt < maxRetries-1 {
-				continue // concurrent update — retry
-			}
-			return fmt.Errorf("addSnapshotRef: failed to update KV key %s: %w", volumeID, err)
-		}
-
-		slog.Info("addSnapshotRef: added snapshot ref", "volumeId", volumeID, "snapshotId", snapshotID)
-		return nil
+	err := s.snapRefs.Upsert(ctx, volumeID, func(snapshots *[]string) (bool, error) {
+		*snapshots = append(*snapshots, snapshotID)
+		return true, nil
+	})
+	if err != nil {
+		return fmt.Errorf("addSnapshotRef: failed to record %s against volume %s: %w", snapshotID, volumeID, err)
 	}
 
-	return fmt.Errorf("addSnapshotRef: exhausted retries for KV key %s", volumeID)
+	slog.Info("addSnapshotRef: added snapshot ref", "volumeId", volumeID, "snapshotId", snapshotID)
+	return nil
 }
 
 // removeSnapshotRefForCleanup removes a reference with a bounded context that
@@ -847,77 +829,89 @@ func (s *SnapshotServiceImpl) removeSnapshotRefForCleanup(ctx context.Context, v
 // Deletes the key if the list becomes empty.
 // Uses CAS (Update with revision) to prevent lost updates under concurrency.
 func (s *SnapshotServiceImpl) removeSnapshotRef(ctx context.Context, volumeID, snapshotID string) error {
-	if s.snapKV == nil {
+	if s.snapRefs == nil {
 		slog.Debug("removeSnapshotRef: snapshotKV is nil, skipping", "volumeId", volumeID, "snapshotId", snapshotID)
 		return nil
 	}
 
-	const maxRetries = 5
-	for attempt := range maxRetries {
-		entry, err := s.snapKV.Get(ctx, volumeID)
-		if err != nil {
-			if errors.Is(err, jetstream.ErrKeyNotFound) {
-				return nil
-			}
-			return fmt.Errorf("removeSnapshotRef: failed to get KV key %s: %w", volumeID, err)
-		}
-
-		var snapshots []string
-		if err := json.Unmarshal(entry.Value(), &snapshots); err != nil {
-			return fmt.Errorf("removeSnapshotRef: failed to unmarshal KV value for %s: %w", volumeID, err)
-		}
-
-		filtered := snapshots[:0]
-		for _, snap := range snapshots {
+	// A removal that finds nothing commits nothing: an identical record would
+	// still bump the revision, waking every watcher and costing the concurrent
+	// writers a CAS attempt each.
+	var removed, nowEmpty bool
+	err := s.snapRefs.Mutate(ctx, volumeID, func(snapshots *[]string) (bool, error) {
+		before := len(*snapshots)
+		filtered := (*snapshots)[:0]
+		for _, snap := range *snapshots {
 			if snap != snapshotID {
 				filtered = append(filtered, snap)
 			}
 		}
-
-		if len(filtered) == 0 {
-			if err := s.snapKV.Delete(ctx, volumeID); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
-				return fmt.Errorf("removeSnapshotRef: failed to delete KV key %s: %w", volumeID, err)
-			}
-		} else {
-			data, err := json.Marshal(filtered)
-			if err != nil {
-				return fmt.Errorf("removeSnapshotRef: failed to marshal snapshot list: %w", err)
-			}
-			if _, err := s.snapKV.Update(ctx, volumeID, data, entry.Revision()); err != nil {
-				if attempt < maxRetries-1 {
-					continue // concurrent update — retry
-				}
-				return fmt.Errorf("removeSnapshotRef: failed to update KV key %s: %w", volumeID, err)
-			}
-		}
-
-		slog.Info("removeSnapshotRef: removed snapshot ref", "volumeId", volumeID, "snapshotId", snapshotID)
+		*snapshots = filtered
+		removed = len(filtered) != before
+		nowEmpty = len(filtered) == 0
+		return removed, nil
+	})
+	if errors.Is(err, kvstore.ErrNotFound) {
 		return nil
 	}
+	if err != nil {
+		return fmt.Errorf("removeSnapshotRef: failed to drop %s from volume %s: %w", snapshotID, volumeID, err)
+	}
 
-	return fmt.Errorf("removeSnapshotRef: exhausted retries for KV key %s", volumeID)
+	// Still runs when nothing was removed, so a key left empty by a lost drop
+	// race is not stranded until the volume gains another snapshot.
+	if nowEmpty {
+		if err := s.dropEmptySnapshotRefs(ctx, volumeID); err != nil {
+			return err
+		}
+	}
+
+	if removed {
+		slog.Info("removeSnapshotRef: removed snapshot ref", "volumeId", volumeID, "snapshotId", snapshotID)
+	}
+	return nil
+}
+
+// dropEmptySnapshotRefs deletes a volume's index key once nothing references it,
+// so a volume that has lost its last snapshot leaves no key behind.
+//
+// The delete is guarded on the revision the emptiness was read at, so a
+// CreateSnapshot landing in between keeps its key rather than having it deleted
+// out from under it.
+func (s *SnapshotServiceImpl) dropEmptySnapshotRefs(ctx context.Context, volumeID string) error {
+	snapshots, rev, err := s.snapRefs.Get(ctx, volumeID)
+	if err != nil {
+		if errors.Is(err, kvstore.ErrNotFound) {
+			return nil
+		}
+		return fmt.Errorf("dropEmptySnapshotRefs: failed to get KV key %s: %w", volumeID, err)
+	}
+	if len(*snapshots) > 0 {
+		return nil
+	}
+	// A lost race means a snapshot was added back, so the key has to stay.
+	if err := s.snapRefs.CompareAndDelete(ctx, volumeID, rev); err != nil &&
+		!errors.Is(err, kvstore.ErrConflict) {
+		return fmt.Errorf("dropEmptySnapshotRefs: failed to delete KV key %s: %w", volumeID, err)
+	}
+	return nil
 }
 
 // volumeHasSnapshots returns true if the volume has any snapshots in KV.
 func (s *SnapshotServiceImpl) volumeHasSnapshots(ctx context.Context, volumeID string) (bool, error) {
-	if s.snapKV == nil {
+	if s.snapRefs == nil {
 		return false, nil
 	}
 
-	entry, err := s.snapKV.Get(ctx, volumeID)
+	snapshots, _, err := s.snapRefs.Get(ctx, volumeID)
 	if err != nil {
-		if errors.Is(err, jetstream.ErrKeyNotFound) {
+		if errors.Is(err, kvstore.ErrNotFound) {
 			return false, nil
 		}
 		return false, fmt.Errorf("volumeHasSnapshots: failed to get KV key %s: %w", volumeID, err)
 	}
 
-	var snapshots []string
-	if err := json.Unmarshal(entry.Value(), &snapshots); err != nil {
-		return false, fmt.Errorf("volumeHasSnapshots: failed to unmarshal KV value for %s: %w", volumeID, err)
-	}
-
-	return len(snapshots) > 0, nil
+	return len(*snapshots) > 0, nil
 }
 
 // ApplyRecordTags mirrors CreateTags into the owning snapshot metadata so

spinifex/handlers/ec2/snapshot/service_impl_test.go

@@ -5,7 +5,9 @@ import (
 	"encoding/json"
 	"errors"
 	"path"
+	"strconv"
 	"strings"
+	"sync"
 	"testing"
 
 	"github.com/aws/aws-sdk-go/aws"
@@ -15,6 +17,7 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/config"
 	"github.com/mulgadc/spinifex/spinifex/ebsmetadata"
 	"github.com/mulgadc/spinifex/spinifex/ebsprovider"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 	"github.com/mulgadc/spinifex/spinifex/objectstore"
 	"github.com/mulgadc/spinifex/spinifex/testutil"
 	"github.com/nats-io/nats.go/jetstream"
@@ -594,9 +597,18 @@ func setupTestNATSKV(t *testing.T) jetstream.KeyValue {
 	return kv
 }
 
+// newRefService builds a service whose only wiring is the volume-snapshot
+// index, which is all the ref tests touch. A nil kv is the disabled case.
+func newRefService(kv jetstream.KeyValue) *SnapshotServiceImpl {
+	if kv == nil {
+		return &SnapshotServiceImpl{}
+	}
+	return &SnapshotServiceImpl{snapRefs: kvstore.Over[[]string](nil, kv, snapshotRefsConfig())}
+}
+
 func TestAddSnapshotRef(t *testing.T) {
 	kv := setupTestNATSKV(t)
-	svc := &SnapshotServiceImpl{snapKV: kv}
+	svc := newRefService(kv)
 
 	require.NoError(t, svc.addSnapshotRef(t.Context(), "vol-1", "snap-a"))
 	require.NoError(t, svc.addSnapshotRef(t.Context(), "vol-1", "snap-b"))
@@ -610,7 +622,7 @@ func TestAddSnapshotRef(t *testing.T) {
 
 func TestRemoveSnapshotRef(t *testing.T) {
 	kv := setupTestNATSKV(t)
-	svc := &SnapshotServiceImpl{snapKV: kv}
+	svc := newRefService(kv)
 
 	require.NoError(t, svc.addSnapshotRef(t.Context(), "vol-1", "snap-a"))
 	require.NoError(t, svc.addSnapshotRef(t.Context(), "vol-1", "snap-b"))
@@ -633,7 +645,7 @@ func TestRemoveSnapshotRef(t *testing.T) {
 
 func TestRemoveSnapshotRef_NonExistentKey(t *testing.T) {
 	kv := setupTestNATSKV(t)
-	svc := &SnapshotServiceImpl{snapKV: kv}
+	svc := newRefService(kv)
 
 	// Should not error on non-existent key
 	require.NoError(t, svc.removeSnapshotRef(t.Context(), "vol-nonexistent", "snap-x"))
@@ -641,7 +653,7 @@ func TestRemoveSnapshotRef_NonExistentKey(t *testing.T) {
 
 func TestRemoveSnapshotRefForCleanupSurvivesCancellation(t *testing.T) {
 	kv := setupTestNATSKV(t)
-	svc := &SnapshotServiceImpl{snapKV: kv}
+	svc := newRefService(kv)
 	require.NoError(t, svc.addSnapshotRef(t.Context(), "vol-1", "snap-a"))
 
 	ctx, cancel := context.WithCancel(t.Context())
@@ -654,7 +666,7 @@ func TestRemoveSnapshotRefForCleanupSurvivesCancellation(t *testing.T) {
 
 func TestVolumeHasSnapshots(t *testing.T) {
 	kv := setupTestNATSKV(t)
-	svc := &SnapshotServiceImpl{snapKV: kv}
+	svc := newRefService(kv)
 
 	// No entry → false
 	has, err := svc.volumeHasSnapshots(t.Context(), "vol-1")
@@ -675,7 +687,7 @@ func TestVolumeHasSnapshots(t *testing.T) {
 }
 
 func TestKVNilFallback(t *testing.T) {
-	svc := &SnapshotServiceImpl{snapKV: nil}
+	svc := newRefService(nil)
 
 	// All methods should be no-ops when KV is nil
 	require.NoError(t, svc.addSnapshotRef(t.Context(), "vol-1", "snap-a"))
@@ -1273,3 +1285,96 @@ type copyFailingProvider struct {
 func (p *copyFailingProvider) CopySnapshot(context.Context, ebsprovider.CopySnapshotRequest) (*ebsprovider.Snapshot, error) {
 	return nil, errors.New("provider copy failed")
 }
+
+// Concurrent adds against one volume must not lose a reference, which is what
+// the CAS loop exists to prevent.
+func TestAddSnapshotRefConcurrentNoLostUpdates(t *testing.T) {
+	kv := setupTestNATSKV(t)
+	svc := newRefService(kv)
+
+	const refs = 16
+	var wg sync.WaitGroup
+	errs := make([]error, refs)
+	for i := range refs {
+		wg.Go(func() {
+			errs[i] = svc.addSnapshotRef(t.Context(), "vol-1", "snap-"+strconv.Itoa(i))
+		})
+	}
+	wg.Wait()
+	require.NoError(t, errors.Join(errs...))
+
+	entry, err := kv.Get(t.Context(), "vol-1")
+	require.NoError(t, err)
+	var snapshots []string
+	require.NoError(t, json.Unmarshal(entry.Value(), &snapshots))
+	assert.Len(t, snapshots, refs)
+}
+
+// Removing the last reference takes the key with it, so a volume that has lost
+// every snapshot leaves nothing behind for volumeHasSnapshots to read.
+func TestRemoveSnapshotRefDropsTheEmptiedKey(t *testing.T) {
+	kv := setupTestNATSKV(t)
+	svc := newRefService(kv)
+
+	require.NoError(t, svc.addSnapshotRef(t.Context(), "vol-1", "snap-a"))
+	require.NoError(t, svc.removeSnapshotRef(t.Context(), "vol-1", "snap-a"))
+
+	_, err := kv.Get(t.Context(), "vol-1")
+	require.ErrorIs(t, err, jetstream.ErrKeyNotFound, "the emptied index key must be gone")
+
+	has, err := svc.volumeHasSnapshots(t.Context(), "vol-1")
+	require.NoError(t, err)
+	assert.False(t, has)
+}
+
+// dropEmptySnapshotRefs leaves a key that is no longer empty alone, so a volume
+// that gained a snapshot after the check keeps its index.
+func TestDropEmptySnapshotRefsLeavesANonEmptyKey(t *testing.T) {
+	kv := setupTestNATSKV(t)
+	svc := newRefService(kv)
+
+	require.NoError(t, svc.addSnapshotRef(t.Context(), "vol-1", "snap-b"))
+	require.NoError(t, svc.dropEmptySnapshotRefs(t.Context(), "vol-1"))
+
+	entry, err := kv.Get(t.Context(), "vol-1")
+	require.NoError(t, err)
+	var snapshots []string
+	require.NoError(t, json.Unmarshal(entry.Value(), &snapshots))
+	assert.Equal(t, []string{"snap-b"}, snapshots)
+}
+
+// The empty-key delete is guarded on the revision the emptiness was read at, so
+// an add landing between that read and the delete keeps its reference. An
+// unguarded delete drops it, and volumeHasSnapshots then reports no snapshots
+// for a volume that has one — which is what DeleteVolume gates on.
+func TestDropEmptySnapshotRefsDoesNotDropAConcurrentAdd(t *testing.T) {
+	kv := setupTestNATSKV(t)
+	svc := newRefService(kv)
+
+	// Repeated because the losing interleaving is a window, not a certainty.
+	for i := range 200 {
+		volume := "vol-" + strconv.Itoa(i)
+		_, err := kv.Put(t.Context(), volume, []byte("[]"))
+		require.NoError(t, err)
+
+		start := make(chan struct{})
+		var wg sync.WaitGroup
+		var dropErr, addErr error
+		wg.Go(func() {
+			<-start
+			dropErr = svc.dropEmptySnapshotRefs(t.Context(), volume)
+		})
+		wg.Go(func() {
+			<-start
+			addErr = svc.addSnapshotRef(t.Context(), volume, "snap-b")
+		})
+		close(start)
+		wg.Wait()
+		require.NoError(t, dropErr)
+		require.NoError(t, addErr)
+
+		has, err := svc.volumeHasSnapshots(t.Context(), volume)
+		require.NoError(t, err)
+		assert.True(t, has, "iteration %d: the concurrent add was dropped by the empty-key delete", i)
+	}
+}

spinifex/handlers/ec2/vpc/external_ipam.go

@@ -50,7 +50,9 @@ func NewExternalIPAM(ctx context.Context, js jetstream.JetStream, pools []extern
 		if err != nil {
 			return nil, err
 		}
-		kv = alloc.KV()
+		if kv, err = alloc.KV(ctx); err != nil {
+			return nil, err
+		}
 	}
 	return &ExternalIPAM{kv: kv, pools: pools, static: alloc, perPool: map[string]external.Allocator{}}, nil
 }

spinifex/handlers/ec2/vpc/ipam.go

@@ -2,13 +2,12 @@ package handlers_ec2_vpc
 
 import (
 	"context"
-	"encoding/json"
 	"errors"
 	"fmt"
 	"log/slog"
 	"net/netip"
 
-	"github.com/mulgadc/spinifex/spinifex/kvutil"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 	"github.com/mulgadc/spinifex/spinifex/migrate"
 	"github.com/nats-io/nats.go/jetstream"
 )
@@ -16,8 +15,15 @@ import (
 const (
 	KVBucketIPAM        = "spinifex-vpc-ipam"
 	KVBucketIPAMVersion = 2
+	KVBucketIPAMHistory = 5
 )
 
+// ipamCASAttempts bounds one allocation's retry on a revision conflict. Every
+// address in a subnet shares one record, so a bulk RunInstances contends with
+// itself once per instance; the default five is below a single realistic launch
+// burst and returns exhaustion where the caller expects an address.
+const ipamCASAttempts = 50
+
 // ErrIPOutOfRange means the requested address is not a usable host address in
 // the subnet: outside the CIDR, in the reserved head, or the broadcast address.
 var ErrIPOutOfRange = errors.New("ip address not usable in subnet")
@@ -42,72 +48,74 @@ type IPAMRecord struct {
 
 // IPAM manages IP address allocation for VPC subnets using NATS KV with CAS.
 type IPAM struct {
-	kv jetstream.KeyValue
+	store *kvstore.Store[IPAMRecord]
+}
+
+// ipamConfig describes the IPAM bucket. js is captured for the migration, which
+// resolves the stream itself and so needs the client rather than the handle; a
+// nil js is a caller supplying its own already-migrated handle.
+func ipamConfig(js jetstream.JetStream) kvstore.Config {
+	cfg := kvstore.Config{
+		Name:     KVBucketIPAM,
+		History:  KVBucketIPAMHistory,
+		Missing:  "IPAM KV bucket not initialized",
+		Attempts: ipamCASAttempts,
+		Exhausted: func(key string, attempts int) error {
+			return fmt.Errorf("IPAM CAS exhausted for subnet %s after %d attempts", key, attempts)
+		},
+	}
+	if js != nil {
+		cfg.OnOpen = func(ctx context.Context, kv jetstream.KeyValue) error {
+			return migrate.DefaultRegistry.RunKVWithJetStream(ctx, KVBucketIPAM, kv, js, KVBucketIPAMVersion)
+		}
+	}
+	return cfg
 }
 
 // NewIPAM creates a new IPAM instance backed by NATS JetStream KV.
 func NewIPAM(ctx context.Context, js jetstream.JetStream) (*IPAM, error) {
-	kv, err := kvutil.GetOrCreateBucket(ctx, js, KVBucketIPAM, 5)
-	if err != nil {
+	store := kvstore.New[IPAMRecord](js, ipamConfig(js))
+	// Opened eagerly: a bucket that cannot be created must fail construction
+	// rather than the first allocation that needs it.
+	if _, err := store.KV(ctx); err != nil {
 		return nil, fmt.Errorf("failed to create IPAM KV bucket: %w", err)
 	}
-	if err := migrate.DefaultRegistry.RunKVWithJetStream(ctx, KVBucketIPAM, kv, js, KVBucketIPAMVersion); err != nil {
-		return nil, fmt.Errorf("migrate %s: %w", KVBucketIPAM, err)
-	}
-	return &IPAM{kv: kv}, nil
+	return &IPAM{store: store}, nil
 }
 
 // NewIPAMWithKV creates an IPAM with an existing KV bucket (for testing).
 func NewIPAMWithKV(kv jetstream.KeyValue) *IPAM {
-	return &IPAM{kv: kv}
+	return &IPAM{store: kvstore.Over[IPAMRecord](nil, kv, ipamConfig(nil))}
 }
 
 // AllocateIP allocates an IP from the subnet, reserving the first 4 and last
 // addresses per AWS convention. Uses CAS for conflict-free multi-node allocation.
 func (m *IPAM) AllocateIP(ctx context.Context, subnetId, cidrBlock, purpose, ownerID string) (string, error) {
-	for attempt := range 5 {
-		record, revision, err := m.getRecord(ctx, subnetId)
-		if err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
-			return "", fmt.Errorf("get IPAM record: %w", err)
+	var ip string
+	err := m.store.Upsert(ctx, subnetId, func(record *IPAMRecord) (bool, error) {
+		// Only an absent record takes the caller's CIDR. A stored one keeps its
+		// own, so a caller passing a stale subnet CIDR cannot re-base a pool
+		// that already has addresses handed out of it.
+		if record.CidrBlock == "" {
+			record.SubnetId = subnetId
+			record.CidrBlock = cidrBlock
 		}
 
-		if record == nil {
-			record = &IPAMRecord{
-				SubnetId:  subnetId,
-				CidrBlock: cidrBlock,
-			}
-		}
-
-		// Find next available IP
-		ip, err := m.nextAvailableIP(record)
+		next, err := m.nextAvailableIP(record)
 		if err != nil {
-			return "", err
+			return false, err
 		}
 
-		record.Allocated = append(record.Allocated, IPEntry{IP: ip, Purpose: purpose, OwnerID: ownerID})
-
-		// CAS write
-		data, err := json.Marshal(record)
-		if err != nil {
-			return "", fmt.Errorf("marshal IPAM record: %w", err)
-		}
-
-		if revision == 0 {
-			_, err = m.kv.Create(ctx, subnetId, data)
-		} else {
-			_, err = m.kv.Update(ctx, subnetId, data, revision)
-		}
-
-		if err != nil {
-			slog.Debug("IPAM CAS conflict, retrying", "subnet", subnetId, "attempt", attempt)
-			continue // CAS conflict, retry
-		}
-
-		slog.Info("IPAM allocated IP", "subnet", subnetId, "ip", ip, "purpose", purpose, "owner", ownerID)
-		return ip, nil
+		record.Allocated = append(record.Allocated, IPEntry{IP: next, Purpose: purpose, OwnerID: ownerID})
+		ip = next
+		return true, nil
+	})
+	if err != nil {
+		return "", err
 	}
 
-	return "", fmt.Errorf("IPAM allocation failed after CAS retries for subnet %s", subnetId)
+	slog.Info("IPAM allocated IP", "subnet", subnetId, "ip", ip, "purpose", purpose, "owner", ownerID)
+	return ip, nil
 }
 
 // ClaimIP reserves a caller-specified IP in the subnet, for a pinned address
@@ -132,48 +140,30 @@ func (m *IPAM) ClaimIP(ctx context.Context, subnetId, cidrBlock, purpose, ownerI
 
 	ip := addr.String()
 
-	for attempt := range 5 {
-		record, revision, err := m.getRecord(ctx, subnetId)
-		if err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
-			return "", fmt.Errorf("get IPAM record: %w", err)
+	err = m.store.Upsert(ctx, subnetId, func(record *IPAMRecord) (bool, error) {
+		// Only an absent record takes the caller's CIDR, matching AllocateIP: a
+		// stale subnet CIDR must not re-base a pool with addresses already out.
+		if record.CidrBlock == "" {
+			record.SubnetId = subnetId
+			record.CidrBlock = cidrBlock
 		}
-
-		if record == nil {
-			record = &IPAMRecord{
-				SubnetId:  subnetId,
-				CidrBlock: cidrBlock,
-			}
-		}
-
+		// Re-checked on every attempt rather than once up front, so a claim that
+		// lost the race to the same address reports it in use instead of
+		// appending a duplicate entry.
 		for _, entry := range record.Allocated {
 			if entry.IP == ip {
-				return "", fmt.Errorf("%w: %s already allocated in subnet %s", ErrIPInUse, ip, subnetId)
+				return false, fmt.Errorf("%w: %s already allocated in subnet %s", ErrIPInUse, ip, subnetId)
 			}
 		}
-
 		record.Allocated = append(record.Allocated, IPEntry{IP: ip, Purpose: purpose, OwnerID: ownerID})
-
-		data, err := json.Marshal(record)
-		if err != nil {
-			return "", fmt.Errorf("marshal IPAM record: %w", err)
-		}
-
-		if revision == 0 {
-			_, err = m.kv.Create(ctx, subnetId, data)
-		} else {
-			_, err = m.kv.Update(ctx, subnetId, data, revision)
-		}
-
-		if err != nil {
-			slog.Debug("IPAM claim CAS conflict, retrying", "subnet", subnetId, "attempt", attempt)
-			continue // CAS conflict, retry: a concurrent writer may have taken this IP
-		}
-
-		slog.Info("IPAM claimed IP", "subnet", subnetId, "ip", ip, "purpose", purpose, "owner", ownerID)
-		return ip, nil
+		return true, nil
+	})
+	if err != nil {
+		return "", err
 	}
 
-	return "", fmt.Errorf("IPAM claim failed after CAS retries for subnet %s", subnetId)
+	slog.Info("IPAM claimed IP", "subnet", subnetId, "ip", ip, "purpose", purpose, "owner", ownerID)
+	return ip, nil
 }
 
 // isClaimableAddr reports whether addr is a usable host address in prefix:
@@ -198,76 +188,38 @@ func isClaimableAddr(prefix netip.Prefix, addr netip.Addr) bool {
 
 // ReleaseIP releases a previously allocated IP address back to the subnet pool.
 func (m *IPAM) ReleaseIP(ctx context.Context, subnetId, ip string) error {
-	for attempt := range 5 {
-		record, revision, err := m.getRecord(ctx, subnetId)
-		if err != nil {
-			return fmt.Errorf("get IPAM record for release: %w", err)
-		}
-		if record == nil {
-			return fmt.Errorf("no IPAM record for subnet %s", subnetId)
-		}
-
-		found := false
+	err := m.store.Mutate(ctx, subnetId, func(record *IPAMRecord) (bool, error) {
 		for i, entry := range record.Allocated {
 			if entry.IP == ip {
 				record.Allocated = append(record.Allocated[:i], record.Allocated[i+1:]...)
-				found = true
-				break
+				return true, nil
 			}
 		}
-		if !found {
-			return fmt.Errorf("IP %s not allocated in subnet %s", ip, subnetId)
-		}
-
-		data, err := json.Marshal(record)
-		if err != nil {
-			return fmt.Errorf("marshal IPAM record: %w", err)
-		}
-
-		if _, err := m.kv.Update(ctx, subnetId, data, revision); err != nil {
-			slog.Debug("IPAM release CAS conflict, retrying", "subnet", subnetId, "attempt", attempt)
-			continue
-		}
-
-		slog.Info("IPAM released IP", "subnet", subnetId, "ip", ip)
-		return nil
+		return false, fmt.Errorf("IP %s not allocated in subnet %s", ip, subnetId)
+	})
+	if errors.Is(err, kvstore.ErrNotFound) {
+		return fmt.Errorf("no IPAM record for subnet %s", subnetId)
+	}
+	if err != nil {
+		return err
 	}
 
-	return fmt.Errorf("IPAM release failed after CAS retries for subnet %s", subnetId)
+	slog.Info("IPAM released IP", "subnet", subnetId, "ip", ip)
+	return nil
 }
 
 // AllocatedIPs returns the list of allocated IP entries for a subnet.
 func (m *IPAM) AllocatedIPs(ctx context.Context, subnetId string) ([]IPEntry, error) {
-	record, _, err := m.getRecord(ctx, subnetId)
+	record, _, err := m.store.Get(ctx, subnetId)
 	if err != nil {
-		if errors.Is(err, jetstream.ErrKeyNotFound) {
+		if errors.Is(err, kvstore.ErrNotFound) {
 			return nil, nil
 		}
 		return nil, err
 	}
-	if record == nil {
-		return nil, nil
-	}
 	return record.Allocated, nil
 }
 
-func (m *IPAM) getRecord(ctx context.Context, subnetId string) (*IPAMRecord, uint64, error) {
-	entry, err := m.kv.Get(ctx, subnetId)
-	if err != nil {
-		if errors.Is(err, jetstream.ErrKeyNotFound) {
-			return nil, 0, jetstream.ErrKeyNotFound
-		}
-		return nil, 0, err
-	}
-
-	var record IPAMRecord
-	if err := json.Unmarshal(entry.Value(), &record); err != nil {
-		return nil, 0, fmt.Errorf("unmarshal IPAM record: %w", err)
-	}
-
-	return &record, entry.Revision(), nil
-}
-
 // nextAvailableIP finds the next available IP in the subnet, skipping reserved addresses.
 func (m *IPAM) nextAvailableIP(record *IPAMRecord) (string, error) {
 	prefix, err := netip.ParsePrefix(record.CidrBlock)

spinifex/handlers/ec2/vpc/ipam_test.go

@@ -347,3 +347,39 @@ func TestCompareIPs_NilSortsFirst(t *testing.T) {
 func itoa(i int) string {
 	return strconv.Itoa(i)
 }
+
+// Concurrent allocations against one subnet must not hand the same address to
+// two callers, which is what the CAS loop exists to prevent.
+func TestIPAM_AllocateConcurrentNoDuplicates(t *testing.T) {
+	t.Parallel()
+	ipam := setupTestIPAM(t)
+
+	const allocations = 16
+	var (
+		wg   sync.WaitGroup
+		mu   sync.Mutex
+		ips  = make(map[string]bool, allocations)
+		errs []error
+	)
+	for i := range allocations {
+		wg.Go(func() {
+			ip, err := ipam.AllocateIP(t.Context(), "subnet-1", "10.0.1.0/24",
+				PurposeENIPrimary, "eni-"+strconv.Itoa(i))
+			mu.Lock()
+			defer mu.Unlock()
+			if err != nil {
+				errs = append(errs, err)
+				return
+			}
+			ips[ip] = true
+		})
+	}
+	wg.Wait()
+
+	require.Empty(t, errs)
+	assert.Len(t, ips, allocations)
+
+	entries, err := ipam.AllocatedIPs(t.Context(), "subnet-1")
+	require.NoError(t, err)
+	assert.Len(t, entries, allocations)
+}

spinifex/handlers/ecs/store.go

@@ -4,12 +4,11 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"log/slog"
 	"time"
 
+	"github.com/mulgadc/spinifex/spinifex/kvlease"
 	"github.com/mulgadc/spinifex/spinifex/kvutil"
 	"github.com/mulgadc/spinifex/spinifex/migrate"
-	"github.com/mulgadc/spinifex/spinifex/otelsetup"
 	"github.com/nats-io/nats.go"
 	"github.com/nats-io/nats.go/jetstream"
 )
@@ -218,28 +217,11 @@ func GetOrCreateAccountBucket(ctx context.Context, js jetstream.JetStream, accou
 }
 
 // InitLeaderBucket creates (or attaches to) the shared spinifex-ecs-leader bucket
-// used for per-cluster scheduler leader-lease CAS locks. The bucket is configured
-// with History=1 and a 60s TTL so stale leases expire on their own when a leader
-// dies mid-cycle. kvutil.GetOrCreateBucket doesn't expose a TTL knob, so this
-// attaches directly and creates only when absent.
+// used for per-cluster scheduler leader-lease CAS locks.
 func InitLeaderBucket(ctx context.Context, js jetstream.JetStream) (jetstream.KeyValue, error) {
-	// Attach before create. This runs on every scheduler tick, and CreateKeyValue
-	// against a bucket that exists with any other config is a STREAM.CREATE the
-	// meta leader answers with an error, so creating first bills one per tick.
-	kv, err := js.KeyValue(ctx, KVBucketECSLeader)
-	if errors.Is(err, jetstream.ErrBucketNotFound) {
-		kv, err = js.CreateKeyValue(ctx, jetstream.KeyValueConfig{
-			Bucket:  KVBucketECSLeader,
-			History: 1,
-			TTL:     KVBucketECSLeaderTTL,
-		})
-	}
-	if err != nil {
-		return nil, fmt.Errorf("failed to create or open ECS leader bucket %s: %w", KVBucketECSLeader, err)
-	}
-	if err := migrate.DefaultRegistry.RunKV(ctx, KVBucketECSLeader, kv, KVBucketECSLeaderVersion); err != nil {
-		return nil, fmt.Errorf("migrate %s: %w", KVBucketECSLeader, err)
-	}
-	slog.Info("ECS leader bucket initialized", "bucket", KVBucketECSLeader, "ttl_ms", otelsetup.Millis(KVBucketECSLeaderTTL))
-	return kv, nil
+	return kvlease.OpenBucket(ctx, js, kvlease.BucketConfig{
+		Name:    KVBucketECSLeader,
+		TTL:     KVBucketECSLeaderTTL,
+		Version: KVBucketECSLeaderVersion,
+	})
 }

spinifex/handlers/eks/store.go

@@ -6,14 +6,13 @@ import (
 	"encoding/hex"
 	"errors"
 	"fmt"
-	"log/slog"
 	"strings"
 	"time"
 
+	"github.com/mulgadc/spinifex/spinifex/kvlease"
 	"github.com/mulgadc/spinifex/spinifex/kvstore"
 	"github.com/mulgadc/spinifex/spinifex/kvutil"
 	"github.com/mulgadc/spinifex/spinifex/migrate"
-	"github.com/mulgadc/spinifex/spinifex/otelsetup"
 	"github.com/nats-io/nats.go"
 	"github.com/nats-io/nats.go/jetstream"
 )
@@ -247,30 +246,12 @@ func GetOrCreateAccountBucket(ctx context.Context, js jetstream.JetStream, accou
 
 // InitLeaderBucket creates (or attaches to) the shared spinifex-eks-leader
 // bucket used for per-cluster reconciler leader-lease CAS locks, at the given
-// replica count (clamped to a minimum of 1). The bucket is configured with
-// History=1 and a 60s TTL so stale leases expire on their own when a leader
-// dies mid-cycle. kvutil.GetOrCreateBucketWithReplicas doesn't expose a TTL
-// knob, so this function sets Replicas on its own create call, which runs only
-// when the bucket is absent.
+// replica count (clamped to a minimum of 1).
 func InitLeaderBucket(ctx context.Context, js jetstream.JetStream, replicas int) (jetstream.KeyValue, error) {
-	// Attach before create, so replicas applies to a genuine first creation only.
-	// CreateKeyValue against a bucket that exists with any other config is a
-	// STREAM.CREATE the meta leader answers with an error rather than a no-op.
-	kv, err := js.KeyValue(ctx, KVBucketEKSLeader)
-	if errors.Is(err, jetstream.ErrBucketNotFound) {
-		kv, err = js.CreateKeyValue(ctx, jetstream.KeyValueConfig{
-			Bucket:   KVBucketEKSLeader,
-			History:  1,
-			TTL:      KVBucketEKSLeaderTTL,
-			Replicas: max(replicas, 1),
-		})
-	}
-	if err != nil {
-		return nil, fmt.Errorf("failed to create or open EKS leader bucket %s: %w", KVBucketEKSLeader, err)
-	}
-	if err := migrate.DefaultRegistry.RunKV(ctx, KVBucketEKSLeader, kv, KVBucketEKSLeaderVersion); err != nil {
-		return nil, fmt.Errorf("migrate %s: %w", KVBucketEKSLeader, err)
-	}
-	slog.Info("EKS leader bucket initialized", "bucket", KVBucketEKSLeader, "ttl_ms", otelsetup.Millis(KVBucketEKSLeaderTTL))
-	return kv, nil
+	return kvlease.OpenBucket(ctx, js, kvlease.BucketConfig{
+		Name:     KVBucketEKSLeader,
+		TTL:      KVBucketEKSLeaderTTL,
+		Replicas: max(replicas, 1),
+		Version:  KVBucketEKSLeaderVersion,
+	})
 }

spinifex/handlers/ochrevector/appliance.go

@@ -367,7 +367,7 @@ func (a *Appliance) promote(ctx context.Context, rec ApplianceRecord, rev uint64
 	rec.Endpoint = endpoint
 	rec.Port = port
 	rec.UpdatedAt = time.Now().UTC()
-	if err := a.store.CompareAndSet(ctx, appliancePostgresKey, &rec, rev); err != nil {
+	if _, err := a.store.CompareAndSet(ctx, appliancePostgresKey, &rec, rev); err != nil {
 		if current, _, getErr := a.getRecord(ctx); getErr == nil && current != nil && current.State == ApplianceStateAvailable {
 			return connInfo(*current), nil
 		}

spinifex/handlers/quota/quota.go

@@ -5,8 +5,10 @@
 package handlers_quota
 
 import (
+	"fmt"
 	"time"
 
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 	"github.com/mulgadc/spinifex/spinifex/utils"
 	"github.com/nats-io/nats.go/jetstream"
 )
@@ -69,10 +71,10 @@ type Limits struct {
 // derive their usage from account-filtered Describe* calls and need no state.
 type Service struct {
 	limits Limits
-	// usage holds the per-account vCPU counters (key {accountID}). Nil when
-	// quotas are disabled, in which case Exempt short-circuits every check
-	// before the counter is touched.
-	usage jetstream.KeyValue
+	// usage holds the per-account vCPU counters (key {accountID}). Backed by a
+	// nil bucket when quotas are disabled, in which case Exempt short-circuits
+	// every check before the counter is touched.
+	usage *kvstore.Store[int]
 
 	// bedrockUsage resolves the stream-fed Bedrock token counter
 	// Nil until SetBedrockUsage is called, in which case
@@ -93,8 +95,25 @@ type Service struct {
 // New constructs a quota Service from the configured limits and the gateway-owned
 // account-usage KV bucket. usage may be nil when quotas are disabled; Exempt then
 // short-circuits every check before the counter is read.
+//
+// The store is built over the handle rather than over a JetStream client, so a
+// bucket the gateway could not open stays unopenable here instead of being
+// quietly recreated under a counter that enforces a cap.
 func New(limits Limits, usage jetstream.KeyValue) *Service {
-	return &Service{limits: limits, usage: usage, rpm: newRPMLimiter()}
+	return &Service{limits: limits, usage: kvstore.Over[int](nil, usage, usageConfig()), rpm: newRPMLimiter()}
+}
+
+// usageConfig describes the account-usage bucket as the vCPU counter reads it.
+// The bucket itself is created by the gateway, so only the CAS policy matters.
+func usageConfig() kvstore.Config {
+	return kvstore.Config{
+		Name:     KVBucketAccountUsage,
+		Attempts: vcpuCASRetries,
+		Missing:  "account usage KV bucket not configured",
+		Exhausted: func(key string, attempts int) error {
+			return fmt.Errorf("vcpu counter CAS exhausted for %s after %d attempts", key, attempts)
+		},
+	}
 }
 
 // SetOverrides attaches the per-account override bucket. Called once at wiring

spinifex/handlers/quota/reconcile_test.go

@@ -81,9 +81,7 @@ func TestReconcileSkipsSystemAccount(t *testing.T) {
 	s := newVCPUService(t, Limits{Enabled: true, VCPUs: 100})
 
 	// Park a non-zero value directly under the system key; AddVCPU would no-op it.
-	if _, err := s.usage.PutString(t.Context(), utils.GlobalAccountID, "42"); err != nil {
-		t.Fatalf("seed system counter: %v", err)
-	}
+	seedCounter(t, s, utils.GlobalAccountID, 42)
 
 	if err := s.Reconcile(context.Background(), accountList(utils.GlobalAccountID, testAccount),
 		staticTotals(map[string]int{utils.GlobalAccountID: 99, testAccount: 2}, true, nil)); err != nil {
@@ -188,9 +186,7 @@ func TestReconcileAccountTouchesOnlyThatAccount(t *testing.T) {
 func TestReconcileAccountSkipsSystemAccount(t *testing.T) {
 	s := newVCPUService(t, Limits{Enabled: true, VCPUs: 100})
 
-	if _, err := s.usage.PutString(t.Context(), utils.GlobalAccountID, "42"); err != nil {
-		t.Fatalf("seed system counter: %v", err)
-	}
+	seedCounter(t, s, utils.GlobalAccountID, 42)
 	if err := s.ReconcileAccount(context.Background(), utils.GlobalAccountID,
 		staticTotals(map[string]int{utils.GlobalAccountID: 0}, true, nil)); err != nil {
 		t.Fatalf("ReconcileAccount: %v", err)
@@ -217,6 +213,15 @@ func TestReconcileDisabledNoop(t *testing.T) {
 	}
 }
 
+// seedCounter writes a counter straight to the bucket, bypassing the CAS path
+// under test.
+func seedCounter(t *testing.T, s *Service, accountID string, value int) {
+	t.Helper()
+	if err := s.usage.Set(t.Context(), accountID, &value); err != nil {
+		t.Fatalf("seed counter %s: %v", accountID, err)
+	}
+}
+
 func assertCounter(t *testing.T, s *Service, accountID string, want int) {
 	t.Helper()
 	got, _, err := s.readVCPU(t.Context(), accountID)

spinifex/handlers/quota/vcpu.go

@@ -2,25 +2,26 @@ package handlers_quota
 
 import (
 	"context"
-	"encoding/json"
 	"errors"
-	"fmt"
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
 	gateway_ec2_instance "github.com/mulgadc/spinifex/spinifex/gateway/ec2/instance"
 	"github.com/mulgadc/spinifex/spinifex/instancetypes"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 	"github.com/nats-io/nats.go"
-	"github.com/nats-io/nats.go/jetstream"
 )
 
 // vcpuCASRetries bounds AddVCPU's retry on a revision conflict. Each retry
 // implies another writer committed, so a single grow can conflict at most as
-// many times as there are concurrent grows on the account; the bound sits well
-// above any realistic in-flight launch burst and trips only as a circuit
-// breaker, never on genuine contention.
-const vcpuCASRetries = 100
+// many times as there are concurrent grows on the account, and this covers a
+// realistic in-flight launch burst.
+//
+// Sized against the backoff, not just the contending set: the retry pauses
+// between attempts, so the budget is a wall-clock one and the counter write
+// sits on the RunInstances response path.
+const vcpuCASRetries = 25
 
 // CheckVCPU rejects with ResourceLimitExceeded when charging want more vCPUs to
 // accountID would exceed the configured cap. It only reads the counter; the
@@ -43,35 +44,14 @@ func (s *Service) CheckVCPU(ctx context.Context, accountID string, want int) err
 // value and a skip flag; skip leaves the counter untouched and never creates a
 // key. A revision conflict is retried; exhausting the bound is a hard error.
 func (s *Service) casVCPU(ctx context.Context, accountID string, next func(current int) (value int, skip bool)) error {
-	for range vcpuCASRetries {
-		current, revision, err := s.readVCPU(ctx, accountID)
-		if err != nil {
-			return err
-		}
-		value, skip := next(current)
+	return s.usage.Upsert(ctx, accountID, func(current *int) (bool, error) {
+		value, skip := next(*current)
 		if skip {
-			return nil
+			return false, nil
 		}
-		data, err := json.Marshal(value)
-		if err != nil {
-			return err
-		}
-		if revision == 0 {
-			_, err = s.usage.Create(ctx, accountID, data)
-		} else {
-			_, err = s.usage.Update(ctx, accountID, 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) {
-			return err
-		}
-	}
-	return fmt.Errorf("vcpu counter CAS exhausted for %s after %d attempts", accountID, vcpuCASRetries)
+		*current = value
+		return true, nil
+	})
 }
 
 // AddVCPU adds delta vCPUs to accountID's counter under CAS so concurrent grows
@@ -111,18 +91,17 @@ func (s *Service) setVCPUAt(ctx context.Context, accountID string, value int, si
 	if revision != sinceRevision || current == value {
 		return nil
 	}
-	data, err := json.Marshal(value)
-	if err != nil {
-		return err
-	}
+	// Single-shot rather than Upsert: the guard is the revision the caller
+	// observed, so a retry against the winner's revision is the write this
+	// deliberately drops.
 	if revision == 0 {
-		_, err = s.usage.Create(ctx, accountID, data)
+		_, err = s.usage.Create(ctx, accountID, &value)
 	} else {
-		_, err = s.usage.Update(ctx, accountID, data, revision)
+		_, err = s.usage.CompareAndSet(ctx, accountID, &value, revision)
 	}
 	// Losing the race is the guard working, not a failure: a charge got there
 	// first and the next pass recomputes against it.
-	if errors.Is(err, jetstream.ErrKeyExists) || errors.Is(err, jetstream.ErrKeyRevisionMismatch) {
+	if errors.Is(err, kvstore.ErrExists) || errors.Is(err, kvstore.ErrConflict) {
 		return nil
 	}
 	return err
@@ -272,15 +251,12 @@ func instanceTypeFromReservations(reservations []*ec2.Reservation, instanceID st
 // it was read at. A missing key is the zero counter at revision 0, which AddVCPU
 // treats as a create.
 func (s *Service) readVCPU(ctx context.Context, accountID string) (count int, revision uint64, err error) {
-	entry, err := s.usage.Get(ctx, accountID)
+	value, revision, err := s.usage.Get(ctx, accountID)
 	if err != nil {
-		if errors.Is(err, jetstream.ErrKeyNotFound) {
+		if errors.Is(err, kvstore.ErrNotFound) {
 			return 0, 0, nil
 		}
 		return 0, 0, err
 	}
-	if err := json.Unmarshal(entry.Value(), &count); err != nil {
-		return 0, 0, fmt.Errorf("unmarshal vcpu counter for %s: %w", accountID, err)
-	}
-	return count, entry.Revision(), nil
+	return *value, revision, nil
 }

spinifex/handlers/rds/agent.go

@@ -7,9 +7,8 @@ import (
 	"log/slog"
 	"time"
 
-	"github.com/nats-io/nats.go/jetstream"
-
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 	"github.com/mulgadc/spinifex/spinifex/vm"
 )
 
@@ -416,7 +415,7 @@ func (s *Service) GetDBBootstrapConfig(ctx context.Context, input *GetDBBootstra
 // A failure to persist the backfill is not a failure to serve it: the value is
 // re-derivable and the next fetch records it, whereas failing here would hold a
 // bootable instance down for a bookkeeping write.
-func (s *Service) resolveClientCIDR(ctx context.Context, kv jetstream.KeyValue, rec *DBInstanceRecord) (string, error) {
+func (s *Service) resolveClientCIDR(ctx context.Context, kv *kvstore.Bucket, rec *DBInstanceRecord) (string, error) {
 	if rec.VpcCIDR != "" {
 		return rec.VpcCIDR, nil
 	}
@@ -453,7 +452,7 @@ func (s *Service) resolveClientCIDR(ctx context.Context, kv jetstream.KeyValue,
 // rather than returned. Propagating it would discard an attach response this
 // call has already built and reinstate the retry loop it exists to avoid; the
 // verdict is idempotent, so the next fetch records it.
-func (s *Service) applyStagedBootstrap(ctx context.Context, kv jetstream.KeyValue, accountID string,
+func (s *Service) applyStagedBootstrap(ctx context.Context, kv *kvstore.Bucket, accountID string,
 	rec *DBInstanceRecord, generationMatches bool, out *GetDBBootstrapConfigOutput) error {
 	envelope, _, err := readBootstrapPayload(ctx, kv, rec.DBInstanceIdentifier)
 	if err != nil {
@@ -500,7 +499,7 @@ func (s *Service) applyStagedBootstrap(ctx context.Context, kv jetstream.KeyValu
 // was never consumed also never bootstrapped, so there is no datadir to save:
 // the password is removed on sight and the instance is marked for recreation
 // rather than left holding a secret at rest for an operator who may never act.
-func (s *Service) scrubLegacyBootstrap(ctx context.Context, kv jetstream.KeyValue,
+func (s *Service) scrubLegacyBootstrap(ctx context.Context, kv *kvstore.Bucket,
 	accountID string, rec *DBInstanceRecord) error {
 	if rec.Bootstrap.MasterUserPassword == "" {
 		return nil
@@ -579,7 +578,7 @@ func (s *Service) AcknowledgeDBBootstrap(ctx context.Context, input *Acknowledge
 	}); err != nil {
 		return nil, err
 	}
-	if err := deleteBootstrapPayload(ctx, kv, input.DBInstanceIdentifier, jetstream.LastRevision(rev)); err != nil {
+	if err := deleteBootstrapPayloadAt(ctx, kv, input.DBInstanceIdentifier, rev); err != nil {
 		return nil, err
 	}
 

spinifex/handlers/rds/agent_test.go

@@ -128,7 +128,9 @@ func readRecord(t *testing.T, svc *Service) (DBInstanceRecord, string) {
 	t.Helper()
 	kv, err := svc.bucket(context.Background(), testAccountID)
 	require.NoError(t, err)
-	entry, err := kv.Get(context.Background(), DBInstanceKey(testDBID))
+	raw, err := kv.KV(context.Background())
+	require.NoError(t, err)
+	entry, err := raw.Get(context.Background(), DBInstanceKey(testDBID))
 	require.NoError(t, err)
 	var rec DBInstanceRecord
 	require.NoError(t, json.Unmarshal(entry.Value(), &rec))

spinifex/handlers/rds/backup.go

@@ -12,7 +12,7 @@ import (
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/rds"
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
-	"github.com/nats-io/nats.go/jetstream"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 )
 
 // Automated backups: one snapshot per instance per day inside its
@@ -203,7 +203,7 @@ var backupCapableStatuses = []Status{StatusAvailable, StatusStopped}
 // succeeded since the window opened, reporting whether it took one. Called from
 // the leader-elected reconciler pass, which is what makes it cluster-singular;
 // the per-instance in-flight guard is what makes it safe anyway.
-func (s *Service) runBackupWindow(ctx context.Context, kv jetstream.KeyValue, rev uint64,
+func (s *Service) runBackupWindow(ctx context.Context, kv *kvstore.Bucket, rev uint64,
 	accountID string, rec *DBInstanceRecord) (bool, error) {
 	now := time.Now().UTC()
 	window, err := s.resolvedBackupWindow(rec)
@@ -267,7 +267,7 @@ func backupRetryDelay(failures int) time.Duration {
 // costs one duplicate snapshot at worst and is swept on schedule, while a stamped
 // backup with no index would be invisible to retention for the life of the
 // instance.
-func (s *Service) takeAutomatedBackup(ctx context.Context, kv jetstream.KeyValue, rev uint64,
+func (s *Service) takeAutomatedBackup(ctx context.Context, kv *kvstore.Bucket, rev uint64,
 	accountID string, rec *DBInstanceRecord, now time.Time) error {
 	record, err := s.snapshotDBInstance(ctx, kv, rev, accountID, rec, &validatedSnapshot{
 		DBSnapshotIdentifier: AutomatedSnapshotIdentifier(rec.DBInstanceIdentifier, now),
@@ -298,7 +298,7 @@ func (s *Service) takeAutomatedBackup(ctx context.Context, kv jetstream.KeyValue
 // Counts the failure and reports it against the DB instance, leaving its status
 // alone. The stamp is what paces the retry inside the window; the count is what
 // makes a backup that has been failing for days visible and testable.
-func (s *Service) recordBackupFailure(ctx context.Context, kv jetstream.KeyValue, accountID string,
+func (s *Service) recordBackupFailure(ctx context.Context, kv *kvstore.Bucket, accountID string,
 	rec *DBInstanceRecord, cause error) error {
 	now := time.Now().UTC()
 	slog.WarnContext(ctx, "rds: an automated backup failed; the DB instance is unaffected",
@@ -316,7 +316,7 @@ func (s *Service) recordBackupFailure(ctx context.Context, kv jetstream.KeyValue
 // A window this instance could not be backed up in, evented once and then paced
 // by the same backoff a failure is. Reported without incrementing the failure
 // count: nothing about the backup failed, the instance was busy elsewhere.
-func (s *Service) recordBackupDeferred(ctx context.Context, kv jetstream.KeyValue, accountID string,
+func (s *Service) recordBackupDeferred(ctx context.Context, kv *kvstore.Bucket, accountID string,
 	rec *DBInstanceRecord, opened time.Time, message string) error {
 	now := time.Now().UTC()
 	if rec.LastAutomatedBackupFailureAt == nil || rec.LastAutomatedBackupFailureAt.Before(opened) {
@@ -337,7 +337,7 @@ func (s *Service) recordBackupDeferred(ctx context.Context, kv jetstream.KeyValu
 // single drain through applyPendingModifications — so a deferred change is
 // applied by exactly the code an immediate one uses, and one instance's VM
 // replace does not hold up the whole fleet's pass.
-func (s *Service) runMaintenanceWindow(ctx context.Context, kv jetstream.KeyValue, rev uint64,
+func (s *Service) runMaintenanceWindow(ctx context.Context, kv *kvstore.Bucket, rev uint64,
 	accountID string, rec *DBInstanceRecord) (bool, error) {
 	pending := rec.PendingModifiedValues
 	if pending.empty() || pending.growingFilesystem() || rec.Status != StatusAvailable {
@@ -364,7 +364,7 @@ func (s *Service) runMaintenanceWindow(ctx context.Context, kv jetstream.KeyValu
 	rec.Status = StatusModifying
 	rec.UpdatedAt = now
 	if err := updateJSON(ctx, kv, DBInstanceKey(rec.DBInstanceIdentifier), rev, rec); err != nil {
-		if errors.Is(err, jetstream.ErrKeyRevisionMismatch) {
+		if errors.Is(err, kvstore.ErrConflict) {
 			// Something else moved the record between the read and here; the next
 			// pass re-reads and the window is still open.
 			return false, nil

spinifex/handlers/rds/backup_sweep.go

@@ -10,8 +10,8 @@ import (
 	"time"
 
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 	"github.com/mulgadc/spinifex/spinifex/vm"
-	"github.com/nats-io/nats.go/jetstream"
 )
 
 // The retention sweep is a cluster-wide vm.Reaper rather than a reconciler tick,
@@ -49,7 +49,7 @@ func (r *BackupRetentionReaper) Sweep(ctx context.Context) (int, error) {
 	}
 	// A truncated listing must never read as "no accounts": that would leave every
 	// tenant's retention unenforced while looking like a clean pass.
-	buckets, err := AccountBucketNames(ctx, js)
+	buckets, err := AccountBuckets(ctx, js)
 	if err != nil {
 		return 0, fmt.Errorf("rds: enumerate account buckets: %w", err)
 	}
@@ -63,16 +63,10 @@ func (r *BackupRetentionReaper) Sweep(ctx context.Context) (int, error) {
 		if err := ctx.Err(); err != nil {
 			return reaped, err
 		}
-		kv, err := js.KeyValue(ctx, bucket)
-		if err != nil {
-			failures = append(failures, fmt.Errorf("open %s: %w", bucket, err))
-			continue
-		}
-		accountID := AccountIDFromBucketName(bucket)
-		swept, err := r.sweepAccount(ctx, kv, accountID, r.limit-reaped)
+		swept, err := r.sweepAccount(ctx, bucket, AccountIDFromBucketName(bucket.Name()), r.limit-reaped)
 		reaped += swept
 		if err != nil {
-			failures = append(failures, fmt.Errorf("sweep %s: %w", bucket, err))
+			failures = append(failures, fmt.Errorf("sweep %s: %w", bucket.Name(), err))
 		}
 	}
 	return reaped, errors.Join(failures...)
@@ -80,7 +74,7 @@ func (r *BackupRetentionReaper) Sweep(ctx context.Context) (int, error) {
 
 // One account: the over-retention automated snapshots of every instance, then the
 // retained volumes whose last snapshot is gone.
-func (r *BackupRetentionReaper) sweepAccount(ctx context.Context, kv jetstream.KeyValue,
+func (r *BackupRetentionReaper) sweepAccount(ctx context.Context, kv *kvstore.Bucket,
 	accountID string, budget int) (int, error) {
 	indexed, err := ListAutomatedBackups(ctx, kv)
 	if err != nil {
@@ -115,7 +109,7 @@ func (r *BackupRetentionReaper) sweepAccount(ctx context.Context, kv jetstream.K
 // available snapshot is never deleted while automated backups are on — if backup
 // creation has been failing for a week, strict retention would delete the whole
 // backup set at the exact moment a backup matters most.
-func (s *Service) sweepInstanceBackups(ctx context.Context, kv jetstream.KeyValue,
+func (s *Service) sweepInstanceBackups(ctx context.Context, kv *kvstore.Bucket,
 	accountID, dbInstanceIdentifier string, stamps []string, budget int) (int, error) {
 	var rec DBInstanceRecord
 	found, err := getJSON(ctx, kv, DBInstanceKey(dbInstanceIdentifier), &rec)
@@ -170,7 +164,7 @@ func (s *Service) sweepInstanceBackups(ctx context.Context, kv jetstream.KeyValu
 // replica reads the same way — and that branch removes a live instance's whole
 // backup set, newest included. The bucket listing is a second, independent read
 // path, so an identifier it still names fails this instance's sweep instead.
-func confirmDBInstanceGone(ctx context.Context, kv jetstream.KeyValue, dbInstanceIdentifier string) error {
+func confirmDBInstanceGone(ctx context.Context, kv *kvstore.Bucket, dbInstanceIdentifier string) error {
 	ids, err := ListDBInstanceIDs(ctx, kv)
 	if err != nil {
 		return err
@@ -187,7 +181,7 @@ func confirmDBInstanceGone(ctx context.Context, kv jetstream.KeyValue, dbInstanc
 // The index entries of one instance, oldest first. An entry whose snapshot record
 // is gone is a half-finished delete: the entry is removed so the index stops
 // naming it, and it counts as nothing reaped because no data went with it.
-func (s *Service) loadAutomatedBackups(ctx context.Context, kv jetstream.KeyValue,
+func (s *Service) loadAutomatedBackups(ctx context.Context, kv *kvstore.Bucket,
 	dbInstanceIdentifier string, stamps []string) ([]automatedBackup, error) {
 	slices.Sort(stamps)
 	entries := make([]automatedBackup, 0, len(stamps))
@@ -248,7 +242,7 @@ func newestAvailableBackup(entries []automatedBackup) string {
 // snapshot is still being read from — an instance restored from it is still
 // running — which is a skip rather than a failure: the next pass retries, and the
 // entry stays so it is not lost from the index.
-func (s *Service) deleteAutomatedBackup(ctx context.Context, kv jetstream.KeyValue,
+func (s *Service) deleteAutomatedBackup(ctx context.Context, kv *kvstore.Bucket,
 	accountID string, entry automatedBackup) (bool, error) {
 	err := s.removeDBSnapshot(ctx, kv, accountID, entry.snapshot)
 	switch {
@@ -276,10 +270,10 @@ func (s *Service) deleteAutomatedBackup(ctx context.Context, kv jetstream.KeyVal
 
 // Last, always: while the entry exists the snapshot is still findable by the
 // sweep, and every step above tolerates work it has already done.
-func (s *Service) dropAutomatedBackupIndex(ctx context.Context, kv jetstream.KeyValue,
+func (s *Service) dropAutomatedBackupIndex(ctx context.Context, kv *kvstore.Bucket,
 	dbInstanceIdentifier, stamp string) error {
 	key := AutomatedBackupKey(dbInstanceIdentifier, stamp)
-	if err := kv.Delete(ctx, key); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
+	if err := kv.Delete(ctx, key); err != nil {
 		return fmt.Errorf("rds: clear the automated-backup index entry %s: %w", key, err)
 	}
 	return nil
@@ -290,7 +284,7 @@ func (s *Service) dropAutomatedBackupIndex(ctx context.Context, kv jetstream.Key
 // BackupRetentionPeriod=0 modify, which turns the feature off, and the teardown
 // of the instance itself, whose automated backups would otherwise pin its data
 // volume after the instance is gone.
-func (s *Service) purgeAutomatedBackups(ctx context.Context, kv jetstream.KeyValue,
+func (s *Service) purgeAutomatedBackups(ctx context.Context, kv *kvstore.Bucket,
 	accountID, dbInstanceIdentifier string) error {
 	stamps, err := listNames(ctx, kv, AutomatedBackupsPrefix(dbInstanceIdentifier))
 	if err != nil {
@@ -313,7 +307,7 @@ func (s *Service) purgeAutomatedBackups(ctx context.Context, kv jetstream.KeyVal
 // deletes the volume inline on the last DeleteDBSnapshot, so this only ever fires
 // for a crash between those two steps — which is precisely what a KV-health-gated
 // cluster-wide sweep is for.
-func (s *Service) reclaimOrphanedVolumes(ctx context.Context, kv jetstream.KeyValue) (int, error) {
+func (s *Service) reclaimOrphanedVolumes(ctx context.Context, kv *kvstore.Bucket) (int, error) {
 	ids, err := listNames(ctx, kv, RetainedVolumesPrefix())
 	if err != nil {
 		return 0, err

spinifex/handlers/rds/backup_sweep_test.go

@@ -8,6 +8,7 @@ import (
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/rds"
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 	"github.com/mulgadc/spinifex/spinifex/vm"
 	"github.com/nats-io/nats.go/jetstream"
 	"github.com/stretchr/testify/assert"
@@ -191,7 +192,10 @@ func TestSweep_LeavesTheBackupSetWhenTheRecordCannotBeRead(t *testing.T) {
 
 	kv, err := h.svc.bucket(t.Context(), testAccountID)
 	require.NoError(t, err)
-	stale := missingKey{KeyValue: kv, key: DBInstanceKey(testDBID)}
+	raw, err := kv.KV(t.Context())
+	require.NoError(t, err)
+	stale := kvstore.NewOpenBucket(nil, missingKey{KeyValue: raw, key: DBInstanceKey(testDBID)},
+		AccountBucketConfig(testAccountID))
 
 	reaped, err := h.svc.sweepInstanceBackups(t.Context(), stale, testAccountID, testDBID,
 		h.automatedStamps(t, testDBID), defaultSweepDeleteLimit)

spinifex/handlers/rds/bootstrap.go

@@ -11,8 +11,8 @@ import (
 
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
 	handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 	"github.com/mulgadc/spinifex/spinifex/utils"
-	"github.com/nats-io/nats.go/jetstream"
 )
 
 // The staged master password is encrypted at rest under bootstrap-payloads/{id}
@@ -79,7 +79,7 @@ func bootstrapKeyID(key []byte) string {
 // Stages the master password for rec's current VM generation. The key's presence
 // is what makes the bootstrap pending, so it is written once and only ever
 // deleted by an acknowledgement, a delete, or a generation bump.
-func (s *Service) writeBootstrapPayload(ctx context.Context, kv jetstream.KeyValue,
+func (s *Service) writeBootstrapPayload(ctx context.Context, kv *kvstore.Bucket,
 	accountID string, rec *DBInstanceRecord, masterPassword string) (string, error) {
 	key, err := s.masterKey()
 	if err != nil {
@@ -126,7 +126,7 @@ func (s *Service) writeBootstrapPayload(ctx context.Context, kv jetstream.KeyVal
 
 // Returns (nil, 0, nil) when nothing is staged, which is what an acknowledged,
 // restored or legacy instance reads as.
-func readBootstrapPayload(ctx context.Context, kv jetstream.KeyValue,
+func readBootstrapPayload(ctx context.Context, kv *kvstore.Bucket,
 	dbInstanceIdentifier string) (*BootstrapPayloadEnvelope, uint64, error) {
 	var envelope BootstrapPayloadEnvelope
 	rev, found, err := getJSONRevision(ctx, kv, BootstrapPayloadKey(dbInstanceIdentifier), &envelope)
@@ -138,10 +138,17 @@ func readBootstrapPayload(ctx context.Context, kv jetstream.KeyValue,
 
 // An already-absent key is the state this is trying to reach, so it is not an
 // error.
-func deleteBootstrapPayload(ctx context.Context, kv jetstream.KeyValue,
-	dbInstanceIdentifier string, opts ...jetstream.KVDeleteOpt) error {
-	err := kv.Delete(ctx, BootstrapPayloadKey(dbInstanceIdentifier), opts...)
-	if err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
+func deleteBootstrapPayload(ctx context.Context, kv *kvstore.Bucket, dbInstanceIdentifier string) error {
+	if err := kv.Delete(ctx, BootstrapPayloadKey(dbInstanceIdentifier)); err != nil {
+		return fmt.Errorf("rds bootstrap: delete the payload for %s: %w", dbInstanceIdentifier, err)
+	}
+	return nil
+}
+
+// deleteBootstrapPayloadAt is deleteBootstrapPayload guarded on rev, for the
+// consumer that must remove only the payload it read.
+func deleteBootstrapPayloadAt(ctx context.Context, kv *kvstore.Bucket, dbInstanceIdentifier string, rev uint64) error {
+	if err := kv.CompareAndDelete(ctx, BootstrapPayloadKey(dbInstanceIdentifier), rev); err != nil {
 		return fmt.Errorf("rds bootstrap: delete the payload for %s: %w", dbInstanceIdentifier, err)
 	}
 	return nil
@@ -150,7 +157,7 @@ func deleteBootstrapPayload(ctx context.Context, kv jetstream.KeyValue,
 // Whether a staged payload is still waiting to be applied. Callers that only
 // need the yes/no — the modify gate, the reconciler's stale-payload sweep — use
 // this rather than decrypting.
-func bootstrapPending(ctx context.Context, kv jetstream.KeyValue, dbInstanceIdentifier string) (bool, error) {
+func bootstrapPending(ctx context.Context, kv *kvstore.Bucket, dbInstanceIdentifier string) (bool, error) {
 	envelope, _, err := readBootstrapPayload(ctx, kv, dbInstanceIdentifier)
 	return envelope != nil, err
 }
@@ -222,7 +229,7 @@ func (s *Service) openBootstrapPayload(envelope *BootstrapPayloadEnvelope, accou
 // Records that the staged password can no longer reach this datadir. Written
 // rather than inferred so the reason rides DescribeDBInstances, and idempotent
 // so a repeated fetch does not churn the record.
-func (s *Service) markBootstrapUnrecoverable(ctx context.Context, kv jetstream.KeyValue,
+func (s *Service) markBootstrapUnrecoverable(ctx context.Context, kv *kvstore.Bucket,
 	accountID, dbInstanceIdentifier, reason string) error {
 	changed := false
 	if err := s.updateInstance(ctx, kv, dbInstanceIdentifier, func(stored *DBInstanceRecord) {
@@ -250,7 +257,7 @@ func (s *Service) markBootstrapUnrecoverable(ctx context.Context, kv jetstream.K
 // password is not what stops the replacement working — replaceInstanceVM
 // revokes the data-volume format grant, and only the initial create can hold
 // one — so rebinding would move the failure rather than fix it.
-func (s *Service) discardPendingBootstrap(ctx context.Context, kv jetstream.KeyValue,
+func (s *Service) discardPendingBootstrap(ctx context.Context, kv *kvstore.Bucket,
 	accountID string, rec *DBInstanceRecord) error {
 	pending, err := bootstrapPending(ctx, kv, rec.DBInstanceIdentifier)
 	if err != nil || !pending {

spinifex/handlers/rds/bootstrap_test.go

@@ -45,7 +45,9 @@ func storedPayload(t *testing.T, svc *Service, id string) (*BootstrapPayloadEnve
 	t.Helper()
 	kv, err := svc.bucket(t.Context(), testAccountID)
 	require.NoError(t, err)
-	entry, err := kv.Get(t.Context(), BootstrapPayloadKey(id))
+	raw, err := kv.KV(t.Context())
+	require.NoError(t, err)
+	entry, err := raw.Get(t.Context(), BootstrapPayloadKey(id))
 	if err != nil {
 		return nil, ""
 	}
@@ -417,16 +419,18 @@ func TestCreateDBInstance_StagesThePasswordEncrypted(t *testing.T) {
 	keys, err := bucketKeys(t.Context(), kv)
 	require.NoError(t, err)
 	require.Contains(t, keys, BootstrapPayloadKey(testDBInstanceID))
+	raw, err := kv.KV(t.Context())
+	require.NoError(t, err)
 
 	// Nothing anywhere in the account bucket may hold the supplied password.
 	for _, key := range keys {
-		entry, err := kv.Get(t.Context(), key)
+		entry, err := raw.Get(t.Context(), key)
 		require.NoError(t, err)
 		assert.NotContains(t, string(entry.Value()), "Sup3rSecret!", "cleartext at rest under %s", key)
 	}
 
 	var envelope BootstrapPayloadEnvelope
-	entry, err := kv.Get(t.Context(), BootstrapPayloadKey(testDBInstanceID))
+	entry, err := raw.Get(t.Context(), BootstrapPayloadKey(testDBInstanceID))
 	require.NoError(t, err)
 	require.NoError(t, json.Unmarshal(entry.Value(), &envelope))
 	plaintext, err := handlers_iam.DecryptSecret(envelope.EncryptedPayload, testMasterKey)

spinifex/handlers/rds/create.go

@@ -9,9 +9,9 @@ import (
 
 	"github.com/aws/aws-sdk-go/service/rds"
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 	"github.com/mulgadc/spinifex/spinifex/utils"
 	"github.com/mulgadc/spinifex/spinifex/vm"
-	"github.com/nats-io/nats.go/jetstream"
 )
 
 // The first VM behind a DB instance. Replacement and recovery increment it,
@@ -61,7 +61,7 @@ func (s *Service) CreateDBInstance(ctx context.Context, input *rds.CreateDBInsta
 	rec := newDBInstanceRecord(accountID, req, placement, parameters)
 	rollbackRev, createErr := createJSONRevision(ctx, kv, key, &rec)
 	if createErr != nil {
-		if errors.Is(createErr, jetstream.ErrKeyExists) {
+		if errors.Is(createErr, kvstore.ErrExists) {
 			return nil, awserrors.Errorf(awserrors.ErrorDBInstanceAlreadyExists,
 				"DB instance %s already exists", req.Identifier)
 		}
@@ -192,7 +192,7 @@ func newDBInstanceRecord(accountID string, req *validatedCreate, placement *endp
 // Folds the launch results into the reserved record and returns it as written.
 // The endpoint is settled here because the ENI IP — and so the vanity name — is
 // only known now.
-func (s *Service) recordLaunch(ctx context.Context, kv jetstream.KeyValue, key, accountID,
+func (s *Service) recordLaunch(ctx context.Context, kv *kvstore.Bucket, key, accountID,
 	expectedResourceID string, launched *LaunchOutput, initialCreate bool) (*DBInstanceRecord, uint64, error) {
 	var rec DBInstanceRecord
 	rev, found, err := getJSONRevision(ctx, kv, key, &rec)
@@ -243,7 +243,7 @@ func (s *Service) recordLaunch(ctx context.Context, kv jetstream.KeyValue, key,
 	return &rec, updatedRev, nil
 }
 
-func (s *Service) rollbackDBInstanceReservation(ctx context.Context, kv jetstream.KeyValue,
+func (s *Service) rollbackDBInstanceReservation(ctx context.Context, kv *kvstore.Bucket,
 	key, identifier, resourceID string, rev uint64) {
 	rbCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), rollbackTimeout)
 	defer cancel()
@@ -264,11 +264,11 @@ func (s *Service) rollbackDBInstanceReservation(ctx context.Context, kv jetstrea
 
 	var rollbackErr error
 	for range rollbackDeleteAttempts {
-		rollbackErr = kv.Delete(rbCtx, key, jetstream.LastRevision(rev))
-		if rollbackErr == nil || errors.Is(rollbackErr, jetstream.ErrKeyNotFound) {
+		rollbackErr = kv.CompareAndDelete(rbCtx, key, rev)
+		if rollbackErr == nil {
 			return
 		}
-		if !errors.Is(rollbackErr, jetstream.ErrKeyRevisionMismatch) {
+		if !errors.Is(rollbackErr, kvstore.ErrConflict) {
 			break
 		}
 

spinifex/handlers/rds/create_test.go

@@ -12,7 +12,7 @@ import (
 	"github.com/aws/aws-sdk-go/service/rds"
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
 	handlers_dns "github.com/mulgadc/spinifex/spinifex/handlers/dns"
-	"github.com/mulgadc/spinifex/spinifex/handlers/ec2/vpc"
+	handlers_ec2_vpc "github.com/mulgadc/spinifex/spinifex/handlers/ec2/vpc"
 	iammock "github.com/mulgadc/spinifex/spinifex/handlers/iam/mock"
 	"github.com/mulgadc/spinifex/spinifex/testutil"
 	"github.com/mulgadc/spinifex/spinifex/utils"

spinifex/handlers/rds/delete.go

@@ -12,9 +12,9 @@ import (
 	"github.com/aws/aws-sdk-go/service/rds"
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
 	"github.com/mulgadc/spinifex/spinifex/handlers/sysinstance"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 	"github.com/mulgadc/spinifex/spinifex/tags"
 	"github.com/mulgadc/spinifex/spinifex/utils"
-	"github.com/nats-io/nats.go/jetstream"
 )
 
 var errFinalSnapshotInProgress = errors.New("rds: final DB snapshot creation is still in progress")
@@ -83,7 +83,7 @@ func (s *Service) DeleteDBInstance(ctx context.Context, input *rds.DeleteDBInsta
 			if reservation.creator {
 				s.rollbackFinalSnapshotReservation(ctx, kv, finalSnapshot, reservation.revision)
 			}
-			if errors.Is(err, jetstream.ErrKeyRevisionMismatch) {
+			if errors.Is(err, kvstore.ErrConflict) {
 				return nil, awserrors.Errorf(awserrors.ErrorDBInstanceInvalidState,
 					"DB instance %s changed state concurrently; retry the delete", id)
 			}
@@ -153,7 +153,7 @@ func finalSnapshotIsOurs(existing *DBSnapshotRecord, rec *DBInstanceRecord) bool
 // Rejects a taken identifier before anything is torn down, as AWS does. Leaving
 // it to the snapshot step would terminate the VM first and then fail on an
 // error no retry can clear.
-func (s *Service) checkFinalSnapshotAvailable(ctx context.Context, kv jetstream.KeyValue, rec *DBInstanceRecord, identifier string) error {
+func (s *Service) checkFinalSnapshotAvailable(ctx context.Context, kv *kvstore.Bucket, rec *DBInstanceRecord, identifier string) error {
 	if identifier == "" {
 		return nil
 	}
@@ -175,7 +175,7 @@ func (s *Service) checkFinalSnapshotAvailable(ctx context.Context, kv jetstream.
 //
 // Every step treats a missing resource as done, so this is safe to re-run from
 // any point it stopped at.
-func (s *Service) teardownDBInstance(ctx context.Context, kv jetstream.KeyValue, accountID string,
+func (s *Service) teardownDBInstance(ctx context.Context, kv *kvstore.Bucket, accountID string,
 	rec *DBInstanceRecord, finalSnapshotCreator bool) error {
 	// Asked to stop cleanly even though the VM is about to go: it is what makes
 	// the final snapshot a clean checkpoint rather than one needing WAL replay.
@@ -203,7 +203,7 @@ func (s *Service) teardownDBInstance(ctx context.Context, kv jetstream.KeyValue,
 	s.deleteInstanceENIs(ctx, accountID, rec)
 
 	if rec.InstanceID != "" {
-		if err := s.DeleteInstanceIndex(ctx, rec.InstanceID); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
+		if err := s.DeleteInstanceIndex(ctx, rec.InstanceID); err != nil {
 			return fmt.Errorf("rds: delete the instance index entry for %s: %w", rec.DBInstanceIdentifier, err)
 		}
 	}
@@ -214,7 +214,7 @@ func (s *Service) teardownDBInstance(ctx context.Context, kv jetstream.KeyValue,
 	}
 	// Last: while it exists the instance is still nameable, and everything above
 	// is reachable only through it.
-	if err := kv.Delete(ctx, DBInstanceKey(rec.DBInstanceIdentifier)); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
+	if err := kv.Delete(ctx, DBInstanceKey(rec.DBInstanceIdentifier)); err != nil {
 		return fmt.Errorf("rds: delete the DB instance record for %s: %w", rec.DBInstanceIdentifier, err)
 	}
 
@@ -244,7 +244,7 @@ func (s *Service) terminateInstanceVM(ctx context.Context, instanceID string) er
 // Reserves the final snapshot name before deletion becomes destructive. The
 // creator is the only worker allowed to cut new data; retries may adopt data it
 // already cut, but cannot race it into making a duplicate snapshot.
-func (s *Service) reserveFinalSnapshot(ctx context.Context, kv jetstream.KeyValue, accountID string,
+func (s *Service) reserveFinalSnapshot(ctx context.Context, kv *kvstore.Bucket, accountID string,
 	rec *DBInstanceRecord, identifier string) (finalSnapshotReservation, error) {
 	if identifier == "" || rec.DataVolumeID == "" {
 		return finalSnapshotReservation{}, nil
@@ -262,7 +262,7 @@ func (s *Service) reserveFinalSnapshot(ctx context.Context, kv jetstream.KeyValu
 	if err == nil {
 		return finalSnapshotReservation{creator: true, revision: rev}, nil
 	}
-	if !errors.Is(err, jetstream.ErrKeyExists) {
+	if !errors.Is(err, kvstore.ErrExists) {
 		return finalSnapshotReservation{}, err
 	}
 
@@ -281,15 +281,14 @@ func (s *Service) reserveFinalSnapshot(ctx context.Context, kv jetstream.KeyValu
 // Removes a reservation when the instance never entered deleting or the EC2
 // create failed. The revision guard prevents cleanup from deleting a record
 // another worker has already completed.
-func (s *Service) rollbackFinalSnapshotReservation(ctx context.Context, kv jetstream.KeyValue,
+func (s *Service) rollbackFinalSnapshotReservation(ctx context.Context, kv *kvstore.Bucket,
 	identifier string, rev uint64) {
 	if identifier == "" || rev == 0 {
 		return
 	}
 	rbCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), rollbackTimeout)
 	defer cancel()
-	if err := kv.Delete(rbCtx, DBSnapshotKey(identifier), jetstream.LastRevision(rev)); err != nil &&
-		!errors.Is(err, jetstream.ErrKeyNotFound) {
+	if err := kv.CompareAndDelete(rbCtx, DBSnapshotKey(identifier), rev); err != nil {
 		slog.WarnContext(rbCtx, "rds: rollback of a final snapshot reservation failed",
 			"dbSnapshot", identifier, "err", err)
 	}
@@ -298,7 +297,7 @@ func (s *Service) rollbackFinalSnapshotReservation(ctx context.Context, kv jetst
 // Records the final snapshot under the reservation written before teardown. A
 // retry first adopts matching EC2 data left by a dead creator; it never cuts a
 // second snapshot while that creator may still be running.
-func (s *Service) takeFinalSnapshot(ctx context.Context, kv jetstream.KeyValue, accountID string,
+func (s *Service) takeFinalSnapshot(ctx context.Context, kv *kvstore.Bucket, accountID string,
 	rec *DBInstanceRecord, creator bool) error {
 	if rec.FinalSnapshotIdentifier == "" || rec.DataVolumeID == "" {
 		return nil
@@ -393,7 +392,7 @@ func (s *Service) takeFinalSnapshot(ctx context.Context, kv jetstream.KeyValue,
 
 // Completes only the snapshot fields under CAS, preserving tags changed while
 // the EC2 snapshot was being cut.
-func (s *Service) completeFinalSnapshot(ctx context.Context, kv jetstream.KeyValue,
+func (s *Service) completeFinalSnapshot(ctx context.Context, kv *kvstore.Bucket,
 	source *DBInstanceRecord, snapshotID string) (bool, error) {
 	key := DBSnapshotKey(source.FinalSnapshotIdentifier)
 	for range tagWriteAttempts {
@@ -426,7 +425,7 @@ func (s *Service) completeFinalSnapshot(ctx context.Context, kv jetstream.KeyVal
 		record.Status = SnapshotStatusAvailable
 		if err := updateJSON(ctx, kv, key, rev, &record); err == nil {
 			return true, nil
-		} else if !errors.Is(err, jetstream.ErrKeyRevisionMismatch) {
+		} else if !errors.Is(err, kvstore.ErrConflict) {
 			return false, err
 		}
 	}
@@ -447,7 +446,7 @@ func (s *Service) recordFinalSnapshotCreated(ctx context.Context, accountID stri
 // copying them, so the volume cannot be deleted while any snapshot survives —
 // including the final one just taken. It is retained instead, recorded with the
 // snapshots holding it so the last DeleteDBSnapshot can release it.
-func (s *Service) releaseDataVolume(ctx context.Context, kv jetstream.KeyValue, accountID string, rec *DBInstanceRecord) error {
+func (s *Service) releaseDataVolume(ctx context.Context, kv *kvstore.Bucket, accountID string, rec *DBInstanceRecord) error {
 	if rec.DataVolumeID == "" {
 		return nil
 	}
@@ -482,7 +481,7 @@ func (s *Service) releaseDataVolume(ctx context.Context, kv jetstream.KeyValue,
 // the last DeleteDBSnapshot can release it. holdersUnresolved marks the
 // case where the volume store refused the delete but named nobody, so a release
 // has to re-check rather than trust an empty list.
-func (s *Service) retainDataVolume(ctx context.Context, kv jetstream.KeyValue, accountID string,
+func (s *Service) retainDataVolume(ctx context.Context, kv *kvstore.Bucket, accountID string,
 	rec *DBInstanceRecord, holders []string, holdersUnresolved bool) error {
 	retained := RetainedVolumeRecord{
 		VolumeID:             rec.DataVolumeID,

spinifex/handlers/rds/describe.go

@@ -8,7 +8,7 @@ import (
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/rds"
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
-	"github.com/nats-io/nats.go/jetstream"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 )
 
 // The prefix AWS gives a DB instance's immutable resource ID, and the two filter
@@ -116,7 +116,7 @@ func dbInstanceFilterMatcher(filters []*rds.Filter) (func(*DBInstanceRecord) boo
 }
 
 // Returns the record plus its revision, for callers that follow with a CAS.
-func (s *Service) getDBInstance(ctx context.Context, kv jetstream.KeyValue, id string) (*DBInstanceRecord, uint64, error) {
+func (s *Service) getDBInstance(ctx context.Context, kv *kvstore.Bucket, id string) (*DBInstanceRecord, uint64, error) {
 	var rec DBInstanceRecord
 	rev, found, err := getJSONRevision(ctx, kv, DBInstanceKey(id), &rec)
 	if err != nil {

spinifex/handlers/rds/dns.go

@@ -5,7 +5,7 @@ import (
 	"log/slog"
 
 	handlers_dns "github.com/mulgadc/spinifex/spinifex/handlers/dns"
-	"github.com/nats-io/nats.go/jetstream"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 )
 
 // The vanity hostname for a DB instance, or "" on a deployment with no base
@@ -31,20 +31,15 @@ func (s *Service) DesiredDNSChanges() (changes []handlers_dns.Change, ok bool) {
 	if err != nil {
 		return nil, false
 	}
-	buckets, err := AccountBucketNames(ctx, js)
+	buckets, err := AccountBuckets(ctx, js)
 	if err != nil {
 		slog.Warn("rds DesiredDNSChanges: enumerate account buckets", "err", err)
 		return nil, false
 	}
 	for _, bucket := range buckets {
-		kv, err := js.KeyValue(ctx, bucket)
+		bucketChanges, err := desiredBucketDNSChanges(ctx, bucket, s.deps.BaseDomain)
 		if err != nil {
-			slog.Warn("rds DesiredDNSChanges: open account bucket", "bucket", bucket, "err", err)
-			return nil, false
-		}
-		bucketChanges, err := desiredBucketDNSChanges(ctx, kv, s.deps.BaseDomain)
-		if err != nil {
-			slog.Warn("rds DesiredDNSChanges: read DB instances", "bucket", bucket, "err", err)
+			slog.Warn("rds DesiredDNSChanges: read DB instances", "bucket", bucket.Name(), "err", err)
 			return nil, false
 		}
 		changes = append(changes, bucketChanges...)
@@ -55,7 +50,7 @@ func (s *Service) DesiredDNSChanges() (changes []handlers_dns.Change, ok bool) {
 // A deleted instance's record is gone, so it contributes nothing and the
 // reconcile prunes its record. Anything still holding an ENI IP keeps its name
 // resolvable, including a failed instance an operator is still investigating.
-func desiredBucketDNSChanges(ctx context.Context, kv jetstream.KeyValue, baseDomain string) ([]handlers_dns.Change, error) {
+func desiredBucketDNSChanges(ctx context.Context, kv *kvstore.Bucket, baseDomain string) ([]handlers_dns.Change, error) {
 	ids, err := ListDBInstanceIDs(ctx, kv)
 	if err != nil {
 		return nil, err

spinifex/handlers/rds/engine_postgres.go

@@ -4,6 +4,7 @@ import (
 	"strconv"
 	"strings"
 	"time"
+
 	// The zone database Go's own LoadLocation reads, vendored into the binary so
 	// the timezone parameter validates identically wherever the daemon runs.
 	_ "time/tzdata"

spinifex/handlers/rds/events.go

@@ -12,7 +12,7 @@ import (
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/rds"
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
-	"github.com/nats-io/nats.go/jetstream"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 )
 
 // The event ring is the only channel some facts have to the customer — a backup
@@ -118,7 +118,7 @@ func (s *Service) appendEvent(ctx context.Context, accountID, sourceType, source
 		if !found {
 			if err := createJSON(ctx, kv, key, &ring); err == nil {
 				return nil
-			} else if !errors.Is(err, jetstream.ErrKeyExists) {
+			} else if !errors.Is(err, kvstore.ErrExists) {
 				return err
 			}
 			// Another writer created the ring first; re-read and append to it.
@@ -130,7 +130,7 @@ func (s *Service) appendEvent(ctx context.Context, accountID, sourceType, source
 		}
 		// A revision mismatch is a lost race with another appender, not a
 		// duplicate: re-read the ring and append to what they wrote.
-		if !errors.Is(err, jetstream.ErrKeyRevisionMismatch) {
+		if !errors.Is(err, kvstore.ErrConflict) {
 			return err
 		}
 	}
@@ -277,7 +277,7 @@ func eventRecordLimit(input *rds.DescribeEventsInput) int {
 
 // A fully qualified read is one Get. Anything broader has to enumerate, because
 // a ring outlives the record that would otherwise name it.
-func eventRingKeys(ctx context.Context, kv jetstream.KeyValue, sourceType, sourceIdentifier string) ([]string, error) {
+func eventRingKeys(ctx context.Context, kv *kvstore.Bucket, sourceType, sourceIdentifier string) ([]string, error) {
 	if sourceType != "" && sourceIdentifier != "" {
 		return []string{EventRingKey(sourceType, sourceIdentifier)}, nil
 	}
@@ -285,11 +285,8 @@ func eventRingKeys(ctx context.Context, kv jetstream.KeyValue, sourceType, sourc
 	if sourceType != "" {
 		prefix += sourceType + "/"
 	}
-	keys, err := kv.Keys(ctx)
+	keys, err := bucketKeys(ctx, kv)
 	if err != nil {
-		if errors.Is(err, jetstream.ErrNoKeysFound) {
-			return nil, nil
-		}
 		return nil, fmt.Errorf("rds: list event rings: %w", err)
 	}
 	matched := make([]string, 0, len(keys))

spinifex/handlers/rds/lifecycle.go

@@ -13,10 +13,10 @@ import (
 	"github.com/aws/aws-sdk-go/service/rds"
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
 	handlers_ec2_instance "github.com/mulgadc/spinifex/spinifex/handlers/ec2/instance"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 	"github.com/mulgadc/spinifex/spinifex/types"
 	"github.com/mulgadc/spinifex/spinifex/utils"
 	"github.com/nats-io/nats.go"
-	"github.com/nats-io/nats.go/jetstream"
 )
 
 // The DB VM's power state is driven over the per-instance EC2 command path
@@ -272,7 +272,7 @@ func uncleanStopMessage(ctx context.Context, engineName, operation string) strin
 // as written. from lists the statuses the operation is legal from, so a stop of
 // an already-stopping instance is rejected as a state error rather than racing
 // the stop already running.
-func (s *Service) beginTransition(ctx context.Context, accountID, id string, to Status, from ...Status) (*DBInstanceRecord, jetstream.KeyValue, error) {
+func (s *Service) beginTransition(ctx context.Context, accountID, id string, to Status, from ...Status) (*DBInstanceRecord, *kvstore.Bucket, error) {
 	if id == "" {
 		return nil, nil, awserrors.Errorf(awserrors.ErrorInvalidParameterValue, "DBInstanceIdentifier is required")
 	}
@@ -310,7 +310,7 @@ func (s *Service) beginTransition(ctx context.Context, accountID, id string, to
 	rec.TransitionStartedAt = &now
 	rec.UpdatedAt = now
 	if err := updateJSON(ctx, kv, DBInstanceKey(id), rev, rec); err != nil {
-		if errors.Is(err, jetstream.ErrKeyRevisionMismatch) {
+		if errors.Is(err, kvstore.ErrConflict) {
 			return nil, nil, awserrors.Errorf(awserrors.ErrorDBInstanceInvalidState,
 				"DB instance %s changed state concurrently; retry the operation", id)
 		}
@@ -321,7 +321,7 @@ func (s *Service) beginTransition(ctx context.Context, accountID, id string, to
 
 // Lands the instance in its post-transition state and clears the transition
 // stamp, returning the record as written.
-func (s *Service) completeTransition(ctx context.Context, kv jetstream.KeyValue, id string, to Status) (*DBInstanceRecord, error) {
+func (s *Service) completeTransition(ctx context.Context, kv *kvstore.Bucket, id string, to Status) (*DBInstanceRecord, error) {
 	var stored *DBInstanceRecord
 	err := s.updateInstance(ctx, kv, id, func(rec *DBInstanceRecord) {
 		rec.Status = to
@@ -336,7 +336,7 @@ func (s *Service) completeTransition(ctx context.Context, kv jetstream.KeyValue,
 
 // Records why an operation failed and returns the error the caller surfaces, so
 // a failed transition cannot be left looking like one still in progress.
-func (s *Service) failTransition(ctx context.Context, kv jetstream.KeyValue, accountID string, rec *DBInstanceRecord, reason string) error {
+func (s *Service) failTransition(ctx context.Context, kv *kvstore.Bucket, accountID string, rec *DBInstanceRecord, reason string) error {
 	if err := s.updateInstance(ctx, kv, rec.DBInstanceIdentifier, func(stored *DBInstanceRecord) {
 		stored.Status = StatusFailed
 		stored.FailureReason = reason
@@ -351,7 +351,7 @@ func (s *Service) failTransition(ctx context.Context, kv jetstream.KeyValue, acc
 
 // A read-modify-write under CAS, replayed on contention so a concurrent agent
 // heartbeat cannot make a lifecycle write disappear.
-func (s *Service) updateInstance(ctx context.Context, kv jetstream.KeyValue, id string, mutate func(*DBInstanceRecord)) error {
+func (s *Service) updateInstance(ctx context.Context, kv *kvstore.Bucket, id string, mutate func(*DBInstanceRecord)) error {
 	_, err := s.updateInstanceIf(ctx, kv, id, func(rec *DBInstanceRecord) bool {
 		mutate(rec)
 		return true
@@ -363,7 +363,7 @@ func (s *Service) updateInstance(ctx context.Context, kv jetstream.KeyValue, id
 // seen the record it would overwrite — a lease claim, whose whole question is
 // whether someone else holds it. Reports whether the write happened; a mutate
 // that returns false leaves the record untouched and is not an error.
-func (s *Service) updateInstanceIf(ctx context.Context, kv jetstream.KeyValue, id string, mutate func(*DBInstanceRecord) bool) (bool, error) {
+func (s *Service) updateInstanceIf(ctx context.Context, kv *kvstore.Bucket, id string, mutate func(*DBInstanceRecord) bool) (bool, error) {
 	key := DBInstanceKey(id)
 	for range tagWriteAttempts {
 		rec, rev, err := s.getDBInstance(ctx, kv, id)
@@ -379,7 +379,7 @@ func (s *Service) updateInstanceIf(ctx context.Context, kv jetstream.KeyValue, i
 		if err == nil {
 			return true, nil
 		}
-		if !errors.Is(err, jetstream.ErrKeyRevisionMismatch) {
+		if !errors.Is(err, kvstore.ErrConflict) {
 			return false, err
 		}
 	}

spinifex/handlers/rds/modify.go

@@ -13,7 +13,7 @@ import (
 	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/aws/aws-sdk-go/service/rds"
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
-	"github.com/nats-io/nats.go/jetstream"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 )
 
 // The resolved effect of a ModifyDBInstance request: what changes now and what
@@ -353,7 +353,7 @@ func (s *Service) planBackupSettings(input *rds.ModifyDBInstanceInput, rec *DBIn
 // rotation, an ENI security-group re-association, and record settings. Each is
 // applied before any disruptive work, so a modify that carries both does not
 // lose the cheap half to a failure in the expensive one.
-func (s *Service) applyImmediateModify(ctx context.Context, kv jetstream.KeyValue, accountID string, rec *DBInstanceRecord, plan *modifyPlan) error {
+func (s *Service) applyImmediateModify(ctx context.Context, kv *kvstore.Bucket, accountID string, rec *DBInstanceRecord, plan *modifyPlan) error {
 	if !plan.immediate() {
 		return nil
 	}
@@ -427,7 +427,7 @@ func (s *Service) applyImmediateModify(ctx context.Context, kv jetstream.KeyValu
 // The retention change itself has already landed, so a sweep that cannot finish
 // is reported rather than allowed to fail the modify: the retention reaper reads
 // the same zero retention and removes whatever is left, including the newest.
-func (s *Service) disableAutomatedBackups(ctx context.Context, kv jetstream.KeyValue, accountID, id string) {
+func (s *Service) disableAutomatedBackups(ctx context.Context, kv *kvstore.Bucket, accountID, id string) {
 	if err := s.purgeAutomatedBackups(ctx, kv, accountID, id); err != nil {
 		slog.WarnContext(ctx, "rds: sweeping the automated backups of a disabled instance failed; the retention reaper will finish it",
 			"dbInstance", id, "accountId", accountID, "err", err)

spinifex/handlers/rds/modify_lease.go

@@ -9,7 +9,7 @@ import (
 	"sync"
 	"time"
 
-	"github.com/nats-io/nats.go/jetstream"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 )
 
 // Applying PendingModifiedValues is the one control-plane operation with two
@@ -31,7 +31,7 @@ var errModifyLeaseLost = errors.New("rds: modify lease lost while applying pendi
 // this change and the caller has nothing to do.
 func (s *Service) withModifyLease(
 	ctx context.Context,
-	kv jetstream.KeyValue,
+	kv *kvstore.Bucket,
 	id string,
 	apply func(context.Context) error,
 ) (bool, error) {
@@ -92,7 +92,7 @@ func (s *Service) newModifyLeaseHolder() string {
 // is allowed, so a caller that already holds it is not deadlocked by itself.
 func (s *Service) claimModifyLease(
 	ctx context.Context,
-	kv jetstream.KeyValue,
+	kv *kvstore.Bucket,
 	id, holder string,
 ) (bool, time.Time, error) {
 	var expiresAt time.Time
@@ -113,7 +113,7 @@ func (s *Service) claimModifyLease(
 // work continue after its last successful renewal expires.
 func (s *Service) renewModifyLease(
 	ctx context.Context,
-	kv jetstream.KeyValue,
+	kv *kvstore.Bucket,
 	id, holder string,
 	expiresAt time.Time,
 	cancelApply context.CancelCauseFunc,

spinifex/handlers/rds/modify_lease_test.go

@@ -218,7 +218,7 @@ func TestWithModifyLease_CancelsWhenRenewalsFailUntilExpiry(t *testing.T) {
 	finished := make(chan error, 1)
 	go func() {
 		_, applyErr := h.svc.withModifyLease(t.Context(), kv, testDBID, func(ctx context.Context) error {
-			bucketDeleted <- js.DeleteKeyValue(ctx, kv.Bucket())
+			bucketDeleted <- js.DeleteKeyValue(ctx, kv.Name())
 			<-ctx.Done()
 			return context.Cause(ctx)
 		})

spinifex/handlers/rds/modify_test.go

@@ -9,9 +9,9 @@ import (
 	"github.com/aws/aws-sdk-go/service/rds"
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
 	iammock "github.com/mulgadc/spinifex/spinifex/handlers/iam/mock"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 	"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"
 )
@@ -95,7 +95,7 @@ func (h *modifyHarness) record(t *testing.T) DBInstanceRecord {
 	return rec
 }
 
-func (h *modifyHarness) kv(t *testing.T) jetstream.KeyValue {
+func (h *modifyHarness) kv(t *testing.T) *kvstore.Bucket {
 	t.Helper()
 	kv, err := h.svc.bucket(t.Context(), testAccountID)
 	require.NoError(t, err)

spinifex/handlers/rds/network.go

@@ -8,7 +8,7 @@ import (
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
-	"github.com/nats-io/nats.go/jetstream"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 )
 
 // The customer-account VPC surface CreateDBInstance reads to place the endpoint
@@ -33,7 +33,7 @@ type endpointPlacement struct {
 // back to the account's default VPC, mirroring AWS's own behaviour. The security
 // groups are validated against whichever VPC that resolved to, because an ENI
 // cannot carry a group from another one.
-func (s *Service) resolvePlacement(ctx context.Context, kv jetstream.KeyValue, accountID string, req *validatedCreate) (*endpointPlacement, error) {
+func (s *Service) resolvePlacement(ctx context.Context, kv *kvstore.Bucket, accountID string, req *validatedCreate) (*endpointPlacement, error) {
 	if s.deps.Network == nil {
 		return nil, awserrors.Errorf(awserrors.ErrorServerInternal, "RDS networking is not wired on this node")
 	}

spinifex/handlers/rds/network_test.go

@@ -4,7 +4,7 @@ import (
 	"testing"
 
 	"github.com/aws/aws-sdk-go/aws"
-	"github.com/mulgadc/spinifex/spinifex/handlers/ec2/vpc"
+	handlers_ec2_vpc "github.com/mulgadc/spinifex/spinifex/handlers/ec2/vpc"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )

spinifex/handlers/rds/parametergroup.go

@@ -13,7 +13,7 @@ import (
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/rds"
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
-	"github.com/nats-io/nats.go/jetstream"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 )
 
 // AWS's cap on one ModifyDBParameterGroup call.
@@ -68,7 +68,7 @@ func (s *Service) CreateDBParameterGroup(ctx context.Context, input *rds.CreateD
 		UpdatedAt:   now,
 	}
 	if err := createJSON(ctx, kv, DBParameterGroupMetaKey(name), &rec); err != nil {
-		if errors.Is(err, jetstream.ErrKeyExists) {
+		if errors.Is(err, kvstore.ErrExists) {
 			return nil, awserrors.Errorf(awserrors.ErrorDBParameterGroupAlreadyExists,
 				"DB parameter group %s already exists", name)
 		}
@@ -198,7 +198,7 @@ func (s *Service) ModifyDBParameterGroup(ctx context.Context, input *rds.ModifyD
 // Applies the complete effective set to every instance currently attached to
 // the group. Each record is re-read before the command so a concurrent detach
 // or delete does not receive parameters for a group it no longer uses.
-func (s *Service) propagateParameterGroup(ctx context.Context, kv jetstream.KeyValue, accountID, name string) error {
+func (s *Service) propagateParameterGroup(ctx context.Context, kv *kvstore.Bucket, accountID, name string) error {
 	ids, err := instancesUsingGroup(ctx, kv, func(rec *DBInstanceRecord) bool {
 		return rec.DBParameterGroupName == name
 	})
@@ -403,7 +403,7 @@ func resolveApplyMethod(spec ParameterSpec, requested string) (string, error) {
 // is synthesised rather than written on read: the record carries nothing a write
 // would preserve, and materialising it on a describe would make a read path a
 // writer for no gain.
-func getDBParameterGroup(ctx context.Context, kv jetstream.KeyValue, accountID, name string) (*DBParameterGroupRecord, uint64, error) {
+func getDBParameterGroup(ctx context.Context, kv *kvstore.Bucket, accountID, name string) (*DBParameterGroupRecord, uint64, error) {
 	var rec DBParameterGroupRecord
 	rev, found, err := getJSONRevision(ctx, kv, DBParameterGroupMetaKey(name), &rec)
 	if err != nil {
@@ -513,7 +513,7 @@ func (s *Service) projectParameterGroupRecord(rec *DBParameterGroupRecord) *rds.
 // Every path that binds a group to an instance comes through here — create,
 // modify, restore, the deferred apply and group propagation — so the
 // cross-engine refusal is one check rather than five.
-func (s *Service) resolveGroupParameters(ctx context.Context, kv jetstream.KeyValue, accountID string, engine Engine, group, instanceClass string) ([]Parameter, error) {
+func (s *Service) resolveGroupParameters(ctx context.Context, kv *kvstore.Bucket, accountID string, engine Engine, group, instanceClass string) ([]Parameter, error) {
 	rec, _, err := getDBParameterGroup(ctx, kv, accountID, group)
 	if err != nil {
 		return nil, err

spinifex/handlers/rds/pending.go

@@ -7,7 +7,7 @@ import (
 	"log/slog"
 
 	"github.com/aws/aws-sdk-go/aws"
-	"github.com/nats-io/nats.go/jetstream"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 )
 
 // The single drain for everything a modify recorded but has not delivered. Both
@@ -25,7 +25,7 @@ import (
 // The caller has already moved the instance into modifying; the record is left
 // there, because the engine has to come back and report healthy before the
 // reconciler calls it available.
-func (s *Service) applyPendingModifications(ctx context.Context, kv jetstream.KeyValue, accountID string, rec *DBInstanceRecord) error {
+func (s *Service) applyPendingModifications(ctx context.Context, kv *kvstore.Bucket, accountID string, rec *DBInstanceRecord) error {
 	pending := rec.PendingModifiedValues
 	if pending.empty() {
 		return nil
@@ -146,7 +146,7 @@ func (s *Service) applyPendingModifications(ctx context.Context, kv jetstream.Ke
 // A re-resolve, never a merge: the whole set is recomputed from the catalog and
 // the new group's overrides, so a parameter the old group set and the new one
 // does not reverts to its default rather than lingering.
-func (s *Service) applyParameterGroup(ctx context.Context, kv jetstream.KeyValue, accountID string, rec *DBInstanceRecord, group, instanceClass string, tolerateUnreachableAgent bool) error {
+func (s *Service) applyParameterGroup(ctx context.Context, kv *kvstore.Bucket, accountID string, rec *DBInstanceRecord, group, instanceClass string, tolerateUnreachableAgent bool) error {
 	engine, err := LookupEngine(rec.Engine)
 	if err != nil {
 		return err
@@ -187,7 +187,7 @@ func (s *Service) applyParameterGroup(ctx context.Context, kv jetstream.KeyValue
 // here in place of the dead agent's reply, with nothing pending a reboot.
 // Persistence uses a detached ctx: the budget that just found the agent
 // unreachable must not also fail the write that unblocks recovery from it.
-func (s *Service) deferParameterApplyToReplacement(ctx context.Context, kv jetstream.KeyValue, rec *DBInstanceRecord, resolved []Parameter) error {
+func (s *Service) deferParameterApplyToReplacement(ctx context.Context, kv *kvstore.Bucket, rec *DBInstanceRecord, resolved []Parameter) error {
 	slog.WarnContext(ctx, "rds: the instance agent is unreachable; deferring the parameter apply to the replacement VM's fresh agent",
 		"dbInstance", rec.DBInstanceIdentifier)
 	rec.Bootstrap.ResolvedParameters = resolved
@@ -211,7 +211,7 @@ func (s *Service) deferParameterApplyToReplacement(ctx context.Context, kv jetst
 // The event is deduped off the stored flag because the reconciler retries a
 // failed modify every pass. A failure to persist joins the cause rather than
 // replacing it — the apply is what the caller asked about.
-func (s *Service) recordParameterApplyFailure(ctx context.Context, kv jetstream.KeyValue, accountID string,
+func (s *Service) recordParameterApplyFailure(ctx context.Context, kv *kvstore.Bucket, accountID string,
 	rec *DBInstanceRecord, group string, cause error,
 ) error {
 	first := false
@@ -237,7 +237,7 @@ func (s *Service) recordParameterApplyFailure(ctx context.Context, kv jetstream.
 // the control plane has already grown the volume, and this extends the guest's
 // filesystem onto the capacity that is now there. Both ext4 and XFS grow while
 // mounted, so it needs no ordering against the engine start.
-func (s *Service) finishFilesystemGrow(ctx context.Context, kv jetstream.KeyValue, accountID string, rec *DBInstanceRecord) error {
+func (s *Service) finishFilesystemGrow(ctx context.Context, kv *kvstore.Bucket, accountID string, rec *DBInstanceRecord) error {
 	if err := s.growFilesystem(ctx, accountID, rec.DBInstanceIdentifier); err != nil {
 		return fmt.Errorf("extend the filesystem of %s onto its grown volume: %w", rec.DBInstanceIdentifier, err)
 	}

spinifex/handlers/rds/reconciler.go

@@ -195,28 +195,23 @@ func (r *Reconciler) reconcileOnce(ctx context.Context) (time.Duration, error) {
 	if err != nil {
 		return 0, err
 	}
-	buckets, err := AccountBucketNames(ctx, js)
+	buckets, err := AccountBuckets(ctx, js)
 	if err != nil {
 		return 0, fmt.Errorf("rds reconciler: enumerate account buckets: %w", err)
 	}
 	var revisit time.Duration
 	var failures []error
 	for _, bucket := range buckets {
-		kv, err := js.KeyValue(ctx, bucket)
-		if err != nil {
-			failures = append(failures, fmt.Errorf("open %s: %w", bucket, err))
-			continue
-		}
-		due, err := r.reconcileAccount(ctx, kv, AccountIDFromBucketName(bucket))
+		due, err := r.reconcileAccount(ctx, bucket, AccountIDFromBucketName(bucket.Name()))
 		revisit = reconciler.Earliest(revisit, due)
 		if err != nil {
-			failures = append(failures, fmt.Errorf("reconcile %s: %w", bucket, err))
+			failures = append(failures, fmt.Errorf("reconcile %s: %w", bucket.Name(), err))
 		}
 	}
 	return revisit, errors.Join(failures...)
 }
 
-func (r *Reconciler) reconcileAccount(ctx context.Context, kv jetstream.KeyValue, accountID string) (time.Duration, error) {
+func (r *Reconciler) reconcileAccount(ctx context.Context, kv *kvstore.Bucket, accountID string) (time.Duration, error) {
 	ids, err := ListDBInstanceIDs(ctx, kv)
 	if err != nil {
 		return 0, err
@@ -244,7 +239,7 @@ func (r *Reconciler) reconcileAccount(ctx context.Context, kv jetstream.KeyValue
 // The reconciler owns every transitional state that no single API call can
 // finish: the one it drives itself (creating), and the ones whose caller may
 // have died partway through. A settled instance is left alone.
-func (r *Reconciler) reconcileInstance(ctx context.Context, kv jetstream.KeyValue, accountID, id string) (time.Duration, error) {
+func (r *Reconciler) reconcileInstance(ctx context.Context, kv *kvstore.Bucket, accountID, id string) (time.Duration, error) {
 	var rec DBInstanceRecord
 	rev, found, err := getJSONRevision(ctx, kv, DBInstanceKey(id), &rec)
 	if err != nil || !found {
@@ -350,7 +345,7 @@ func (r *Reconciler) failureRevisit(rec *DBInstanceRecord) time.Duration {
 // vpcd applies the requested group list declaratively, so this removes the ENI
 // from the system VPC's default group rather than merely adding the new one
 // alongside it — and that default group is the whole of the exposure.
-func (r *Reconciler) remediateSystemENISG(ctx context.Context, kv jetstream.KeyValue,
+func (r *Reconciler) remediateSystemENISG(ctx context.Context, kv *kvstore.Bucket,
 	rev uint64, rec *DBInstanceRecord) (bool, error) {
 	if rec.SystemENIID == "" || rec.Status == StatusDeleting {
 		return false, nil
@@ -409,7 +404,7 @@ func (r *Reconciler) ensuredSystemSG(ctx context.Context) (string, error) {
 // operator acts. Only reported: deleting the payload on the grounds that a
 // healthy engine implies a completed bootstrap is exactly the inference this
 // protocol refuses to make.
-func (r *Reconciler) reportStalePendingBootstrap(ctx context.Context, kv jetstream.KeyValue,
+func (r *Reconciler) reportStalePendingBootstrap(ctx context.Context, kv *kvstore.Bucket,
 	accountID string, rec *DBInstanceRecord) error {
 	key := accountID + "/" + rec.DBInstanceIdentifier
 	stale := rec.Status == StatusAvailable && time.Since(rec.CreatedAt) > r.svc.bootstrapTimeout()
@@ -454,7 +449,7 @@ func (r *Reconciler) forgetPendingReport(key string) {
 	delete(r.reportedPending, key)
 }
 
-func (r *Reconciler) reconcileCreating(ctx context.Context, kv jetstream.KeyValue, rev uint64, accountID string, rec *DBInstanceRecord) error {
+func (r *Reconciler) reconcileCreating(ctx context.Context, kv *kvstore.Bucket, rev uint64, accountID string, rec *DBInstanceRecord) error {
 	// No lower bound on the heartbeat: the VM is new, so any beat naming it is
 	// necessarily this instance's.
 	healthy, err := r.engineReady(ctx, accountID, rec, time.Time{})
@@ -485,7 +480,7 @@ func withAgentReason(rec *DBInstanceRecord, reason string) string {
 // Reboot and start both end the same way: the engine comes back and says so.
 // The API call that began them returns before that happens, so this is what
 // actually lands the instance in available.
-func (r *Reconciler) reconcileRestarting(ctx context.Context, kv jetstream.KeyValue, rev uint64, accountID string, rec *DBInstanceRecord) error {
+func (r *Reconciler) reconcileRestarting(ctx context.Context, kv *kvstore.Bucket, rev uint64, accountID string, rec *DBInstanceRecord) error {
 	// The VM keeps its instance ID across a restart, so only a beat sent after
 	// the transition began proves the engine came back rather than that it was
 	// up before it went down.
@@ -516,7 +511,7 @@ func (r *Reconciler) reconcileRestarting(ctx context.Context, kv jetstream.KeyVa
 // half-applied, and the in-guest filesystem grow, which can only run once the
 // agent is up again. Both are driven from here so a customer's modify completes
 // without them, not just without them watching.
-func (r *Reconciler) reconcileModifying(ctx context.Context, kv jetstream.KeyValue, rev uint64, accountID string, rec *DBInstanceRecord) error {
+func (r *Reconciler) reconcileModifying(ctx context.Context, kv *kvstore.Bucket, rev uint64, accountID string, rec *DBInstanceRecord) error {
 	pending := rec.PendingModifiedValues
 	overrun := time.Since(transitionStarted(rec)) > transitionTimeout
 
@@ -590,7 +585,7 @@ func (r *Reconciler) reconcileModifying(ctx context.Context, kv jetstream.KeyVal
 // A stop whose caller died leaves the VM possibly still running, so the stop is
 // re-issued rather than assumed: it is idempotent, and a VM no node holds is
 // confirmed down against the fleet before the record calls it stopped.
-func (r *Reconciler) reconcileStopping(ctx context.Context, kv jetstream.KeyValue, rev uint64, accountID string, rec *DBInstanceRecord) error {
+func (r *Reconciler) reconcileStopping(ctx context.Context, kv *kvstore.Bucket, rev uint64, accountID string, rec *DBInstanceRecord) error {
 	if r.svc.deps.Instances == nil {
 		return errors.New("rds reconciler: no instance command path configured")
 	}
@@ -617,7 +612,7 @@ func (r *Reconciler) reconcileStopping(ctx context.Context, kv jetstream.KeyValu
 // resource, so replaying it converges rather than failing on what it already
 // did; only a teardown still stuck past the bound is called failed, which is
 // what lets the customer retry the delete.
-func (r *Reconciler) reconcileDeleting(ctx context.Context, kv jetstream.KeyValue, rev uint64, accountID string, rec *DBInstanceRecord) error {
+func (r *Reconciler) reconcileDeleting(ctx context.Context, kv *kvstore.Bucket, rev uint64, accountID string, rec *DBInstanceRecord) error {
 	err := r.svc.teardownDBInstance(ctx, kv, accountID, rec, false)
 	if err == nil || errors.Is(err, errFinalSnapshotInProgress) {
 		return nil
@@ -636,7 +631,7 @@ func (r *Reconciler) reconcileDeleting(ctx context.Context, kv jetstream.KeyValu
 // died. It is returned to where the snapshot found it; the quiesce needs no
 // undoing, because the agent's own hold deadline is far shorter than this bound
 // and has already released the engine.
-func (r *Reconciler) reconcileBackingUp(ctx context.Context, kv jetstream.KeyValue, rev uint64, accountID string, rec *DBInstanceRecord) error {
+func (r *Reconciler) reconcileBackingUp(ctx context.Context, kv *kvstore.Bucket, rev uint64, accountID string, rec *DBInstanceRecord) error {
 	if time.Since(transitionStarted(rec)) <= transitionTimeout {
 		return nil
 	}
@@ -660,7 +655,7 @@ func (r *Reconciler) reconcileBackingUp(ctx context.Context, kv jetstream.KeyVal
 // The record is re-read rather than carried over from the status pass above,
 // which may have transitioned the instance — a backup must not be started against
 // a status that has since moved.
-func (r *Reconciler) reconcileWindows(ctx context.Context, kv jetstream.KeyValue, accountID, id string) error {
+func (r *Reconciler) reconcileWindows(ctx context.Context, kv *kvstore.Bucket, accountID, id string) error {
 	var rec DBInstanceRecord
 	rev, found, err := getJSONRevision(ctx, kv, DBInstanceKey(id), &rec)
 	if err != nil || !found {
@@ -680,7 +675,7 @@ func (r *Reconciler) reconcileWindows(ctx context.Context, kv jetstream.KeyValue
 // snapshot is the authority: it was tagged with the DB snapshot identifier before
 // the record was flipped, so its presence says the data exists and its absence
 // says the create died before cutting it.
-func (r *Reconciler) reconcileSnapshots(ctx context.Context, kv jetstream.KeyValue, accountID string) (time.Duration, error) {
+func (r *Reconciler) reconcileSnapshots(ctx context.Context, kv *kvstore.Bucket, accountID string) (time.Duration, error) {
 	ids, err := ListDBSnapshotIDs(ctx, kv)
 	if err != nil {
 		return 0, fmt.Errorf("list DB snapshots: %w", err)
@@ -713,16 +708,16 @@ func (r *Reconciler) reconcileSnapshots(ctx context.Context, kv jetstream.KeyVal
 // Either adopts the EC2 snapshot the dead worker cut, or withdraws the record so
 // its identifier is usable again. A record left in creating would otherwise hold
 // the name forever while naming nothing a customer can restore.
-func (r *Reconciler) resolveCreatingSnapshot(ctx context.Context, kv jetstream.KeyValue, rev uint64,
+func (r *Reconciler) resolveCreatingSnapshot(ctx context.Context, kv *kvstore.Bucket, rev uint64,
 	accountID string, rec *DBSnapshotRecord) error {
 	snapshotID, err := r.svc.findEC2SnapshotFor(ctx, accountID, rec.DBSnapshotIdentifier)
 	if err != nil {
 		return err
 	}
 	if snapshotID == "" {
-		if err := kv.Delete(ctx, DBSnapshotKey(rec.DBSnapshotIdentifier), jetstream.LastRevision(rev)); err != nil {
+		if err := kv.CompareAndDelete(ctx, DBSnapshotKey(rec.DBSnapshotIdentifier), rev); err != nil {
 			switch {
-			case errors.Is(err, jetstream.ErrKeyNotFound), errors.Is(err, jetstream.ErrKeyRevisionMismatch):
+			case errors.Is(err, kvstore.ErrConflict):
 				// A concurrent completion or delete owns the newer revision.
 				return nil
 			default:
@@ -743,7 +738,7 @@ func (r *Reconciler) resolveCreatingSnapshot(ctx context.Context, kv jetstream.K
 	// conservative reading is the one that never overstates the snapshot.
 	rec.CrashConsistent = true
 	if err := updateJSON(ctx, kv, DBSnapshotKey(rec.DBSnapshotIdentifier), rev, rec); err != nil {
-		if errors.Is(err, jetstream.ErrKeyRevisionMismatch) {
+		if errors.Is(err, kvstore.ErrConflict) {
 			return nil
 		}
 		return err
@@ -829,7 +824,7 @@ func (r *Reconciler) engineReady(ctx context.Context, accountID string, rec *DBI
 
 // A CAS write, so a transition raced by an agent report or a lifecycle op is
 // dropped rather than clobbering the newer state; the next pass re-reads.
-func (r *Reconciler) transition(ctx context.Context, kv jetstream.KeyValue, rev uint64, rec *DBInstanceRecord, to Status, reason string) error {
+func (r *Reconciler) transition(ctx context.Context, kv *kvstore.Bucket, rev uint64, rec *DBInstanceRecord, to Status, reason string) error {
 	if !CanTransition(rec.Status, to) {
 		return fmt.Errorf("illegal transition %s -> %s", rec.Status, to)
 	}
@@ -845,7 +840,7 @@ func (r *Reconciler) transition(ctx context.Context, kv jetstream.KeyValue, rev
 	rec.UpdatedAt = time.Now().UTC()
 
 	if err := updateJSON(ctx, kv, DBInstanceKey(rec.DBInstanceIdentifier), rev, rec); err != nil {
-		if errors.Is(err, jetstream.ErrKeyRevisionMismatch) {
+		if errors.Is(err, kvstore.ErrConflict) {
 			slog.DebugContext(ctx, "rds reconciler: transition lost a revision race; retrying next pass",
 				"dbInstance", rec.DBInstanceIdentifier, "to", to)
 			return nil
@@ -855,7 +850,7 @@ func (r *Reconciler) transition(ctx context.Context, kv jetstream.KeyValue, rev
 	slog.InfoContext(ctx, "rds reconciler: DB instance transitioned",
 		"dbInstance", rec.DBInstanceIdentifier, "from", from, "to", to, "reason", reason)
 	if from == StatusCreating && to == StatusAvailable {
-		r.svc.RecordEvent(ctx, AccountIDFromBucketName(kv.Bucket()), EventSourceTypeDBInstance,
+		r.svc.RecordEvent(ctx, AccountIDFromBucketName(kv.Name()), EventSourceTypeDBInstance,
 			rec.DBInstanceIdentifier, "DB instance is available.", EventCategoryAvailability)
 	}
 	return nil
@@ -863,7 +858,7 @@ func (r *Reconciler) transition(ctx context.Context, kv jetstream.KeyValue, rev
 
 // transition against a freshly read revision, for a caller that has written the
 // record itself since the pass opened and so cannot use the revision it read.
-func (r *Reconciler) transitionFresh(ctx context.Context, kv jetstream.KeyValue, id string, to Status, reason string) error {
+func (r *Reconciler) transitionFresh(ctx context.Context, kv *kvstore.Bucket, id string, to Status, reason string) error {
 	var rec DBInstanceRecord
 	rev, found, err := getJSONRevision(ctx, kv, DBInstanceKey(id), &rec)
 	if err != nil || !found {

spinifex/handlers/rds/recovery.go

@@ -7,7 +7,7 @@ import (
 	"log/slog"
 	"time"
 
-	"github.com/nats-io/nats.go/jetstream"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 )
 
 // How long an instance may be observed dark before it is called failed, unless
@@ -73,7 +73,7 @@ func classifyHealth(obs healthObservation) healthVerdict {
 // engine reports healthy again — the recovery the AMI's in-guest
 // Restart=on-failure and EC2's VM auto-restart provide underneath. v1.0 does no
 // repair of its own.
-func (r *Reconciler) reconcileHealth(ctx context.Context, kv jetstream.KeyValue, rev uint64, accountID string, rec *DBInstanceRecord) error {
+func (r *Reconciler) reconcileHealth(ctx context.Context, kv *kvstore.Bucket, rev uint64, accountID string, rec *DBInstanceRecord) error {
 	obs := r.observeAgent(accountID, rec)
 
 	// The VM lookup is a fleet-wide describe fan-out, so it is issued only where
@@ -101,7 +101,7 @@ func (r *Reconciler) reconcileHealth(ctx context.Context, kv jetstream.KeyValue,
 // A healthy heartbeat is the only thing that resets the failure clock. VM state
 // alone must not: a VM that boots and immediately wedges would otherwise reset
 // it every pass and mask a persistent fault indefinitely.
-func (r *Reconciler) clearFailure(ctx context.Context, kv jetstream.KeyValue, rev uint64, accountID string, rec *DBInstanceRecord) error {
+func (r *Reconciler) clearFailure(ctx context.Context, kv *kvstore.Bucket, rev uint64, accountID string, rec *DBInstanceRecord) error {
 	if rec.Status == StatusFailed {
 		rec.UnhealthySince = nil
 		if err := r.transition(ctx, kv, rev, rec, StatusAvailable, ""); err != nil {
@@ -122,7 +122,7 @@ func (r *Reconciler) clearFailure(ctx context.Context, kv jetstream.KeyValue, re
 // Starts the failure clock on the first dark pass and fails the instance once
 // the grace window has elapsed. Both writes carry the clock, so the timestamp a
 // later leader measures against is the one the original observation stamped.
-func (r *Reconciler) recordFailure(ctx context.Context, kv jetstream.KeyValue, rev uint64, accountID string, rec *DBInstanceRecord, obs healthObservation) error {
+func (r *Reconciler) recordFailure(ctx context.Context, kv *kvstore.Bucket, rev uint64, accountID string, rec *DBInstanceRecord, obs healthObservation) error {
 	// Terminal for the control plane in v1.0: nothing retries, so re-recording a
 	// failure that is already recorded would only churn the record.
 	if rec.Status == StatusFailed {
@@ -151,10 +151,10 @@ func (r *Reconciler) recordFailure(ctx context.Context, kv jetstream.KeyValue, r
 // Writes the clock without a status change. A lost revision race is dropped
 // rather than retried: the record moved under this pass, and the next one
 // re-reads and re-stamps against whatever it moved to.
-func (r *Reconciler) persistFailureClock(ctx context.Context, kv jetstream.KeyValue, rev uint64, rec *DBInstanceRecord) error {
+func (r *Reconciler) persistFailureClock(ctx context.Context, kv *kvstore.Bucket, rev uint64, rec *DBInstanceRecord) error {
 	rec.UpdatedAt = time.Now().UTC()
 	err := updateJSON(ctx, kv, DBInstanceKey(rec.DBInstanceIdentifier), rev, rec)
-	if errors.Is(err, jetstream.ErrKeyRevisionMismatch) {
+	if errors.Is(err, kvstore.ErrConflict) {
 		slog.DebugContext(ctx, "rds reconciler: failure clock lost a revision race; retrying next pass",
 			"dbInstance", rec.DBInstanceIdentifier)
 		return nil

spinifex/handlers/rds/replace.go

@@ -6,8 +6,8 @@ import (
 	"fmt"
 	"log/slog"
 
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 	"github.com/mulgadc/spinifex/spinifex/utils"
-	"github.com/nats-io/nats.go/jetstream"
 )
 
 // The one primitive a class change, a storage grow that rides the same outage,
@@ -40,7 +40,7 @@ type replaceInput struct {
 // Only the VM's identity is written here — the class, the storage size and the
 // pending values belong to the caller that asked for the change, so a failure
 // part-way through leaves the request still recorded and retryable.
-func (s *Service) replaceInstanceVM(ctx context.Context, kv jetstream.KeyValue, accountID string, rec *DBInstanceRecord, in replaceInput) error {
+func (s *Service) replaceInstanceVM(ctx context.Context, kv *kvstore.Bucket, accountID string, rec *DBInstanceRecord, in replaceInput) error {
 	if rec.ENIID == "" || rec.DataVolumeID == "" {
 		return fmt.Errorf("rds: DB instance %s has no persisted ENI and data volume to replace onto",
 			rec.DBInstanceIdentifier)
@@ -199,7 +199,7 @@ func (s *Service) launchReplacementVM(ctx context.Context, accountID string, rec
 // agent's IMDS credentials still resolve to this DB instance.
 func (s *Service) rewriteInstanceIndex(ctx context.Context, accountID string, rec *DBInstanceRecord, oldInstanceID, newInstanceID string) error {
 	if oldInstanceID != "" && oldInstanceID != newInstanceID {
-		if err := s.DeleteInstanceIndex(ctx, oldInstanceID); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
+		if err := s.DeleteInstanceIndex(ctx, oldInstanceID); err != nil {
 			return fmt.Errorf("rds: drop the superseded instance index entry for %s: %w", rec.DBInstanceIdentifier, err)
 		}
 	}
@@ -226,7 +226,7 @@ func (s *Service) rollbackReplacementLaunch(
 	defer cancel()
 
 	var cleanupErrs []error
-	if err := s.DeleteInstanceIndex(cleanupCtx, launched.InstanceID); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
+	if err := s.DeleteInstanceIndex(cleanupCtx, launched.InstanceID); err != nil {
 		cleanupErrs = append(cleanupErrs, fmt.Errorf("delete replacement instance index %s: %w", launched.InstanceID, err))
 	}
 	if oldInstanceID != "" {

spinifex/handlers/rds/restore.go

@@ -11,9 +11,9 @@ import (
 	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/aws/aws-sdk-go/service/rds"
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 	"github.com/mulgadc/spinifex/spinifex/tags"
 	"github.com/mulgadc/spinifex/spinifex/utils"
-	"github.com/nats-io/nats.go/jetstream"
 )
 
 // Builds a new DB instance on a volume created from the snapshot. It is a fresh
@@ -76,7 +76,7 @@ func (s *Service) RestoreDBInstanceFromDBSnapshot(ctx context.Context, input *rd
 	rec := newRestoredDBInstanceRecord(accountID, req, placement, parameters, snapshot)
 	rollbackRev, createErr := createJSONRevision(ctx, kv, key, &rec)
 	if createErr != nil {
-		if errors.Is(createErr, jetstream.ErrKeyExists) {
+		if errors.Is(createErr, kvstore.ErrExists) {
 			return nil, awserrors.Errorf(awserrors.ErrorDBInstanceAlreadyExists,
 				"DB instance %s already exists", req.Identifier)
 		}

spinifex/handlers/rds/service.go

@@ -11,6 +11,7 @@ import (
 	"time"
 
 	"github.com/mulgadc/spinifex/spinifex/admin"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 	"github.com/nats-io/nats.go"
 	"github.com/nats-io/nats.go/jetstream"
 )
@@ -86,6 +87,13 @@ type Service struct {
 	// persisted only on change or on the slower floor.
 	livenessMu sync.Mutex
 	liveness   map[string]*agentLiveness
+
+	// buckets memoises one Bucket per account, plus the single system bucket.
+	// Resolution is per request, so without this every RDS call pays a bucket
+	// create round trip before it does any work.
+	bucketsMu sync.Mutex
+	buckets   map[string]*kvstore.Bucket
+	systemB   *kvstore.Bucket
 }
 
 type agentLiveness struct {
@@ -158,20 +166,48 @@ func (s *Service) js() (jetstream.JetStream, error) {
 	return jetstream.New(s.nc)
 }
 
-func (s *Service) bucket(ctx context.Context, accountID string) (jetstream.KeyValue, error) {
+// bucket resolves accountID's bucket, memoised: the resolution used to bill a
+// stream create per RDS call, and a Bucket also carries the recovery a raw
+// handle cannot — a lost stream is reopened under the next read rather than
+// failing every call until the process restarts.
+func (s *Service) bucket(ctx context.Context, accountID string) (*kvstore.Bucket, error) {
 	js, err := s.js()
 	if err != nil {
 		return nil, err
 	}
-	return GetOrCreateAccountBucket(ctx, js, accountID)
+	s.bucketsMu.Lock()
+	b, ok := s.buckets[accountID]
+	if !ok {
+		b = kvstore.NewBucket(js, AccountBucketConfig(accountID))
+		if s.buckets == nil {
+			s.buckets = make(map[string]*kvstore.Bucket)
+		}
+		s.buckets[accountID] = b
+	}
+	s.bucketsMu.Unlock()
+	// Opened outside the lock, under the Bucket's own, so one account's slow
+	// first open does not hold every other account's resolution behind it.
+	if _, err := b.KV(ctx); err != nil {
+		return nil, err
+	}
+	return b, nil
 }
 
-func (s *Service) systemBucket(ctx context.Context) (jetstream.KeyValue, error) {
+func (s *Service) systemBucket(ctx context.Context) (*kvstore.Bucket, error) {
 	js, err := s.js()
 	if err != nil {
 		return nil, err
 	}
-	return GetOrCreateSystemBucket(ctx, js)
+	s.bucketsMu.Lock()
+	if s.systemB == nil {
+		s.systemB = kvstore.NewBucket(js, SystemBucketConfig())
+	}
+	b := s.systemB
+	s.bucketsMu.Unlock()
+	if _, err := b.KV(ctx); err != nil {
+		return nil, err
+	}
+	return b, nil
 }
 
 // Both paths empty deliberately disables TLS; a partial configuration is an
@@ -200,83 +236,79 @@ func (s *Service) tlsAvailable() (bool, error) {
 	return caCert != nil && caKey != nil, nil
 }
 
+// The eight helpers below are the account bucket's record layer. They are
+// adapters over kvstore rather than a codec of their own: the JSON, the error
+// wrapping and the stream recovery belong to kvstore, and what stays here is
+// the not-found convention its callers are written against.
+
 // Returns (false, nil) when the key is absent.
-func getJSON(ctx context.Context, kv jetstream.KeyValue, key string, out any) (bool, error) {
-	entry, err := kv.Get(ctx, key)
-	if errors.Is(err, jetstream.ErrKeyNotFound) {
+func getJSON[T any](ctx context.Context, b *kvstore.Bucket, key string, out *T) (bool, error) {
+	v, _, err := kvstore.On[T](b).Get(ctx, key)
+	if errors.Is(err, kvstore.ErrNotFound) {
 		return false, nil
 	}
 	if err != nil {
 		return false, err
 	}
-	if err := json.Unmarshal(entry.Value(), out); err != nil {
-		return false, fmt.Errorf("unmarshal %s: %w", key, err)
-	}
+	*out = *v
 	return true, nil
 }
 
 // getJSON plus the entry revision, for callers that follow with a CAS update.
-func getJSONRevision(ctx context.Context, kv jetstream.KeyValue, key string, out any) (uint64, bool, error) {
-	entry, err := kv.Get(ctx, key)
-	if errors.Is(err, jetstream.ErrKeyNotFound) {
+func getJSONRevision[T any](ctx context.Context, b *kvstore.Bucket, key string, out *T) (uint64, bool, error) {
+	v, rev, err := kvstore.On[T](b).Get(ctx, key)
+	if errors.Is(err, kvstore.ErrNotFound) {
 		return 0, false, nil
 	}
 	if err != nil {
 		return 0, false, err
 	}
-	if err := json.Unmarshal(entry.Value(), out); err != nil {
-		return 0, false, fmt.Errorf("unmarshal %s: %w", key, err)
-	}
-	return entry.Revision(), true, nil
+	*out = *v
+	return rev, true, nil
 }
 
-func putJSON(ctx context.Context, kv jetstream.KeyValue, key string, v any) error {
-	data, err := json.Marshal(v)
+// getJSONRevisionAny is getJSONRevision for the one caller whose record type is
+// chosen at runtime from the ARN, so no type parameter can name it. The bytes
+// still come through kvstore; only the decode is dynamic.
+func getJSONRevisionAny(ctx context.Context, b *kvstore.Bucket, key string, out any) (uint64, bool, error) {
+	raw, rev, err := kvstore.On[json.RawMessage](b).Get(ctx, key)
+	if errors.Is(err, kvstore.ErrNotFound) {
+		return 0, false, nil
+	}
 	if err != nil {
-		return err
+		return 0, false, err
 	}
-	if _, err := kv.Put(ctx, key, data); err != nil {
-		return fmt.Errorf("put %s: %w", key, err)
+	if err := json.Unmarshal(*raw, out); err != nil {
+		return 0, false, fmt.Errorf("unmarshal %s: %w", key, err)
 	}
-	return nil
+	return rev, true, nil
+}
+
+func putJSON[T any](ctx context.Context, b *kvstore.Bucket, key string, v T) error {
+	return kvstore.On[T](b).Set(ctx, key, &v)
 }
 
 // Writes v at key only if nothing is stored there, so the key doubles as a
-// cluster-wide reservation. Returns jetstream.ErrKeyExists when it is taken.
-func createJSON(ctx context.Context, kv jetstream.KeyValue, key string, v any) error {
-	_, err := createJSONRevision(ctx, kv, key, v)
+// cluster-wide reservation. Returns kvstore.ErrExists when it is taken.
+func createJSON[T any](ctx context.Context, b *kvstore.Bucket, key string, v T) error {
+	_, err := createJSONRevision(ctx, b, key, v)
 	return err
 }
 
 // createJSON plus the created entry's revision, for callers that must undo only
 // the reservation they created rather than a concurrent replacement.
-func createJSONRevision(ctx context.Context, kv jetstream.KeyValue, key string, v any) (uint64, error) {
-	data, err := json.Marshal(v)
-	if err != nil {
-		return 0, err
-	}
-	rev, err := kv.Create(ctx, key, data)
-	if err != nil {
-		return 0, err
-	}
-	return rev, nil
+func createJSONRevision[T any](ctx context.Context, b *kvstore.Bucket, key string, v T) (uint64, error) {
+	return kvstore.On[T](b).Create(ctx, key, &v)
 }
 
-// Writes v at key only if the stored entry is still at rev.
-func updateJSON(ctx context.Context, kv jetstream.KeyValue, key string, rev uint64, v any) error {
-	_, err := updateJSONRevision(ctx, kv, key, rev, v)
+// Writes v at key only if the stored entry is still at rev. Returns
+// kvstore.ErrConflict when another writer got there first.
+func updateJSON[T any](ctx context.Context, b *kvstore.Bucket, key string, rev uint64, v T) error {
+	_, err := updateJSONRevision(ctx, b, key, rev, v)
 	return err
 }
 
 // updateJSON plus the revision written by the successful CAS update.
-func updateJSONRevision(ctx context.Context, kv jetstream.KeyValue, key string, rev uint64, v any) (uint64, error) {
-	data, err := json.Marshal(v)
-	if err != nil {
-		return 0, err
-	}
-	updatedRev, err := kv.Update(ctx, key, data, rev)
-	if err != nil {
-		return 0, fmt.Errorf("update %s: %w", key, err)
-	}
-	return updatedRev, nil
+func updateJSONRevision[T any](ctx context.Context, b *kvstore.Bucket, key string, rev uint64, v T) (uint64, error) {
+	return kvstore.On[T](b).CompareAndSet(ctx, key, &v, rev)
 }

spinifex/handlers/rds/snapshot.go

@@ -13,9 +13,9 @@ import (
 	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/aws/aws-sdk-go/service/rds"
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 	"github.com/mulgadc/spinifex/spinifex/tags"
 	"github.com/mulgadc/spinifex/spinifex/utils"
-	"github.com/nats-io/nats.go/jetstream"
 )
 
 // The EC2 snapshot surface the RDS control plane drives. A DB snapshot is an
@@ -78,7 +78,7 @@ func (s *Service) CreateDBSnapshot(ctx context.Context, input *rds.CreateDBSnaps
 // rev is the revision the caller read the DB instance at: the CAS that moves it
 // into backing-up is what serialises the two kinds of snapshot against each other
 // and against a lifecycle op.
-func (s *Service) snapshotDBInstance(ctx context.Context, kv jetstream.KeyValue, rev uint64,
+func (s *Service) snapshotDBInstance(ctx context.Context, kv *kvstore.Bucket, rev uint64,
 	accountID string, rec *DBInstanceRecord, req *validatedSnapshot) (*DBSnapshotRecord, error) {
 	if s.deps.Snapshots == nil {
 		return nil, errors.New("rds: no snapshot service configured")
@@ -107,7 +107,7 @@ func (s *Service) snapshotDBInstance(ctx context.Context, kv jetstream.KeyValue,
 	record := newDBSnapshotRecord(accountID, rec, req)
 	key := DBSnapshotKey(req.DBSnapshotIdentifier)
 	if err := createJSON(ctx, kv, key, &record); err != nil {
-		if errors.Is(err, jetstream.ErrKeyExists) {
+		if errors.Is(err, kvstore.ErrExists) {
 			return nil, awserrors.Errorf(awserrors.ErrorDBSnapshotAlreadyExists,
 				"DB snapshot %s already exists", req.DBSnapshotIdentifier)
 		}
@@ -308,7 +308,7 @@ func (s *Service) DeleteDBSnapshot(ctx context.Context, input *rds.DeleteDBSnaps
 // data volume behind it. Shared with the retention sweep, which cannot go
 // through DeleteDBSnapshot: that rejects the rds: namespace automated snapshots
 // live in, so a customer can never delete one by hand.
-func (s *Service) removeDBSnapshot(ctx context.Context, kv jetstream.KeyValue, accountID string, rec *DBSnapshotRecord) error {
+func (s *Service) removeDBSnapshot(ctx context.Context, kv *kvstore.Bucket, accountID string, rec *DBSnapshotRecord) error {
 	if s.deps.Snapshots == nil {
 		return errors.New("rds: no snapshot service configured")
 	}
@@ -330,7 +330,7 @@ func (s *Service) removeDBSnapshot(ctx context.Context, kv jetstream.KeyValue, a
 	}
 	// Last: while it exists the snapshot is still nameable, and a retry re-runs
 	// the steps above, each of which tolerates work it has already done.
-	if err := kv.Delete(ctx, DBSnapshotKey(id)); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
+	if err := kv.Delete(ctx, DBSnapshotKey(id)); err != nil {
 		return fmt.Errorf("rds: delete the DB snapshot record for %s: %w", id, err)
 	}
 
@@ -343,7 +343,7 @@ func (s *Service) removeDBSnapshot(ctx context.Context, kv jetstream.KeyValue, a
 // the case where an instance restored from it is still running. That is
 // translated into an RDS fault naming the instance rather than surfaced as the
 // raw EC2 code, which says nothing a customer can act on.
-func (s *Service) deleteEC2Snapshot(ctx context.Context, kv jetstream.KeyValue, accountID string, rec *DBSnapshotRecord) error {
+func (s *Service) deleteEC2Snapshot(ctx context.Context, kv *kvstore.Bucket, accountID string, rec *DBSnapshotRecord) error {
 	if rec.SnapshotID == "" {
 		return nil
 	}
@@ -365,7 +365,7 @@ func (s *Service) deleteEC2Snapshot(ctx context.Context, kv jetstream.KeyValue,
 // What the customer has to remove before the snapshot can go. A restored
 // instance is named; anything else is a volume outside RDS's own bookkeeping,
 // which is reported as such rather than guessed at.
-func (s *Service) describeSnapshotDependents(ctx context.Context, kv jetstream.KeyValue, dbSnapshotIdentifier string) string {
+func (s *Service) describeSnapshotDependents(ctx context.Context, kv *kvstore.Bucket, dbSnapshotIdentifier string) string {
 	ids, err := ListDBInstanceIDs(ctx, kv)
 	if err != nil {
 		slog.WarnContext(ctx, "rds: listing the instances restored from a snapshot failed",
@@ -392,7 +392,7 @@ func (s *Service) describeSnapshotDependents(ctx context.Context, kv jetstream.K
 // rather than taken from the retained record: that index is what DeleteVolume
 // enforces against, and a record written before another snapshot was taken would
 // otherwise read as "nothing holds it".
-func (s *Service) releaseRetainedVolume(ctx context.Context, kv jetstream.KeyValue, rec *DBSnapshotRecord) error {
+func (s *Service) releaseRetainedVolume(ctx context.Context, kv *kvstore.Bucket, rec *DBSnapshotRecord) error {
 	if rec.SourceVolumeID == "" {
 		return nil
 	}
@@ -408,7 +408,7 @@ func (s *Service) releaseRetainedVolume(ctx context.Context, kv jetstream.KeyVal
 // Deletes a retained data volume once nothing holds it, or records the holders it
 // still has. Reports whether the volume actually went, which is what lets the
 // reaper count the ones it reclaimed.
-func (s *Service) reclaimRetainedVolume(ctx context.Context, kv jetstream.KeyValue,
+func (s *Service) reclaimRetainedVolume(ctx context.Context, kv *kvstore.Bucket,
 	retained *RetainedVolumeRecord) (bool, error) {
 	holders, err := s.snapshotsHolding(ctx, retained.VolumeID)
 	if err != nil {
@@ -438,7 +438,7 @@ func (s *Service) reclaimRetainedVolume(ctx context.Context, kv jetstream.KeyVal
 	default:
 		return false, fmt.Errorf("rds: delete the retained data volume %s: %w", retained.VolumeID, err)
 	}
-	if err := kv.Delete(ctx, RetainedVolumeKey(retained.VolumeID)); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
+	if err := kv.Delete(ctx, RetainedVolumeKey(retained.VolumeID)); err != nil {
 		return false, fmt.Errorf("rds: clear the retained-volume record for %s: %w", retained.VolumeID, err)
 	}
 	slog.InfoContext(ctx, "rds: retained data volume reclaimed; its last snapshot is gone",
@@ -586,7 +586,7 @@ func newDBSnapshotRecord(accountID string, rec *DBInstanceRecord, req *validated
 }
 
 // Returns the record plus its revision, for callers that follow with a CAS.
-func (s *Service) getDBSnapshot(ctx context.Context, kv jetstream.KeyValue, id string) (*DBSnapshotRecord, uint64, error) {
+func (s *Service) getDBSnapshot(ctx context.Context, kv *kvstore.Bucket, id string) (*DBSnapshotRecord, uint64, error) {
 	var rec DBSnapshotRecord
 	rev, found, err := getJSONRevision(ctx, kv, DBSnapshotKey(id), &rec)
 	if err != nil {
@@ -601,7 +601,7 @@ func (s *Service) getDBSnapshot(ctx context.Context, kv jetstream.KeyValue, id s
 // Rejects a taken identifier before any work starts. A snapshot record outlives
 // the instance it came from, so a name can be held by a snapshot of an instance
 // that no longer exists.
-func (s *Service) checkDBSnapshotAvailable(ctx context.Context, kv jetstream.KeyValue, id string) error {
+func (s *Service) checkDBSnapshotAvailable(ctx context.Context, kv *kvstore.Bucket, id string) error {
 	var existing DBSnapshotRecord
 	found, err := getJSON(ctx, kv, DBSnapshotKey(id), &existing)
 	if err != nil {
@@ -615,10 +615,10 @@ func (s *Service) checkDBSnapshotAvailable(ctx context.Context, kv jetstream.Key
 
 // Withdraws a creating record whose snapshot never happened, on a context
 // detached from the caller's so an expired request deadline still clears it.
-func (s *Service) discardSnapshotRecord(ctx context.Context, kv jetstream.KeyValue, id string) {
+func (s *Service) discardSnapshotRecord(ctx context.Context, kv *kvstore.Bucket, id string) {
 	rbCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), rollbackTimeout)
 	defer cancel()
-	if err := kv.Delete(rbCtx, DBSnapshotKey(id)); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
+	if err := kv.Delete(rbCtx, DBSnapshotKey(id)); err != nil {
 		slog.WarnContext(rbCtx, "rds: withdrawing the record of a snapshot that failed left it behind",
 			"dbSnapshot", id, "err", err)
 	}
@@ -628,7 +628,7 @@ func (s *Service) discardSnapshotRecord(ctx context.Context, kv jetstream.KeyVal
 // CAS, returning the status the instance goes back to. That write is the
 // per-instance guard: a second snapshot, an automated one or a lifecycle
 // op finds the instance already backing-up and is rejected rather than queued.
-func (s *Service) beginSnapshotOperation(ctx context.Context, kv jetstream.KeyValue, rev uint64,
+func (s *Service) beginSnapshotOperation(ctx context.Context, kv *kvstore.Bucket, rev uint64,
 	rec *DBInstanceRecord, dbSnapshotIdentifier string) (Status, error) {
 	legal := []Status{StatusAvailable, StatusStopped}
 	if !slices.Contains(legal, rec.Status) {
@@ -649,7 +649,7 @@ func (s *Service) beginSnapshotOperation(ctx context.Context, kv jetstream.KeyVa
 	rec.UpdatedAt = now
 
 	if err := updateJSON(ctx, kv, DBInstanceKey(rec.DBInstanceIdentifier), rev, rec); err != nil {
-		if errors.Is(err, jetstream.ErrKeyRevisionMismatch) {
+		if errors.Is(err, kvstore.ErrConflict) {
 			return "", awserrors.Errorf(awserrors.ErrorDBInstanceInvalidState,
 				"DB instance %s changed state concurrently; retry the snapshot", rec.DBInstanceIdentifier)
 		}
@@ -662,7 +662,7 @@ func (s *Service) beginSnapshotOperation(ctx context.Context, kv jetstream.KeyVa
 // leaves a usable instance rather than one stuck in backing-up. A delete
 // accepted while the snapshot ran owns the record by now, so the resume only
 // applies to an instance still in backing-up.
-func (s *Service) endSnapshotOperation(ctx context.Context, kv jetstream.KeyValue, id string, resume Status) {
+func (s *Service) endSnapshotOperation(ctx context.Context, kv *kvstore.Bucket, id string, resume Status) {
 	endCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), defaultTimeout)
 	defer cancel()
 	err := s.updateInstance(endCtx, kv, id, func(stored *DBInstanceRecord) {

spinifex/handlers/rds/store.go

@@ -11,6 +11,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/mulgadc/spinifex/spinifex/kvlease"
 	"github.com/mulgadc/spinifex/spinifex/kvstore"
 	"github.com/mulgadc/spinifex/spinifex/kvutil"
 	"github.com/mulgadc/spinifex/spinifex/migrate"
@@ -190,15 +191,51 @@ func AccountBucketName(accountID string) string {
 	return KVBucketRDSAccountPrefix + accountID
 }
 
+// AccountBucketConfig describes one tenant's bucket. The name carries the
+// account, so a Bucket is per-tenant and the Service holds one per account
+// rather than one for the package.
+func AccountBucketConfig(accountID string) kvstore.Config {
+	bucket := AccountBucketName(accountID)
+	return kvstore.Config{
+		Name:    bucket,
+		History: KVBucketRDSAccountHistory,
+		Missing: "rds service: nil nats connection",
+		OnOpen: func(ctx context.Context, kv jetstream.KeyValue) error {
+			return migrate.DefaultRegistry.RunKV(ctx, bucket, kv, KVBucketRDSAccountVersion)
+		},
+	}
+}
+
+// AccountBucketReadConfig describes one tenant's bucket for a cross-tenant pass
+// that only reads it. No OnOpen: a reconcile pass is not where a schema
+// migration should be discovered, and it would bill a version read per account.
+func AccountBucketReadConfig(accountID string) kvstore.Config {
+	return kvstore.Config{
+		Name:       AccountBucketName(accountID),
+		History:    KVBucketRDSAccountHistory,
+		Missing:    "rds service: nil nats connection",
+		AttachOnly: true,
+	}
+}
+
+// SystemBucketConfig describes the cross-tenant bucket holding the instanceID →
+// DB instance reverse index.
+func SystemBucketConfig() kvstore.Config {
+	return kvstore.Config{
+		Name:    KVBucketRDSSystem,
+		History: KVBucketRDSSystemHistory,
+		Missing: "rds service: nil nats connection",
+		OnOpen: func(ctx context.Context, kv jetstream.KeyValue) error {
+			return migrate.DefaultRegistry.RunKV(ctx, KVBucketRDSSystem, kv, KVBucketRDSSystemVersion)
+		},
+	}
+}
+
 // Creates the bucket on first use; subsequent calls return the existing handle.
 func GetOrCreateAccountBucket(ctx context.Context, js jetstream.JetStream, accountID string) (jetstream.KeyValue, error) {
-	bucket := AccountBucketName(accountID)
-	kv, err := kvutil.GetOrCreateBucket(ctx, js, bucket, KVBucketRDSAccountHistory)
+	kv, err := kvstore.NewBucket(js, AccountBucketConfig(accountID)).KV(ctx)
 	if err != nil {
-		return nil, fmt.Errorf("failed to create RDS per-account KV bucket %s: %w", bucket, err)
-	}
-	if err := migrate.DefaultRegistry.RunKV(ctx, bucket, kv, KVBucketRDSAccountVersion); err != nil {
-		return nil, fmt.Errorf("migrate %s: %w", bucket, err)
+		return nil, fmt.Errorf("failed to create RDS per-account KV bucket %s: %w", AccountBucketName(accountID), err)
 	}
 	return kv, nil
 }
@@ -206,37 +243,21 @@ func GetOrCreateAccountBucket(ctx context.Context, js jetstream.JetStream, accou
 // Created lazily alongside the first DB instance rather than at daemon boot, so
 // a cluster with no RDS usage carries no bucket.
 func GetOrCreateSystemBucket(ctx context.Context, js jetstream.JetStream) (jetstream.KeyValue, error) {
-	kv, err := kvutil.GetOrCreateBucket(ctx, js, KVBucketRDSSystem, KVBucketRDSSystemHistory)
+	kv, err := kvstore.NewBucket(js, SystemBucketConfig()).KV(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create RDS system KV bucket %s: %w", KVBucketRDSSystem, err)
 	}
-	if err := migrate.DefaultRegistry.RunKV(ctx, KVBucketRDSSystem, kv, KVBucketRDSSystemVersion); err != nil {
-		return nil, fmt.Errorf("migrate %s: %w", KVBucketRDSSystem, err)
-	}
 	return kv, nil
 }
 
-// kvutil.GetOrCreateBucket exposes no TTL knob, so the lease bucket attaches
-// directly and creates only when absent.
+// InitLeaderBucket creates (or attaches to) the shared spinifex-rds-leader
+// bucket used for the reconciler's leader-lease CAS lock.
 func InitLeaderBucket(ctx context.Context, js jetstream.JetStream) (jetstream.KeyValue, error) {
-	// Attach before create. This runs on every reconcile tick, and CreateKeyValue
-	// against a bucket that exists with any other config is a STREAM.CREATE the
-	// meta leader answers with an error, so creating first bills one per tick.
-	kv, err := js.KeyValue(ctx, KVBucketRDSLeader)
-	if errors.Is(err, jetstream.ErrBucketNotFound) {
-		kv, err = js.CreateKeyValue(ctx, jetstream.KeyValueConfig{
-			Bucket:  KVBucketRDSLeader,
-			History: 1,
-			TTL:     KVBucketRDSLeaderTTL,
-		})
-	}
-	if err != nil {
-		return nil, fmt.Errorf("failed to create or open RDS leader bucket %s: %w", KVBucketRDSLeader, err)
-	}
-	if err := migrate.DefaultRegistry.RunKV(ctx, KVBucketRDSLeader, kv, KVBucketRDSLeaderVersion); err != nil {
-		return nil, fmt.Errorf("migrate %s: %w", KVBucketRDSLeader, err)
-	}
-	return kv, nil
+	return kvlease.OpenBucket(ctx, js, kvlease.BucketConfig{
+		Name:    KVBucketRDSLeader,
+		TTL:     KVBucketRDSLeaderTTL,
+		Version: KVBucketRDSLeaderVersion,
+	})
 }
 
 // Every RDS per-account bucket in the cluster. The reconciler and the DNS
@@ -255,22 +276,33 @@ func AccountBucketNames(ctx context.Context, js jetstream.JetStream) ([]string,
 	return names, nil
 }
 
-// AccountWatchBuckets returns every per-account bucket as a watchable handle,
-// for a reconciler that must be woken by a cluster change rather than poll for
-// it. The set is re-read on each call because a new account's bucket appears
+// AccountBuckets returns every per-account bucket as a typed handle, for the
+// cross-tenant passes — the reconciler, the backup sweep and the DNS desired
+// set. The set is re-read on each call because a new account's bucket appears
 // without notice: JetStream publishes no bucket-created event.
-func AccountWatchBuckets(ctx context.Context, js jetstream.JetStream) ([]*kvstore.Bucket, error) {
+//
+// The handles are unopened, so a bucket that cannot be reached surfaces from
+// the pass that reads it rather than costing every caller an open it may not
+// need. They attach and never create: these passes enumerate the buckets they
+// read, so an absent one means it was deleted under them, not that it is new.
+func AccountBuckets(ctx context.Context, js jetstream.JetStream) ([]*kvstore.Bucket, error) {
 	names, err := AccountBucketNames(ctx, js)
 	if err != nil {
 		return nil, err
 	}
 	buckets := make([]*kvstore.Bucket, 0, len(names))
 	for _, name := range names {
-		buckets = append(buckets, kvstore.NewBucket(js, kvstore.Config{Name: name, History: 1}))
+		buckets = append(buckets, kvstore.NewBucket(js, AccountBucketReadConfig(AccountIDFromBucketName(name))))
 	}
 	return buckets, nil
 }
 
+// AccountWatchBuckets is AccountBuckets for a reconciler that must be woken by
+// a cluster change rather than poll for it.
+func AccountWatchBuckets(ctx context.Context, js jetstream.JetStream) ([]*kvstore.Bucket, error) {
+	return AccountBuckets(ctx, js)
+}
+
 // The account a per-account bucket belongs to, for callers that enumerated
 // buckets rather than starting from an account.
 func AccountIDFromBucketName(bucket string) string {
@@ -279,13 +311,13 @@ func AccountIDFromBucketName(bucket string) string {
 
 // The DB instance identifiers held in one account bucket. An empty bucket
 // yields no names rather than an error.
-func ListDBInstanceIDs(ctx context.Context, kv jetstream.KeyValue) ([]string, error) {
+func ListDBInstanceIDs(ctx context.Context, kv *kvstore.Bucket) ([]string, error) {
 	return listNames(ctx, kv, DBInstancesPrefix())
 }
 
 // Manual and automated snapshots alike: the type is on the record, so a listing
 // filtered by it still has to read every one.
-func ListDBSnapshotIDs(ctx context.Context, kv jetstream.KeyValue) ([]string, error) {
+func ListDBSnapshotIDs(ctx context.Context, kv *kvstore.Bucket) ([]string, error) {
 	names, err := listNames(ctx, kv, DBSnapshotsPrefix())
 	if err != nil {
 		return nil, err
@@ -298,7 +330,7 @@ func ListDBSnapshotIDs(ctx context.Context, kv jetstream.KeyValue) ([]string, er
 
 // Walks the .../meta keys, which is what makes a group's own record findable
 // among the per-parameter keys hanging off the same prefix.
-func ListDBParameterGroupNames(ctx context.Context, kv jetstream.KeyValue) ([]string, error) {
+func ListDBParameterGroupNames(ctx context.Context, kv *kvstore.Bucket) ([]string, error) {
 	keys, err := bucketKeys(ctx, kv)
 	if err != nil {
 		return nil, err
@@ -320,7 +352,7 @@ func ListDBParameterGroupNames(ctx context.Context, kv jetstream.KeyValue) ([]st
 }
 
 // The stored overrides of one parameter group, keyed by parameter name.
-func ListDBParameterOverrides(ctx context.Context, kv jetstream.KeyValue, group string) (map[string]DBParameterRecord, error) {
+func ListDBParameterOverrides(ctx context.Context, kv *kvstore.Bucket, group string) (map[string]DBParameterRecord, error) {
 	prefix := DBParameterGroupParamsPrefix(group)
 	names, err := listNames(ctx, kv, prefix)
 	if err != nil {
@@ -345,7 +377,7 @@ func ListDBParameterOverrides(ctx context.Context, kv jetstream.KeyValue, group
 // Every account's automated-backup index, grouped by DB instance, from one
 // bucket listing: the retention sweep needs every instance's set, and a Keys call
 // per instance would cost one listing each.
-func ListAutomatedBackups(ctx context.Context, kv jetstream.KeyValue) (map[string][]string, error) {
+func ListAutomatedBackups(ctx context.Context, kv *kvstore.Bucket) (map[string][]string, error) {
 	keys, err := bucketKeys(ctx, kv)
 	if err != nil {
 		return nil, err
@@ -382,7 +414,7 @@ func splitAutomatedBackupKey(key string) (string, string, bool) {
 
 // The leaf names directly under prefix. A nested key belongs to a sub-space and
 // is skipped rather than reported as a malformed name.
-func listNames(ctx context.Context, kv jetstream.KeyValue, prefix string) ([]string, error) {
+func listNames(ctx context.Context, kv *kvstore.Bucket, prefix string) ([]string, error) {
 	keys, err := bucketKeys(ctx, kv)
 	if err != nil {
 		return nil, err
@@ -403,7 +435,11 @@ func listNames(ctx context.Context, kv jetstream.KeyValue, prefix string) ([]str
 
 // An empty bucket yields no keys rather than an error, so a first-ever describe
 // answers with an empty list instead of failing.
-func bucketKeys(ctx context.Context, kv jetstream.KeyValue) ([]string, error) {
+func bucketKeys(ctx context.Context, b *kvstore.Bucket) ([]string, error) {
+	kv, err := b.KV(ctx)
+	if err != nil {
+		return nil, err
+	}
 	keys, err := kv.Keys(ctx)
 	if err != nil {
 		if errors.Is(err, jetstream.ErrNoKeysFound) {

spinifex/handlers/rds/subnetgroup.go

@@ -13,7 +13,7 @@ import (
 	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/aws/aws-sdk-go/service/rds"
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
-	"github.com/nats-io/nats.go/jetstream"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 )
 
 // AWS's own name rules. The name is a KV key rather than a DNS label, so the
@@ -76,7 +76,7 @@ func (s *Service) CreateDBSubnetGroup(ctx context.Context, input *rds.CreateDBSu
 	// Create rather than Put, so two concurrent creates of one name have exactly
 	// one winner instead of both reporting success over each other's subnets.
 	if err := createJSON(ctx, kv, DBSubnetGroupKey(name), &rec); err != nil {
-		if errors.Is(err, jetstream.ErrKeyExists) {
+		if errors.Is(err, kvstore.ErrExists) {
 			return nil, awserrors.Errorf(awserrors.ErrorDBSubnetGroupAlreadyExists,
 				"DB subnet group %s already exists", name)
 		}
@@ -234,7 +234,7 @@ func (s *Service) resolveGroupSubnets(ctx context.Context, accountID string, req
 
 // The record plus its revision. A missing group raises AWS's own fault, so a
 // well-formed name that resolves to nothing is distinguishable from a bad one.
-func getDBSubnetGroup(ctx context.Context, kv jetstream.KeyValue, name string) (*DBSubnetGroupRecord, uint64, error) {
+func getDBSubnetGroup(ctx context.Context, kv *kvstore.Bucket, name string) (*DBSubnetGroupRecord, uint64, error) {
 	var rec DBSubnetGroupRecord
 	rev, found, err := getJSONRevision(ctx, kv, DBSubnetGroupKey(name), &rec)
 	if err != nil {
@@ -249,7 +249,7 @@ func getDBSubnetGroup(ctx context.Context, kv jetstream.KeyValue, name string) (
 // The identifiers of every instance in the account matching uses, sorted. Both
 // group deletes share it: an in-use guard that missed one instance would strand
 // a live database's configuration.
-func instancesUsingGroup(ctx context.Context, kv jetstream.KeyValue, uses func(*DBInstanceRecord) bool) ([]string, error) {
+func instancesUsingGroup(ctx context.Context, kv *kvstore.Bucket, uses func(*DBInstanceRecord) bool) ([]string, error) {
 	ids, err := ListDBInstanceIDs(ctx, kv)
 	if err != nil {
 		return nil, err

spinifex/handlers/rds/tags.go

@@ -11,7 +11,7 @@ import (
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/rds"
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
-	"github.com/nats-io/nats.go/jetstream"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 )
 
 // AWS's own tag limits. The reserved prefix is refused rather than dropped: a
@@ -180,9 +180,9 @@ func (s *Service) resolveTaggable(resourceName, accountID string) (taggableResou
 // The record plus its revision. A missing record raises the resource's own
 // not-found fault, so an ARN that parses but names nothing is distinguishable
 // from one that does not parse.
-func readTagged(ctx context.Context, kv jetstream.KeyValue, resource taggableResource, identifier string) (TaggedRecord, uint64, error) {
+func readTagged(ctx context.Context, kv *kvstore.Bucket, resource taggableResource, identifier string) (TaggedRecord, uint64, error) {
 	rec := resource.newRecord()
-	rev, found, err := getJSONRevision(ctx, kv, resource.key(identifier), rec)
+	rev, found, err := getJSONRevisionAny(ctx, kv, resource.key(identifier), rec)
 	if err != nil {
 		return nil, 0, err
 	}
@@ -195,7 +195,7 @@ func readTagged(ctx context.Context, kv jetstream.KeyValue, resource taggableRes
 // Applies mutate to the record's tags under CAS. Losing the CAS means another
 // tag write landed first, so the mutation is replayed against its result rather
 // than overwriting it.
-func mutateTags(ctx context.Context, kv jetstream.KeyValue, resource taggableResource, identifier string, mutate func(map[string]string) (map[string]string, error)) error {
+func mutateTags(ctx context.Context, kv *kvstore.Bucket, resource taggableResource, identifier string, mutate func(map[string]string) (map[string]string, error)) error {
 	key := resource.key(identifier)
 	for range tagWriteAttempts {
 		rec, rev, err := readTagged(ctx, kv, resource, identifier)
@@ -214,7 +214,7 @@ func mutateTags(ctx context.Context, kv jetstream.KeyValue, resource taggableRes
 		}
 		// A revision mismatch is a lost race with another writer, not a
 		// duplicate: re-read the record and re-apply the tag mutation.
-		if !errors.Is(err, jetstream.ErrKeyRevisionMismatch) {
+		if !errors.Is(err, kvstore.ErrConflict) {
 			return err
 		}
 	}

spinifex/handlers/sts/assume_role.go

@@ -19,10 +19,10 @@ import (
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/iam"
 	"github.com/aws/aws-sdk-go/service/sts"
-	"github.com/nats-io/nats.go/jetstream"
 
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
 	handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 	"github.com/mulgadc/spinifex/spinifex/utils"
 )
 
@@ -553,8 +553,8 @@ func (s *STSServiceImpl) mintSession(ctx context.Context, env sessionEnvelope, d
 			ExpiresAt:         expiresAt,
 			CreatedAt:         now,
 		}
-		if err := putSessionCredential(ctx, s.sessionsBucket, cred); err != nil {
-			if errors.Is(err, jetstream.ErrKeyExists) {
+		if err := putSessionCredential(ctx, s.sessions, cred); err != nil {
+			if errors.Is(err, kvstore.ErrExists) {
 				continue
 			}
 			return nil, "", "", fmt.Errorf("persist session credential: %w", err)

spinifex/handlers/sts/service_impl.go

@@ -7,6 +7,7 @@ import (
 	"log/slog"
 
 	handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 	"github.com/nats-io/nats.go"
 	"github.com/nats-io/nats.go/jetstream"
 )
@@ -26,11 +27,11 @@ const masterKeySize = 32 // AES-256, must match handlers_iam.
 // change. RunJanitor is the exception: it already owns a real context and
 // threads it into the sweep.
 type STSServiceImpl struct {
-	natsConn       *nats.Conn
-	js             jetstream.JetStream
-	sessionsBucket jetstream.KeyValue
-	iamSvc         handlers_iam.IAMService
-	masterKey      []byte
+	natsConn  *nats.Conn
+	js        jetstream.JetStream
+	sessions  *kvstore.Store[SessionCredential]
+	iamSvc    handlers_iam.IAMService
+	masterKey []byte
 }
 
 var _ STSService = (*STSServiceImpl)(nil)
@@ -57,7 +58,7 @@ func NewSTSServiceImpl(ctx context.Context, natsConn *nats.Conn, iamSvc handlers
 		return nil, fmt.Errorf("get JetStream context: %w", err)
 	}
 
-	sessionsBucket, err := initSessionCredentialsBucket(ctx, js, replicas)
+	sessions, err := initSessionCredentialsStore(ctx, js, replicas)
 	if err != nil {
 		return nil, fmt.Errorf("init session credentials bucket: %w", err)
 	}
@@ -67,10 +68,10 @@ func NewSTSServiceImpl(ctx context.Context, natsConn *nats.Conn, iamSvc handlers
 		"replicas", replicas)
 
 	return &STSServiceImpl{
-		natsConn:       natsConn,
-		js:             js,
-		sessionsBucket: sessionsBucket,
-		iamSvc:         iamSvc,
-		masterKey:      masterKey,
+		natsConn:  natsConn,
+		js:        js,
+		sessions:  sessions,
+		iamSvc:    iamSvc,
+		masterKey: masterKey,
 	}, nil
 }

spinifex/handlers/sts/service_impl_test.go

@@ -60,8 +60,8 @@ func TestNewSTSServiceImpl_RejectsWrongMasterKeySize(t *testing.T) {
 
 func TestNewSTSServiceImpl_InitializesBucket(t *testing.T) {
 	svc, _ := newTestSetup(t)
-	require.NotNil(t, svc.sessionsBucket)
-	assert.Equal(t, KVBucketSessionCredentials, svc.sessionsBucket.Bucket())
+	require.NotNil(t, svc.sessions)
+	assert.Equal(t, KVBucketSessionCredentials, svc.sessions.Name())
 }
 
 func TestNewSTSServiceImpl_NormalisesNegativeClusterSize(t *testing.T) {
@@ -116,7 +116,7 @@ func TestLookupSessionCredential_HitRoundTrips(t *testing.T) {
 		ExpiresAt:         now.Add(time.Hour),
 		CreatedAt:         now,
 	}
-	require.NoError(t, putSessionCredential(t.Context(), svc.sessionsBucket, cred))
+	require.NoError(t, putSessionCredential(t.Context(), svc.sessions, cred))
 
 	got, err := svc.LookupSessionCredential(cred.AccessKeyID)
 	require.NoError(t, err)
@@ -134,13 +134,13 @@ func TestLookupSessionCredential_UnmarshalFailureSurfacesError(t *testing.T) {
 	// the prefix is valid (so the lookup reaches the bucket) but the JSON
 	// body is garbage. This guards against the "lookup returns nil silently
 	// on parse failure" silent-failure mode.
-	_, err := svc.sessionsBucket.Put(t.Context(), akid, []byte("not json"))
+	_, err := sessionsKV(t, svc).Put(t.Context(), akid, []byte("not json"))
 	require.NoError(t, err)
 
 	got, err := svc.LookupSessionCredential(akid)
 	require.Error(t, err)
 	assert.Nil(t, got)
-	assert.Contains(t, err.Error(), "unmarshal session credential")
+	assert.Contains(t, err.Error(), "decode")
 }
 
 // Sanity check: the production marshaller round-trips through the lookup

spinifex/handlers/sts/session_credentials.go

@@ -11,6 +11,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 	"github.com/mulgadc/spinifex/spinifex/kvutil"
 	"github.com/mulgadc/spinifex/spinifex/migrate"
 	"github.com/mulgadc/spinifex/spinifex/otelsetup"
@@ -63,27 +64,35 @@ type SessionCredential struct {
 	CreatedAt      time.Time `json:"created_at"`
 }
 
-// initSessionCredentialsBucket opens (or creates) the session-credentials KV bucket.
-// History is fixed at 1: credentials are write-once at mint and delete-once at expiry.
-func initSessionCredentialsBucket(ctx context.Context, js jetstream.JetStream, replicas int) (jetstream.KeyValue, error) {
-	// kvutil clamps replicas to a minimum of 1, so a zero clusterSize still creates.
-	kv, err := kvutil.GetOrCreateBucketWithReplicas(ctx, js, KVBucketSessionCredentials, 1, replicas)
-	if err != nil {
+// sessionCredentialsConfig describes the session-credentials bucket. History is
+// fixed at 1: credentials are write-once at mint and delete-once at expiry.
+func sessionCredentialsConfig(replicas int) kvstore.Config {
+	return kvstore.Config{
+		Name:     KVBucketSessionCredentials,
+		History:  1,
+		Replicas: replicas,
+		Missing:  "session credentials KV bucket not initialized",
+		OnOpen: func(ctx context.Context, kv jetstream.KeyValue) error {
+			return migrate.DefaultRegistry.RunKV(ctx, KVBucketSessionCredentials, kv, KVBucketSessionCredentialsVersion)
+		},
+	}
+}
+
+// initSessionCredentialsStore opens (or creates) the session-credentials bucket.
+func initSessionCredentialsStore(ctx context.Context, js jetstream.JetStream, replicas int) (*kvstore.Store[SessionCredential], error) {
+	// kvstore clamps replicas to a minimum of 1, so a zero clusterSize still creates.
+	store := kvstore.New[SessionCredential](js, sessionCredentialsConfig(replicas))
+	// Opened eagerly: a bucket that cannot be created must fail service
+	// construction, not the first AssumeRole that needs to mint into it.
+	if _, err := store.KV(ctx); err != nil {
 		return nil, fmt.Errorf("open session credentials bucket: %w", err)
 	}
-
-	if err := migrate.DefaultRegistry.RunKV(
-		ctx, KVBucketSessionCredentials, kv, KVBucketSessionCredentialsVersion,
-	); err != nil {
-		return nil, fmt.Errorf("migrate %s: %w", KVBucketSessionCredentials, err)
-	}
-
-	return kv, nil
+	return store, nil
 }
 
 // putSessionCredential persists a SessionCredential via CAS create, enforcing the ASIA-prefix
-// invariant. Returns jetstream.ErrKeyExists on collision so callers can retry.
-func putSessionCredential(ctx context.Context, bucket jetstream.KeyValue, cred *SessionCredential) error {
+// invariant. Returns kvstore.ErrExists on collision so callers can retry.
+func putSessionCredential(ctx context.Context, store *kvstore.Store[SessionCredential], cred *SessionCredential) error {
 	if cred == nil {
 		return errors.New("nil session credential")
 	}
@@ -91,11 +100,7 @@ func putSessionCredential(ctx context.Context, bucket jetstream.KeyValue, cred *
 		return fmt.Errorf("session AKID must start with %q, got %q",
 			SessionAccessKeyIDPrefix, cred.AccessKeyID)
 	}
-	data, err := json.Marshal(cred)
-	if err != nil {
-		return fmt.Errorf("marshal session credential: %w", err)
-	}
-	if _, err := bucket.Create(ctx, cred.AccessKeyID, data); err != nil {
+	if _, err := store.Create(ctx, cred.AccessKeyID, cred); err != nil {
 		return fmt.Errorf("store session credential: %w", err)
 	}
 	return nil
@@ -127,18 +132,14 @@ func (s *STSServiceImpl) LookupSessionCredential(accessKeyID string) (*SessionCr
 	if !strings.HasPrefix(accessKeyID, SessionAccessKeyIDPrefix) {
 		return nil, nil
 	}
-	entry, err := s.sessionsBucket.Get(ctx, accessKeyID)
+	cred, _, err := s.sessions.Get(ctx, accessKeyID)
 	if err != nil {
-		if errors.Is(err, jetstream.ErrKeyNotFound) {
+		if errors.Is(err, kvstore.ErrNotFound) {
 			return nil, nil
 		}
 		return nil, fmt.Errorf("get session credential: %w", err)
 	}
-	var cred SessionCredential
-	if err := json.Unmarshal(entry.Value(), &cred); err != nil {
-		return nil, fmt.Errorf("unmarshal session credential: %w", err)
-	}
-	return &cred, nil
+	return cred, nil
 }
 
 // RunJanitor periodically removes expired session credentials. Idempotent; multiple
@@ -164,10 +165,20 @@ func (s *STSServiceImpl) RunJanitor(ctx context.Context) {
 
 // sweepExpired deletes all records whose ExpiresAt is past the grace period.
 // Per-key errors are logged and skipped; returns the delete count.
+//
+// Deliberately on the raw handle rather than Store.List: the sweep needs the key
+// to delete it, and one undecodable record must not stop it expiring every
+// other credential.
 func (s *STSServiceImpl) sweepExpired(ctx context.Context, now time.Time) int {
 	cutoff := now.Add(-janitorGracePeriod)
 
-	keys, err := kvutil.Keys(ctx, s.sessionsBucket)
+	bucket, err := s.sessions.KV(ctx)
+	if err != nil {
+		slog.Warn("STS janitor: open session credential bucket failed", "err", err)
+		return 0
+	}
+
+	keys, err := kvutil.Keys(ctx, bucket)
 	if err != nil {
 		if errors.Is(err, jetstream.ErrNoKeysFound) {
 			return 0
@@ -181,7 +192,7 @@ func (s *STSServiceImpl) sweepExpired(ctx context.Context, now time.Time) int {
 		if key == utils.VersionKey {
 			continue
 		}
-		entry, err := s.sessionsBucket.Get(ctx, key)
+		entry, err := bucket.Get(ctx, key)
 		if err != nil {
 			if errors.Is(err, jetstream.ErrKeyNotFound) {
 				continue
@@ -202,7 +213,7 @@ func (s *STSServiceImpl) sweepExpired(ctx context.Context, now time.Time) int {
 			continue
 		}
 
-		if err := s.sessionsBucket.Delete(ctx, key); err != nil {
+		if err := bucket.Delete(ctx, key); err != nil {
 			slog.Warn("STS janitor: delete expired session credential failed",
 				"key", key, "err", err)
 			continue

spinifex/handlers/sts/session_credentials_test.go

@@ -2,11 +2,11 @@ package handlers_sts
 
 import (
 	"context"
-	"encoding/json"
 	"errors"
 	"testing"
 	"time"
 
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 	"github.com/mulgadc/spinifex/spinifex/kvutil"
 	"github.com/mulgadc/spinifex/spinifex/testutil"
 	"github.com/mulgadc/spinifex/spinifex/utils"
@@ -15,10 +15,19 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func setupBucket(t *testing.T) jetstream.KeyValue {
+func setupBucket(t *testing.T) *kvstore.Store[SessionCredential] {
 	t.Helper()
 	_, nc, _ := testutil.StartTestJetStream(t)
-	kv, err := initSessionCredentialsBucket(t.Context(), testutil.NewJetStream(t, nc), 1)
+	store, err := initSessionCredentialsStore(t.Context(), testutil.NewJetStream(t, nc), 1)
+	require.NoError(t, err)
+	return store
+}
+
+// sessionsKV reaches the raw handle under the store, for the tests that seed a
+// record the codec would refuse or assert on a key without decoding it.
+func sessionsKV(t *testing.T, svc *STSServiceImpl) jetstream.KeyValue {
+	t.Helper()
+	kv, err := svc.sessions.KV(t.Context())
 	require.NoError(t, err)
 	return kv
 }
@@ -44,7 +53,9 @@ func TestInitSessionCredentialsBucket_StampsVersion(t *testing.T) {
 	_, nc, _ := testutil.StartTestJetStream(t)
 	js := testutil.NewJetStream(t, nc)
 
-	kv, err := initSessionCredentialsBucket(t.Context(), js, 1)
+	store, err := initSessionCredentialsStore(t.Context(), js, 1)
+	require.NoError(t, err)
+	kv, err := store.KV(t.Context())
 	require.NoError(t, err)
 
 	version, err := kvutil.ReadVersion(t.Context(), kv)
@@ -56,13 +67,13 @@ func TestInitSessionCredentialsBucket_Idempotent(t *testing.T) {
 	_, nc, _ := testutil.StartTestJetStream(t)
 	js := testutil.NewJetStream(t, nc)
 
-	kv1, err := initSessionCredentialsBucket(t.Context(), js, 1)
+	store1, err := initSessionCredentialsStore(t.Context(), js, 1)
 	require.NoError(t, err)
 
 	// Reopen — must return the same bucket without error.
-	kv2, err := initSessionCredentialsBucket(t.Context(), js, 1)
+	store2, err := initSessionCredentialsStore(t.Context(), js, 1)
 	require.NoError(t, err)
-	assert.Equal(t, kv1.Bucket(), kv2.Bucket())
+	assert.Equal(t, store1.Name(), store2.Name())
 }
 
 func TestPutSessionCredential_RoundTrip(t *testing.T) {
@@ -71,11 +82,9 @@ func TestPutSessionCredential_RoundTrip(t *testing.T) {
 
 	require.NoError(t, putSessionCredential(t.Context(), bucket, cred))
 
-	entry, err := bucket.Get(t.Context(), cred.AccessKeyID)
+	got, _, err := bucket.Get(t.Context(), cred.AccessKeyID)
 	require.NoError(t, err)
 
-	var got SessionCredential
-	require.NoError(t, json.Unmarshal(entry.Value(), &got))
 	assert.Equal(t, cred.AccessKeyID, got.AccessKeyID)
 	assert.Equal(t, cred.AssumedRoleARN, got.AssumedRoleARN)
 	assert.Equal(t, cred.SessionTokenHMAC, got.SessionTokenHMAC)
@@ -108,7 +117,9 @@ func TestPutSessionCredential_RejectsNonASIAPrefix(t *testing.T) {
 
 	// Bucket must contain no credentials — every put failed before reaching
 	// bucket.Create, so no AKID should be present.
-	keys, err := bucket.Keys(t.Context())
+	kv, err := bucket.KV(t.Context())
+	require.NoError(t, err)
+	keys, err := kv.Keys(t.Context())
 	if !errors.Is(err, jetstream.ErrNoKeysFound) {
 		require.NoError(t, err)
 		for _, k := range keys {
@@ -132,11 +143,11 @@ func TestPutSessionCredential_CollisionReturnsKeyExists(t *testing.T) {
 
 	require.NoError(t, putSessionCredential(t.Context(), bucket, cred))
 
-	// Second create with the same AKID must surface jetstream.ErrKeyExists so the
+	// Second create with the same AKID must surface kvstore.ErrExists so the
 	// mint helper can retry with a freshly generated AKID.
 	err := putSessionCredential(t.Context(), bucket, cred)
 	require.Error(t, err)
-	assert.ErrorIs(t, err, jetstream.ErrKeyExists)
+	assert.ErrorIs(t, err, kvstore.ErrExists)
 }
 
 func TestVerifySessionToken_MatchAndMismatch(t *testing.T) {
@@ -176,7 +187,7 @@ func putCredWithExpiry(t *testing.T, svc *STSServiceImpl, akid string, expiresAt
 	t.Helper()
 	cred := newTestSessionCredential(akid)
 	cred.ExpiresAt = expiresAt
-	require.NoError(t, putSessionCredential(t.Context(), svc.sessionsBucket, cred))
+	require.NoError(t, putSessionCredential(t.Context(), svc.sessions, cred))
 }
 
 func TestSweepExpired_DeletesPastGraceOnly(t *testing.T) {
@@ -192,14 +203,14 @@ func TestSweepExpired_DeletesPastGraceOnly(t *testing.T) {
 	assert.Equal(t, 2, deleted)
 
 	// Live and within-grace must still exist; past-grace must be gone.
-	_, err := svc.sessionsBucket.Get(t.Context(), "ASIALIVE000000000001")
+	_, err := sessionsKV(t, svc).Get(t.Context(), "ASIALIVE000000000001")
 	require.NoError(t, err)
-	_, err = svc.sessionsBucket.Get(t.Context(), "ASIAJUSTEXPIRED00002")
+	_, err = sessionsKV(t, svc).Get(t.Context(), "ASIAJUSTEXPIRED00002")
 	require.NoError(t, err)
 
-	_, err = svc.sessionsBucket.Get(t.Context(), "ASIAPASTGRACE0000003")
+	_, err = sessionsKV(t, svc).Get(t.Context(), "ASIAPASTGRACE0000003")
 	require.ErrorIs(t, err, jetstream.ErrKeyNotFound)
-	_, err = svc.sessionsBucket.Get(t.Context(), "ASIAANCIENT000000004")
+	_, err = sessionsKV(t, svc).Get(t.Context(), "ASIAANCIENT000000004")
 	require.ErrorIs(t, err, jetstream.ErrKeyNotFound)
 }
 
@@ -215,18 +226,18 @@ func TestSweepExpired_SkipsCorruptRecord(t *testing.T) {
 	svc, _ := newTestSetup(t)
 	now := time.Now().UTC()
 
-	_, err := svc.sessionsBucket.Put(t.Context(), "ASIACORRUPT000000001", []byte("not json"))
+	_, err := sessionsKV(t, svc).Put(t.Context(), "ASIACORRUPT000000001", []byte("not json"))
 	require.NoError(t, err)
 	putCredWithExpiry(t, svc, "ASIAEXPIRED000000002", now.Add(-24*time.Hour))
 
 	deleted := svc.sweepExpired(t.Context(), now)
 	assert.Equal(t, 1, deleted)
 
-	_, err = svc.sessionsBucket.Get(t.Context(), "ASIAEXPIRED000000002")
+	_, err = sessionsKV(t, svc).Get(t.Context(), "ASIAEXPIRED000000002")
 	require.ErrorIs(t, err, jetstream.ErrKeyNotFound)
 	// Corrupt record is left in place — janitor is not authorised to delete
 	// data it cannot interpret; an operator must inspect it.
-	_, err = svc.sessionsBucket.Get(t.Context(), "ASIACORRUPT000000001")
+	_, err = sessionsKV(t, svc).Get(t.Context(), "ASIACORRUPT000000001")
 	require.NoError(t, err)
 }
 
@@ -237,7 +248,7 @@ func TestSweepExpired_IgnoresVersionKey(t *testing.T) {
 	svc, _ := newTestSetup(t)
 	assert.Equal(t, 0, svc.sweepExpired(t.Context(), time.Now().UTC()))
 
-	_, err := svc.sessionsBucket.Get(t.Context(), utils.VersionKey)
+	_, err := sessionsKV(t, svc).Get(t.Context(), utils.VersionKey)
 	require.NoError(t, err, "version key must survive the sweep")
 }
 

spinifex/kvlease/bucket.go

@@ -4,8 +4,11 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"log/slog"
 	"time"
 
+	"github.com/mulgadc/spinifex/spinifex/migrate"
+	"github.com/mulgadc/spinifex/spinifex/otelsetup"
 	"github.com/nats-io/nats.go"
 	"github.com/nats-io/nats.go/jetstream"
 )
@@ -16,9 +19,57 @@ var (
 	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.
+// BucketConfig describes a lease bucket. History is not a field: a lease key has
+// no history worth keeping, so every lease bucket is created with one.
+type BucketConfig struct {
+	Name string
+
+	// TTL expires a lease its holder died without releasing, so a crashed leader
+	// does not park the key forever. It must match the lease's own Config.TTL.
+	TTL time.Duration
+
+	// Replicas is the replica count applied on a genuine first creation only.
+	// Zero leaves it to the server, which means one.
+	Replicas int
+
+	// Version runs the registered schema migrations for the bucket after it is
+	// opened. Zero skips them, for a bucket holding nothing but lease keys.
+	Version int
+}
+
+// OpenBucket attaches to a lease bucket, creating it only when genuinely absent,
+// and runs any schema migration registered for it.
+//
+// Attach before create, because callers re-open this on every reconcile tick:
+// CreateKeyValue against a bucket that already exists is a STREAM.CREATE the
+// meta leader answers with an error, so creating first bills one per tick. It is
+// also what makes Replicas mean "on first creation" rather than a config change
+// against a live bucket.
+func OpenBucket(ctx context.Context, js jetstream.JetStream, cfg BucketConfig) (jetstream.KeyValue, error) {
+	kv, err := js.KeyValue(ctx, cfg.Name)
+	if errors.Is(err, jetstream.ErrBucketNotFound) {
+		kv, err = js.CreateKeyValue(ctx, jetstream.KeyValueConfig{
+			Bucket:   cfg.Name,
+			History:  1,
+			TTL:      cfg.TTL,
+			Replicas: cfg.Replicas,
+		})
+	}
+	if err != nil {
+		return nil, fmt.Errorf("kvlease: open or create lease bucket %s: %w", cfg.Name, err)
+	}
+	if cfg.Version > 0 {
+		if err := migrate.DefaultRegistry.RunKV(ctx, cfg.Name, kv, cfg.Version); err != nil {
+			return nil, fmt.Errorf("migrate %s: %w", cfg.Name, err)
+		}
+	}
+	slog.DebugContext(ctx, "kvlease: lease bucket ready", "bucket", cfg.Name, "ttl_ms", otelsetup.Millis(cfg.TTL))
+	return kv, nil
+}
+
+// NATSBucket is OpenBucket as a BucketFunc, retrying while JetStream is still
+// forming quorum on a cold multi-node start. 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)
@@ -27,16 +78,7 @@ func NATSBucket(nc *nats.Conn, bucket string, ttl time.Duration) BucketFunc {
 		}
 		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,
-				})
-			}
+			kv, err := OpenBucket(ctx, js, BucketConfig{Name: bucket, TTL: ttl})
 			if err == nil {
 				return kv, nil
 			}

spinifex/kvstore/bucket.go

@@ -40,6 +40,11 @@ type Config struct {
 	// as instance state is by the owning node's next write.
 	RecreateIfMissing bool
 
+	// AttachOnly opens an existing bucket and never creates one, for a caller
+	// reading buckets it did not make: one deleted between a listing and the
+	// read is then an error rather than a silently recreated empty bucket.
+	AttachOnly bool
+
 	// OnOpen runs after every successful open, including a recovery reopen, so
 	// a recreated bucket is re-stamped rather than left unversioned. Owners
 	// use it to run their schema migrations.
@@ -175,13 +180,16 @@ func (b *Bucket) Reopen(ctx context.Context) (jetstream.KeyValue, error) {
 	b.mu.Lock()
 	defer b.mu.Unlock()
 
-	b.kv = nil
+	// Checked before the handle is discarded. A Bucket over a fixed handle has
+	// nothing to resolve against, so dropping it would turn one transient
+	// outage into a permanent failure for every later call.
 	if b.js == nil {
 		if b.cfg.Missing == "" {
 			return nil, errors.New("kvstore: cannot reopen without a JetStream client")
 		}
 		return nil, errors.New(b.cfg.Missing)
 	}
+	b.kv = nil
 
 	// Another goroutine may already have repaired it, so reconnecting comes
 	// first and is the only step that is always safe.
@@ -231,6 +239,82 @@ func (b *Bucket) withKV(ctx context.Context, op func(jetstream.KeyValue) error)
 	return op(kv)
 }
 
+// Delete removes a key. Idempotent: an already-absent key is success.
+//
+// On Bucket rather than Store because it decodes nothing, so every typed view
+// over one bucket shares it and a codec-free caller needs no type parameter.
+func (b *Bucket) Delete(ctx context.Context, key string) error {
+	err := b.withKV(ctx, func(kv jetstream.KeyValue) error {
+		if err := kv.Delete(ctx, key); err != nil {
+			return fmt.Errorf("kvstore: delete %s: %w", key, err)
+		}
+		return nil
+	})
+	if err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
+		return err
+	}
+	return nil
+}
+
+// CompareAndDelete removes a key only if it is still at rev, returning
+// ErrConflict when it is not.
+//
+// Unlike Delete this is not idempotent: a guarded delete carries an expected
+// last sequence, and an absent, tombstoned or purged key fails that check, so
+// an already-gone key reports ErrConflict rather than success.
+//
+// This is the delete half of CompareAndSet, for a caller undoing a reservation
+// it made — it must remove its own write and not a replacement that landed on
+// top of it.
+func (b *Bucket) CompareAndDelete(ctx context.Context, key string, rev uint64) error {
+	kv, err := b.KV(ctx)
+	if err != nil {
+		return err
+	}
+	// Not through withKV: the revision is the caller's, so a replay against a
+	// reopened bucket would be guarded on a revision from the bucket it lost.
+	if err := kv.Delete(ctx, key, jetstream.LastRevision(rev)); err != nil {
+		if errors.Is(err, jetstream.ErrKeyRevisionMismatch) {
+			return fmt.Errorf("%w: %s", ErrConflict, key)
+		}
+		return fmt.Errorf("kvstore: delete %s: %w", key, err)
+	}
+	return nil
+}
+
+// Purge is Delete for a key whose history must go with it, so a later Create
+// sees a key that never existed rather than one with a delete marker on top.
+func (b *Bucket) Purge(ctx context.Context, key string) error {
+	err := b.withKV(ctx, func(kv jetstream.KeyValue) error {
+		if err := kv.Purge(ctx, key); err != nil {
+			return fmt.Errorf("kvstore: purge %s: %w", key, err)
+		}
+		return nil
+	})
+	if err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
+		return err
+	}
+	return nil
+}
+
+// Exists reports whether a key is present without decoding its value, so a
+// record that cannot be unmarshalled is still reported as present.
+func (b *Bucket) Exists(ctx context.Context, key string) (bool, error) {
+	err := b.withKV(ctx, func(kv jetstream.KeyValue) error {
+		if _, err := kv.Get(ctx, key); err != nil {
+			return fmt.Errorf("kvstore: get %s: %w", key, err)
+		}
+		return nil
+	})
+	if err != nil {
+		if errors.Is(err, jetstream.ErrKeyNotFound) {
+			return false, nil
+		}
+		return false, err
+	}
+	return true, nil
+}
+
 // onOpen runs the configured open hook, naming the bucket on failure so a
 // migration error is not mistaken for a connection one.
 func (b *Bucket) onOpen(ctx context.Context, kv jetstream.KeyValue) error {
@@ -298,6 +382,12 @@ func (b *Bucket) WatchFrom(ctx context.Context, filter string, revision uint64)
 // open picks the kvutil helper matching the configured TTL and replica count.
 func (b *Bucket) open(ctx context.Context) (jetstream.KeyValue, error) {
 	switch {
+	case b.cfg.AttachOnly:
+		kv, err := b.js.KeyValue(ctx, b.cfg.Name)
+		if err != nil {
+			return nil, fmt.Errorf("open KV bucket %s: %w", b.cfg.Name, err)
+		}
+		return kv, nil
 	case b.cfg.TTL > 0:
 		return kvutil.GetOrCreateBucketWithOptions(ctx, b.js, kvutil.BucketOptions{
 			Name:        b.cfg.Name,

spinifex/kvstore/recovery_test.go

@@ -191,11 +191,25 @@ func TestStore_CompareAndSetDoesNotReRun(t *testing.T) {
 	require.NoError(t, err)
 	loseStream(t, js)
 
-	err = store.CompareAndSet(t.Context(), "acct-a/one", &record{Name: "two"}, rev)
+	_, err = store.CompareAndSet(t.Context(), "acct-a/one", &record{Name: "two"}, rev)
 	require.Error(t, err, "a revision-guarded write must not be replayed onto a reopened bucket")
 	assert.ErrorContains(t, err, "acct-a/one")
 }
 
+// CompareAndDelete guards on the caller's revision, so replaying it against a
+// reopened bucket would delete under a revision from the bucket it lost.
+func TestBucket_CompareAndDeleteDoesNotReRun(t *testing.T) {
+	t.Parallel()
+	_, js, store := newRecoverableStore(t, kvstore.Config{RecreateIfMissing: true})
+	rev, err := store.Create(t.Context(), "acct-a/one", &record{Name: "one"})
+	require.NoError(t, err)
+	loseStream(t, js)
+
+	err = store.CompareAndDelete(t.Context(), "acct-a/one", rev)
+	require.Error(t, err, "a revision-guarded delete must not be replayed onto a reopened bucket")
+	assert.ErrorContains(t, err, "acct-a/one")
+}
+
 // TestBucket_OnOpenRunsOnEveryOpen covers the hook a recreated bucket depends
 // on: an unstamped, unmigrated bucket is not a recovered one.
 func TestBucket_OnOpenRunsOnEveryOpen(t *testing.T) {
@@ -403,3 +417,31 @@ func TestBucket_DescriptionReachesTheCreatedBucket(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, "a bucket that says what it is", status.(*jetstream.KeyValueBucketStatus).StreamInfo().Config.Description)
 }
+
+// A Store over a fixed handle cannot reopen, so a transient outage must leave
+// the handle it already has. Discarding it would end every later call with the
+// configured Missing reason, long after the stream came back.
+func TestBucket_FixedHandleSurvivesATransientOutage(t *testing.T) {
+	t.Parallel()
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
+
+	const bucket = "kvstore-fixed-handle-test"
+	cfg := jetstream.KeyValueConfig{Bucket: bucket, History: 1}
+	kv, err := js.CreateKeyValue(t.Context(), cfg)
+	require.NoError(t, err)
+
+	store := kvstore.Over[record](nil, kv, kvstore.Config{Name: bucket, Missing: "test: no JetStream client"})
+	require.NoError(t, store.Set(t.Context(), "acct-a/one", &record{Name: "one"}))
+
+	loseStream := func() { require.NoError(t, js.DeleteKeyValue(t.Context(), bucket)) }
+	loseStream()
+	_, _, err = store.Get(t.Context(), "acct-a/one")
+	require.Error(t, err, "the call meeting the lost stream still fails")
+
+	_, err = js.CreateKeyValue(t.Context(), cfg)
+	require.NoError(t, err)
+
+	err = store.Set(t.Context(), "acct-a/one", &record{Name: "two"})
+	require.NoError(t, err, "the store must work again once the stream is back")
+}

spinifex/kvstore/store.go

@@ -99,35 +99,6 @@ func (s *Store[T]) Upsert(ctx context.Context, key string, fn func(*T) (bool, er
 	})
 }
 
-// Delete removes a key. Idempotent: an already-absent key is success.
-func (s *Store[T]) Delete(ctx context.Context, key string) error {
-	err := s.withKV(ctx, func(kv jetstream.KeyValue) error {
-		if err := kv.Delete(ctx, key); err != nil {
-			return fmt.Errorf("kvstore: delete %s: %w", key, err)
-		}
-		return nil
-	})
-	if err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
-		return err
-	}
-	return nil
-}
-
-// Purge is Delete for a key whose history must go with it, so a later Create
-// sees a key that never existed rather than one with a delete marker on top.
-func (s *Store[T]) Purge(ctx context.Context, key string) error {
-	err := s.withKV(ctx, func(kv jetstream.KeyValue) error {
-		if err := kv.Purge(ctx, key); err != nil {
-			return fmt.Errorf("kvstore: purge %s: %w", key, err)
-		}
-		return nil
-	})
-	if err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
-		return err
-	}
-	return nil
-}
-
 // List decodes every record whose key carries the given prefix. An empty
 // prefix matches everything and an empty bucket is not an error. A key that
 // disappears between the listing and the read is skipped.
@@ -316,43 +287,29 @@ func (s *Store[T]) Set(ctx context.Context, key string, v *T) error {
 	})
 }
 
-// Exists reports whether a key is present without decoding its value, so a
-// record that cannot be unmarshalled is still reported as present.
-func (s *Store[T]) Exists(ctx context.Context, key string) (bool, error) {
-	err := s.withKV(ctx, func(kv jetstream.KeyValue) error {
-		if _, err := kv.Get(ctx, key); err != nil {
-			return fmt.Errorf("kvstore: get %s: %w", key, err)
-		}
-		return nil
-	})
-	if err != nil {
-		if errors.Is(err, jetstream.ErrKeyNotFound) {
-			return false, nil
-		}
-		return false, err
-	}
-	return true, nil
-}
-
 // CompareAndSet writes v only if key is still at rev, returning ErrConflict
 // when it is not. For callers whose lost race is a decision rather than a
 // failure; callers that should retry want Mutate.
-func (s *Store[T]) CompareAndSet(ctx context.Context, key string, v *T, rev uint64) error {
+//
+// The revision the write landed at is returned, for a caller chaining a second
+// CAS onto its own first one.
+func (s *Store[T]) CompareAndSet(ctx context.Context, key string, v *T, rev uint64) (uint64, error) {
 	kv, err := s.KV(ctx)
 	if err != nil {
-		return err
+		return 0, err
 	}
 	data, err := json.Marshal(v)
 	if err != nil {
-		return fmt.Errorf("kvstore: encode %s: %w", key, err)
+		return 0, fmt.Errorf("kvstore: encode %s: %w", key, err)
 	}
-	if _, err := kv.Update(ctx, key, data, rev); err != nil {
+	next, err := kv.Update(ctx, key, data, rev)
+	if err != nil {
 		if errors.Is(err, jetstream.ErrKeyRevisionMismatch) || errors.Is(err, jetstream.ErrKeyExists) {
-			return fmt.Errorf("%w: %s", ErrConflict, key)
+			return 0, fmt.Errorf("%w: %s", ErrConflict, key)
 		}
-		return fmt.Errorf("kvstore: update %s: %w", key, err)
+		return 0, fmt.Errorf("kvstore: update %s: %w", key, err)
 	}
-	return nil
+	return next, nil
 }
 
 // keys lists an already-opened bucket, treating an empty bucket as an empty

spinifex/kvstore/store_test.go

@@ -316,19 +316,58 @@ func TestStore_CompareAndSetRejectsAStaleRevision(t *testing.T) {
 	// Another writer commits first, moving the revision on.
 	seedRaw(t, store, "acct-a/one", record{Name: "one", Count: 2})
 
-	err := store.CompareAndSet(t.Context(), "acct-a/one", &record{Name: "one", Count: 99}, stale)
+	_, err := store.CompareAndSet(t.Context(), "acct-a/one", &record{Name: "one", Count: 99}, stale)
 	require.ErrorIs(t, err, kvstore.ErrConflict)
 
 	got, fresh, err := store.Get(t.Context(), "acct-a/one")
 	require.NoError(t, err)
 	assert.Equal(t, 2, got.Count, "a rejected write must leave the winner's value in place")
 
-	require.NoError(t, store.CompareAndSet(t.Context(), "acct-a/one", &record{Name: "one", Count: 3}, fresh))
+	_, err = store.CompareAndSet(t.Context(), "acct-a/one", &record{Name: "one", Count: 3}, fresh)
+	require.NoError(t, err)
 	got, _, err = store.Get(t.Context(), "acct-a/one")
 	require.NoError(t, err)
 	assert.Equal(t, 3, got.Count, "the same call commits at the current revision")
 }
 
+// TestBucket_CompareAndDeleteWithdrawsOnlyItsOwnWrite pins the reservation
+// rollback: the holder removes the record it wrote, and a replacement that
+// landed on top of it survives.
+func TestBucket_CompareAndDeleteWithdrawsOnlyItsOwnWrite(t *testing.T) {
+	t.Parallel()
+	store := newStore(t)
+	own := mustCreate(t, store, "acct-a/mine", record{Name: "mine", Count: 1})
+
+	require.NoError(t, store.CompareAndDelete(t.Context(), "acct-a/mine", own))
+	_, _, err := store.Get(t.Context(), "acct-a/mine")
+	require.ErrorIs(t, err, kvstore.ErrNotFound, "the holder's own write is withdrawn")
+
+	stale := mustCreate(t, store, "acct-a/theirs", record{Name: "theirs", Count: 1})
+	// A replacement lands on top of the reservation.
+	seedRaw(t, store, "acct-a/theirs", record{Name: "theirs", Count: 2})
+
+	err = store.CompareAndDelete(t.Context(), "acct-a/theirs", stale)
+	require.ErrorIs(t, err, kvstore.ErrConflict)
+	got, _, err := store.Get(t.Context(), "acct-a/theirs")
+	require.NoError(t, err, "a lost guard must leave the replacement in place")
+	assert.Equal(t, 2, got.Count)
+}
+
+// Unlike Delete, a guarded delete carries an expected last sequence, so a key
+// that is already gone fails the guard rather than reporting success.
+func TestBucket_CompareAndDeleteReportsAnAbsentKeyAsConflict(t *testing.T) {
+	t.Parallel()
+	store := newStore(t)
+	rev := mustCreate(t, store, "acct-a/one", record{Name: "one"})
+	require.NoError(t, store.CompareAndDelete(t.Context(), "acct-a/one", rev))
+
+	err := store.CompareAndDelete(t.Context(), "acct-a/one", rev)
+	require.ErrorIs(t, err, kvstore.ErrConflict, "the spent revision does not delete a second time")
+
+	err = store.CompareAndDelete(t.Context(), "never-existed", 1)
+	require.ErrorIs(t, err, kvstore.ErrConflict)
+}
+
 // TestStore_CreateReturnsTheClaimRevision pins the claim-then-CAS sequence a
 // single-writer state machine runs: the winner of the create holds a revision
 // good enough to commit its own first update, with no read in between.
@@ -344,11 +383,11 @@ func TestStore_CreateReturnsTheClaimRevision(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, read, claimed, "the claim revision is the record's current revision")
 
-	require.NoError(t, store.CompareAndSet(t.Context(), "singleton", &record{Name: "available"}, claimed),
-		"the winner promotes its own claim without re-reading")
+	_, err = store.CompareAndSet(t.Context(), "singleton", &record{Name: "available"}, claimed)
+	require.NoError(t, err, "the winner promotes its own claim without re-reading")
 
 	// The spent revision must not commit a second time.
-	err = store.CompareAndSet(t.Context(), "singleton", &record{Name: "impostor"}, claimed)
+	_, err = store.CompareAndSet(t.Context(), "singleton", &record{Name: "impostor"}, claimed)
 	require.ErrorIs(t, err, kvstore.ErrConflict)
 
 	got, _, err := store.Get(t.Context(), "singleton")

spinifex/kvutil/cas.go

@@ -70,8 +70,11 @@ func retryCAS(ctx context.Context, cfg CASConfig, op, key string, attempt func()
 			return ctx.Err()
 		}
 	}
+	// The caller's error is wrapped around the last conflict rather than
+	// replacing it: the API error naming the condition is what tells a stale
+	// replica read apart from a genuine race.
 	if cfg.Exhausted != nil {
-		return cfg.Exhausted(key, attempts)
+		return fmt.Errorf("%w: %w", cfg.Exhausted(key, attempts), lastErr)
 	}
 	return fmt.Errorf("CAS %s exhausted %d attempts for key %s: %w", op, attempts, key, lastErr)
 }

spinifex/network/external/dhcp/manager.go

@@ -5,7 +5,6 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"github.com/mulgadc/spinifex/spinifex/otelsetup"
 	"log/slog"
 	"math/rand/v2"
 	"net"
@@ -13,6 +12,8 @@ import (
 	"sync"
 	"time"
 
+	"github.com/mulgadc/spinifex/spinifex/otelsetup"
+
 	"github.com/nats-io/nats.go"
 	"github.com/nats-io/nats.go/jetstream"
 	"golang.org/x/sync/singleflight"

spinifex/network/external/static_pool.go

@@ -2,7 +2,6 @@ package external
 
 import (
 	"context"
-	"encoding/json"
 	"errors"
 	"fmt"
 	"log/slog"
@@ -10,7 +9,7 @@ import (
 	"net/netip"
 
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
-	"github.com/mulgadc/spinifex/spinifex/kvutil"
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
 	"github.com/mulgadc/spinifex/spinifex/migrate"
 	"github.com/nats-io/nats.go/jetstream"
 )
@@ -21,8 +20,12 @@ const (
 	// carries over with no manual migration.
 	KVBucketStaticPool        = "spinifex-external-ipam"
 	KVBucketStaticPoolVersion = 2
+	KVBucketStaticPoolHistory = 5
 
-	staticPoolCASRetries = 5
+	// One pool record is one key, so every concurrent allocation in a pool
+	// contends on it. A launch asking for sixteen public IPs is sixteen
+	// writers, and each round has exactly one winner.
+	staticPoolCASRetries = 25
 
 	// purposeIGWLRP duplicates handlers/ec2/vpc.PurposeIGWLRP so the
 	// allocator can reserve the gateway slot without importing handlers.
@@ -58,11 +61,34 @@ type PoolRecord struct {
 	Cursor string `json:"cursor,omitempty"`
 }
 
+// errPoolContended means a pool key stayed contended for the whole retry
+// budget, which the boot-time drift reconcile tolerates and a caller asking for
+// an address does not.
+var errPoolContended = errors.New("external IPAM: pool contended")
+
+// staticPoolConfig describes the pool bucket. The name and history are
+// unchanged from the pre-Q1 ExternalIPAM so existing cluster data carries over,
+// and the migration runs on every open rather than only the first.
+func staticPoolConfig() kvstore.Config {
+	return kvstore.Config{
+		Name:     KVBucketStaticPool,
+		History:  KVBucketStaticPoolHistory,
+		Attempts: staticPoolCASRetries,
+		Missing:  "external IPAM: no JetStream client configured",
+		OnOpen: func(ctx context.Context, kv jetstream.KeyValue) error {
+			return migrate.DefaultRegistry.RunKV(ctx, KVBucketStaticPool, kv, KVBucketStaticPoolVersion)
+		},
+		Exhausted: func(key string, attempts int) error {
+			return fmt.Errorf("%w: %s after %d attempts", errPoolContended, key, attempts)
+		},
+	}
+}
+
 // StaticPoolAllocator implements Allocator backed by a NATS JetStream KV
 // bucket. Bucket schema and CAS semantics are unchanged from the pre-Q1
 // ExternalIPAM.
 type StaticPoolAllocator struct {
-	kv    jetstream.KeyValue
+	store *kvstore.Store[PoolRecord]
 	pools []ExternalPoolConfig
 }
 
@@ -71,14 +97,7 @@ var _ Allocator = (*StaticPoolAllocator)(nil)
 // NewStaticPoolAllocator creates the KV bucket (if missing), runs pending
 // migrations, and seeds each pool's record.
 func NewStaticPoolAllocator(ctx context.Context, js jetstream.JetStream, pools []ExternalPoolConfig) (*StaticPoolAllocator, error) {
-	kv, err := kvutil.GetOrCreateBucket(ctx, js, KVBucketStaticPool, 5)
-	if err != nil {
-		return nil, fmt.Errorf("create external IPAM KV bucket: %w", err)
-	}
-	if err := migrate.DefaultRegistry.RunKV(ctx, KVBucketStaticPool, kv, KVBucketStaticPoolVersion); err != nil {
-		return nil, fmt.Errorf("migrate %s: %w", KVBucketStaticPool, err)
-	}
-	a := &StaticPoolAllocator{kv: kv, pools: pools}
+	a := &StaticPoolAllocator{store: kvstore.New[PoolRecord](js, staticPoolConfig()), pools: pools}
 	if err := a.initPools(ctx); err != nil {
 		return nil, fmt.Errorf("init external IPAM pools: %w", err)
 	}
@@ -88,12 +107,14 @@ func NewStaticPoolAllocator(ctx context.Context, js jetstream.JetStream, pools [
 // NewStaticPoolAllocatorWithKV is the test constructor — skips bucket
 // creation and migrations.
 func NewStaticPoolAllocatorWithKV(kv jetstream.KeyValue, pools []ExternalPoolConfig) *StaticPoolAllocator {
-	return &StaticPoolAllocator{kv: kv, pools: pools}
+	return &StaticPoolAllocator{store: kvstore.Over[PoolRecord](nil, kv, staticPoolConfig()), pools: pools}
 }
 
-// KV exposes the underlying bucket so callers that share the bucket
-// (notably ExternalIPAM facade tests) can construct sibling allocators.
-func (a *StaticPoolAllocator) KV() jetstream.KeyValue { return a.kv }
+// KV exposes the underlying bucket so callers that share it (the ExternalIPAM
+// facade and its tests) can construct sibling allocators over the same handle.
+func (a *StaticPoolAllocator) KV(ctx context.Context) (jetstream.KeyValue, error) {
+	return a.store.KV(ctx)
+}
 
 // Pools returns the allocator's pool list. Used by the ExternalIPAM
 // facade to satisfy pool-by-region lookups.
@@ -101,97 +122,86 @@ func (a *StaticPoolAllocator) Pools() []ExternalPoolConfig { return a.pools }
 
 // Allocate implements Allocator.
 func (a *StaticPoolAllocator) Allocate(ctx context.Context, req AllocateRequest) (netip.Addr, error) {
-	for attempt := range staticPoolCASRetries {
-		record, revision, err := a.getRecord(ctx, req.PoolName)
-		if err != nil {
-			return netip.Addr{}, fmt.Errorf("get external IPAM record: %w", err)
+	var ip string
+	err := a.store.Mutate(ctx, req.PoolName, func(record *PoolRecord) (bool, error) {
+		// A record written before the field existed decodes to a nil map, which
+		// the assignment below would panic on.
+		if record.Allocated == nil {
+			record.Allocated = map[string]ExternalIPAllocation{}
 		}
-
-		ip, err := nextAvailableIP(record)
+		next, err := nextAvailableIP(record)
 		if err != nil {
-			return netip.Addr{}, err
+			return false, err
 		}
-
-		record.Allocated[ip] = ExternalIPAllocation{
+		record.Allocated[next] = ExternalIPAllocation{
 			Purpose:      req.Purpose,
 			AllocationID: req.AllocationID,
 			ENIId:        req.ENIID,
 			InstanceId:   req.InstanceID,
 		}
-		record.Cursor = ip
-
-		data, err := json.Marshal(record)
-		if err != nil {
-			return netip.Addr{}, fmt.Errorf("marshal external IPAM record: %w", err)
-		}
-
-		if _, err := a.kv.Update(ctx, req.PoolName, data, revision); err != nil {
-			slog.Debug("external IPAM CAS conflict, retrying", "pool", req.PoolName, "attempt", attempt)
-			continue
-		}
-
-		addr, ok := netip.AddrFromSlice(net.ParseIP(ip).To4())
-		if !ok {
-			return netip.Addr{}, fmt.Errorf("parse allocated ip %q", ip)
-		}
-		slog.Info("external IPAM allocated IP", "pool", req.PoolName, "ip", ip, "purpose", req.Purpose)
-		return addr.Unmap(), nil
+		record.Cursor = next
+		ip = next
+		return true, nil
+	})
+	if err != nil {
+		return netip.Addr{}, err
 	}
-	return netip.Addr{}, fmt.Errorf("external IPAM allocation failed after CAS retries for pool %s", req.PoolName)
+
+	addr, ok := netip.AddrFromSlice(net.ParseIP(ip).To4())
+	if !ok {
+		return netip.Addr{}, fmt.Errorf("parse allocated ip %q", ip)
+	}
+	slog.InfoContext(ctx, "external IPAM allocated IP", "pool", req.PoolName, "ip", ip, "purpose", req.Purpose)
+	return addr.Unmap(), nil
 }
 
 // Release implements Allocator.
 func (a *StaticPoolAllocator) Release(ctx context.Context, poolName string, ip netip.Addr, ownerENIID string) error {
 	target := ip.String()
-	for attempt := range staticPoolCASRetries {
-		record, revision, err := a.getRecord(ctx, poolName)
-		if err != nil {
-			return fmt.Errorf("get external IPAM record for release: %w", err)
-		}
+	var released bool
+	err := a.store.Mutate(ctx, poolName, func(record *PoolRecord) (bool, error) {
+		// Reset per attempt: a retry that finds the IP already gone must not
+		// inherit the previous attempt's verdict.
+		released = false
 
 		alloc, ok := record.Allocated[target]
 		if !ok {
 			// An owner-scoped release of an already-freed IP is an idempotent
 			// no-op (a duplicated or stale teardown), not an error.
 			if ownerENIID != "" {
-				return nil
+				return false, nil
 			}
-			return fmt.Errorf("IP %s not allocated in pool %s", target, poolName)
+			return false, fmt.Errorf("IP %s not allocated in pool %s", target, poolName)
 		}
 		// Ownership guard: when a caller names an owner, only free the lease if
 		// it still belongs to that ENI. External IPs are recycled and GC/teardown
 		// sweeps re-emit releases, so a release for a prior owner must not free an
 		// IP since reassigned to a live instance (that would double-allocate).
 		if ownerENIID != "" && alloc.ENIId != "" && alloc.ENIId != ownerENIID {
-			slog.Info("external IPAM release skip — IP reassigned to a different ENI (stale release)",
+			slog.InfoContext(ctx, "external IPAM release skip — IP reassigned to a different ENI (stale release)",
 				"pool", poolName, "ip", target, "stale_owner_eni", ownerENIID, "current_owner_eni", alloc.ENIId)
-			return nil
+			return false, nil
 		}
 		if alloc.Purpose == purposeIGWLRP {
-			return fmt.Errorf("cannot release gateway IP %s in pool %s", target, poolName)
+			return false, fmt.Errorf("cannot release gateway IP %s in pool %s", target, poolName)
 		}
 
 		delete(record.Allocated, target)
-
-		data, err := json.Marshal(record)
-		if err != nil {
-			return fmt.Errorf("marshal external IPAM record: %w", err)
-		}
-
-		if _, err := a.kv.Update(ctx, poolName, data, revision); err != nil {
-			slog.Debug("external IPAM release CAS conflict, retrying", "pool", poolName, "attempt", attempt)
-			continue
-		}
-
-		slog.Info("external IPAM released IP", "pool", poolName, "ip", target)
-		return nil
+		released = true
+		return true, nil
+	})
+	if err != nil {
+		return err
 	}
-	return fmt.Errorf("external IPAM release failed after CAS retries for pool %s", poolName)
+	if released {
+		slog.InfoContext(ctx, "external IPAM released IP", "pool", poolName, "ip", target)
+	}
+	return nil
 }
 
 // GetPoolRecord returns the current pool record.
 func (a *StaticPoolAllocator) GetPoolRecord(ctx context.Context, poolName string) (*PoolRecord, error) {
-	rec, _, err := a.getRecord(ctx, poolName)
+	rec, _, err := a.store.Get(ctx, poolName)
 	return rec, err
 }
 
@@ -204,43 +214,63 @@ func (a *StaticPoolAllocator) initPools(ctx context.Context) error {
 	return nil
 }
 
-func (a *StaticPoolAllocator) initPool(ctx context.Context, pool ExternalPoolConfig) error {
-	chk, revision, err := a.getRecord(ctx, pool.Name)
+// rangesMatch reports whether rec already carries the ranges pool configures.
+func rangesMatch(rec *PoolRecord, pool ExternalPoolConfig) bool {
+	return rec.RangeStart == pool.RangeStart && rec.RangeEnd == pool.RangeEnd &&
+		rec.GwLrpRangeStart == pool.GwLrpRangeStart && rec.GwLrpRangeEnd == pool.GwLrpRangeEnd
+}
 
-	if err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
+// reconcileDrift commits the corrected ranges under the store's retry budget. A
+// pool is one key, so an ordinary Allocate or Release wins this race routinely;
+// conceding to it would leave the drift unfixed until some later boot won.
+func (a *StaticPoolAllocator) reconcileDrift(ctx context.Context, pool ExternalPoolConfig) error {
+	return a.store.Mutate(ctx, pool.Name, func(rec *PoolRecord) (bool, error) {
+		// Re-checked per attempt: the winner of a lost race may already have
+		// stored the ranges this node wanted, which is nothing left to do.
+		if rangesMatch(rec, pool) {
+			return false, nil
+		}
+		rec.RangeStart = pool.RangeStart
+		rec.RangeEnd = pool.RangeEnd
+		rec.GwLrpRangeStart = pool.GwLrpRangeStart
+		rec.GwLrpRangeEnd = pool.GwLrpRangeEnd
+		return true, nil
+	})
+}
+
+func (a *StaticPoolAllocator) initPool(ctx context.Context, pool ExternalPoolConfig) error {
+	chk, _, err := a.store.Get(ctx, pool.Name)
+
+	if err != nil && !errors.Is(err, kvstore.ErrNotFound) {
 		return err
 	}
 
 	if err == nil {
-		if chk.RangeStart != pool.RangeStart || chk.RangeEnd != pool.RangeEnd ||
-			chk.GwLrpRangeStart != pool.GwLrpRangeStart || chk.GwLrpRangeEnd != pool.GwLrpRangeEnd {
+		if !rangesMatch(chk, pool) {
 			slog.Info("external IPAM pool config drift, reconciling KV",
 				"pool", pool.Name,
 				"old_range", chk.RangeStart+"-"+chk.RangeEnd, "new_range", pool.RangeStart+"-"+pool.RangeEnd,
 				"old_gw_lrp_range", chk.GwLrpRangeStart+"-"+chk.GwLrpRangeEnd,
 				"new_gw_lrp_range", pool.GwLrpRangeStart+"-"+pool.GwLrpRangeEnd)
 
-			chk.RangeStart = pool.RangeStart
-			chk.RangeEnd = pool.RangeEnd
-			chk.GwLrpRangeStart = pool.GwLrpRangeStart
-			chk.GwLrpRangeEnd = pool.GwLrpRangeEnd
-
-			data, err := json.Marshal(chk)
-			if err != nil {
-				return fmt.Errorf("marshal external IPAM record: %w", err)
-			}
-
-			if _, err := a.kv.Update(ctx, pool.Name, data, revision); err != nil {
-				slog.Warn("external IPAM update failed", "pool", pool.Name, "err", err)
-				return err
+			if err := a.reconcileDrift(ctx, pool); err != nil {
+				if !errors.Is(err, errPoolContended) {
+					return err
+				}
+				// Failing the boot over a key that stayed contended would take
+				// this node down, so an unreconciled drift is reported instead.
+				slog.WarnContext(ctx, "external IPAM drift reconcile stayed contended; the stored ranges still differ",
+					"pool", pool.Name,
+					"want_range", pool.RangeStart+"-"+pool.RangeEnd,
+					"want_gw_lrp_range", pool.GwLrpRangeStart+"-"+pool.GwLrpRangeEnd, "err", err)
 			}
 			return nil
 		}
-		slog.Debug("external IPAM pool already initialized", "pool", pool.Name)
+		slog.DebugContext(ctx, "external IPAM pool already initialized", "pool", pool.Name)
 		return nil
 	}
 
-	slog.Info("external IPAM pool not found, creating", "pool", pool.Name)
+	slog.InfoContext(ctx, "external IPAM pool not found, creating", "pool", pool.Name)
 
 	gwIP := pool.GatewayIP
 	if gwIP == "" {
@@ -263,35 +293,17 @@ func (a *StaticPoolAllocator) initPool(ctx context.Context, pool ExternalPoolCon
 		},
 	}
 
-	data, err := json.Marshal(record)
-	if err != nil {
-		return fmt.Errorf("marshal pool record: %w", err)
-	}
-
-	if _, err := a.kv.Create(ctx, pool.Name, data); err != nil {
-		if errors.Is(err, jetstream.ErrKeyExists) {
+	if _, err := a.store.Create(ctx, pool.Name, record); err != nil {
+		if errors.Is(err, kvstore.ErrExists) {
 			return nil
 		}
 		return fmt.Errorf("create pool KV entry: %w", err)
 	}
 
-	slog.Info("external IPAM pool initialized", "pool", pool.Name, "gateway_ip", gwIP)
+	slog.InfoContext(ctx, "external IPAM pool initialized", "pool", pool.Name, "gateway_ip", gwIP)
 	return nil
 }
 
-func (a *StaticPoolAllocator) getRecord(ctx context.Context, poolName string) (*PoolRecord, uint64, error) {
-	entry, err := a.kv.Get(ctx, poolName)
-	if err != nil {
-		return nil, 0, err
-	}
-
-	var record PoolRecord
-	if err := json.Unmarshal(entry.Value(), &record); err != nil {
-		return nil, 0, fmt.Errorf("unmarshal external IPAM record: %w", err)
-	}
-	return &record, entry.Revision(), nil
-}
-
 // nextAvailableIP picks the next unallocated address, skipping [GwLrpRangeStart,
 // GwLrpRangeEnd]. Resumes one past Cursor, wrapping at RangeEnd.
 func nextAvailableIP(record *PoolRecord) (string, error) {

spinifex/network/external/static_pool_test.go

@@ -144,6 +144,71 @@ func TestStaticPool_CASConflict(t *testing.T) {
 	}
 }
 
+// A launch asking for many public IPs at once is that many writers on one pool
+// key. The retry bound has to cover the whole contending set, not a guess at
+// timing: at the inherited five attempts this drops callers under load.
+func TestStaticPool_AllocateUnderHeavyContention(t *testing.T) {
+	pool := wanPool()
+	pool.RangeEnd = "192.168.1.200"
+	a := newStaticAllocator(t, []ExternalPoolConfig{pool})
+
+	const writers = 16
+	results := make([]netip.Addr, writers)
+	errs := make([]error, writers)
+	var wg sync.WaitGroup
+	for i := range writers {
+		wg.Go(func() {
+			results[i], errs[i] = a.Allocate(t.Context(), AllocateRequest{PoolName: "wan"})
+		})
+	}
+	wg.Wait()
+
+	seen := make(map[string]bool, writers)
+	for i, err := range errs {
+		require.NoError(t, err, "concurrent allocation %d must not be dropped", i)
+		ip := results[i].String()
+		assert.False(t, seen[ip], "duplicate IP %s handed to two callers", ip)
+		seen[ip] = true
+	}
+}
+
+// A pool is one key, so the boot-time drift reconcile contends with every live
+// allocation on it. It must retry onto the winner rather than concede, which
+// would leave the configured ranges unstored until some later boot won the race.
+func TestStaticPool_DriftReconcileConvergesUnderContention(t *testing.T) {
+	narrow := wanPool()
+	narrow.RangeEnd = "192.168.1.180"
+	a := newStaticAllocator(t, []ExternalPoolConfig{narrow})
+
+	widened := narrow
+	widened.RangeEnd = "192.168.1.200"
+
+	const writers = 16
+	var wg sync.WaitGroup
+	for range writers {
+		wg.Go(func() {
+			_, _ = a.Allocate(t.Context(), AllocateRequest{PoolName: "wan"})
+		})
+	}
+	var driftErr error
+	wg.Go(func() { driftErr = a.reconcileDrift(t.Context(), widened) })
+	wg.Wait()
+	require.NoError(t, driftErr)
+
+	rec, err := a.GetPoolRecord(t.Context(), "wan")
+	require.NoError(t, err)
+	assert.Equal(t, "192.168.1.200", rec.RangeEnd, "the reconcile must land the configured range on the winner")
+}
+
+// A write that is not a lost race must surface with its cause rather than being
+// retried into a generic exhaustion, which is what the old loop did.
+func TestStaticPool_ReleaseSurfacesANonConflictError(t *testing.T) {
+	a := newStaticAllocator(t, []ExternalPoolConfig{wanPool()})
+	err := a.Release(t.Context(), "no-such-pool", mustAddr(t, "192.168.1.151"), "")
+	require.Error(t, err)
+	assert.NotErrorIs(t, err, errPoolContended, "an absent pool is not contention")
+}
+
 func TestStaticPool_ReleaseUnknown(t *testing.T) {
 	a := newStaticAllocator(t, []ExternalPoolConfig{wanPool()})
 	err := a.Release(context.Background(), "wan", mustAddr(t, "192.168.1.200"), "")