feat: recover kvstore buckets from stream loss (#915)

spinifex/handlers/acm/store.go

@@ -187,7 +187,9 @@ func NewStore(ctx context.Context, nc *nats.Conn, masterKey []byte) (*Store, err
 
 	slog.Info("ACM store initialised", "bucket", KVBucketACM)
 	return &Store{
-		store: kvstore.Over[CertRecord](kv, kvstore.Config{
+		store: kvstore.Over[CertRecord](js, kv, kvstore.Config{
+			Name:     KVBucketACM,
+			History:  KVBucketACMVersion,
 			Attempts: maxInUseByCASAttempts,
 			Exhausted: func(key string, attempts int) error {
 				return fmt.Errorf("acm store: exceeded %d CAS attempts updating %s", attempts, key)

spinifex/handlers/elbv2/store.go

@@ -48,6 +48,9 @@ const (
 type Store struct {
 	kv jetstream.KeyValue
 
+	// js is what a watcher's bucket reopens against after the stream is lost.
+	js jetstream.JetStream
+
 	// now and claimTTL back ClaimLBName's pending-vs-orphan check. Both are
 	// overridden in tests (same package) so the crash-orphan reclaim path never
 	// needs a real sleep past the TTL.
@@ -93,7 +96,7 @@ func NewStore(ctx context.Context, nc *nats.Conn) (*Store, error) {
 	}
 
 	slog.Info("ELBv2 store initialized", "bucket", KVBucketELBv2)
-	return &Store{kv: kv, now: time.Now, claimTTL: lbNameClaimTTL}, nil
+	return &Store{kv: kv, js: js, now: time.Now, claimTTL: lbNameClaimTTL}, nil
 }
 
 // --- Load Balancer CRUD ---
@@ -216,7 +219,7 @@ func (s *Store) ListLoadBalancersStrict(ctx context.Context) ([]*LoadBalancerRec
 // groups, listeners and rules, so a caller wanting only one of those filters on
 // the corresponding key prefix.
 func (s *Store) WatchBucket() *kvstore.Bucket {
-	return kvstore.NewOpenBucket(s.kv, kvstore.Config{Name: KVBucketELBv2})
+	return kvstore.NewOpenBucket(s.js, s.kv, kvstore.Config{Name: KVBucketELBv2})
 }
 
 // GetLoadBalancerByArn finds a load balancer by the short ID in the ARN's final

spinifex/kvstore/bucket.go

@@ -7,6 +7,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"log/slog"
 	"sync"
 	"time"
 
@@ -23,6 +24,22 @@ type Config struct {
 	Missing   string
 	Attempts  int
 	Exhausted func(key string, attempts int) error
+
+	// RecreateIfMissing allows recovery to create the bucket again when the
+	// stream is not merely unreachable but gone. It governs recovery only —
+	// the first open is always get-or-create, because bootstrapping an absent
+	// bucket is the intent there.
+	//
+	// Default false, and deliberately so: a bucket that existed and now does
+	// not has lost its records, and recreating it turns that into a success
+	// the caller cannot see. Set it only where the data is reconstructible,
+	// as instance state is by the owning node's next write.
+	RecreateIfMissing 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.
+	OnOpen func(ctx context.Context, kv jetstream.KeyValue) error
 }
 
 // Bucket memoises a lazily created KV bucket. Construct it with NewBucket.
@@ -42,9 +59,10 @@ func NewBucket(js jetstream.JetStream, cfg Config) *Bucket {
 
 // NewOpenBucket returns a Bucket over an already-open handle, for callers that
 // resolve their bucket at construction so a bad connection fails at startup
-// rather than on first use.
-func NewOpenBucket(kv jetstream.KeyValue, cfg Config) *Bucket {
-	return &Bucket{kv: kv, cfg: cfg}
+// rather than on first use. js is what recovery reopens against; pass nil to
+// keep the handle fixed for a caller not ready to recover.
+func NewOpenBucket(js jetstream.JetStream, kv jetstream.KeyValue, cfg Config) *Bucket {
+	return &Bucket{js: js, kv: kv, cfg: cfg}
 }
 
 // Name returns the configured bucket name, for callers holding several buckets
@@ -76,10 +94,88 @@ func (b *Bucket) KV(ctx context.Context) (jetstream.KeyValue, error) {
 	if err != nil {
 		return nil, err
 	}
+	if err := b.onOpen(ctx, kv); err != nil {
+		return nil, err
+	}
 	b.kv = kv
 	return kv, nil
 }
 
+// Reopen discards the memoised handle and resolves the bucket again, for a
+// caller that saw the stream go out from under it. It reconnects; whether it
+// may also recreate an absent bucket is Config.RecreateIfMissing.
+func (b *Bucket) Reopen(ctx context.Context) (jetstream.KeyValue, error) {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+
+	b.kv = nil
+	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)
+	}
+
+	// Another goroutine may already have repaired it, so reconnecting comes
+	// first and is the only step that is always safe.
+	kv, err := b.js.KeyValue(ctx, b.cfg.Name)
+	if err != nil {
+		if !b.cfg.RecreateIfMissing {
+			return nil, fmt.Errorf("kvstore: reopen %s: %w", b.cfg.Name, err)
+		}
+		if !errors.Is(err, jetstream.ErrBucketNotFound) && !kvutil.IsStreamUnavailable(err) {
+			return nil, fmt.Errorf("kvstore: reopen %s: %w", b.cfg.Name, err)
+		}
+		slog.WarnContext(ctx, "kvstore: bucket stream lost, recreating", "bucket", b.cfg.Name)
+		if kv, err = b.open(ctx); err != nil {
+			return nil, fmt.Errorf("kvstore: recreate %s: %w", b.cfg.Name, err)
+		}
+	}
+
+	if err := b.onOpen(ctx, kv); err != nil {
+		return nil, err
+	}
+	b.kv = kv
+	slog.InfoContext(ctx, "kvstore: bucket reopened", "bucket", b.cfg.Name)
+	return kv, nil
+}
+
+// withKV runs op against the bucket, reopening and running it once more if the
+// stream was lost underneath it. A failure to reopen surfaces the original
+// error, not the reopen's — the first one is what the caller was doing.
+//
+// op therefore has to be safe to run twice. Anything carrying a revision the
+// caller supplied is not, and must not use this.
+func (b *Bucket) withKV(ctx context.Context, op func(jetstream.KeyValue) error) error {
+	kv, err := b.KV(ctx)
+	if err != nil {
+		return err
+	}
+	err = op(kv)
+	if err == nil || !kvutil.IsStreamUnavailable(err) {
+		return err
+	}
+
+	kv, reopenErr := b.Reopen(ctx)
+	if reopenErr != nil {
+		slog.WarnContext(ctx, "kvstore: recovery failed", "bucket", b.cfg.Name, "err", reopenErr)
+		return err
+	}
+	return op(kv)
+}
+
+// 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 {
+	if b.cfg.OnOpen == nil {
+		return nil
+	}
+	if err := b.cfg.OnOpen(ctx, kv); err != nil {
+		return fmt.Errorf("kvstore: on-open %s: %w", b.cfg.Name, err)
+	}
+	return nil
+}
+
 // Watch returns a watcher over the keys matching filter, which takes NATS
 // subject wildcards ("node.*", "lb.*", ">" for everything).
 //
@@ -87,15 +183,24 @@ func (b *Bucket) KV(ctx context.Context) (jetstream.KeyValue, error) {
 // existing key would fire a redundant pass, and replaying them again on every
 // re-establish would defeat the debounce. The cost is that a disconnect gap is
 // invisible, which is what the caller's periodic resync exists to cover.
+// A watcher that dies with the stream is not resurrected by a later reopen;
+// only establishing a new one recovers, which is what the caller's periodic
+// resync drives.
 func (b *Bucket) Watch(ctx context.Context, filter string) (jetstream.KeyWatcher, error) {
-	kv, err := b.KV(ctx)
+	var w jetstream.KeyWatcher
+	// Wrapped inside the closure, so a bucket that could not be opened at all
+	// reports the configured reason rather than a watch failure on top of it.
+	err := b.withKV(ctx, func(kv jetstream.KeyValue) error {
+		var err error
+		w, err = kv.Watch(ctx, filter, jetstream.UpdatesOnly())
+		if err != nil {
+			return fmt.Errorf("kvstore: watch %s %s: %w", b.cfg.Name, filter, err)
+		}
+		return nil
+	})
 	if err != nil {
 		return nil, err
 	}
-	w, err := kv.Watch(ctx, filter, jetstream.UpdatesOnly())
-	if err != nil {
-		return nil, fmt.Errorf("kvstore: watch %s %s: %w", b.cfg.Name, filter, err)
-	}
 	return w, nil
 }
 

spinifex/kvstore/recovery_test.go

@@ -0,0 +1,245 @@
+package kvstore_test
+
+import (
+	"context"
+	"errors"
+	"sync/atomic"
+	"testing"
+
+	"github.com/mulgadc/spinifex/spinifex/kvstore"
+	"github.com/mulgadc/spinifex/spinifex/kvutil"
+	"github.com/mulgadc/spinifex/spinifex/testutil"
+	"github.com/nats-io/nats-server/v2/server"
+	"github.com/nats-io/nats.go/jetstream"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+const recoveryBucket = "kvstore-recovery-test"
+
+// newRecoverableStore returns a store whose bucket may be recreated by recovery,
+// already open so the next operation is the one that meets the lost stream.
+func newRecoverableStore(t *testing.T, cfg kvstore.Config) (*server.Server, jetstream.JetStream, *kvstore.Store[record]) {
+	t.Helper()
+	ns, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
+
+	cfg.Name = recoveryBucket
+	cfg.History = 1
+	store := kvstore.New[record](js, cfg)
+	_, err := store.KV(t.Context())
+	require.NoError(t, err, "the bucket must be open before the stream is taken away")
+	return ns, js, store
+}
+
+// loseStream deletes the bucket out from under the store's memoised handle,
+// which is what cluster formation does to a low-replication stream.
+func loseStream(t *testing.T, js jetstream.JetStream) {
+	t.Helper()
+	require.NoError(t, js.DeleteKeyValue(t.Context(), recoveryBucket))
+}
+
+// TestStore_RecoversAfterStreamLost runs every operation against a bucket whose
+// stream has gone, and asserts each one reaches the recreated bucket. The
+// recreated bucket is empty, so a read reporting ErrNotFound has recovered:
+// what must not survive is a stream-unavailable error reaching the caller.
+func TestStore_RecoversAfterStreamLost(t *testing.T) {
+	tests := []struct {
+		name    string
+		op      func(context.Context, *kvstore.Store[record]) error
+		wantErr error
+	}{
+		{"Get", func(ctx context.Context, s *kvstore.Store[record]) error {
+			_, _, err := s.Get(ctx, "acct-a/one")
+			return err
+		}, kvstore.ErrNotFound},
+		{"Create", func(ctx context.Context, s *kvstore.Store[record]) error {
+			_, err := s.Create(ctx, "acct-a/one", &record{Name: "one"})
+			return err
+		}, nil},
+		{"Set", func(ctx context.Context, s *kvstore.Store[record]) error {
+			return s.Set(ctx, "acct-a/one", &record{Name: "one"})
+		}, nil},
+		{"Mutate", func(ctx context.Context, s *kvstore.Store[record]) error {
+			return s.Mutate(ctx, "acct-a/one", func(*record) (bool, error) { return true, nil })
+		}, kvstore.ErrNotFound},
+		{"Upsert", func(ctx context.Context, s *kvstore.Store[record]) error {
+			return s.Upsert(ctx, "counter", func(r *record) (bool, error) {
+				r.Count++
+				return true, nil
+			})
+		}, nil},
+		{"Delete", func(ctx context.Context, s *kvstore.Store[record]) error {
+			return s.Delete(ctx, "acct-a/one")
+		}, nil},
+		{"Purge", func(ctx context.Context, s *kvstore.Store[record]) error {
+			return s.Purge(ctx, "acct-a/one")
+		}, nil},
+		{"Exists", func(ctx context.Context, s *kvstore.Store[record]) error {
+			_, err := s.Exists(ctx, "acct-a/one")
+			return err
+		}, nil},
+		{"List", func(ctx context.Context, s *kvstore.Store[record]) error {
+			_, err := s.List(ctx, "")
+			return err
+		}, nil},
+		{"DeletePrefix", func(ctx context.Context, s *kvstore.Store[record]) error {
+			return s.DeletePrefix(ctx, "acct-a/")
+		}, nil},
+		{"Watch", func(ctx context.Context, s *kvstore.Store[record]) error {
+			w, err := s.Watch(ctx, ">")
+			if err == nil {
+				_ = w.Stop()
+			}
+			return err
+		}, nil},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			_, js, store := newRecoverableStore(t, kvstore.Config{RecreateIfMissing: true})
+			require.NoError(t, store.Set(t.Context(), "acct-a/seed", &record{Name: "seed"}))
+			loseStream(t, js)
+
+			err := tt.op(t.Context(), store)
+			assert.False(t, kvutil.IsStreamUnavailable(err),
+				"the operation must reach the reopened bucket, not surface the loss: %v", err)
+			if tt.wantErr != nil {
+				require.ErrorIs(t, err, tt.wantErr)
+			} else {
+				require.NoError(t, err)
+			}
+
+			// The recreated bucket is the one the store now holds, so a
+			// subsequent write lands without a second recovery.
+			require.NoError(t, store.Set(t.Context(), "acct-a/after", &record{Name: "after"}))
+			got, _, err := store.Get(t.Context(), "acct-a/after")
+			require.NoError(t, err)
+			assert.Equal(t, "after", got.Name)
+		})
+	}
+}
+
+// TestStore_RecreateIfMissingDefaultsOff is the guard on the data-loss half of
+// recovery: a bucket that existed and now does not has lost its records, and a
+// store that has not opted in must say so rather than hand back an empty one.
+func TestStore_RecreateIfMissingDefaultsOff(t *testing.T) {
+	_, js, store := newRecoverableStore(t, kvstore.Config{})
+	require.NoError(t, store.Set(t.Context(), "acct-a/one", &record{Name: "one"}))
+	loseStream(t, js)
+
+	err := store.Set(t.Context(), "acct-a/one", &record{Name: "two"})
+	require.Error(t, err, "recovery must not silently recreate a bucket the caller did not opt into")
+
+	_, err = js.KeyValue(t.Context(), recoveryBucket)
+	require.Error(t, err, "the bucket must still be absent")
+}
+
+// TestStore_ReconnectsWithoutRecreating covers the first half of Reopen: a
+// bucket that is merely unreachable through a stale handle is reconnected, and
+// its records are still there. RecreateIfMissing is off, so nothing else could
+// have produced this result.
+func TestStore_ReconnectsWithoutRecreating(t *testing.T) {
+	_, js, store := newRecoverableStore(t, kvstore.Config{})
+	require.NoError(t, store.Set(t.Context(), "acct-a/one", &record{Name: "one"}))
+
+	kv, err := store.Reopen(t.Context())
+	require.NoError(t, err)
+	require.NotNil(t, kv)
+
+	got, _, err := store.Get(t.Context(), "acct-a/one")
+	require.NoError(t, err)
+	assert.Equal(t, "one", got.Name, "a reconnect keeps the records a recreate would have dropped")
+
+	_, err = js.KeyValue(t.Context(), recoveryBucket)
+	require.NoError(t, err)
+}
+
+// TestStore_RecoveryFailureSurfacesTheOriginalError pins which of the two errors
+// the caller sees. With the server gone both the operation and the reopen fail,
+// and the operation's error is the one describing what the caller was doing.
+func TestStore_RecoveryFailureSurfacesTheOriginalError(t *testing.T) {
+	ns, _, store := newRecoverableStore(t, kvstore.Config{RecreateIfMissing: true})
+	ns.Shutdown()
+	ns.WaitForShutdown()
+
+	err := store.Set(t.Context(), "acct-a/one", &record{Name: "one"})
+	require.Error(t, err)
+	assert.ErrorContains(t, err, "put acct-a/one",
+		"the surfaced error must name the operation, not the failed reopen")
+}
+
+// TestStore_CompareAndSetDoesNotReRun is the documented exception to the retry.
+// The revision is the caller's, so a second attempt against a reopened bucket
+// would be committing a precondition that no longer means anything.
+func TestStore_CompareAndSetDoesNotReRun(t *testing.T) {
+	_, 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.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")
+}
+
+// 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) {
+	var opens atomic.Int32
+	_, js, store := newRecoverableStore(t, kvstore.Config{
+		RecreateIfMissing: true,
+		OnOpen: func(ctx context.Context, kv jetstream.KeyValue) error {
+			opens.Add(1)
+			return kvutil.WriteVersion(ctx, kv, 7)
+		},
+	})
+
+	// newRecoverableStore opened it once already; the count is what the reopen
+	// adds to that.
+	before := opens.Load()
+	require.Positive(t, before, "the hook must run on the first open")
+
+	loseStream(t, js)
+	require.NoError(t, store.Set(t.Context(), "acct-a/one", &record{Name: "one"}))
+	assert.Equal(t, before+1, opens.Load(), "the hook must run again on a recovery reopen")
+
+	kv, err := store.KV(t.Context())
+	require.NoError(t, err)
+	version, err := kvutil.ReadVersion(t.Context(), kv)
+	require.NoError(t, err)
+	assert.Equal(t, 7, version, "a recreated bucket must be re-stamped, not left unversioned")
+}
+
+// TestBucket_OnOpenFailureFailsTheOpen keeps a half-opened bucket from being
+// memoised: a migration that could not run leaves records the caller's decoder
+// will not understand, so the open is the right place to stop.
+func TestBucket_OnOpenFailureFailsTheOpen(t *testing.T) {
+	_, nc, _ := testutil.StartTestJetStream(t)
+	boom := errors.New("migration refused")
+	store := kvstore.New[record](testutil.NewJetStream(t, nc), kvstore.Config{
+		Name:    "kvstore-onopen-fail-test",
+		History: 1,
+		OnOpen: func(context.Context, jetstream.KeyValue) error {
+			return boom
+		},
+	})
+
+	_, _, err := store.Get(t.Context(), "acct-a/one")
+	require.ErrorIs(t, err, boom)
+
+	_, _, err = store.Get(t.Context(), "acct-a/one")
+	require.ErrorIs(t, err, boom, "a failed open must not be memoised as a good one")
+}
+
+// TestBucket_ReopenWithoutAJetStreamClientReportsTheConfiguredMessage covers the
+// eager caller that passed nil: it cannot recover, and must say why.
+func TestBucket_ReopenWithoutAJetStreamClientReportsTheConfiguredMessage(t *testing.T) {
+	bucket := kvstore.NewOpenBucket(nil, newOpenBucket(t), kvstore.Config{
+		Name:    "kvstore-over-test",
+		Missing: "kvstore test: no JetStream client configured",
+	})
+
+	_, err := bucket.Reopen(t.Context())
+	require.ErrorContains(t, err, "no JetStream client configured")
+}

spinifex/kvstore/store.go

@@ -32,31 +32,41 @@ func New[T any](js jetstream.JetStream, cfg Config) *Store[T] {
 // Get reads and decodes one record, returning ErrNotFound when the key is
 // absent. The revision is returned for callers doing their own CAS.
 func (s *Store[T]) Get(ctx context.Context, key string) (*T, uint64, error) {
-	kv, err := s.KV(ctx)
-	if err != nil {
-		return nil, 0, err
-	}
-	return s.get(ctx, kv, key)
+	var (
+		v   *T
+		rev uint64
+	)
+	err := s.withKV(ctx, func(kv jetstream.KeyValue) error {
+		var err error
+		v, rev, err = s.get(ctx, kv, key)
+		return err
+	})
+	return v, rev, err
 }
 
 // Create writes a record only if the key is absent, returning ErrExists when
 // it is not. The create-only write is the single-writer claim, and the
 // returned revision is the CAS precondition for the claim's own first update.
 func (s *Store[T]) Create(ctx context.Context, key string, v *T) (uint64, error) {
-	kv, err := s.KV(ctx)
-	if err != nil {
-		return 0, err
-	}
 	data, err := json.Marshal(v)
 	if err != nil {
 		return 0, fmt.Errorf("kvstore: encode %s: %w", key, err)
 	}
-	rev, err := kv.Create(ctx, key, data)
+	// Wrapped inside the closure, so a bucket that could not be opened at all
+	// reports the configured reason rather than a create failure on top of it.
+	var rev uint64
+	err = s.withKV(ctx, func(kv jetstream.KeyValue) error {
+		var createErr error
+		if rev, createErr = kv.Create(ctx, key, data); createErr != nil {
+			return fmt.Errorf("kvstore: create %s: %w", key, createErr)
+		}
+		return nil
+	})
 	if err != nil {
 		if errors.Is(err, jetstream.ErrKeyExists) {
 			return 0, fmt.Errorf("%w: %s", ErrExists, key)
 		}
-		return 0, fmt.Errorf("kvstore: create %s: %w", key, err)
+		return 0, err
 	}
 	return rev, nil
 }
@@ -64,55 +74,57 @@ func (s *Store[T]) Create(ctx context.Context, key string, v *T) (uint64, error)
 // Mutate applies fn under a revision-guarded retry loop. fn reports whether it
 // changed anything; false commits no write. An absent key is ErrNotFound.
 func (s *Store[T]) Mutate(ctx context.Context, key string, fn func(*T) (bool, error)) error {
-	kv, err := s.KV(ctx)
-	if err != nil {
+	// Safe to re-run: Update re-reads the record inside its own CAS loop, so
+	// the second attempt starts from whatever the reopened bucket holds.
+	return s.withKV(ctx, func(kv jetstream.KeyValue) error {
+		_, err := kvutil.Update(ctx, kv, key, kvutil.CASConfig{
+			Attempts:  s.cfg.Attempts,
+			NotFound:  fmt.Errorf("%w: %s", ErrNotFound, key),
+			Exhausted: s.cfg.Exhausted,
+		}, fn)
 		return err
-	}
-	_, err = kvutil.Update(ctx, kv, key, kvutil.CASConfig{
-		Attempts:  s.cfg.Attempts,
-		NotFound:  fmt.Errorf("%w: %s", ErrNotFound, key),
-		Exhausted: s.cfg.Exhausted,
-	}, fn)
-	return err
+	})
 }
 
 // Upsert is Mutate for a counter-style record whose first write creates it: an
 // absent key starts fn from the zero value rather than reporting ErrNotFound.
 func (s *Store[T]) Upsert(ctx context.Context, key string, fn func(*T) (bool, error)) error {
-	kv, err := s.KV(ctx)
-	if err != nil {
+	return s.withKV(ctx, func(kv jetstream.KeyValue) error {
+		_, err := kvutil.Update(ctx, kv, key, kvutil.CASConfig{
+			Attempts:       s.cfg.Attempts,
+			CreateIfAbsent: true,
+			Exhausted:      s.cfg.Exhausted,
+		}, fn)
 		return err
-	}
-	_, err = kvutil.Update(ctx, kv, key, kvutil.CASConfig{
-		Attempts:       s.cfg.Attempts,
-		CreateIfAbsent: true,
-		Exhausted:      s.cfg.Exhausted,
-	}, fn)
-	return err
+	})
 }
 
 // Delete removes a key. Idempotent: an already-absent key is success.
 func (s *Store[T]) Delete(ctx context.Context, key string) error {
-	kv, err := s.KV(ctx)
-	if err != nil {
+	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
 	}
-	if err := kv.Delete(ctx, key); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
-		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 (s *Store[T]) Purge(ctx context.Context, key string) error {
-	kv, err := s.KV(ctx)
-	if err != nil {
+	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
 	}
-	if err := kv.Purge(ctx, key); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
-		return fmt.Errorf("kvstore: purge %s: %w", key, err)
-	}
 	return nil
 }
 
@@ -120,79 +132,94 @@ func (s *Store[T]) Purge(ctx context.Context, key string) error {
 // prefix matches everything and an empty bucket is not an error. A key that
 // disappears between the listing and the read is skipped.
 func (s *Store[T]) List(ctx context.Context, prefix string) ([]T, error) {
-	kv, keys, err := s.keys(ctx)
+	var out []T
+	// Whole listing sits inside the wrapper, not just the key enumeration: a
+	// stream lost partway through the reads has to restart from the listing,
+	// so out is reset rather than appended to twice.
+	err := s.withKV(ctx, func(kv jetstream.KeyValue) error {
+		out = nil
+		keys, err := s.keys(ctx, kv)
+		if err != nil {
+			return err
+		}
+		for _, key := range keys {
+			if !strings.HasPrefix(key, prefix) {
+				continue
+			}
+			v, _, err := s.get(ctx, kv, key)
+			if errors.Is(err, ErrNotFound) {
+				continue
+			}
+			if err != nil {
+				return err
+			}
+			out = append(out, *v)
+		}
+		return nil
+	})
 	if err != nil {
 		return nil, err
 	}
-	var out []T
-	for _, key := range keys {
-		if !strings.HasPrefix(key, prefix) {
-			continue
-		}
-		v, _, err := s.get(ctx, kv, key)
-		if errors.Is(err, ErrNotFound) {
-			continue
-		}
-		if err != nil {
-			return nil, err
-		}
-		out = append(out, *v)
-	}
 	return out, nil
 }
 
 // DeletePrefix removes every key carrying the given prefix. An empty prefix
 // purges the bucket. Idempotent, like Delete.
 func (s *Store[T]) DeletePrefix(ctx context.Context, prefix string) error {
-	kv, keys, err := s.keys(ctx)
-	if err != nil {
-		return err
-	}
-	for _, key := range keys {
-		if !strings.HasPrefix(key, prefix) {
-			continue
+	// Safe to re-run: a key deleted before the stream was lost is already
+	// absent on the retry, which Delete treats as success.
+	return s.withKV(ctx, func(kv jetstream.KeyValue) error {
+		keys, err := s.keys(ctx, kv)
+		if err != nil {
+			return err
 		}
-		if err := kv.Delete(ctx, key); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
-			return fmt.Errorf("kvstore: delete %s: %w", key, err)
+		for _, key := range keys {
+			if !strings.HasPrefix(key, prefix) {
+				continue
+			}
+			if err := kv.Delete(ctx, key); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
+				return fmt.Errorf("kvstore: delete %s: %w", key, err)
+			}
 		}
-	}
-	return nil
+		return nil
+	})
 }
 
-// Over returns a Store for records of type T over an already-open bucket.
-func Over[T any](kv jetstream.KeyValue, cfg Config) *Store[T] {
-	return &Store[T]{Bucket: NewOpenBucket(kv, cfg)}
+// Over returns a Store for records of type T over an already-open bucket. js is
+// what recovery reopens against; pass nil to keep the handle fixed.
+func Over[T any](js jetstream.JetStream, kv jetstream.KeyValue, cfg Config) *Store[T] {
+	return &Store[T]{Bucket: NewOpenBucket(js, kv, cfg)}
 }
 
 // Set writes a record whether or not the key exists, for callers whose write
 // is a replacement rather than a claim or a read-modify-write.
 func (s *Store[T]) Set(ctx context.Context, key string, v *T) error {
-	kv, err := s.KV(ctx)
-	if err != nil {
-		return err
-	}
 	data, err := json.Marshal(v)
 	if err != nil {
 		return fmt.Errorf("kvstore: encode %s: %w", key, err)
 	}
-	if _, err := kv.Put(ctx, key, data); err != nil {
-		return fmt.Errorf("kvstore: put %s: %w", key, err)
-	}
-	return nil
+	return s.withKV(ctx, func(kv jetstream.KeyValue) error {
+		if _, err := kv.Put(ctx, key, data); err != nil {
+			return fmt.Errorf("kvstore: put %s: %w", key, 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 (s *Store[T]) Exists(ctx context.Context, key string) (bool, error) {
-	kv, err := s.KV(ctx)
+	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 {
-		return false, err
-	}
-	if _, err := kv.Get(ctx, key); err != nil {
 		if errors.Is(err, jetstream.ErrKeyNotFound) {
 			return false, nil
 		}
-		return false, fmt.Errorf("kvstore: get %s: %w", key, err)
+		return false, err
 	}
 	return true, nil
 }
@@ -218,20 +245,17 @@ func (s *Store[T]) CompareAndSet(ctx context.Context, key string, v *T, rev uint
 	return nil
 }
 
-// keys lists the bucket, treating an empty bucket as an empty listing.
-func (s *Store[T]) keys(ctx context.Context) (jetstream.KeyValue, []string, error) {
-	kv, err := s.KV(ctx)
-	if err != nil {
-		return nil, nil, err
-	}
+// keys lists an already-opened bucket, treating an empty bucket as an empty
+// listing. Recovery is the caller's, so the raw error passes through.
+func (s *Store[T]) keys(ctx context.Context, kv jetstream.KeyValue) ([]string, error) {
 	keys, err := kv.Keys(ctx)
 	if err != nil {
 		if errors.Is(err, jetstream.ErrNoKeysFound) {
-			return kv, nil, nil
+			return nil, nil
 		}
-		return nil, nil, fmt.Errorf("kvstore: list keys: %w", err)
+		return nil, err
 	}
-	return kv, keys, nil
+	return keys, nil
 }
 
 // get is Get against an already-opened bucket, for loops that opened it once.

spinifex/kvstore/store_test.go

@@ -233,7 +233,7 @@ func TestOver_ServesAPreOpenedBucket(t *testing.T) {
 	_, err = kv.Put(t.Context(), "acct-a/one", data)
 	require.NoError(t, err)
 
-	store := kvstore.Over[record](kv, kvstore.Config{})
+	store := kvstore.Over[record](nil, kv, kvstore.Config{})
 
 	got, _, err := store.Get(t.Context(), "acct-a/one")
 	require.NoError(t, err)
@@ -441,7 +441,7 @@ func TestBucket_ConfiguredReportsWhetherThereIsAnythingToOpen(t *testing.T) {
 
 	assert.False(t, kvstore.NewBucket(nil, kvstore.Config{Name: "b"}).Configured())
 	assert.True(t, kvstore.NewBucket(testutil.NewJetStream(t, nc), kvstore.Config{Name: "b"}).Configured())
-	assert.True(t, kvstore.NewOpenBucket(newOpenBucket(t), kvstore.Config{Name: "b"}).Configured())
+	assert.True(t, kvstore.NewOpenBucket(nil, newOpenBucket(t), kvstore.Config{Name: "b"}).Configured())
 }
 
 // TestBucket_TTLOpensABucketWithThatMaxAge covers Config.TTL, the branch of