From 286e11da03617fb6f092031eb47247a6fe9033aa Mon Sep 17 00:00:00 2001 From: Josh Siviour Date: Thu, 27 Aug 2026 16:19:41 +1000 Subject: [PATCH] feat: recover kvstore buckets from stream loss (#915) --- spinifex/handlers/acm/store.go | 4 +- spinifex/handlers/elbv2/store.go | 7 +- spinifex/kvstore/bucket.go | 121 ++++++++++++++- spinifex/kvstore/recovery_test.go | 245 ++++++++++++++++++++++++++++++ spinifex/kvstore/store.go | 204 ++++++++++++++----------- spinifex/kvstore/store_test.go | 4 +- 6 files changed, 482 insertions(+), 103 deletions(-) create mode 100644 spinifex/kvstore/recovery_test.go diff --git a/spinifex/handlers/acm/store.go b/spinifex/handlers/acm/store.go index 1e267d753..5f2f4c8fb 100644 --- a/spinifex/handlers/acm/store.go +++ b/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) diff --git a/spinifex/handlers/elbv2/store.go b/spinifex/handlers/elbv2/store.go index 327d42837..e304a8338 100644 --- a/spinifex/handlers/elbv2/store.go +++ b/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 diff --git a/spinifex/kvstore/bucket.go b/spinifex/kvstore/bucket.go index 5df9686e1..098088716 100644 --- a/spinifex/kvstore/bucket.go +++ b/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 } diff --git a/spinifex/kvstore/recovery_test.go b/spinifex/kvstore/recovery_test.go new file mode 100644 index 000000000..982c33bbc --- /dev/null +++ b/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") +} diff --git a/spinifex/kvstore/store.go b/spinifex/kvstore/store.go index 7f35a5622..fcdd5c852 100644 --- a/spinifex/kvstore/store.go +++ b/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. diff --git a/spinifex/kvstore/store_test.go b/spinifex/kvstore/store_test.go index d4b1ad6f8..e279f9d7a 100644 --- a/spinifex/kvstore/store_test.go +++ b/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