From 6960ef6620f28f065a95172473583772e8312fc2 Mon Sep 17 00:00:00 2001 From: Josh Siviour Date: Mon, 31 Aug 2026 10:46:57 +1000 Subject: [PATCH] fix: never lower a quota counter from a stale instance view (#943) --- spinifex/handlers/quota/records.go | 53 ++++++++++--- spinifex/handlers/quota/records_test.go | 99 +++++++++++++++++++++++++ spinifex/kvstore/bucket.go | 64 ++++++++++++++++ spinifex/kvstore/snapshot_test.go | 98 ++++++++++++++++++++++++ spinifex/kvstore/store.go | 60 +++++++++++++++ 5 files changed, 365 insertions(+), 9 deletions(-) create mode 100644 spinifex/kvstore/snapshot_test.go diff --git a/spinifex/handlers/quota/records.go b/spinifex/handlers/quota/records.go index 5d3f8d1e9..87505bfb4 100644 --- a/spinifex/handlers/quota/records.go +++ b/spinifex/handlers/quota/records.go @@ -3,6 +3,7 @@ package handlers_quota import ( "context" "encoding/json" + "log/slog" "github.com/mulgadc/spinifex/spinifex/instancetypes" "github.com/mulgadc/spinifex/spinifex/kvstore" @@ -18,27 +19,47 @@ import ( // which is what zeroes an account that has just released everything. // // complete reports whether the view is good enough to lower a counter. It is -// false when the read failed, so a counter may still be raised from whatever -// was seen but never lowered from a view known to be short. +// false when the read failed and when the read cannot prove it was current, so +// a counter may still be raised from whatever was seen but never lowered from +// a view known to be short or suspected of being stale. type InstanceVCPULister func(ctx context.Context) (perAccount map[string]int, complete bool, err error) +// recordSource is the pair of reads the sweep needs: the stream position it +// must reach, and the snapshot it judges against that position. It exists so a +// test can hand the sweep a snapshot that is deliberately behind, which a live +// single-node server never produces. +type recordSource interface { + LastSequence(ctx context.Context, filter string) (uint64, error) + Snapshot(ctx context.Context, filter string) ([]kvstore.Item[vm.InstanceRecord], uint64, error) +} + +var _ recordSource = (*kvstore.Store[vm.InstanceRecord])(nil) + // RecordVCPULister totals vCPUs from the per-resource instance record space. // The account is on the record and vCPUs derive from the instance type, which -// is spec, so the whole sweep is one KV list where it used to be a NATS +// is spec, so the whole sweep is one KV read where it used to be a NATS // fan-out of describes costing passes x accounts x (nodes + 2). // // There is no index by account, so a recompute of one account still reads the -// whole prefix. That is one list per settled burst, which the debounce already +// whole prefix. That is one read per settled burst, which the debounce already // coalesces, against a fan-out that ran per account on every tick. -func RecordVCPULister(records *kvstore.Store[vm.InstanceRecord], prefix string) InstanceVCPULister { +// +// The watermark is taken before the snapshot, not after. Every record write +// that had completed when the pass began sits at or below it, so a snapshot +// reaching it has seen all of them and may lower a counter. A write landing +// after it belongs to the pass its own watch update triggers. +func RecordVCPULister(records recordSource, prefix string) InstanceVCPULister { + filter := prefix + "*" return func(ctx context.Context) (map[string]int, bool, error) { - list, err := records.List(ctx, prefix) + watermark, watermarkErr := records.LastSequence(ctx, filter) + items, highWater, err := records.Snapshot(ctx, filter) if err != nil { return nil, false, err } - totals := make(map[string]int, len(list)) - for i := range list { - record := &list[i] + + totals := make(map[string]int, len(items)) + for i := range items { + record := &items[i].Value accountID := record.Metadata.AccountID if accountID == "" || accountID == utils.GlobalAccountID { continue @@ -52,6 +73,20 @@ func RecordVCPULister(records *kvstore.Store[vm.InstanceRecord], prefix string) totals[accountID] += vcpus } } + + // A snapshot short of the watermark was served by a replica that has + // not caught up, and lowering from it would undo a charge for an + // instance it never saw. Raising is still safe and still happens. + if watermarkErr != nil { + slog.WarnContext(ctx, "quota: instance watermark unavailable, counters may not be lowered", + "err", watermarkErr) + return totals, false, nil + } + if highWater < watermark { + slog.WarnContext(ctx, "quota: instance snapshot behind the stream, counters may not be lowered", + "high_water", highWater, "watermark", watermark) + return totals, false, nil + } return totals, true, nil } } diff --git a/spinifex/handlers/quota/records_test.go b/spinifex/handlers/quota/records_test.go index e7315778b..c0ba852c4 100644 --- a/spinifex/handlers/quota/records_test.go +++ b/spinifex/handlers/quota/records_test.go @@ -1,7 +1,9 @@ package handlers_quota_test import ( + "context" "encoding/json" + "errors" "testing" handlers_quota "github.com/mulgadc/spinifex/spinifex/handlers/quota" @@ -89,6 +91,103 @@ func TestRecordVCPULister_UnknownInstanceTypeContributesNothing(t *testing.T) { require.Equal(t, map[string]int{account: 2}, totals) } +// stubSource stands in for the record store so a snapshot can be made +// deliberately stale, which a single-node server never does on its own. +type stubSource struct { + items []kvstore.Item[vm.InstanceRecord] + highWater uint64 + watermark uint64 + watermarkErr error + snapshotCalls int +} + +func (s *stubSource) LastSequence(context.Context, string) (uint64, error) { + return s.watermark, s.watermarkErr +} + +func (s *stubSource) Snapshot(context.Context, string) ([]kvstore.Item[vm.InstanceRecord], uint64, error) { + s.snapshotCalls++ + return s.items, s.highWater, nil +} + +func stubItem(id, accountID, instanceType string, state vm.InstanceState) kvstore.Item[vm.InstanceRecord] { + return kvstore.Item[vm.InstanceRecord]{ + Key: recordPrefix + id, + Revision: 1, + Value: vm.InstanceRecord{ + Metadata: resource.Metadata{Name: id, AccountID: accountID}, + Spec: vm.InstanceSpec{InstanceType: instanceType}, + Status: vm.InstanceStatus{Status: state}, + }, + } +} + +// A snapshot that reached the watermark saw every write that had completed +// when the pass began, so it is good enough to lower a counter. +func TestRecordVCPULister_CompleteWhenTheSnapshotReachesTheWatermark(t *testing.T) { + const account = "111111111111" + source := &stubSource{ + items: []kvstore.Item[vm.InstanceRecord]{stubItem("i-1", account, "m5.xlarge", vm.StateRunning)}, + highWater: 12, + watermark: 12, + } + + totals, complete, err := handlers_quota.RecordVCPULister(source, recordPrefix)(t.Context()) + require.NoError(t, err) + require.True(t, complete) + require.Equal(t, map[string]int{account: 4}, totals) + require.Equal(t, 1, source.snapshotCalls) +} + +// A snapshot behind the watermark was served by a replica that has not caught +// up. The totals are still returned so the pass can raise; complete is false +// so it cannot lower and undo a charge for an instance it never saw. +func TestRecordVCPULister_IncompleteWhenTheSnapshotIsBehindTheWatermark(t *testing.T) { + const account = "111111111111" + source := &stubSource{ + items: []kvstore.Item[vm.InstanceRecord]{stubItem("i-1", account, "m5.xlarge", vm.StateRunning)}, + highWater: 11, + watermark: 12, + } + + totals, complete, err := handlers_quota.RecordVCPULister(source, recordPrefix)(t.Context()) + require.NoError(t, err) + require.False(t, complete) + require.Equal(t, map[string]int{account: 4}, totals) +} + +// A watermark that cannot be read is not a reason to fail the pass, but it is +// a reason not to trust the snapshot's currency. +func TestRecordVCPULister_IncompleteWhenTheWatermarkIsUnavailable(t *testing.T) { + const account = "111111111111" + source := &stubSource{ + items: []kvstore.Item[vm.InstanceRecord]{stubItem("i-1", account, "t3.micro", vm.StateRunning)}, + highWater: 12, + watermarkErr: errors.New("stream leader unavailable"), + } + + totals, complete, err := handlers_quota.RecordVCPULister(source, recordPrefix)(t.Context()) + require.NoError(t, err) + require.False(t, complete) + require.Equal(t, map[string]int{account: 2}, totals) +} + +// The live pairing is what makes complete reachable at all: against a real +// bucket the snapshot and the watermark agree, including when the newest +// event is the delete that releases the quota. +func TestRecordVCPULister_CompleteAgainstALiveBucketAfterADelete(t *testing.T) { + store := recordStore(t) + const account = "111111111111" + putRecord(t, store, "i-1", account, "m5.xlarge", vm.StateRunning) + putRecord(t, store, "i-2", account, "t3.micro", vm.StateRunning) + require.NoError(t, store.Delete(t.Context(), recordPrefix+"i-2")) + + totals, complete, err := handlers_quota.RecordVCPULister(store, recordPrefix)(t.Context()) + require.NoError(t, err) + require.True(t, complete) + require.Equal(t, map[string]int{account: 4}, totals) +} + func TestAccountForRecord(t *testing.T) { valid, err := json.Marshal(vm.InstanceRecord{ Metadata: resource.Metadata{Name: "i-1", AccountID: "111111111111"}, diff --git a/spinifex/kvstore/bucket.go b/spinifex/kvstore/bucket.go index e6d94d272..6aac39f6c 100644 --- a/spinifex/kvstore/bucket.go +++ b/spinifex/kvstore/bucket.go @@ -5,6 +5,7 @@ package kvstore import ( "context" + "encoding/json" "errors" "fmt" "log/slog" @@ -104,6 +105,69 @@ func (b *Bucket) KV(ctx context.Context) (jetstream.KeyValue, error) { return kv, nil } +// jsAPIStreamMsgGet is the non-direct stream read, spelled out rather than +// reached through nats.go. Every convenience wrapper switches to +// $JS.API.DIRECT.GET on a stream with AllowDirect, and KV buckets always +// have it. The default API prefix holds because every caller builds its +// client with plain jetstream.New, so no domain is in play. +const jsAPIStreamMsgGet = "$JS.API.STREAM.MSG.GET.KV_%s" + +// jsErrCodeMessageNotFound is returned for a filter matching nothing, which is +// an empty prefix rather than a failure. +const jsErrCodeMessageNotFound = 10037 + +// LastSequence reports the stream sequence of the newest message matching +// filter, or zero when nothing matches. A tombstone is a message, so the +// sequence follows the stream rather than the surviving key set. +// +// The stream leader answers, which is the entire point of the hand-rolled +// request: nats.go's GetLastMsgForSubject would serve this from any replica, +// so a reader asking "am I behind?" could be answered by a replica that is +// equally behind. +func (b *Bucket) LastSequence(ctx context.Context, filter string) (uint64, error) { + b.mu.Lock() + js := b.js + b.mu.Unlock() + if js == nil { + if b.cfg.Missing == "" { + return 0, errors.New("kvstore: no JetStream client configured") + } + return 0, errors.New(b.cfg.Missing) + } + + req, err := json.Marshal(map[string]string{"last_by_subj": "$KV." + b.cfg.Name + "." + filter}) + if err != nil { + return 0, fmt.Errorf("kvstore: encode last-sequence request for %s: %w", b.cfg.Name, err) + } + msg, err := js.Conn().RequestWithContext(ctx, fmt.Sprintf(jsAPIStreamMsgGet, b.cfg.Name), req) + if err != nil { + return 0, fmt.Errorf("kvstore: last sequence %s %s: %w", b.cfg.Name, filter, err) + } + + var resp struct { + Message *struct { + Sequence uint64 `json:"seq"` + } `json:"message"` + Error *struct { + ErrCode int `json:"err_code"` + Description string `json:"description"` + } `json:"error"` + } + if err := json.Unmarshal(msg.Data, &resp); err != nil { + return 0, fmt.Errorf("kvstore: decode last sequence %s: %w", b.cfg.Name, err) + } + if resp.Error != nil { + if resp.Error.ErrCode == jsErrCodeMessageNotFound { + return 0, nil + } + return 0, fmt.Errorf("kvstore: last sequence %s %s: %s", b.cfg.Name, filter, resp.Error.Description) + } + if resp.Message == nil { + return 0, fmt.Errorf("kvstore: last sequence %s %s: empty response", b.cfg.Name, filter) + } + return resp.Message.Sequence, 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. diff --git a/spinifex/kvstore/snapshot_test.go b/spinifex/kvstore/snapshot_test.go new file mode 100644 index 000000000..b76cea5f1 --- /dev/null +++ b/spinifex/kvstore/snapshot_test.go @@ -0,0 +1,98 @@ +package kvstore_test + +import ( + "testing" + + "github.com/mulgadc/spinifex/spinifex/kvstore" + "github.com/mulgadc/spinifex/spinifex/testutil" + "github.com/stretchr/testify/require" +) + +type snapRecord struct { + Name string `json:"name"` +} + +func snapshotStore(t *testing.T, name string) *kvstore.Store[snapRecord] { + t.Helper() + _, nc, _ := testutil.StartTestJetStream(t) + return kvstore.New[snapRecord](testutil.NewJetStream(t, nc), kvstore.Config{Name: name, History: 1}) +} + +func TestSnapshotReturnsMatchingRecordsWithTheirRevisions(t *testing.T) { + store := snapshotStore(t, "snapshot-records") + require.NoError(t, store.Set(t.Context(), "i.a", &snapRecord{Name: "a"})) + require.NoError(t, store.Set(t.Context(), "i.b", &snapRecord{Name: "b"})) + require.NoError(t, store.Set(t.Context(), "node.1", &snapRecord{Name: "elsewhere"})) + + items, highWater, err := store.Snapshot(t.Context(), "i.*") + require.NoError(t, err) + require.Len(t, items, 2) + + byKey := map[string]kvstore.Item[snapRecord]{} + for _, item := range items { + byKey[item.Key] = item + require.NotZero(t, item.Revision) + } + require.Equal(t, "a", byKey["i.a"].Value.Name) + require.Equal(t, "b", byKey["i.b"].Value.Name) + require.NotContains(t, byKey, "node.1") + require.Equal(t, byKey["i.b"].Revision, highWater) +} + +// The high-water mark tracks the reader's position in the stream, not what +// survives in it. A delete is a message, so a snapshot whose newest event is a +// delete still reaches the watermark and can still lower a counter. +func TestSnapshotCountsATombstoneTowardsTheHighWaterMark(t *testing.T) { + store := snapshotStore(t, "snapshot-tombstone") + require.NoError(t, store.Set(t.Context(), "i.a", &snapRecord{Name: "a"})) + require.NoError(t, store.Set(t.Context(), "i.b", &snapRecord{Name: "b"})) + require.NoError(t, store.Delete(t.Context(), "i.b")) + + items, highWater, err := store.Snapshot(t.Context(), "i.*") + require.NoError(t, err) + require.Len(t, items, 1) + require.Equal(t, "i.a", items[0].Key) + + watermark, err := store.LastSequence(t.Context(), "i.*") + require.NoError(t, err) + require.Equal(t, watermark, highWater) + require.Greater(t, highWater, items[0].Revision) +} + +func TestSnapshotOfNothingIsNotAnError(t *testing.T) { + store := snapshotStore(t, "snapshot-empty") + require.NoError(t, store.Set(t.Context(), "node.1", &snapRecord{Name: "elsewhere"})) + + items, highWater, err := store.Snapshot(t.Context(), "i.*") + require.NoError(t, err) + require.Empty(t, items) + require.Zero(t, highWater) +} + +func TestLastSequenceIsTheNewestMessageMatchingTheFilter(t *testing.T) { + store := snapshotStore(t, "watermark-newest") + require.NoError(t, store.Set(t.Context(), "i.a", &snapRecord{Name: "a"})) + items, highWater, err := store.Snapshot(t.Context(), "i.*") + require.NoError(t, err) + require.Len(t, items, 1) + + watermark, err := store.LastSequence(t.Context(), "i.*") + require.NoError(t, err) + require.Equal(t, highWater, watermark) + + // A write outside the filter must not move it, which is what makes the + // mark comparable to a snapshot of the same filter. + require.NoError(t, store.Set(t.Context(), "node.1", &snapRecord{Name: "elsewhere"})) + after, err := store.LastSequence(t.Context(), "i.*") + require.NoError(t, err) + require.Equal(t, watermark, after) +} + +func TestLastSequenceIsZeroWhenNothingMatches(t *testing.T) { + store := snapshotStore(t, "watermark-empty") + require.NoError(t, store.Set(t.Context(), "node.1", &snapRecord{Name: "elsewhere"})) + + watermark, err := store.LastSequence(t.Context(), "i.*") + require.NoError(t, err) + require.Zero(t, watermark) +} diff --git a/spinifex/kvstore/store.go b/spinifex/kvstore/store.go index 909131046..70bf1ef83 100644 --- a/spinifex/kvstore/store.go +++ b/spinifex/kvstore/store.go @@ -163,6 +163,66 @@ func (s *Store[T]) List(ctx context.Context, prefix string) ([]T, error) { return out, nil } +// Item is one record from a Snapshot, carrying the revision it was read at so +// the caller can CAS against it. +type Item[T any] struct { + Key string + Value T + Revision uint64 +} + +// Snapshot reads every record matching filter in one ordered pass and reports +// the highest stream sequence it saw. Filter takes NATS subject wildcards. +// +// One consumer answers the whole read, so the keys and the values are +// mutually consistent. List is not: it enumerates keys through a consumer and +// then fetches each value through a direct get, and each of those gets is +// routed independently, so no two of them need come from the same replica. +// +// A deleted key is not returned, but its revision still counts towards +// highWater. That is deliberate: the mark says how far this reader got +// through the stream, not what survives in it, and a delete is how far a +// reader gets when the newest thing that happened was a delete. +func (s *Store[T]) Snapshot(ctx context.Context, filter string) (items []Item[T], highWater uint64, err error) { + // Pure read, so it is safe for withKV to run it twice; the results are + // reset on entry rather than appended to across a retry. + err = s.withKV(ctx, func(kv jetstream.KeyValue) error { + items, highWater = nil, 0 + + watcher, err := kv.Watch(ctx, filter) + if err != nil { + return fmt.Errorf("kvstore: snapshot %s %s: %w", s.cfg.Name, filter, err) + } + defer func() { _ = watcher.Stop() }() + + for entry := range watcher.Updates() { + // nil marks the end of the existing values, which is the whole + // snapshot: anything after it is a live update, not state. + if entry == nil { + return nil + } + if entry.Revision() > highWater { + highWater = entry.Revision() + } + if entry.Operation() != jetstream.KeyValuePut { + continue + } + var v T + if err := json.Unmarshal(entry.Value(), &v); err != nil { + return fmt.Errorf("kvstore: decode %s: %w", entry.Key(), err) + } + items = append(items, Item[T]{Key: entry.Key(), Value: v, Revision: entry.Revision()}) + } + // The channel closed without the marker, so the watcher died mid-pass + // and what was collected is a partial view rather than a snapshot. + return fmt.Errorf("kvstore: snapshot %s %s: watcher closed early", s.cfg.Name, filter) + }) + if err != nil { + return nil, 0, err + } + return items, highWater, 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 {