From 8dddf9d55ef394d04f8e336590e4468e95f83cda Mon Sep 17 00:00:00 2001 From: Josh Siviour Date: Wed, 26 Aug 2026 23:55:58 +1000 Subject: [PATCH] refactor: convert the bedrock access, credential and usage stores to kvstore (#898) --- spinifex/gateway/bedrock/access.go | 45 ++---- spinifex/gateway/bedrock/access_test.go | 21 +++ spinifex/gateway/bedrock/credentials.go | 41 ++--- spinifex/gateway/bedrock/credentials_test.go | 11 ++ spinifex/gateway/bedrock/usage.go | 149 ++++++------------- spinifex/gateway/bedrock/usage_test.go | 20 +++ spinifex/kvstore/bucket.go | 8 + spinifex/kvstore/store.go | 15 ++ spinifex/kvstore/store_test.go | 85 +++++++++++ 9 files changed, 231 insertions(+), 164 deletions(-) diff --git a/spinifex/gateway/bedrock/access.go b/spinifex/gateway/bedrock/access.go index 7b1c269ff..3e4878fa3 100644 --- a/spinifex/gateway/bedrock/access.go +++ b/spinifex/gateway/bedrock/access.go @@ -6,9 +6,8 @@ import ( "errors" "fmt" "strings" - "sync" - "github.com/mulgadc/spinifex/spinifex/kvutil" + "github.com/mulgadc/spinifex/spinifex/kvstore" "github.com/mulgadc/spinifex/spinifex/utils" "github.com/nats-io/nats.go/jetstream" ) @@ -57,13 +56,10 @@ type AccessResolver interface { } // ModelAccessStore resolves per-account model grants from the -// bedrock-model-access JetStream KV bucket. +// bedrock-model-access JetStream KV bucket. A grant carries no value, so this +// is a Bucket rather than a Store[T]: presence of the key is the whole record. type ModelAccessStore struct { - js jetstream.JetStream - replicas int - - mu sync.Mutex - kv jetstream.KeyValue + bucket *kvstore.Bucket } var _ AccessResolver = (*ModelAccessStore)(nil) @@ -71,23 +67,12 @@ var _ AccessResolver = (*ModelAccessStore)(nil) // NewModelAccessStore constructs a ModelAccessStore over the cluster's // JetStream client, replicated across replicas nodes. func NewModelAccessStore(js jetstream.JetStream, replicas int) *ModelAccessStore { - return &ModelAccessStore{js: js, replicas: replicas} -} - -// bucket lazily opens (or creates) the bedrock-model-access KV bucket, -// caching the handle, mirroring WeightsStore.bucket. -func (s *ModelAccessStore) bucket(ctx context.Context) (jetstream.KeyValue, error) { - s.mu.Lock() - defer s.mu.Unlock() - if s.kv != nil { - return s.kv, nil - } - kv, err := kvutil.GetOrCreateBucketWithReplicas(ctx, s.js, modelAccessBucket, modelAccessHistory, s.replicas) - if err != nil { - return nil, err - } - s.kv = kv - return kv, nil + return &ModelAccessStore{bucket: kvstore.NewBucket(js, kvstore.Config{ + Name: modelAccessBucket, + History: modelAccessHistory, + Replicas: replicas, + Missing: "bedrock: model access store has no JetStream client configured", + })} } // Granted reports whether accountID holds a grant on modelID. The system @@ -97,7 +82,7 @@ func (s *ModelAccessStore) Granted(ctx context.Context, accountID, modelID strin if accountID == utils.GlobalAccountID { return true, nil } - kv, err := s.bucket(ctx) + kv, err := s.bucket.KV(ctx) if err != nil { return false, err } @@ -114,7 +99,7 @@ func (s *ModelAccessStore) Granted(ctx context.Context, accountID, modelID strin // Grant gives accountID access to modelID. It is idempotent: re-granting an // existing grant is a no-op rather than an error. func (s *ModelAccessStore) Grant(ctx context.Context, accountID, modelID string) error { - kv, err := s.bucket(ctx) + kv, err := s.bucket.KV(ctx) if err != nil { return err } @@ -127,7 +112,7 @@ func (s *ModelAccessStore) Grant(ctx context.Context, accountID, modelID string) // Revoke removes accountID's access to modelID. Revoking a grant that does // not exist succeeds, so callers need not check first. func (s *ModelAccessStore) Revoke(ctx context.Context, accountID, modelID string) error { - kv, err := s.bucket(ctx) + kv, err := s.bucket.KV(ctx) if err != nil { return err } @@ -148,7 +133,7 @@ func (s *ModelAccessStore) Revoke(ctx context.Context, accountID, modelID string // leaving a half-seeded account. The marker itself is a conditional create: // every node runs this at startup and only the first need win. func (s *ModelAccessStore) SeedAccountGrants(ctx context.Context, accountID string, modelIDs []string) (bool, error) { - kv, err := s.bucket(ctx) + kv, err := s.bucket.KV(ctx) if err != nil { return false, err } @@ -178,7 +163,7 @@ func (s *ModelAccessStore) SeedAccountGrants(ctx context.Context, accountID stri // order. Keys that do not decode are skipped rather than failing the call, so // one malformed key cannot hide an account's whole grant set. func (s *ModelAccessStore) List(ctx context.Context, accountID string) ([]string, error) { - kv, err := s.bucket(ctx) + kv, err := s.bucket.KV(ctx) if err != nil { return nil, err } diff --git a/spinifex/gateway/bedrock/access_test.go b/spinifex/gateway/bedrock/access_test.go index fb50fdec5..f65e8a880 100644 --- a/spinifex/gateway/bedrock/access_test.go +++ b/spinifex/gateway/bedrock/access_test.go @@ -242,3 +242,24 @@ func TestSeedAccountGrants_MarkerIsNotAGrant(t *testing.T) { require.NoError(t, err) assert.Empty(t, models, "the seed marker must not be listed as a grant") } + +// TestModelAccessStore_RequiresJetStream covers a store constructed with no +// JetStream client: every accessor must report the misconfiguration rather +// than panic on a nil handle. +func TestModelAccessStore_RequiresJetStream(t *testing.T) { + store := NewModelAccessStore(nil, 1) + ctx := context.Background() + + _, err := store.Granted(ctx, "000000000001", "anthropic.claude-3-5-sonnet-20240620-v1:0") + require.Error(t, err) + assert.Contains(t, err.Error(), "no JetStream client configured") + + require.Error(t, store.Grant(ctx, "000000000001", "m")) + require.Error(t, store.Revoke(ctx, "000000000001", "m")) + + _, err = store.List(ctx, "000000000001") + require.Error(t, err) + + _, err = store.SeedAccountGrants(ctx, "000000000001", []string{"m"}) + require.Error(t, err) +} diff --git a/spinifex/gateway/bedrock/credentials.go b/spinifex/gateway/bedrock/credentials.go index 442629fce..bc3c4277b 100644 --- a/spinifex/gateway/bedrock/credentials.go +++ b/spinifex/gateway/bedrock/credentials.go @@ -4,10 +4,9 @@ import ( "context" "errors" "fmt" - "sync" handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam" - "github.com/mulgadc/spinifex/spinifex/kvutil" + "github.com/mulgadc/spinifex/spinifex/kvstore" "github.com/nats-io/nats.go/jetstream" ) @@ -26,14 +25,12 @@ func credentialKey(accountID, vendor string) string { // CredentialStore resolves per-account provider API keys from the // bedrock-credentials JetStream KV bucket, falling back to a platform-wide // default when an account has none. Keys are encrypted at rest. +// Values are IAM-encrypted ciphertext rather than JSON, so this is a Bucket: +// a Store[string] would wrap the ciphertext in a second encoding. type CredentialStore struct { - js jetstream.JetStream + bucket *kvstore.Bucket masterKey []byte - replicas int platformDefaults map[string]string - - mu sync.Mutex - kv jetstream.KeyValue } var _ CredentialResolver = (*CredentialStore)(nil) @@ -43,36 +40,24 @@ var _ CredentialResolver = (*CredentialStore)(nil) // own; pass an empty map to disable platform defaults entirely. func NewCredentialStore(js jetstream.JetStream, masterKey []byte, replicas int, platformDefaults map[string]string) *CredentialStore { return &CredentialStore{ - js: js, + bucket: kvstore.NewBucket(js, kvstore.Config{ + Name: bedrockCredentialsBucket, + History: bedrockCredentialsHistory, + Replicas: replicas, + Missing: "bedrock: credential store has no JetStream client configured", + }), masterKey: masterKey, - replicas: replicas, platformDefaults: platformDefaults, } } -// bucket lazily opens (or creates) the cluster-replicated bedrock-credentials -// KV bucket, caching the handle for subsequent calls. -func (s *CredentialStore) bucket(ctx context.Context) (jetstream.KeyValue, error) { - s.mu.Lock() - defer s.mu.Unlock() - if s.kv != nil { - return s.kv, nil - } - kv, err := kvutil.GetOrCreateBucketWithReplicas(ctx, s.js, bedrockCredentialsBucket, bedrockCredentialsHistory, s.replicas) - if err != nil { - return nil, err - } - s.kv = kv - return kv, nil -} - // Resolve returns accountID's vendor API key: a per-account key if one is // stored, else the platform default, else ("", false, nil). A nil js is // tolerated so a store built with only platformDefaults (e.g. in tests) can // still serve the default-only path without touching JetStream. func (s *CredentialStore) Resolve(ctx context.Context, accountID, vendor string) (string, bool, error) { - if s.js != nil { - kv, err := s.bucket(ctx) + if s.bucket.Configured() { + kv, err := s.bucket.KV(ctx) if err != nil { return "", false, err } @@ -96,7 +81,7 @@ func (s *CredentialStore) Resolve(ctx context.Context, accountID, vendor string) // PutCredential encrypts and stores key as accountID's vendor credential. func (s *CredentialStore) PutCredential(ctx context.Context, accountID, vendor, key string) error { - kv, err := s.bucket(ctx) + kv, err := s.bucket.KV(ctx) if err != nil { return err } diff --git a/spinifex/gateway/bedrock/credentials_test.go b/spinifex/gateway/bedrock/credentials_test.go index 362454ed4..489620339 100644 --- a/spinifex/gateway/bedrock/credentials_test.go +++ b/spinifex/gateway/bedrock/credentials_test.go @@ -114,3 +114,14 @@ func TestEncryptDecryptSecret_RoundTrip(t *testing.T) { require.NoError(t, err) assert.Equal(t, plaintext, decrypted) } + +// TestCredentialStore_PutCredential_RequiresJetStream is the counterpart to +// Resolve's tolerated nil: a write has no platform-default fallback to fall +// back to, so it must report the misconfiguration. +func TestCredentialStore_PutCredential_RequiresJetStream(t *testing.T) { + store := NewCredentialStore(nil, fixedMasterKey(), 1, map[string]string{"anthropic": "sk-default"}) + + err := store.PutCredential(context.Background(), "000000000001", "anthropic", "sk-account") + require.Error(t, err) + assert.Contains(t, err.Error(), "no JetStream client configured") +} diff --git a/spinifex/gateway/bedrock/usage.go b/spinifex/gateway/bedrock/usage.go index 28377f39d..59ec3fc24 100644 --- a/spinifex/gateway/bedrock/usage.go +++ b/spinifex/gateway/bedrock/usage.go @@ -7,9 +7,9 @@ import ( "fmt" "log/slog" "strings" - "sync" "time" + "github.com/mulgadc/spinifex/spinifex/kvstore" "github.com/mulgadc/spinifex/spinifex/kvutil" "github.com/nats-io/nats.go/jetstream" ) @@ -89,14 +89,10 @@ type UsageReader interface { // bedrock-usage JetStream KV bucket, and dedupes applied RequestIDs in a // second bucket so at-least-once stream delivery never double-counts. type UsageStore struct { - js jetstream.JetStream - replicas int - - mu sync.Mutex - kv jetstream.KeyValue - - dedupeMu sync.Mutex - dedupe jetstream.KeyValue + store *kvstore.Store[UsageCounters] + // dedupe holds valueless markers under a bucket-wide TTL, so it is a + // Bucket rather than a Store[T]. + dedupe *kvstore.Bucket } var _ UsageReader = (*UsageStore)(nil) @@ -104,39 +100,28 @@ var _ UsageReader = (*UsageStore)(nil) // NewUsageStore constructs a UsageStore over the cluster's JetStream client, // replicated across replicas nodes. func NewUsageStore(js jetstream.JetStream, replicas int) *UsageStore { - return &UsageStore{js: js, replicas: replicas} + return &UsageStore{ + store: kvstore.New[UsageCounters](js, kvstore.Config{ + Name: bedrockUsageBucket, + History: bedrockUsageHistory, + Replicas: replicas, + Missing: "bedrock: usage store has no JetStream client configured", + Attempts: usageCASRetries, + Exhausted: usageCASExhausted, + }), + dedupe: kvstore.NewBucket(js, kvstore.Config{ + Name: bedrockUsageDedupeBucket, + History: bedrockUsageDedupeHistory, + TTL: invocationStreamMaxAge, + Missing: "bedrock: usage store has no JetStream client configured", + }), + } } -// bucket lazily opens (or creates) the cluster-replicated bedrock-usage KV -// bucket, caching the handle for subsequent calls. -func (s *UsageStore) bucket(ctx context.Context) (jetstream.KeyValue, error) { - s.mu.Lock() - defer s.mu.Unlock() - if s.kv != nil { - return s.kv, nil - } - kv, err := kvutil.GetOrCreateBucketWithReplicas(ctx, s.js, bedrockUsageBucket, bedrockUsageHistory, s.replicas) - if err != nil { - return nil, err - } - s.kv = kv - return kv, nil -} - -// dedupeBucket lazily opens (or creates) the cluster-replicated -// bedrock-usage-dedupe KV bucket. -func (s *UsageStore) dedupeBucket(ctx context.Context) (jetstream.KeyValue, error) { - s.dedupeMu.Lock() - defer s.dedupeMu.Unlock() - if s.dedupe != nil { - return s.dedupe, nil - } - kv, err := kvutil.GetOrCreateBucketWithTTL(ctx, s.js, bedrockUsageDedupeBucket, bedrockUsageDedupeHistory, invocationStreamMaxAge) - if err != nil { - return nil, err - } - s.dedupe = kv - return kv, nil +// usageCASExhausted names the counter whose contention outlasted the retry +// budget, so a spent budget is distinguishable from any other write failure. +func usageCASExhausted(key string, attempts int) error { + return fmt.Errorf("usage counter CAS exhausted for %s after %d attempts", key, attempts) } // MarkProcessed atomically claims requestID for the usage consumer. first is @@ -144,7 +129,7 @@ func (s *UsageStore) dedupeBucket(ctx context.Context) (jetstream.KeyValue, erro // (redelivery) returns false with no error, so the caller can skip // re-applying counters without treating the duplicate as a failure. func (s *UsageStore) MarkProcessed(ctx context.Context, requestID string) (first bool, err error) { - kv, err := s.dedupeBucket(ctx) + kv, err := s.dedupe.KV(ctx) if err != nil { return false, err } @@ -160,23 +145,14 @@ func (s *UsageStore) MarkProcessed(ctx context.Context, requestID string) (first // Get returns accountID's usage counters for modelID in period, or // (zero, false, nil) if the account has accrued nothing yet. func (s *UsageStore) Get(ctx context.Context, accountID, modelID, period string) (UsageCounters, bool, error) { - kv, err := s.bucket(ctx) - if err != nil { - return UsageCounters{}, false, err - } - entry, err := kv.Get(ctx, usageKey(accountID, modelID, period)) - switch { - case err == nil: - var counters UsageCounters - if uerr := json.Unmarshal(entry.Value(), &counters); uerr != nil { - return UsageCounters{}, false, fmt.Errorf("decode usage counters for %s/%s/%s: %w", accountID, modelID, period, uerr) - } - return counters, true, nil - case errors.Is(err, jetstream.ErrKeyNotFound): + counters, _, err := s.store.Get(ctx, usageKey(accountID, modelID, period)) + if errors.Is(err, kvstore.ErrNotFound) { return UsageCounters{}, false, nil - default: - return UsageCounters{}, false, fmt.Errorf("kv get usage counters for %s/%s/%s: %w", accountID, modelID, period, err) } + if err != nil { + return UsageCounters{}, false, fmt.Errorf("usage counters for %s/%s/%s: %w", accountID, modelID, period, err) + } + return *counters, true, nil } // microUSDCost returns tokens priced at pricePerMillion integer micro-USD @@ -194,67 +170,28 @@ func microUSDCost(tokens, pricePerMillion int64) int64 { // RequestID has not been applied before — Apply itself has no idempotency of // its own and will double-count a record applied twice. func (s *UsageStore) Apply(ctx context.Context, rec InvocationRecord, price Price) error { - kv, err := s.bucket(ctx) - if err != nil { - return err - } key := usageKey(rec.AccountID, rec.ModelID, billingPeriod(rec.Timestamp)) - - for range usageCASRetries { - var current UsageCounters - var revision uint64 - entry, err := kv.Get(ctx, key) - switch { - case err == nil: - if uerr := json.Unmarshal(entry.Value(), ¤t); uerr != nil { - return fmt.Errorf("decode usage counters for %s: %w", key, uerr) - } - revision = entry.Revision() - case errors.Is(err, jetstream.ErrKeyNotFound): - // Zero value at revision 0: the CAS write below creates it. - default: - return fmt.Errorf("kv get usage counters for %s: %w", key, err) - } - - next := current - next.InputTokens += rec.InputTokens - next.OutputTokens += rec.OutputTokens - next.RequestCount++ + return s.store.Upsert(ctx, key, func(c *UsageCounters) (bool, error) { + c.InputTokens += rec.InputTokens + c.OutputTokens += rec.OutputTokens + c.RequestCount++ if price.Known { - next.CostMicroUSD += microUSDCost(rec.InputTokens, price.InputMicroUSDPerMillion) + microUSDCost(rec.OutputTokens, price.OutputMicroUSDPerMillion) + c.CostMicroUSD += microUSDCost(rec.InputTokens, price.InputMicroUSDPerMillion) + microUSDCost(rec.OutputTokens, price.OutputMicroUSDPerMillion) } else { - next.CostUnknown = true + c.CostUnknown = true } - - data, merr := json.Marshal(next) - if merr != nil { - return fmt.Errorf("encode usage counters for %s: %w", key, merr) - } - - if revision == 0 { - _, err = kv.Create(ctx, key, data) - } else { - _, err = kv.Update(ctx, key, data, revision) - } - if err == nil { - return nil - } - // Create reports a lost race as ErrKeyExists, Update as - // ErrKeyRevisionMismatch — and only the latter holds on a replicated - // bucket, where the conflict carries a different API error code. - if !errors.Is(err, jetstream.ErrKeyExists) && !errors.Is(err, jetstream.ErrKeyRevisionMismatch) { - return fmt.Errorf("kv write usage counters for %s: %w", key, err) - } - } - return fmt.Errorf("usage counter CAS exhausted for %s after %d attempts", key, usageCASRetries) + return true, nil + }) } // TokensThisPeriod sums accountID's input+output tokens across every model // for the current billing period, for the tokens-per-month quota dimension. // A few seconds of staleness against the usage consumer is acceptable for a // monthly cap (see handlers_quota.CheckBedrockTokens). +// It stays on the raw handle rather than Store.List because the period is +// matched on the key's suffix, which List does not surface. func (s *UsageStore) TokensThisPeriod(ctx context.Context, accountID string) (int64, error) { - kv, err := s.bucket(ctx) + kv, err := s.store.KV(ctx) if err != nil { return 0, err } diff --git a/spinifex/gateway/bedrock/usage_test.go b/spinifex/gateway/bedrock/usage_test.go index 44efe4001..bbddcc338 100644 --- a/spinifex/gateway/bedrock/usage_test.go +++ b/spinifex/gateway/bedrock/usage_test.go @@ -288,3 +288,23 @@ func TestBillingPeriod_MonthlyFormat(t *testing.T) { ts := time.Date(2026, time.August, 3, 12, 0, 0, 0, time.UTC) assert.Equal(t, "2026-08", billingPeriod(ts)) } + +// TestUsageStore_RequiresJetStream covers a store constructed with no +// JetStream client: both the counter bucket and the dedupe bucket must report +// the misconfiguration rather than panic on a nil handle. +func TestUsageStore_RequiresJetStream(t *testing.T) { + store := NewUsageStore(nil, 1) + ctx := context.Background() + + _, _, err := store.Get(ctx, "000000000001", "m", "2026-08") + require.Error(t, err) + assert.Contains(t, err.Error(), "no JetStream client configured") + + _, err = store.MarkProcessed(ctx, "req-1") + require.Error(t, err) + + require.Error(t, store.Apply(ctx, InvocationRecord{AccountID: "000000000001", ModelID: "m"}, Price{})) + + _, err = store.TokensThisPeriod(ctx, "000000000001") + require.Error(t, err) +} diff --git a/spinifex/kvstore/bucket.go b/spinifex/kvstore/bucket.go index 1983e3a4e..457553e55 100644 --- a/spinifex/kvstore/bucket.go +++ b/spinifex/kvstore/bucket.go @@ -46,6 +46,14 @@ func NewOpenBucket(kv jetstream.KeyValue, cfg Config) *Bucket { return &Bucket{kv: kv, cfg: cfg} } +// Configured reports whether the bucket has anything to open against, for the +// callers whose absent-KV path is a legitimate fallback rather than an error. +func (b *Bucket) Configured() bool { + b.mu.Lock() + defer b.mu.Unlock() + return b.js != nil || b.kv != nil +} + // KV opens the bucket on first use and returns the cached handle thereafter. func (b *Bucket) KV(ctx context.Context) (jetstream.KeyValue, error) { b.mu.Lock() diff --git a/spinifex/kvstore/store.go b/spinifex/kvstore/store.go index 81584a1ef..3173d913a 100644 --- a/spinifex/kvstore/store.go +++ b/spinifex/kvstore/store.go @@ -76,6 +76,21 @@ func (s *Store[T]) Mutate(ctx context.Context, key string, fn func(*T) (bool, er 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 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) diff --git a/spinifex/kvstore/store_test.go b/spinifex/kvstore/store_test.go index 0f6d79259..746aa1dc5 100644 --- a/spinifex/kvstore/store_test.go +++ b/spinifex/kvstore/store_test.go @@ -4,6 +4,7 @@ import ( "encoding/json" "errors" "testing" + "time" "github.com/mulgadc/spinifex/spinifex/kvstore" "github.com/mulgadc/spinifex/spinifex/kvutil" @@ -375,3 +376,87 @@ func TestStore_ConfigAttemptsBoundsMutate(t *testing.T) { require.NoError(t, err) assert.Equal(t, 2, got.Count, "an exhausted Mutate commits nothing of its own") } + +// TestStore_UpsertCreatesAnAbsentKey covers Upsert's difference from Mutate: +// the first write starts from the zero value rather than reporting ErrNotFound. +func TestStore_UpsertCreatesAnAbsentKey(t *testing.T) { + store := newStore(t) + + err := store.Upsert(t.Context(), "counter", func(r *record) (bool, error) { + r.Name = "first" + r.Count += 5 + return true, nil + }) + require.NoError(t, err) + + got, _, err := store.Get(t.Context(), "counter") + require.NoError(t, err) + assert.Equal(t, record{Name: "first", Count: 5}, *got) +} + +// TestStore_UpsertAccumulatesOntoAnExistingRecord proves the second call reads +// the first one's value rather than starting from zero again. +func TestStore_UpsertAccumulatesOntoAnExistingRecord(t *testing.T) { + store := newStore(t) + + for range 3 { + require.NoError(t, store.Upsert(t.Context(), "counter", func(r *record) (bool, error) { + r.Count += 2 + return true, nil + })) + } + + got, _, err := store.Get(t.Context(), "counter") + require.NoError(t, err) + assert.Equal(t, 6, got.Count) +} + +// TestStore_UpsertRetriesARevisionConflict asserts a write that lands under an +// in-flight Upsert is re-read rather than clobbered, so no increment is lost. +func TestStore_UpsertRetriesARevisionConflict(t *testing.T) { + store := newStore(t) + mustCreate(t, store, "counter", record{Count: 1}) + + var attempts int + err := store.Upsert(t.Context(), "counter", func(r *record) (bool, error) { + attempts++ + if attempts == 1 { + seedRaw(t, store, "counter", record{Count: 10}) + } + r.Count++ + return true, nil + }) + require.NoError(t, err) + assert.Equal(t, 2, attempts, "the first attempt must lose its revision and retry") + + got, _, err := store.Get(t.Context(), "counter") + require.NoError(t, err) + assert.Equal(t, 11, got.Count, "the retry must build on the interloper's value, not the stale one") +} + +// TestBucket_ConfiguredReportsWhetherThereIsAnythingToOpen covers the flag the +// callers whose absent-KV path is a legitimate fallback branch on. +func TestBucket_ConfiguredReportsWhetherThereIsAnythingToOpen(t *testing.T) { + _, nc, _ := testutil.StartTestJetStream(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()) +} + +// TestBucket_TTLOpensABucketWithThatMaxAge covers Config.TTL, the branch of +// Bucket.open no caller reached until the usage dedupe markers. +func TestBucket_TTLOpensABucketWithThatMaxAge(t *testing.T) { + _, nc, _ := testutil.StartTestJetStream(t) + bucket := kvstore.NewBucket(testutil.NewJetStream(t, nc), kvstore.Config{ + Name: "kvstore-ttl-test", + History: 1, + TTL: 90 * time.Minute, + }) + + kv, err := bucket.KV(t.Context()) + require.NoError(t, err) + status, err := kv.Status(t.Context()) + require.NoError(t, err) + assert.Equal(t, 90*time.Minute, status.TTL(), "Config.TTL must reach the created bucket") +}