From 3409d50ec422fd8094c55e4424bbadd3315d15f8 Mon Sep 17 00:00:00 2001 From: Josh Siviour Date: Sun, 30 Aug 2026 13:48:04 +1000 Subject: [PATCH] feat: sweep quota from the instance record space on the per-key reconciler (#940) The vCPU sweep ran on a 30s ticker and called the account-filtered DescribeInstances fan-out once per account per pass, costing passes x accounts x (nodes + 2) and measured at ~62,000 instance describes a day, none of it user traffic and none of it tracking how often anything actually changed. The account is on the instance record and vCPUs derive from the instance type, which is spec, so the whole sweep is one read of the record space. The loop moves onto reconciler.Run: a change to one instance recomputes that one account, and the resync sweeps every account as the backstop. The completeness flag has to stay, for a different reason than the one it was written for. nats.go hardcodes AllowDirect on KV buckets, so a get at R3 may be served by a lagging replica, and ChargeLaunch increments the counter while reconcile overwrites it absolutely. A scan that misses a just-charged launch would lower the counter and lift the cap. The lower is now written only while the counter still stands at the revision read before the scan started, so a charge landing mid-scan wins and the scan is dropped for the next pass. Two behaviours change with the source. A failed sweep was a per-account describe error the pass continued past; it is now one read covering every account, so a failure moves no counter at all. And an account holding nothing is absent from the scan rather than reporting empty, so the caller's account list is what zeroes it. The retype gate and the volume delete guard still use the fan-out; only the sweep comes off it. --- .../gateway/ec2/instance/DescribeInstances.go | 13 +- spinifex/handlers/quota/reconcile.go | 113 +++++---- spinifex/handlers/quota/reconcile_test.go | 235 ++++++++++-------- spinifex/handlers/quota/records.go | 85 +++++++ spinifex/handlers/quota/records_test.go | 126 ++++++++++ spinifex/handlers/quota/vcpu.go | 47 +++- spinifex/handlers/quota/vcpu_retype_test.go | 17 ++ spinifex/services/awsgw/awsgw.go | 77 ++++-- .../services/awsgw/quota_reconcile_test.go | 145 ++++++++--- 9 files changed, 632 insertions(+), 226 deletions(-) create mode 100644 spinifex/handlers/quota/records.go create mode 100644 spinifex/handlers/quota/records_test.go diff --git a/spinifex/gateway/ec2/instance/DescribeInstances.go b/spinifex/gateway/ec2/instance/DescribeInstances.go index 50488d7c5..774a501bc 100644 --- a/spinifex/gateway/ec2/instance/DescribeInstances.go +++ b/spinifex/gateway/ec2/instance/DescribeInstances.go @@ -77,12 +77,13 @@ func DescribeInstancesChecked(ctx context.Context, input *ec2.DescribeInstancesI return &ec2.DescribeInstancesOutput{Reservations: reservations}, nil } -// DescribeInstancesForReconcile is the strict variant the quota reconcile uses. -// complete is true only when the sweep observed every expected node and both -// instance buckets, so reconcile may lower a counter only from a provably -// complete view. A partial sweep — a node down, a timed-out fan-out, or a failed -// bucket query — returns complete=false, and reconcile leaves the counter for the -// next clean pass rather than under-counting usage and lifting the cap. +// DescribeInstancesForReconcile is the strict variant, for a caller that may act +// only on a provably whole view. complete is true only when the sweep observed +// every expected node and both instance buckets; a partial sweep — a node down, +// a timed-out fan-out, or a failed bucket query — returns complete=false so the +// caller can decline to act rather than act on a short count. The quota vCPU +// sweep reads the instance record space instead of this; the retype gate and the +// volume delete guard are what still call it. func DescribeInstancesForReconcile(ctx context.Context, input *ec2.DescribeInstancesInput, natsConn *nats.Conn, expectedNodes int, accountID string) (reservations []*ec2.Reservation, complete bool, err error) { reservations, complete, _, err = gatherInstances(ctx, input, natsConn, expectedNodes, accountID) return reservations, complete, err diff --git a/spinifex/handlers/quota/reconcile.go b/spinifex/handlers/quota/reconcile.go index 19fa45364..9ee162493 100644 --- a/spinifex/handlers/quota/reconcile.go +++ b/spinifex/handlers/quota/reconcile.go @@ -6,49 +6,28 @@ import ( "log/slog" "github.com/aws/aws-sdk-go/service/ec2" - gateway_ec2_instance "github.com/mulgadc/spinifex/spinifex/gateway/ec2/instance" "github.com/mulgadc/spinifex/spinifex/instancetypes" "github.com/mulgadc/spinifex/spinifex/utils" - "github.com/nats-io/nats.go" ) // AccountLister enumerates the account IDs whose vCPU counters reconcile should -// recompute. It is satisfied by the gateway's active-account enumerator, so a -// global describe (which does not exist) is never needed: each account is swept -// through its own account-filtered Describe* call. +// recompute. It is satisfied by the gateway's active-account enumerator, and is +// what tells a whole-set pass about an account holding nothing: an account with +// no instances is absent from the scan, so only this list can zero it. type AccountLister func() ([]string, error) -// InstanceLister returns the reservations one account currently holds and whether -// the sweep was complete. The production implementation is the account-filtered -// DescribeInstances fan-out, which bundles running plus stopped plus terminated; -// reconcile drops the terminal set so only the "existing" (running plus stopped) -// vCPUs are summed. complete is false when a node or instance bucket did not -// answer, so reconcile must not lower a counter from that partial view. -type InstanceLister func(accountID string) (reservations []*ec2.Reservation, complete bool, err error) - -// NATSInstanceLister builds the production InstanceLister from a NATS connection -// and the configured cluster node count. The count is the expected node total, -// not the live-active count, so a node that is down makes the sweep incomplete -// rather than silently dropping its instances; reconcile then leaves the counter -// untouched instead of lowering it. expectedNodes is re-evaluated per call so a -// config change between passes is reflected without a gateway restart. -func NATSInstanceLister(natsConn *nats.Conn, expectedNodes func() int) InstanceLister { - return func(accountID string) ([]*ec2.Reservation, bool, error) { - return gateway_ec2_instance.DescribeInstancesForReconcile( - context.Background(), &ec2.DescribeInstancesInput{}, natsConn, expectedNodes(), accountID) - } -} - // Reconcile recomputes every account's vCPU counter from the live running-plus- // stopped instance set and CAS-overwrites it. It is the only path that lowers a // counter: it corrects the drift an out-of-band termination or a retype leaves // behind and zeroes accounts that now hold nothing. The system account is exempt -// and never charged. A per-account describe or write failure is logged and the -// pass continues; the first such error is returned so the caller can surface it. -// A counter is only lowered when the sweep was complete (every node and instance -// bucket answered); a partial sweep may raise but never lower, so a transient -// node outage cannot silently under-count usage and lift the cap. -func (s *Service) Reconcile(ctx context.Context, accounts AccountLister, list InstanceLister) error { +// and never charged. A per-account write failure is logged and the pass +// continues; the first such error is returned so the caller can surface it. +// +// A counter is only lowered from a view good enough to lower it: the scan +// succeeded, and the counter has not moved since the revision read before the +// scan started. Otherwise the pass may raise but never lower, so an under-count +// cannot silently lift the cap. +func (s *Service) Reconcile(ctx context.Context, accounts AccountLister, list InstanceVCPULister) error { if s == nil || !s.limits.Enabled { return nil } @@ -56,23 +35,31 @@ func (s *Service) Reconcile(ctx context.Context, accounts AccountLister, list In if err != nil { return fmt.Errorf("quota reconcile: list accounts: %w", err) } - var firstErr error + charged := make([]string, 0, len(ids)) for _, accountID := range ids { + if accountID != utils.GlobalAccountID { + charged = append(charged, accountID) + } + } + + // Revisions come before the scan, not after: a charge landing while the + // scan runs moves the revision, and the write then loses to the charge + // rather than overwriting a launch the scan never saw. + before := s.counterRevisions(ctx, charged) + totals, complete, err := list(ctx) + if err != nil { + slog.Warn("quota reconcile: instance scan failed, counters left unchanged", "err", err) + return err + } + + var firstErr error + for _, accountID := range charged { if err := ctx.Err(); err != nil { return err } - if accountID == utils.GlobalAccountID { - continue - } - reservations, complete, err := list(accountID) - if err != nil { - slog.Warn("quota reconcile: describe failed, counter left unchanged", "account", accountID, "err", err) - if firstErr == nil { - firstErr = err - } - continue - } - if err := s.reconcileVCPU(ctx, accountID, sumReservationVCPUs(reservations), complete); err != nil { + // An account absent from the scan holds nothing and is zeroed, which + // is the drift an out-of-band termination leaves behind. + if err := s.reconcileVCPU(ctx, accountID, totals[accountID], complete, before[accountID]); err != nil { slog.Warn("quota reconcile: counter overwrite failed", "account", accountID, "err", err) if firstErr == nil { firstErr = err @@ -82,6 +69,42 @@ func (s *Service) Reconcile(ctx context.Context, accounts AccountLister, list In return firstErr } +// ReconcileAccount recomputes one account's counter. It is the per-key entry +// point: a change to one instance costs that account's recompute rather than a +// sweep of every account. The scan still reads the whole record space, because +// nothing indexes it by account, but it runs once per settled burst instead of +// once per account per tick. +func (s *Service) ReconcileAccount(ctx context.Context, accountID string, list InstanceVCPULister) error { + if s == nil || !s.limits.Enabled || accountID == utils.GlobalAccountID { + return nil + } + before := s.counterRevisions(ctx, []string{accountID}) + totals, complete, err := list(ctx) + if err != nil { + return fmt.Errorf("quota reconcile %s: instance scan: %w", accountID, err) + } + if err := s.reconcileVCPU(ctx, accountID, totals[accountID], complete, before[accountID]); err != nil { + return fmt.Errorf("quota reconcile %s: counter overwrite: %w", accountID, err) + } + return nil +} + +// counterRevisions reads the current revision of each account's counter. A +// read that fails reports revision zero, which reads as "no counter yet" and +// makes the later write a create that a concurrent charge wins. +func (s *Service) counterRevisions(ctx context.Context, accountIDs []string) map[string]uint64 { + out := make(map[string]uint64, len(accountIDs)) + for _, accountID := range accountIDs { + _, revision, err := s.readVCPU(ctx, accountID) + if err != nil { + slog.Warn("quota reconcile: counter revision read failed", "account", accountID, "err", err) + continue + } + out[accountID] = revision + } + return out +} + // sumReservationVCPUs totals the catalog vCPUs of every non-terminal instance // across the reservations. Terminated and shutting-down instances no longer // exist for quota and are skipped; an unknown instance type contributes nothing. diff --git a/spinifex/handlers/quota/reconcile_test.go b/spinifex/handlers/quota/reconcile_test.go index a69127574..d6a5f5532 100644 --- a/spinifex/handlers/quota/reconcile_test.go +++ b/spinifex/handlers/quota/reconcile_test.go @@ -5,36 +5,18 @@ import ( "errors" "testing" - "github.com/aws/aws-sdk-go/aws" - "github.com/aws/aws-sdk-go/service/ec2" "github.com/mulgadc/spinifex/spinifex/utils" ) -// reservation builds a single reservation whose instances carry the given type -// and state, so a test can describe an account's running-plus-stopped shape. -func reservation(instances ...*ec2.Instance) *ec2.Reservation { - return &ec2.Reservation{Instances: instances} -} - -// instance builds one ec2.Instance with a type and state name; an empty state -// leaves State nil (treated as live). -func instance(instanceType, stateName string) *ec2.Instance { - inst := &ec2.Instance{InstanceType: aws.String(instanceType)} - if stateName != "" { - inst.State = &ec2.InstanceState{Name: aws.String(stateName)} - } - return inst -} - -// staticLister returns an InstanceLister serving a fixed per-account map as a -// complete sweep, and records which accounts were described so a test can assert -// the system account is never swept. -func staticLister(byAccount map[string][]*ec2.Reservation, seen *[]string) InstanceLister { - return func(accountID string) ([]*ec2.Reservation, bool, error) { - if seen != nil { - *seen = append(*seen, accountID) +// staticTotals serves a fixed per-account total as one scan, and records how +// many times it was called so a test can assert the sweep reads the record +// space once rather than once per account. +func staticTotals(byAccount map[string]int, complete bool, calls *int) InstanceVCPULister { + return func(context.Context) (map[string]int, bool, error) { + if calls != nil { + *calls++ } - return byAccount[accountID], true, nil + return byAccount, complete, nil } } @@ -50,58 +32,51 @@ func TestReconcileCorrectsOverCount(t *testing.T) { if err := s.AddVCPU(t.Context(), testAccount, 16); err != nil { t.Fatalf("seed AddVCPU: %v", err) } - lister := staticLister(map[string][]*ec2.Reservation{ - testAccount: {reservation( - instance("m5.xlarge", ec2.InstanceStateNameRunning), - instance("t3.micro", ec2.InstanceStateNameStopped), - )}, - }, nil) - if err := s.Reconcile(context.Background(), accountList(testAccount), lister); err != nil { + if err := s.Reconcile(context.Background(), accountList(testAccount), + staticTotals(map[string]int{testAccount: 6}, true, nil)); err != nil { t.Fatalf("Reconcile: %v", err) } assertCounter(t, s, testAccount, 6) } -// A terminated instance still lingering in the terminated KV must not be charged: -// the counter drops to only the surviving running instance's vCPUs. -func TestReconcileCorrectsStaleTermination(t *testing.T) { - s := newVCPUService(t, Limits{Enabled: true, VCPUs: 100}) - - if err := s.AddVCPU(t.Context(), testAccount, 6); err != nil { - t.Fatalf("seed AddVCPU: %v", err) - } - lister := staticLister(map[string][]*ec2.Reservation{ - testAccount: {reservation( - instance("t3.micro", ec2.InstanceStateNameRunning), // counts: 2 - instance("m5.xlarge", ec2.InstanceStateNameTerminated), // excluded - instance("c5.large", ec2.InstanceStateNameShuttingDown), // excluded - )}, - }, nil) - - if err := s.Reconcile(context.Background(), accountList(testAccount), lister); err != nil { - t.Fatalf("Reconcile: %v", err) - } - assertCounter(t, s, testAccount, 2) -} - -// An account that has terminated everything is zeroed. -func TestReconcileZeroesEmptyAccount(t *testing.T) { +// An account that has terminated everything is absent from the scan entirely, +// which is the only signal that it now holds nothing, so it must be zeroed +// rather than skipped. +func TestReconcileZeroesAccountAbsentFromScan(t *testing.T) { s := newVCPUService(t, Limits{Enabled: true, VCPUs: 100}) if err := s.AddVCPU(t.Context(), testAccount, 8); err != nil { t.Fatalf("seed AddVCPU: %v", err) } - lister := staticLister(map[string][]*ec2.Reservation{}, nil) // describes empty - if err := s.Reconcile(context.Background(), accountList(testAccount), lister); err != nil { + if err := s.Reconcile(context.Background(), accountList(testAccount), + staticTotals(map[string]int{}, true, nil)); err != nil { t.Fatalf("Reconcile: %v", err) } assertCounter(t, s, testAccount, 0) } -// The system account is exempt: it is skipped entirely (never described) and any -// counter parked under its key is left untouched. +// One scan covers every account, where the fan-out cost one describe each. +func TestReconcileScansOnceForEveryAccount(t *testing.T) { + s := newVCPUService(t, Limits{Enabled: true, VCPUs: 100}) + + const a, b, c = "111111111111", "222222222222", "333333333333" + calls := 0 + if err := s.Reconcile(context.Background(), accountList(a, b, c), + staticTotals(map[string]int{a: 2, b: 4, c: 8}, true, &calls)); err != nil { + t.Fatalf("Reconcile: %v", err) + } + if calls != 1 { + t.Fatalf("scans = %d, want 1 for three accounts", calls) + } + assertCounter(t, s, a, 2) + assertCounter(t, s, b, 4) + assertCounter(t, s, c, 8) +} + +// The system account is exempt: any counter parked under its key is left +// untouched even when the scan reports a total against it. func TestReconcileSkipsSystemAccount(t *testing.T) { s := newVCPUService(t, Limits{Enabled: true, VCPUs: 100}) @@ -109,93 +84,137 @@ func TestReconcileSkipsSystemAccount(t *testing.T) { if _, err := s.usage.PutString(t.Context(), utils.GlobalAccountID, "42"); err != nil { t.Fatalf("seed system counter: %v", err) } - var seen []string - lister := staticLister(map[string][]*ec2.Reservation{ - testAccount: {reservation(instance("t3.micro", ec2.InstanceStateNameRunning))}, - }, &seen) - if err := s.Reconcile(context.Background(), accountList(utils.GlobalAccountID, testAccount), lister); err != nil { + if err := s.Reconcile(context.Background(), accountList(utils.GlobalAccountID, testAccount), + staticTotals(map[string]int{utils.GlobalAccountID: 99, testAccount: 2}, true, nil)); err != nil { t.Fatalf("Reconcile: %v", err) } - for _, id := range seen { - if id == utils.GlobalAccountID { - t.Fatalf("system account was described, want skipped") - } - } assertCounter(t, s, utils.GlobalAccountID, 42) assertCounter(t, s, testAccount, 2) } -// A describe failure on one account is logged and the pass continues for the -// rest; the first error is returned and the failed account's counter is left -// untouched for the next pass to retry. -func TestReconcileContinuesOnDescribeError(t *testing.T) { +// The scan is one read covering every account, so a failed read is not a +// per-account outage that the pass can continue past: no counter may move, and +// the error is returned for the next pass to retry. +func TestReconcileScanFailureLeavesEveryCounterUntouched(t *testing.T) { s := newVCPUService(t, Limits{Enabled: true, VCPUs: 100}) - const bad, good = "111111111111", "222222222222" - if err := s.AddVCPU(t.Context(), bad, 8); err != nil { - t.Fatalf("seed bad: %v", err) + const a, b = "111111111111", "222222222222" + if err := s.AddVCPU(t.Context(), a, 8); err != nil { + t.Fatalf("seed a: %v", err) } - sentinel := errors.New("describe boom") - lister := func(accountID string) ([]*ec2.Reservation, bool, error) { - if accountID == bad { - return nil, false, sentinel - } - return []*ec2.Reservation{reservation(instance("m5.xlarge", ec2.InstanceStateNameRunning))}, true, nil + if err := s.AddVCPU(t.Context(), b, 4); err != nil { + t.Fatalf("seed b: %v", err) } + sentinel := errors.New("scan boom") + failing := func(context.Context) (map[string]int, bool, error) { return nil, false, sentinel } - err := s.Reconcile(context.Background(), accountList(bad, good), lister) + err := s.Reconcile(context.Background(), accountList(a, b), failing) if !errors.Is(err, sentinel) { t.Fatalf("Reconcile err = %v, want %v", err, sentinel) } - assertCounter(t, s, bad, 8) // unchanged: left for the next pass - assertCounter(t, s, good, 4) + assertCounter(t, s, a, 8) + assertCounter(t, s, b, 4) } -// An incomplete sweep (a node down or a failed bucket query) must never lower a -// counter: a short count is dropped and the counter left for the next clean pass. -// A higher observed count still raises it, since those instances do exist. +// An incomplete sweep must never lower a counter: a short count is dropped and +// the counter left for the next clean pass. A higher observed count still +// raises it, since those instances do exist. func TestReconcileIncompleteSweepDoesNotLower(t *testing.T) { s := newVCPUService(t, Limits{Enabled: true, VCPUs: 100}) - // Counter holds 8 (two m5.xlarge across two nodes); a partial sweep sees only - // one node's 4 vCPUs and reports complete=false. if err := s.AddVCPU(t.Context(), testAccount, 8); err != nil { t.Fatalf("seed AddVCPU: %v", err) } - partial := func(accountID string) ([]*ec2.Reservation, bool, error) { - return []*ec2.Reservation{reservation(instance("m5.xlarge", ec2.InstanceStateNameRunning))}, false, nil - } - if err := s.Reconcile(context.Background(), accountList(testAccount), partial); err != nil { + if err := s.Reconcile(context.Background(), accountList(testAccount), + staticTotals(map[string]int{testAccount: 4}, false, nil)); err != nil { t.Fatalf("Reconcile partial: %v", err) } assertCounter(t, s, testAccount, 8) // unchanged: a partial sweep cannot lower - // A partial sweep that observes more than the counter may still raise it. - higher := func(accountID string) ([]*ec2.Reservation, bool, error) { - return []*ec2.Reservation{reservation( - instance("m5.xlarge", ec2.InstanceStateNameRunning), - instance("m5.xlarge", ec2.InstanceStateNameRunning), - instance("m5.xlarge", ec2.InstanceStateNameRunning), - )}, false, nil - } - if err := s.Reconcile(context.Background(), accountList(testAccount), higher); err != nil { + if err := s.Reconcile(context.Background(), accountList(testAccount), + staticTotals(map[string]int{testAccount: 12}, false, nil)); err != nil { t.Fatalf("Reconcile raise: %v", err) } - assertCounter(t, s, testAccount, 12) // raised to the observed 3 x 4 = 12 + assertCounter(t, s, testAccount, 12) } -// A disabled service never reaches the KV: Reconcile is a no-op even with a nil -// lister and account enumerator that would otherwise panic. +// The guard that makes an absolute overwrite safe: a charge landing between the +// revision read and the write moves the revision, so the charge stands and the +// scan is dropped. Without it the overwrite would silently undo the launch. +func TestReconcileDoesNotClobberAChargeLandingMidScan(t *testing.T) { + s := newVCPUService(t, Limits{Enabled: true, VCPUs: 100}) + + if err := s.AddVCPU(t.Context(), testAccount, 4); err != nil { + t.Fatalf("seed AddVCPU: %v", err) + } + // The scan sees the account holding only the seeded 4 vCPUs; a launch is + // charged while it runs, so the counter it would write is already stale. + racing := func(context.Context) (map[string]int, bool, error) { + if err := s.AddVCPU(context.Background(), testAccount, 2); err != nil { + return nil, false, err + } + return map[string]int{testAccount: 4}, true, nil + } + + if err := s.Reconcile(context.Background(), accountList(testAccount), racing); err != nil { + t.Fatalf("Reconcile: %v", err) + } + assertCounter(t, s, testAccount, 6) // the charge, not the scan's stale 4 +} + +// ReconcileAccount is the per-key entry point: it settles the account it was +// given and leaves every other account alone, which is what makes a change to +// one instance cost one account's work. +func TestReconcileAccountTouchesOnlyThatAccount(t *testing.T) { + s := newVCPUService(t, Limits{Enabled: true, VCPUs: 100}) + + const mine, other = "111111111111", "222222222222" + if err := s.AddVCPU(t.Context(), mine, 16); err != nil { + t.Fatalf("seed mine: %v", err) + } + if err := s.AddVCPU(t.Context(), other, 8); err != nil { + t.Fatalf("seed other: %v", err) + } + + if err := s.ReconcileAccount(context.Background(), mine, + staticTotals(map[string]int{mine: 6, other: 0}, true, nil)); err != nil { + t.Fatalf("ReconcileAccount: %v", err) + } + assertCounter(t, s, mine, 6) + assertCounter(t, s, other, 8) // untouched, though the scan saw it hold nothing +} + +func TestReconcileAccountSkipsSystemAccount(t *testing.T) { + s := newVCPUService(t, Limits{Enabled: true, VCPUs: 100}) + + if _, err := s.usage.PutString(t.Context(), utils.GlobalAccountID, "42"); err != nil { + t.Fatalf("seed system counter: %v", err) + } + if err := s.ReconcileAccount(context.Background(), utils.GlobalAccountID, + staticTotals(map[string]int{utils.GlobalAccountID: 0}, true, nil)); err != nil { + t.Fatalf("ReconcileAccount: %v", err) + } + assertCounter(t, s, utils.GlobalAccountID, 42) +} + +// A disabled service never reaches the KV: both entry points are a no-op even +// with a nil lister and account enumerator that would otherwise panic. func TestReconcileDisabledNoop(t *testing.T) { s := New(Limits{Enabled: false}, nil) if err := s.Reconcile(context.Background(), nil, nil); err != nil { t.Fatalf("disabled Reconcile = %v, want nil", err) } + if err := s.ReconcileAccount(context.Background(), testAccount, nil); err != nil { + t.Fatalf("disabled ReconcileAccount = %v, want nil", err) + } var nilService *Service if err := nilService.Reconcile(context.Background(), nil, nil); err != nil { t.Fatalf("nil Reconcile = %v, want nil", err) } + if err := nilService.ReconcileAccount(context.Background(), testAccount, nil); err != nil { + t.Fatalf("nil ReconcileAccount = %v, want nil", err) + } } func assertCounter(t *testing.T, s *Service, accountID string, want int) { diff --git a/spinifex/handlers/quota/records.go b/spinifex/handlers/quota/records.go new file mode 100644 index 000000000..5d3f8d1e9 --- /dev/null +++ b/spinifex/handlers/quota/records.go @@ -0,0 +1,85 @@ +package handlers_quota + +import ( + "context" + "encoding/json" + + "github.com/mulgadc/spinifex/spinifex/instancetypes" + "github.com/mulgadc/spinifex/spinifex/kvstore" + "github.com/mulgadc/spinifex/spinifex/utils" + "github.com/mulgadc/spinifex/spinifex/vm" +) + +// InstanceVCPULister totals the vCPUs each account currently holds, in one +// read of the whole instance record space rather than one describe per account. +// +// An account with no instances is absent from the map rather than present as +// zero: the caller knows which accounts exist and reads a missing one as zero, +// 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. +type InstanceVCPULister func(ctx context.Context) (perAccount map[string]int, complete bool, err error) + +// 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 +// 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 +// coalesces, against a fan-out that ran per account on every tick. +func RecordVCPULister(records *kvstore.Store[vm.InstanceRecord], prefix string) InstanceVCPULister { + return func(ctx context.Context) (map[string]int, bool, error) { + list, err := records.List(ctx, prefix) + if err != nil { + return nil, false, err + } + totals := make(map[string]int, len(list)) + for i := range list { + record := &list[i] + accountID := record.Metadata.AccountID + if accountID == "" || accountID == utils.GlobalAccountID { + continue + } + if recordIsTerminal(record) { + continue + } + // An unknown instance type contributes nothing, matching the sum + // the fan-out produced. + if vcpus, ok := instancetypes.DefaultVCPUs(record.Spec.InstanceType); ok { + totals[accountID] += vcpus + } + } + return totals, true, nil + } +} + +// recordIsTerminal reports whether a record has left the counted set. Pending, +// running, stopping and stopped all exist and are charged; only shutting-down +// and terminated are excluded, so a reaped instance frees quota. This mirrors +// isTerminalState, which reads the same states off the AWS-shaped projection. +func recordIsTerminal(record *vm.InstanceRecord) bool { + switch record.Status.Status { + case vm.StateShuttingDown, vm.StateTerminated: + return true + default: + return false + } +} + +// AccountForRecord returns the account a raw instance record belongs to, for +// mapping a watch update onto the unit of work quota recomputes. ok is false +// when the value does not decode or carries no account, which the caller reads +// as "cannot attribute" and answers with a whole-set pass. +func AccountForRecord(value []byte) (accountID string, ok bool) { + var record vm.InstanceRecord + if err := json.Unmarshal(value, &record); err != nil { + return "", false + } + if record.Metadata.AccountID == "" || record.Metadata.AccountID == utils.GlobalAccountID { + return "", false + } + return record.Metadata.AccountID, true +} diff --git a/spinifex/handlers/quota/records_test.go b/spinifex/handlers/quota/records_test.go new file mode 100644 index 000000000..e7315778b --- /dev/null +++ b/spinifex/handlers/quota/records_test.go @@ -0,0 +1,126 @@ +package handlers_quota_test + +import ( + "encoding/json" + "testing" + + handlers_quota "github.com/mulgadc/spinifex/spinifex/handlers/quota" + "github.com/mulgadc/spinifex/spinifex/kvstore" + "github.com/mulgadc/spinifex/spinifex/resource" + "github.com/mulgadc/spinifex/spinifex/testutil" + "github.com/mulgadc/spinifex/spinifex/utils" + "github.com/mulgadc/spinifex/spinifex/vm" + "github.com/stretchr/testify/require" +) + +const recordPrefix = "i." + +func recordStore(t *testing.T) *kvstore.Store[vm.InstanceRecord] { + t.Helper() + _, nc, _ := testutil.StartTestJetStream(t) + return kvstore.New[vm.InstanceRecord](testutil.NewJetStream(t, nc), kvstore.Config{ + Name: "spinifex-instance-state", + History: 1, + Replicas: 1, + }) +} + +func putRecord(t *testing.T, store *kvstore.Store[vm.InstanceRecord], id, accountID, instanceType string, state vm.InstanceState) { + t.Helper() + require.NoError(t, store.Set(t.Context(), recordPrefix+id, &vm.InstanceRecord{ + Metadata: resource.Metadata{Name: id, AccountID: accountID}, + Spec: vm.InstanceSpec{InstanceType: instanceType}, + Status: vm.InstanceStatus{Status: state}, + })) +} + +// The sum is per account, from the record space, in one read. +func TestRecordVCPULister_TotalsPerAccount(t *testing.T) { + store := recordStore(t) + const a, b = "111111111111", "222222222222" + putRecord(t, store, "i-1", a, "m5.xlarge", vm.StateRunning) // 4 + putRecord(t, store, "i-2", a, "t3.micro", vm.StateRunning) // 2 + putRecord(t, store, "i-3", b, "m5.xlarge", vm.StateRunning) // 4 + + totals, complete, err := handlers_quota.RecordVCPULister(store, recordPrefix)(t.Context()) + require.NoError(t, err) + require.True(t, complete) + require.Equal(t, map[string]int{a: 6, b: 4}, totals) +} + +// A stopped instance still exists and is charged; only shutting-down and +// terminated leave the counted set, so a reaped instance frees quota. +func TestRecordVCPULister_ChargesStoppedButNotTerminal(t *testing.T) { + store := recordStore(t) + const account = "111111111111" + putRecord(t, store, "i-1", account, "t3.micro", vm.StateRunning) // 2 + putRecord(t, store, "i-2", account, "t3.micro", vm.StateStopped) // 2 + putRecord(t, store, "i-3", account, "t3.micro", vm.StatePending) // 2 + putRecord(t, store, "i-4", account, "m5.xlarge", vm.StateShuttingDown) // excluded + putRecord(t, store, "i-5", account, "m5.xlarge", vm.StateTerminated) // excluded + + totals, _, err := handlers_quota.RecordVCPULister(store, recordPrefix)(t.Context()) + require.NoError(t, err) + require.Equal(t, map[string]int{account: 6}, totals) +} + +// The system account is never charged, and an account holding nothing is absent +// rather than present as zero — the caller's account list is what zeroes it. +func TestRecordVCPULister_SkipsSystemAccountAndReportsNothingAsAbsent(t *testing.T) { + store := recordStore(t) + putRecord(t, store, "i-1", utils.GlobalAccountID, "m5.xlarge", vm.StateRunning) + + totals, complete, err := handlers_quota.RecordVCPULister(store, recordPrefix)(t.Context()) + require.NoError(t, err) + require.True(t, complete) + require.Empty(t, totals) +} + +// An unknown instance type contributes nothing rather than failing the sweep, +// matching the sum the fan-out produced. +func TestRecordVCPULister_UnknownInstanceTypeContributesNothing(t *testing.T) { + store := recordStore(t) + const account = "111111111111" + putRecord(t, store, "i-1", account, "not-a-real-type", vm.StateRunning) + putRecord(t, store, "i-2", account, "t3.micro", vm.StateRunning) + + totals, _, err := handlers_quota.RecordVCPULister(store, recordPrefix)(t.Context()) + require.NoError(t, err) + require.Equal(t, map[string]int{account: 2}, totals) +} + +func TestAccountForRecord(t *testing.T) { + valid, err := json.Marshal(vm.InstanceRecord{ + Metadata: resource.Metadata{Name: "i-1", AccountID: "111111111111"}, + }) + require.NoError(t, err) + system, err := json.Marshal(vm.InstanceRecord{ + Metadata: resource.Metadata{Name: "i-2", AccountID: utils.GlobalAccountID}, + }) + require.NoError(t, err) + unowned, err := json.Marshal(vm.InstanceRecord{Metadata: resource.Metadata{Name: "i-3"}}) + require.NoError(t, err) + + for _, tc := range []struct { + name string + value []byte + want string + wantOK bool + }{ + {"a record names its account", valid, "111111111111", true}, + // The system account is never charged, so attributing to it would queue + // work that reconcile then skips. + {"the system account is not a unit of work", system, "", false}, + {"a record with no account cannot be attributed", unowned, "", false}, + // A delete tombstone carries no value, which is the case the fallback + // to a whole-set pass exists for. + {"an empty value cannot be attributed", nil, "", false}, + {"a corrupt value cannot be attributed", []byte("{not json"), "", false}, + } { + t.Run(tc.name, func(t *testing.T) { + got, ok := handlers_quota.AccountForRecord(tc.value) + require.Equal(t, tc.wantOK, ok) + require.Equal(t, tc.want, got) + }) + } +} diff --git a/spinifex/handlers/quota/vcpu.go b/spinifex/handlers/quota/vcpu.go index 2ea82d259..fae7ce736 100644 --- a/spinifex/handlers/quota/vcpu.go +++ b/spinifex/handlers/quota/vcpu.go @@ -96,15 +96,48 @@ func (s *Service) setVCPU(ctx context.Context, accountID string, value int) erro }) } +// setVCPUAt lowers or raises accountID's counter to value, but only while the +// counter still stands at sinceRevision. +// +// The guard is what stops a scan clobbering a launch it never saw: ChargeLaunch +// landing between the revision read and this write moves the revision, so the +// charge wins and the scan is dropped for the next pass to settle. Without it +// the absolute overwrite would silently undo the charge and lift the cap. +func (s *Service) setVCPUAt(ctx context.Context, accountID string, value int, sinceRevision uint64) error { + current, revision, err := s.readVCPU(ctx, accountID) + if err != nil { + return err + } + if revision != sinceRevision || current == value { + return nil + } + data, err := json.Marshal(value) + if err != nil { + return err + } + if revision == 0 { + _, err = s.usage.Create(ctx, accountID, data) + } else { + _, err = s.usage.Update(ctx, accountID, data, revision) + } + // Losing the race is the guard working, not a failure: a charge got there + // first and the next pass recomputes against it. + if errors.Is(err, jetstream.ErrKeyExists) || errors.Is(err, jetstream.ErrKeyRevisionMismatch) { + return nil + } + return err +} + // reconcileVCPU writes accountID's counter from a reconcile sweep. A complete -// sweep overwrites unconditionally (the only path that may lower a counter). An -// incomplete sweep — a node down or a failed bucket query — may only raise the -// counter: lowering from a partial view would under-count usage and lift the -// cap, so a short count is dropped and left for the next clean pass. Raising is -// always safe, since the instances actually observed do exist. -func (s *Service) reconcileVCPU(ctx context.Context, accountID string, value int, complete bool) error { +// sweep overwrites (the only path that may lower a counter), guarded on the +// counter not having moved since sinceRevision. An incomplete sweep — a failed +// scan — may only raise the counter: lowering from a partial view would +// under-count usage and lift the cap, so a short count is dropped and left for +// the next clean pass. Raising is always safe, since the instances actually +// observed do exist. +func (s *Service) reconcileVCPU(ctx context.Context, accountID string, value int, complete bool, sinceRevision uint64) error { if complete { - return s.setVCPU(ctx, accountID, value) + return s.setVCPUAt(ctx, accountID, value, sinceRevision) } current, _, err := s.readVCPU(ctx, accountID) if err != nil { diff --git a/spinifex/handlers/quota/vcpu_retype_test.go b/spinifex/handlers/quota/vcpu_retype_test.go index 9e4f07921..103ee733b 100644 --- a/spinifex/handlers/quota/vcpu_retype_test.go +++ b/spinifex/handlers/quota/vcpu_retype_test.go @@ -10,6 +10,23 @@ import ( "github.com/mulgadc/spinifex/spinifex/utils" ) +// reservation builds a single reservation whose instances carry the given type +// and state. The retype gate still reads the AWS-shaped projection, so these +// live here now the vCPU sweep sums from the record space instead. +func reservation(instances ...*ec2.Instance) *ec2.Reservation { + return &ec2.Reservation{Instances: instances} +} + +// instance builds one ec2.Instance with a type and state name; an empty state +// leaves State nil (treated as live). +func instance(instanceType, stateName string) *ec2.Instance { + inst := &ec2.Instance{InstanceType: aws.String(instanceType)} + if stateName != "" { + inst.State = &ec2.InstanceState{Name: aws.String(stateName)} + } + return inst +} + // staticResolver returns an InstanceTypeResolver serving one fixed answer, so a // test drives EnforceRetype without a DescribeInstances round trip. func staticResolver(instanceType string, ok bool, err error) InstanceTypeResolver { diff --git a/spinifex/services/awsgw/awsgw.go b/spinifex/services/awsgw/awsgw.go index ede4c70f2..79ae0d2f2 100644 --- a/spinifex/services/awsgw/awsgw.go +++ b/spinifex/services/awsgw/awsgw.go @@ -19,6 +19,7 @@ import ( "github.com/mulgadc/bluebottle/pkg/tlsconfig" "github.com/mulgadc/spinifex/spinifex/admin" "github.com/mulgadc/spinifex/spinifex/config" + "github.com/mulgadc/spinifex/spinifex/daemon" "github.com/mulgadc/spinifex/spinifex/gateway" gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock" gateway_ecr "github.com/mulgadc/spinifex/spinifex/gateway/ecr" @@ -29,9 +30,12 @@ import ( handlers_ochrevector "github.com/mulgadc/spinifex/spinifex/handlers/ochrevector" handlers_quota "github.com/mulgadc/spinifex/spinifex/handlers/quota" handlers_sts "github.com/mulgadc/spinifex/spinifex/handlers/sts" + "github.com/mulgadc/spinifex/spinifex/kvstore" "github.com/mulgadc/spinifex/spinifex/network/reconcile" "github.com/mulgadc/spinifex/spinifex/objectstore" + "github.com/mulgadc/spinifex/spinifex/reconciler" "github.com/mulgadc/spinifex/spinifex/utils" + "github.com/mulgadc/spinifex/spinifex/vm" "github.com/nats-io/nats.go" "github.com/nats-io/nats.go/jetstream" toml "github.com/pelletier/go-toml/v2" @@ -548,11 +552,7 @@ func launchService(config *config.ClusterConfig) error { // terminations free quota. Started only when quotas are enabled so default-off // gateways spin no ticker. if quotaCfg.Enabled { - // Sweep against the configured node total, not the live-active count: a - // node that is down must make the sweep incomplete so reconcile leaves - // the counter alone rather than lowering it from a partial view. - expectedNodes := func() int { return len(config.Nodes) } - go runQuotaReconcile(janitorCtx, gw.Quota, natsConn, activeAccountIDs(iamService), expectedNodes) + go runQuotaReconcile(janitorCtx, gw.Quota, natsConn, js, activeAccountIDs(iamService), len(config.Nodes)) } // Bedrock RPM enforcement is local and immediate (never gated on @@ -593,37 +593,60 @@ func launchService(config *config.ClusterConfig) error { return nil } -// runQuotaReconcile drives the per-account vCPU reconcile: a startup pass plus a -// ReconcileInterval ticker, each guarded by the dedicated quota reconcile leader -// lock so exactly one gateway sweeps at a time across a multi-gateway deployment. -// The lock is distinct from vpcd's network-reconcile lock so the two loops never -// block each other. It runs until ctx is cancelled. -func runQuotaReconcile(ctx context.Context, quota *handlers_quota.Service, natsConn *nats.Conn, accounts handlers_quota.AccountLister, expectedNodes func() int) { +// runQuotaReconcile drives the per-account vCPU reconcile off changes to the +// instance record space rather than a ticker: a change to one instance costs +// that one account's recompute, and the resync sweeps every account as the +// backstop. Each pass is guarded by the dedicated quota reconcile leader lock +// so exactly one gateway sweeps at a time across a multi-gateway deployment. +// The lock is distinct from vpcd's network-reconcile lock so the two loops +// never block each other. It runs until ctx is cancelled. +func runQuotaReconcile(ctx context.Context, quota *handlers_quota.Service, natsConn *nats.Conn, + js jetstream.JetStream, accounts handlers_quota.AccountLister, replicas int) { holder, _ := os.Hostname() - list := handlers_quota.NATSInstanceLister(natsConn, expectedNodes) + cfg := kvstore.Config{Name: daemon.InstanceStateBucket, History: 1, Replicas: replicas} + list := handlers_quota.RecordVCPULister( + kvstore.New[vm.InstanceRecord](js, cfg), daemon.InstanceRecordPrefix) - runPass := func() { + // Leadership is taken per pass rather than held, matching the other + // reconcile loops: the lock keeps two gateways off the same counters, it is + // not what keeps this loop alive. + underLeader := func(run func() error) error { release, elected := reconcile.AcquireLeader(ctx, natsConn, handlers_quota.KVBucketQuotaReconcile, holder) if !elected { - return + return nil } defer release() - if err := quota.Reconcile(ctx, accounts, list); err != nil { - slog.Warn("quota reconcile pass failed", "err", err) - } + return run() } - runPass() - ticker := time.NewTicker(handlers_quota.ReconcileInterval) - defer ticker.Stop() - for { - select { - case <-ctx.Done(): - return - case <-ticker.C: - runPass() - } + reconciler.Run(ctx, reconciler.Config{ + Name: "quota", + Sources: []reconciler.Source{reconciler.Fixed(kvstore.NewBucket(js, cfg), daemon.InstanceRecordPrefix+"*")}, + Reconcile: func(ctx context.Context) error { + return underLeader(func() error { return quota.Reconcile(ctx, accounts, list) }) + }, + ReconcileKey: func(ctx context.Context, accountID string) error { + return underLeader(func() error { return quota.ReconcileAccount(ctx, accountID, list) }) + }, + KeyFor: quotaKeyFor, + Resync: handlers_quota.ReconcileInterval, + }) +} + +// quotaKeyFor maps an instance record update onto the account whose counter it +// dirties, which is quota's unit of work: fifty instances changing in one +// account are one recompute. A delete carries no value to read the account +// from, so it reports ok=false and the loop falls back to a whole-set pass — +// which is also the pass that lowers the counter for the instance that went. +func quotaKeyFor(entry jetstream.KeyValueEntry) ([]string, bool) { + if entry.Operation() != jetstream.KeyValuePut { + return nil, false } + accountID, ok := handlers_quota.AccountForRecord(entry.Value()) + if !ok { + return nil, false + } + return []string{accountID}, true } // accountLister is the slice of IAMService the lifecycle sweeper needs. diff --git a/spinifex/services/awsgw/quota_reconcile_test.go b/spinifex/services/awsgw/quota_reconcile_test.go index 7b058f7a6..c9f770c1b 100644 --- a/spinifex/services/awsgw/quota_reconcile_test.go +++ b/spinifex/services/awsgw/quota_reconcile_test.go @@ -7,28 +7,14 @@ import ( "testing" "time" - "github.com/aws/aws-sdk-go/aws" - "github.com/aws/aws-sdk-go/service/ec2" handlers_quota "github.com/mulgadc/spinifex/spinifex/handlers/quota" + "github.com/mulgadc/spinifex/spinifex/resource" "github.com/mulgadc/spinifex/spinifex/testutil" - "github.com/nats-io/nats.go" + "github.com/mulgadc/spinifex/spinifex/vm" "github.com/nats-io/nats.go/jetstream" "github.com/stretchr/testify/require" ) -// respondJSON subscribes to subject and replies with out to every request, so -// the reconcile describe fan-out and KV bucket queries resolve without daemons. -func respondJSON(t *testing.T, nc *nats.Conn, subject string, out *ec2.DescribeInstancesOutput) { - t.Helper() - data, err := json.Marshal(out) - require.NoError(t, err) - sub, err := nc.Subscribe(subject, func(msg *nats.Msg) { - _ = msg.Respond(data) - }) - require.NoError(t, err) - t.Cleanup(func() { _ = sub.Unsubscribe() }) -} - func TestOpenAccountUsageBucketPreservesExistingConfiguration(t *testing.T) { _, nc, _ := testutil.StartTestJetStream(t) js := testutil.NewJetStream(t, nc) @@ -45,37 +31,60 @@ func TestOpenAccountUsageBucketPreservesExistingConfiguration(t *testing.T) { require.EqualValues(t, 5, status.History()) } -// TestRunQuotaReconcile drives the leader-locked reconcile loop end to end: the -// startup pass elects this gateway, describes the account's instances over NATS -// via the production NATSInstanceLister, and writes the recomputed vCPU counter. -func TestRunQuotaReconcile(t *testing.T) { +// putInstanceRecord writes one instance record at the key the sweep reads, so a +// test can state an account's holdings without a node or a describe responder. +func putInstanceRecord(t *testing.T, records jetstream.KeyValue, id, accountID, instanceType string, state vm.InstanceState) { + t.Helper() + data, err := json.Marshal(vm.InstanceRecord{ + Metadata: resource.Metadata{Name: id, AccountID: accountID}, + Spec: vm.InstanceSpec{InstanceType: instanceType}, + Status: vm.InstanceStatus{Status: state}, + }) + require.NoError(t, err) + _, err = records.Put(t.Context(), "i."+id, data) + require.NoError(t, err) +} + +// instanceRecordBucket creates the shared instance-state bucket the sweep reads +// and the loop watches. +func instanceRecordBucket(t *testing.T, js jetstream.JetStream) jetstream.KeyValue { + t.Helper() + kv, err := js.CreateKeyValue(t.Context(), jetstream.KeyValueConfig{ + Bucket: "spinifex-instance-state", + History: 1, + }) + require.NoError(t, err) + return kv +} + +// The startup pass recomputes the counter from the instance record space. No +// describe responder is registered anywhere in this test: if the loop still +// fanned out over NATS it would have nothing to fan out to. +func TestRunQuotaReconcileCountsFromTheRecordSpace(t *testing.T) { _, nc, _ := testutil.StartTestJetStream(t) - bucket, err := openAccountUsageBucket(t.Context(), testutil.NewJetStream(t, nc), 1) + js := testutil.NewJetStream(t, nc) + bucket, err := openAccountUsageBucket(t.Context(), js, 1) require.NoError(t, err) quota := handlers_quota.New(handlers_quota.Limits{Enabled: true, VCPUs: 100}, bucket) const account = "123456789012" - // Running fan-out reports one m5.xlarge (4 vCPUs); the KV buckets are empty. - respondJSON(t, nc, "ec2.DescribeInstances", &ec2.DescribeInstancesOutput{ - Reservations: []*ec2.Reservation{{Instances: []*ec2.Instance{{ - InstanceType: aws.String("m5.xlarge"), - State: &ec2.InstanceState{Name: aws.String(ec2.InstanceStateNameRunning)}, - }}}}, - }) - respondJSON(t, nc, "ec2.DescribeStoppedInstances", &ec2.DescribeInstancesOutput{}) - respondJSON(t, nc, "ec2.DescribeTerminatedInstances", &ec2.DescribeInstancesOutput{}) + records := instanceRecordBucket(t, js) + putInstanceRecord(t, records, "i-aaa", account, "m5.xlarge", vm.StateRunning) // 4 + putInstanceRecord(t, records, "i-bbb", account, "t3.micro", vm.StateStopped) // 2 + // Terminal instances no longer exist for quota. + putInstanceRecord(t, records, "i-ccc", account, "m5.xlarge", vm.StateTerminated) accounts := func() ([]string, error) { return []string{account}, nil } ctx, cancel := context.WithCancel(context.Background()) done := make(chan struct{}) go func() { - runQuotaReconcile(ctx, quota, nc, accounts, func() int { return 1 }) + runQuotaReconcile(ctx, quota, nc, js, accounts, 1) close(done) }() require.Eventually(t, func() bool { - return readUsageCounter(t, bucket, account) == 4 - }, 5*time.Second, 20*time.Millisecond, "counter should reconcile to 4 vCPUs") + return readUsageCounter(t, bucket, account) == 6 + }, 10*time.Second, 20*time.Millisecond, "counter should reconcile to 6 vCPUs from the records") cancel() select { @@ -85,6 +94,76 @@ func TestRunQuotaReconcile(t *testing.T) { } } +// A launch is a record write, and the loop is woken by it rather than by a +// tick: the counter moves well inside the resync interval. +func TestRunQuotaReconcileFollowsARecordChange(t *testing.T) { + _, nc, _ := testutil.StartTestJetStream(t) + js := testutil.NewJetStream(t, nc) + bucket, err := openAccountUsageBucket(t.Context(), js, 1) + require.NoError(t, err) + quota := handlers_quota.New(handlers_quota.Limits{Enabled: true, VCPUs: 100}, bucket) + + const account = "123456789012" + records := instanceRecordBucket(t, js) + putInstanceRecord(t, records, "i-aaa", account, "m5.xlarge", vm.StateRunning) + + accounts := func() ([]string, error) { return []string{account}, nil } + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + go runQuotaReconcile(ctx, quota, nc, js, accounts, 1) + + require.Eventually(t, func() bool { + return readUsageCounter(t, bucket, account) == 4 + }, 10*time.Second, 20*time.Millisecond, "startup pass should settle at 4") + + putInstanceRecord(t, records, "i-bbb", account, "t3.micro", vm.StateRunning) + require.Eventually(t, func() bool { + return readUsageCounter(t, bucket, account) == 6 + }, 10*time.Second, 20*time.Millisecond, "the record write should wake the loop and raise to 6") + + // A terminate removes the record. The account cannot be read off a + // tombstone, so this is the whole-set fallback rather than a per-key pass. + require.NoError(t, records.Delete(t.Context(), "i.i-bbb")) + require.Eventually(t, func() bool { + return readUsageCounter(t, bucket, account) == 4 + }, 10*time.Second, 20*time.Millisecond, "the delete should lower the counter back to 4") +} + +// A record change recomputes only the account that owns it. The drifted counter +// on the other account is what proves it: a whole-set pass would correct it too, +// and the resync that eventually will is a full interval away. +func TestRunQuotaReconcileRecomputesOnlyTheChangedAccount(t *testing.T) { + _, nc, _ := testutil.StartTestJetStream(t) + js := testutil.NewJetStream(t, nc) + bucket, err := openAccountUsageBucket(t.Context(), js, 1) + require.NoError(t, err) + quota := handlers_quota.New(handlers_quota.Limits{Enabled: true, VCPUs: 100}, bucket) + + const changed, untouched = "111111111111", "222222222222" + records := instanceRecordBucket(t, js) + putInstanceRecord(t, records, "i-aaa", changed, "m5.xlarge", vm.StateRunning) + + accounts := func() ([]string, error) { return []string{changed, untouched}, nil } + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + go runQuotaReconcile(ctx, quota, nc, js, accounts, 1) + + require.Eventually(t, func() bool { + return readUsageCounter(t, bucket, changed) == 4 + }, 10*time.Second, 20*time.Millisecond, "startup pass should settle the first account") + + // Drift the other account after the startup pass has been and gone. + require.NoError(t, quota.AddVCPU(t.Context(), untouched, 32)) + + putInstanceRecord(t, records, "i-bbb", changed, "t3.micro", vm.StateRunning) + require.Eventually(t, func() bool { + return readUsageCounter(t, bucket, changed) == 6 + }, 10*time.Second, 20*time.Millisecond, "the changed account should recompute to 6") + + require.Equal(t, 32, readUsageCounter(t, bucket, untouched), + "a change to one account must not recompute another; the resync is what eventually corrects it") +} + // readUsageCounter reads an account's integer vCPU counter from the usage bucket, // returning -1 while the key is still absent so Eventually keeps polling. func readUsageCounter(t *testing.T, bucket jetstream.KeyValue, account string) int {