fix: never lower a quota counter from a stale instance view (#943)

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
 	}
 }

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"},

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.

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)
+}

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 {