feat(quota): add unwired owner-side quota override cleanup (#1240)

Pin that account teardown removes the usage counter but leaves the per-account quota override, so a reused account ID inherits raised or Unlimited limits. Add an unwired, idempotent CleanupAccount on the quota service: a present override key is removed by a revision-guarded delete and reported deleted with its revision; a missing key is reported only as observed-absent-at-scan.
2026-10-23 20:57:02 +00:00 · 2026-10-23 11:56:08 +10:00
parent bdfa01459c
commit fb0f8e9291
3 changed files with 338 additions and 0 deletions
@@ -0,0 +1,61 @@
+package accountteardown
+
+//test:in-package — drives the real final-stage reapers against the real quota service to pin that teardown leaves the override bucket alone.
+
+import (
+	"testing"
+
+	"github.com/mulgadc/spinifex/internal/testkit"
+	"github.com/mulgadc/spinifex/spinifex/domains/admission/quota"
+	"github.com/mulgadc/spinifex/spinifex/foundation/state/kvutil"
+	handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam"
+	"github.com/nats-io/nats.go/jetstream"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// No reaper in this package references the quota override bucket or calls an
+// override delete. Running the real engine and the real StageAccount reapers
+// against an account with a raised vCPU override, with usage wired, shows the
+// consequence: the override outlives the account record it belongs to, while
+// the usage counter the reapers do own is removed.
+func TestCurrentBehaviour_AccountTermination_QuotaOverridesSurviveTeardown(t *testing.T) {
+	_, _, js := testutil.StartTestJetStream(t)
+	ctx := testCtx(t)
+
+	usage, err := kvutil.GetOrCreateBucket(ctx, js, quota.KVBucketAccountUsage, 1)
+	require.NoError(t, err)
+	_, err = usage.Put(ctx, testAccountID, []byte(`{"vcpus":4}`))
+	require.NoError(t, err)
+
+	overridesKV, err := kvutil.GetOrCreateBucket(ctx, js, quota.KVBucketAccountQuota, 1)
+	require.NoError(t, err)
+	svc := quota.New(quota.Limits{Enabled: true, VCPUs: 16}, usage)
+	svc.SetOverrides(overridesKV)
+	unlimited := quota.Unlimited
+	require.NoError(t, svc.PutAccountQuota(ctx, testAccountID, quota.Overrides{VCPUs: &unlimited}, "operator"))
+
+	accounts := newFakeAccounts()
+	iam := &accountIAM{account: &handlers_iam.Account{AccountID: testAccountID, AccountName: testAccountName}}
+	engine := quietEngine(accounts, AccountReapers(iam, nil, usage)...)
+
+	result, err := engine.Teardown(ctx, request())
+	require.NoError(t, err)
+	assert.True(t, result.AccountDeleted)
+	assert.True(t, accounts.deleted)
+	assert.Zero(t, result.StuckCount())
+
+	// The reaper this package does own removed the usage counter, which shows
+	// which store teardown actually clears.
+	_, err = usage.Get(ctx, testAccountID)
+	assert.ErrorIs(t, err, jetstream.ErrKeyNotFound, "teardown's quota-usage reaper must remove the vCPU counter")
+
+	_, err = overridesKV.Get(ctx, testAccountID)
+	require.NoError(t, err, "the override key must still be present after the account is torn down")
+
+	over, effective, err := svc.GetAccountQuota(ctx, testAccountID)
+	require.NoError(t, err)
+	require.NotNil(t, over.VCPUs)
+	assert.Equal(t, quota.Unlimited, *over.VCPUs)
+	assert.Equal(t, quota.Unlimited, effective.VCPUs)
+}
@@ -0,0 +1,93 @@
+package quota
+
+import (
+	"context"
+	"errors"
+	"fmt"
+
+	awserrors "github.com/mulgadc/spinifex/spinifex/foundation/aws/errors"
+	awsidentifiers "github.com/mulgadc/spinifex/spinifex/foundation/aws/identifiers"
+	"github.com/nats-io/nats.go/jetstream"
+)
+
+// OverrideCleanupDisposition is CleanupAccount's outcome for an account's
+// quota override key.
+type OverrideCleanupDisposition string
+
+const (
+	OverrideCleanupDeleted OverrideCleanupDisposition = "deleted"
+	OverrideCleanupAbsent  OverrideCleanupDisposition = "absent"
+	OverrideCleanupFailed  OverrideCleanupDisposition = "failed"
+)
+
+// evidenceRevisionGuardedDelete and evidenceObservedAbsent are the fixed
+// evidence bases CleanupAccount reports; see kvstore/bucket.go:112-120 for why
+// a plain read cannot support a stronger absence claim.
+const (
+	evidenceRevisionGuardedDelete = "revision-guarded-delete"
+	evidenceObservedAbsent        = "observed-absent-at-scan"
+)
+
+// cleanupMaxAttempts bounds retries against a revision that keeps moving
+// under concurrent writes, so cleanup fails closed instead of spinning.
+const cleanupMaxAttempts = 5
+
+// OverrideCleanup is CleanupAccount's result: what happened to an account's
+// quota override key, and the fixed evidence basis for that claim.
+type OverrideCleanup struct {
+	Disposition OverrideCleanupDisposition
+	Evidence    string
+	Revision    uint64
+}
+
+// CleanupAccount removes accountID's quota override key so a reused account ID
+// does not inherit a deleted account's raised or lowered limits. It is
+// unwired: nothing calls it yet.
+func (s *Service) CleanupAccount(ctx context.Context, accountID string) (*OverrideCleanup, error) {
+	if accountID == "" {
+		return nil, awserrors.Errorf(awserrors.ErrorMissingParameter, "accountID is required for quota override cleanup")
+	}
+	if !awsidentifiers.IsAccountID(accountID) {
+		return nil, awserrors.Errorf(awserrors.ErrorInvalidParameterValue, "accountID %q is not a valid AWS account ID", accountID)
+	}
+	if accountID == awsidentifiers.GlobalAccountID {
+		return nil, awserrors.Errorf(awserrors.ErrorInvalidParameterValue, "quota override cleanup must not target the global account")
+	}
+	if s == nil || s.overrides == nil {
+		return &OverrideCleanup{Disposition: OverrideCleanupFailed}, errors.New("quota: override bucket not configured")
+	}
+
+	// accountID is never resolved through any fallback here: cleanup must act
+	// on exactly accountID and never fall back to the global key.
+	for range cleanupMaxAttempts {
+		entry, err := s.overrides.Get(ctx, accountID)
+		if err != nil {
+			if errors.Is(err, jetstream.ErrKeyNotFound) {
+				return &OverrideCleanup{Disposition: OverrideCleanupAbsent, Evidence: evidenceObservedAbsent}, nil
+			}
+			return &OverrideCleanup{Disposition: OverrideCleanupFailed}, fmt.Errorf("get quota override for cleanup: %w", err)
+		}
+		rev := entry.Revision()
+
+		err = s.overrides.Delete(ctx, accountID, jetstream.LastRevision(rev))
+		switch {
+		case err == nil:
+			return &OverrideCleanup{
+				Disposition: OverrideCleanupDeleted,
+				Evidence:    evidenceRevisionGuardedDelete,
+				Revision:    rev,
+			}, nil
+		case errors.Is(err, jetstream.ErrKeyNotFound):
+			return &OverrideCleanup{Disposition: OverrideCleanupAbsent, Evidence: evidenceObservedAbsent}, nil
+		case errors.Is(err, jetstream.ErrKeyRevisionMismatch):
+			// A concurrent write landed between the read and the guarded
+			// delete; retry against whatever revision now holds.
+			continue
+		default:
+			return &OverrideCleanup{Disposition: OverrideCleanupFailed}, fmt.Errorf("delete quota override for cleanup: %w", err)
+		}
+	}
+
+	return &OverrideCleanup{Disposition: OverrideCleanupFailed},
+		fmt.Errorf("quota override cleanup for account: exhausted %d attempts against concurrent writes", cleanupMaxAttempts)
+}
@@ -0,0 +1,184 @@
+//test:in-package — CleanupAccount deletes through the unexported overrides field, and the retry/race behaviour is only observable from inside the package.
+package quota
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	awsidentifiers "github.com/mulgadc/spinifex/spinifex/foundation/aws/identifiers"
+	"github.com/nats-io/nats.go/jetstream"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// overrideCleanupFakeKV wraps a real KeyValue so a test can inject a race
+// between CleanupAccount's read and its guarded delete, or a store failure,
+// deterministically rather than through goroutine timing.
+type overrideCleanupFakeKV struct {
+	jetstream.KeyValue
+
+	deleteCalls  int
+	beforeDelete func(attempt int) error
+	getErr       error
+}
+
+func (f *overrideCleanupFakeKV) Get(ctx context.Context, key string) (jetstream.KeyValueEntry, error) {
+	if f.getErr != nil {
+		return nil, f.getErr
+	}
+	return f.KeyValue.Get(ctx, key)
+}
+
+func (f *overrideCleanupFakeKV) Delete(ctx context.Context, key string, opts ...jetstream.KVDeleteOpt) error {
+	f.deleteCalls++
+	if f.beforeDelete != nil {
+		if err := f.beforeDelete(f.deleteCalls); err != nil {
+			return err
+		}
+	}
+	return f.KeyValue.Delete(ctx, key, opts...)
+}
+
+func TestCleanupAccount_DeletesTheOverrideKey(t *testing.T) {
+	svc, _ := newOverrideService(t)
+	ctx := t.Context()
+	require.NoError(t, svc.PutAccountQuota(ctx, testAccount, Overrides{VCPUs: ptr(32)}, "operator"))
+
+	result, err := svc.CleanupAccount(ctx, testAccount)
+	require.NoError(t, err)
+	assert.Equal(t, OverrideCleanupDeleted, result.Disposition)
+	assert.Equal(t, evidenceRevisionGuardedDelete, result.Evidence)
+	assert.NotZero(t, result.Revision)
+
+	_, err = svc.overrides.Get(ctx, testAccount)
+	assert.ErrorIs(t, err, jetstream.ErrKeyNotFound)
+}
+
+func TestCleanupAccount_RepeatCallIsIdempotent(t *testing.T) {
+	svc, _ := newOverrideService(t)
+	ctx := t.Context()
+	require.NoError(t, svc.PutAccountQuota(ctx, testAccount, Overrides{VCPUs: ptr(32)}, "operator"))
+
+	first, err := svc.CleanupAccount(ctx, testAccount)
+	require.NoError(t, err)
+	assert.Equal(t, OverrideCleanupDeleted, first.Disposition)
+
+	second, err := svc.CleanupAccount(ctx, testAccount)
+	require.NoError(t, err)
+	assert.Equal(t, OverrideCleanupAbsent, second.Disposition)
+	assert.Equal(t, evidenceObservedAbsent, second.Evidence)
+}
+
+func TestCleanupAccount_AbsentKeyIsIdempotent(t *testing.T) {
+	svc, _ := newOverrideService(t)
+	result, err := svc.CleanupAccount(t.Context(), "333333333333")
+	require.NoError(t, err)
+	assert.Equal(t, OverrideCleanupAbsent, result.Disposition)
+	assert.Equal(t, evidenceObservedAbsent, result.Evidence)
+}
+
+func TestCleanupAccount_OtherAccountsUntouched(t *testing.T) {
+	svc, _ := newOverrideService(t)
+	ctx := t.Context()
+	accountA, accountB := testAccount, "222222222222"
+	require.NoError(t, svc.PutAccountQuota(ctx, accountA, Overrides{VCPUs: ptr(32)}, "operator"))
+	require.NoError(t, svc.PutAccountQuota(ctx, accountB, Overrides{VCPUs: ptr(8)}, "operator"))
+
+	result, err := svc.CleanupAccount(ctx, accountA)
+	require.NoError(t, err)
+	assert.Equal(t, OverrideCleanupDeleted, result.Disposition)
+
+	over, _, err := svc.GetAccountQuota(ctx, accountB)
+	require.NoError(t, err)
+	require.NotNil(t, over.VCPUs)
+	assert.Equal(t, 8, *over.VCPUs)
+}
+
+func TestCleanupAccount_RefusesGlobalAccountID(t *testing.T) {
+	svc, _ := newOverrideService(t)
+	result, err := svc.CleanupAccount(t.Context(), awsidentifiers.GlobalAccountID)
+	assert.Error(t, err)
+	assert.Nil(t, result)
+}
+
+func TestCleanupAccount_RefusesEmptyAccountID(t *testing.T) {
+	svc, _ := newOverrideService(t)
+	result, err := svc.CleanupAccount(t.Context(), "")
+	assert.Error(t, err)
+	assert.Nil(t, result)
+}
+
+func TestCleanupAccount_RefusesInvalidAccountID(t *testing.T) {
+	svc, _ := newOverrideService(t)
+	result, err := svc.CleanupAccount(t.Context(), "not-an-account")
+	assert.Error(t, err)
+	assert.Nil(t, result)
+}
+
+func TestCleanupAccount_WithoutBucketFails(t *testing.T) {
+	svc := New(baseLimits(), nil)
+	result, err := svc.CleanupAccount(t.Context(), testAccount)
+	assert.Error(t, err)
+	require.NotNil(t, result)
+	assert.Equal(t, OverrideCleanupFailed, result.Disposition)
+}
+
+// A concurrent Put landing between CleanupAccount's read and its
+// revision-guarded delete must not be lost silently: the delete is rejected
+// and retried against the revision the write left behind.
+func TestCleanupAccount_ConcurrentPutDuringCleanupRetriesAndDeletes(t *testing.T) {
+	svc, _ := newOverrideService(t)
+	ctx := t.Context()
+	require.NoError(t, svc.PutAccountQuota(ctx, testAccount, Overrides{VCPUs: ptr(32)}, "operator"))
+
+	fake := &overrideCleanupFakeKV{KeyValue: svc.overrides}
+	fake.beforeDelete = func(attempt int) error {
+		if attempt != 1 {
+			return nil
+		}
+		// Simulate a racing override write landing after CleanupAccount's
+		// read but before its guarded delete, through the real KV so the
+		// delete's expected revision is now stale.
+		return svc.PutAccountQuota(ctx, testAccount, Overrides{VCPUs: ptr(8)}, "operator")
+	}
+	svc.overrides = fake
+
+	result, err := svc.CleanupAccount(ctx, testAccount)
+	require.NoError(t, err)
+	assert.Equal(t, OverrideCleanupDeleted, result.Disposition)
+	assert.Equal(t, 2, fake.deleteCalls, "the first delete is rejected by the race and retried once")
+}
+
+// A revision that keeps moving under sustained concurrent writes must not
+// spin forever: cleanup gives up and reports failed, bounded by attempts.
+func TestCleanupAccount_PersistentConcurrentWritesFailsBounded(t *testing.T) {
+	svc, _ := newOverrideService(t)
+	ctx := t.Context()
+	require.NoError(t, svc.PutAccountQuota(ctx, testAccount, Overrides{VCPUs: ptr(32)}, "operator"))
+
+	fake := &overrideCleanupFakeKV{KeyValue: svc.overrides}
+	fake.beforeDelete = func(int) error {
+		return svc.PutAccountQuota(ctx, testAccount, Overrides{VCPUs: ptr(8)}, "operator")
+	}
+	svc.overrides = fake
+
+	result, err := svc.CleanupAccount(ctx, testAccount)
+	assert.Error(t, err)
+	require.NotNil(t, result)
+	assert.Equal(t, OverrideCleanupFailed, result.Disposition)
+	assert.Equal(t, cleanupMaxAttempts, fake.deleteCalls)
+}
+
+// A store error must report failed, never absent: absence is a claim about
+// what the store holds, which an error gives no basis to make.
+func TestCleanupAccount_StoreErrorIsFailedNotAbsent(t *testing.T) {
+	svc, _ := newOverrideService(t)
+	fake := &overrideCleanupFakeKV{KeyValue: svc.overrides, getErr: errors.New("kv unavailable")}
+	svc.overrides = fake
+
+	result, err := svc.CleanupAccount(t.Context(), testAccount)
+	assert.Error(t, err)
+	require.NotNil(t, result)
+	assert.Equal(t, OverrideCleanupFailed, result.Disposition)
+}