From dca7a3835686cec9d23cffe6bcac23508434e6ff Mon Sep 17 00:00:00 2001 From: Ben Duncan <5178243+benduncan@users.noreply.github.com> Date: Thu, 20 Aug 2026 01:52:07 +1000 Subject: [PATCH] feat: resolve service quotas per account from a KV override record --- ...ttings.local.json.tmp.2466042.87acb86176ea | 25 +++ cmd/spinifex/cmd/templates/awsgw.toml | 14 +- spinifex/gateway/elbv2.go | 3 + spinifex/gateway/rds.go | 8 +- spinifex/gateway/rds/handler.go | 7 +- spinifex/gateway/rds/instance.go | 7 +- spinifex/handlers/quota/live.go | 40 ++++- spinifex/handlers/quota/live_test.go | 12 ++ spinifex/handlers/quota/override.go | 131 ++++++++++++++++ spinifex/handlers/quota/override_kv_test.go | 124 +++++++++++++++ spinifex/handlers/quota/override_test.go | 145 ++++++++++++++++++ spinifex/handlers/quota/quota.go | 28 ++++ spinifex/handlers/quota/storage.go | 13 +- spinifex/handlers/quota/vcpu.go | 5 +- spinifex/services/awsgw/awsgw.go | 33 +++- 15 files changed, 577 insertions(+), 18 deletions(-) create mode 100644 .claude/settings.local.json.tmp.2466042.87acb86176ea create mode 100644 spinifex/handlers/quota/override.go create mode 100644 spinifex/handlers/quota/override_kv_test.go create mode 100644 spinifex/handlers/quota/override_test.go diff --git a/.claude/settings.local.json.tmp.2466042.87acb86176ea b/.claude/settings.local.json.tmp.2466042.87acb86176ea new file mode 100644 index 000000000..8d1eb5d6c --- /dev/null +++ b/.claude/settings.local.json.tmp.2466042.87acb86176ea @@ -0,0 +1,25 @@ +{ + "permissions": { + "allow": [ + "Bash(timeout 30 bd ready)", + "Bash(python3 -c ' *)", + "Bash(timeout 30 bd label list)", + "Bash(timeout 30 bd labels)", + "Bash(git fetch *)", + "Bash(git checkout *)", + "Bash(grep -n \"^#\\\\{2,3\\\\} \" docs/development/improvements/new-multi-node-networking.md)", + "Bash(go version *)", + "Read(//home/tf-user/go/pkg/mod/golang.org/**)", + "Read(//usr/local/go/**)", + "Bash(go build *)", + "Bash(go test *)", + "Bash(go vet *)", + "Bash(python3 -)", + "Bash(bash -n scripts/iso-builder/build/spinifex-init.sh)", + "Bash(timeout 900 make preflight)", + "Bash(timeout 300 golangci-lint run ./cmd/installer/...)", + "Bash(timeout 300 golangci-lint run ./...)", + "Bash(git stash *)" + ] + } +} diff --git a/cmd/spinifex/cmd/templates/awsgw.toml b/cmd/spinifex/cmd/templates/awsgw.toml index f23c04369..022c4bb33 100644 --- a/cmd/spinifex/cmd/templates/awsgw.toml +++ b/cmd/spinifex/cmd/templates/awsgw.toml @@ -40,11 +40,17 @@ burst = 5 [signup] max_accounts = 128 -# Per-account service quotas (default off) +# Per-account service quotas (default off). These values are the baseline every +# account inherits; raise an individual account with +# `spx admin account quota set`, which stores an override in NATS KV. +# A dimension set to -1 is uncapped for accounts that inherit it. [quota] enabled = false -vcpus = 8 -vpcs = 8 +vcpus = 16 +vpcs = 4 subnets = 16 eips = 4 -volumes_gib = 100 +volumes = 16 +volumes_gib = 200 +rds_instances = 2 +load_balancers = 2 diff --git a/spinifex/gateway/elbv2.go b/spinifex/gateway/elbv2.go index d46eca715..03db2050f 100644 --- a/spinifex/gateway/elbv2.go +++ b/spinifex/gateway/elbv2.go @@ -44,6 +44,9 @@ func elbv2Handler[In any](handler func(context.Context, *In, *GatewayConfig, str var elbv2Actions = map[string]ELBv2Handler{ "CreateLoadBalancer": elbv2Handler(func(ctx context.Context, input *elbv2.CreateLoadBalancerInput, gw *GatewayConfig, accountID string) (any, error) { + if err := gw.Quota.EnforceLoadBalancers(ctx, gw.NATSConn, accountID, 1); err != nil { + return nil, err + } return gateway_elbv2.CreateLoadBalancer(ctx, input, gw.NATSConn, accountID) }), "DeleteLoadBalancer": elbv2Handler(func(ctx context.Context, input *elbv2.DeleteLoadBalancerInput, gw *GatewayConfig, accountID string) (any, error) { diff --git a/spinifex/gateway/rds.go b/spinifex/gateway/rds.go index f1d69d5d5..b03b9e139 100644 --- a/spinifex/gateway/rds.go +++ b/spinifex/gateway/rds.go @@ -1,6 +1,7 @@ package gateway import ( + "context" "errors" "log/slog" "net/http" @@ -55,7 +56,12 @@ func (gw *GatewayConfig) RDS_Request(w http.ResponseWriter, r *http.Request) err } xmlOutput, err := gateway_rds.Dispatch(r.Context(), action, queryArgs, gw.NATSConn, caller, - gateway_rds.Env{ExpectedNodes: gw.ExpectedNodes}) + gateway_rds.Env{ + ExpectedNodes: gw.ExpectedNodes, + QuotaCheck: func(ctx context.Context, accountID string, want int) error { + return gw.Quota.EnforceRDSInstances(ctx, gw.NATSConn, accountID, want) + }, + }) if err != nil { return err } diff --git a/spinifex/gateway/rds/handler.go b/spinifex/gateway/rds/handler.go index e50f9914e..fb9d7c4f8 100644 --- a/spinifex/gateway/rds/handler.go +++ b/spinifex/gateway/rds/handler.go @@ -36,6 +36,11 @@ type Env struct { // How many nodes a fan-out waits for before it stops waiting. Without it a // gather burns its full timeout on every call instead of early-exiting. ExpectedNodes int + + // QuotaCheck gates a create on the account's DB-instance cap. It is a + // function so this package stays free of the quota dependency; nil skips + // the check, which is what a gateway with quotas disabled supplies. + QuotaCheck func(ctx context.Context, accountID string, want int) error } type Handler func(ctx context.Context, action string, q map[string]string, nc *nats.Conn, caller Caller, env Env) ([]byte, error) @@ -94,7 +99,7 @@ type actionDef struct { // distinct from an unknown one. var actions = map[string]actionDef{ // Instance lifecycle. - "CreateDBInstance": {handler: typed(CreateDBInstance)}, + "CreateDBInstance": {handler: typedEnv(CreateDBInstance)}, "DescribeDBInstances": {handler: typed(DescribeDBInstances)}, "ModifyDBInstance": {handler: typed(ModifyDBInstance), scope: dbInstanceScope}, "DeleteDBInstance": {handler: typed(DeleteDBInstance), scope: dbInstanceScope}, diff --git a/spinifex/gateway/rds/instance.go b/spinifex/gateway/rds/instance.go index f832cba58..d176ee860 100644 --- a/spinifex/gateway/rds/instance.go +++ b/spinifex/gateway/rds/instance.go @@ -10,7 +10,12 @@ import ( // The whole orchestration runs on the daemon side, so this only carries the // caller's account through — the request body never names the owner. -func CreateDBInstance(ctx context.Context, input *rds.CreateDBInstanceInput, nc *nats.Conn, caller Caller) (any, error) { +func CreateDBInstance(ctx context.Context, input *rds.CreateDBInstanceInput, nc *nats.Conn, caller Caller, env Env) (any, error) { + if env.QuotaCheck != nil { + if err := env.QuotaCheck(ctx, caller.AccountID, 1); err != nil { + return nil, err + } + } return handlers_rds.NewNATSService(nc).CreateDBInstance(ctx, input, caller.AccountID) } diff --git a/spinifex/handlers/quota/live.go b/spinifex/handlers/quota/live.go index 0ed68b20f..62ac56f0e 100644 --- a/spinifex/handlers/quota/live.go +++ b/spinifex/handlers/quota/live.go @@ -5,9 +5,13 @@ import ( "errors" "github.com/aws/aws-sdk-go/service/ec2" + "github.com/aws/aws-sdk-go/service/elbv2" + "github.com/aws/aws-sdk-go/service/rds" "github.com/mulgadc/spinifex/spinifex/awserrors" gateway_ec2_eip "github.com/mulgadc/spinifex/spinifex/gateway/ec2/eip" gateway_ec2_vpc "github.com/mulgadc/spinifex/spinifex/gateway/ec2/vpc" + gateway_elbv2 "github.com/mulgadc/spinifex/spinifex/gateway/elbv2" + handlers_rds "github.com/mulgadc/spinifex/spinifex/handlers/rds" "github.com/nats-io/nats.go" ) @@ -16,6 +20,9 @@ import ( // the shared comparison for every live dimension; the per-dimension methods // supply count from the relevant Describe* call. func exceeds(count, want, limit int) error { + if limit == Unlimited { + return nil + } if count+want > limit { return errors.New(awserrors.ErrorResourceLimitExceeded) } @@ -31,7 +38,7 @@ func (s *Service) EnforceVPCs(ctx context.Context, natsConn *nats.Conn, accountI if err != nil { return err } - return exceeds(len(out.Vpcs), want, s.limits.VPCs) + return exceeds(len(out.Vpcs), want, s.limitsFor(ctx, accountID).VPCs) } // EnforceSubnets gates CreateSubnet on the account's live DescribeSubnets count. @@ -44,7 +51,7 @@ func (s *Service) EnforceSubnets(ctx context.Context, natsConn *nats.Conn, accou if err != nil { return err } - return exceeds(len(out.Subnets), want, s.limits.Subnets) + return exceeds(len(out.Subnets), want, s.limitsFor(ctx, accountID).Subnets) } // EnforceEIPs gates AllocateAddress on the account's live DescribeAddresses count. @@ -56,5 +63,32 @@ func (s *Service) EnforceEIPs(ctx context.Context, natsConn *nats.Conn, accountI if err != nil { return err } - return exceeds(len(out.Addresses), want, s.limits.EIPs) + return exceeds(len(out.Addresses), want, s.limitsFor(ctx, accountID).EIPs) +} + +// EnforceLoadBalancers gates CreateLoadBalancer on the account's live count. +// ALBs and NLBs share one cap, as both consume a load-balancer appliance VM. +func (s *Service) EnforceLoadBalancers(ctx context.Context, natsConn *nats.Conn, accountID string, want int) error { + if s.Exempt(accountID) { + return nil + } + out, err := gateway_elbv2.DescribeLoadBalancers(ctx, &elbv2.DescribeLoadBalancersInput{}, natsConn, accountID) + if err != nil { + return err + } + return exceeds(len(out.LoadBalancers), want, s.limitsFor(ctx, accountID).LoadBalancers) +} + +// EnforceRDSInstances gates CreateDBInstance on the account's live count. A DB +// instance is capped by count rather than by vCPU because its VM runs in the +// system account and never reaches the tenant's vCPU counter. +func (s *Service) EnforceRDSInstances(ctx context.Context, natsConn *nats.Conn, accountID string, want int) error { + if s.Exempt(accountID) { + return nil + } + out, err := handlers_rds.NewNATSService(natsConn).DescribeDBInstances(ctx, &rds.DescribeDBInstancesInput{}, accountID) + if err != nil { + return err + } + return exceeds(len(out.DBInstances), want, s.limitsFor(ctx, accountID).RDSInstances) } diff --git a/spinifex/handlers/quota/live_test.go b/spinifex/handlers/quota/live_test.go index b7c2c00da..020ee23e9 100644 --- a/spinifex/handlers/quota/live_test.go +++ b/spinifex/handlers/quota/live_test.go @@ -57,6 +57,18 @@ func TestEnforceLiveExemptShortCircuits(t *testing.T) { {"vpc system account", func() error { return enabled.EnforceVPCs(context.Background(), nil, utils.GlobalAccountID, 1) }}, {"subnet system account", func() error { return enabled.EnforceSubnets(context.Background(), nil, utils.GlobalAccountID, 1) }}, {"eip system account", func() error { return enabled.EnforceEIPs(context.Background(), nil, utils.GlobalAccountID, 1) }}, + {"lb disabled", func() error { + return disabled.EnforceLoadBalancers(context.Background(), nil, normalAccount, 1) + }}, + {"lb system account", func() error { + return enabled.EnforceLoadBalancers(context.Background(), nil, utils.GlobalAccountID, 1) + }}, + {"rds disabled", func() error { + return disabled.EnforceRDSInstances(context.Background(), nil, normalAccount, 1) + }}, + {"rds system account", func() error { + return enabled.EnforceRDSInstances(context.Background(), nil, utils.GlobalAccountID, 1) + }}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { diff --git a/spinifex/handlers/quota/override.go b/spinifex/handlers/quota/override.go new file mode 100644 index 000000000..a9ac5bd3e --- /dev/null +++ b/spinifex/handlers/quota/override.go @@ -0,0 +1,131 @@ +package handlers_quota + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "log/slog" + "time" + + "github.com/nats-io/nats.go/jetstream" +) + +// Overrides raises or lowers individual dimensions for one account. Every field +// is a pointer so an explicit 0 (a limit that denies everything) is distinct +// from an absent field, which inherits the configured value. +type Overrides struct { + VCPUs *int `json:"vcpus,omitempty"` + VPCs *int `json:"vpcs,omitempty"` + Subnets *int `json:"subnets,omitempty"` + EIPs *int `json:"eips,omitempty"` + VolumesGiB *int `json:"volumes_gib,omitempty"` + Volumes *int `json:"volumes,omitempty"` + RDSInstances *int `json:"rds_instances,omitempty"` + LoadBalancers *int `json:"load_balancers,omitempty"` + + // Raising a customer's quota is a commercial act, so it is attributable. + UpdatedAt string `json:"updated_at,omitempty"` + UpdatedBy string `json:"updated_by,omitempty"` +} + +// Empty reports whether no dimension is overridden, in which case the record +// carries nothing but provenance and the account resolves to the config. +func (o Overrides) Empty() bool { + return o.VCPUs == nil && o.VPCs == nil && o.Subnets == nil && o.EIPs == nil && + o.VolumesGiB == nil && o.Volumes == nil && o.RDSInstances == nil && o.LoadBalancers == nil +} + +// apply overlays the set fields onto base, leaving Enabled and the Bedrock +// dimensions alone: enablement is a cluster-wide switch, never per-account. +func (o Overrides) apply(base Limits) Limits { + out := base + for _, f := range []struct { + src *int + dst *int + }{ + {o.VCPUs, &out.VCPUs}, + {o.VPCs, &out.VPCs}, + {o.Subnets, &out.Subnets}, + {o.EIPs, &out.EIPs}, + {o.VolumesGiB, &out.VolumesGiB}, + {o.Volumes, &out.Volumes}, + {o.RDSInstances, &out.RDSInstances}, + {o.LoadBalancers, &out.LoadBalancers}, + } { + if f.src != nil { + *f.dst = *f.src + } + } + return out +} + +// limitsFor resolves the effective limits for accountID: its overrides where +// set, the configured limits everywhere else. A read failure falls back to the +// configured limits rather than to no limit, so a KV outage cannot lift a cap. +func (s *Service) limitsFor(ctx context.Context, accountID string) Limits { + over, err := s.overridesFor(ctx, accountID) + if err != nil { + slog.WarnContext(ctx, "quota: override read failed, using configured limits", + "account", accountID, "err", err) + return s.limits + } + return over.apply(s.limits) +} + +// overridesFor reads accountID's override record. An absent key is the normal +// case for nearly every account and yields the zero value, not an error; only a +// genuine read or decode failure is returned. +func (s *Service) overridesFor(ctx context.Context, accountID string) (Overrides, error) { + if s == nil || s.overrides == nil { + return Overrides{}, nil + } + entry, err := s.overrides.Get(ctx, accountID) + if errors.Is(err, jetstream.ErrKeyNotFound) { + return Overrides{}, nil + } + if err != nil { + return Overrides{}, fmt.Errorf("read quota override for %s: %w", accountID, err) + } + var over Overrides + if err := json.Unmarshal(entry.Value(), &over); err != nil { + return Overrides{}, fmt.Errorf("decode quota override for %s: %w", accountID, err) + } + return over, nil +} + +// GetAccountQuota returns accountID's stored overrides and the limits they +// resolve to. A read failure is surfaced here rather than masked, because an +// operator inspecting a quota must not be shown the fallback as if it were real. +func (s *Service) GetAccountQuota(ctx context.Context, accountID string) (Overrides, Limits, error) { + over, err := s.overridesFor(ctx, accountID) + if err != nil { + return Overrides{}, Limits{}, err + } + return over, over.apply(s.limits), nil +} + +// PutAccountQuota stores accountID's overrides, stamping provenance. An empty +// set deletes the record so the account resolves to the config, which keeps +// "has this account been touched?" a single existence check. +func (s *Service) PutAccountQuota(ctx context.Context, accountID string, over Overrides, updatedBy string) error { + if s == nil || s.overrides == nil { + return errors.New("quota: override bucket not configured") + } + if over.Empty() { + if err := s.overrides.Delete(ctx, accountID); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) { + return fmt.Errorf("clear quota override for %s: %w", accountID, err) + } + return nil + } + over.UpdatedAt = time.Now().UTC().Format(time.RFC3339) + over.UpdatedBy = updatedBy + data, err := json.Marshal(over) + if err != nil { + return fmt.Errorf("encode quota override for %s: %w", accountID, err) + } + if _, err := s.overrides.Put(ctx, accountID, data); err != nil { + return fmt.Errorf("store quota override for %s: %w", accountID, err) + } + return nil +} diff --git a/spinifex/handlers/quota/override_kv_test.go b/spinifex/handlers/quota/override_kv_test.go new file mode 100644 index 000000000..62fb664d9 --- /dev/null +++ b/spinifex/handlers/quota/override_kv_test.go @@ -0,0 +1,124 @@ +package handlers_quota + +import ( + "testing" + + "github.com/mulgadc/spinifex/spinifex/testutil" + "github.com/nats-io/nats.go/jetstream" + "github.com/stretchr/testify/require" +) + +// newOverrideService builds a Service backed by a real JetStream KV bucket, so +// the resolver is exercised against the store it uses in production rather than +// a stand-in that cannot reproduce a read failure. +func newOverrideService(t *testing.T) (*Service, jetstream.KeyValue) { + t.Helper() + _, _, js := testutil.StartTestJetStream(t) + kv, err := js.CreateKeyValue(t.Context(), jetstream.KeyValueConfig{ + Bucket: KVBucketAccountQuota, + History: 1, + }) + require.NoError(t, err) + + svc := New(baseLimits(), nil) + svc.SetOverrides(kv) + return svc, kv +} + +// An account with no record is the normal state, and it must inherit the +// configured baseline without erroring. +func TestLimitsForAbsentKeyInherits(t *testing.T) { + svc, _ := newOverrideService(t) + + got := svc.limitsFor(t.Context(), "123456789012") + require.Equal(t, baseLimits(), got) +} + +// The property the whole tier rests on: two accounts on one gateway resolve to +// different limits. +func TestLimitsForIsPerAccount(t *testing.T) { + svc, _ := newOverrideService(t) + const upgraded, sandbox = "000000000002", "000000000003" + + require.NoError(t, svc.PutAccountQuota(t.Context(), upgraded, + Overrides{VCPUs: ptr(32), RDSInstances: ptr(8)}, "operator")) + + require.Equal(t, 32, svc.limitsFor(t.Context(), upgraded).VCPUs) + require.Equal(t, 8, svc.limitsFor(t.Context(), upgraded).RDSInstances) + require.Equal(t, baseLimits().VCPUs, svc.limitsFor(t.Context(), sandbox).VCPUs) + require.Equal(t, baseLimits().RDSInstances, svc.limitsFor(t.Context(), sandbox).RDSInstances) +} + +// A stored override must round-trip through KV with its provenance, and an +// unset dimension must still inherit. +func TestPutAndGetAccountQuota(t *testing.T) { + svc, _ := newOverrideService(t) + const account = "000000000002" + + require.NoError(t, svc.PutAccountQuota(t.Context(), account, Overrides{EIPs: ptr(12)}, "operator")) + + over, effective, err := svc.GetAccountQuota(t.Context(), account) + require.NoError(t, err) + require.NotNil(t, over.EIPs) + require.Equal(t, 12, *over.EIPs) + require.Equal(t, "operator", over.UpdatedBy) + require.NotEmpty(t, over.UpdatedAt) + require.Equal(t, 12, effective.EIPs) + require.Equal(t, baseLimits().VCPUs, effective.VCPUs, "unset dimension must inherit") +} + +// Clearing must delete the record rather than leave an empty one, so "has this +// account been touched?" stays a single existence check. +func TestPutAccountQuotaEmptyClears(t *testing.T) { + svc, kv := newOverrideService(t) + const account = "000000000002" + + require.NoError(t, svc.PutAccountQuota(t.Context(), account, Overrides{VCPUs: ptr(32)}, "operator")) + require.NoError(t, svc.PutAccountQuota(t.Context(), account, Overrides{}, "operator")) + + _, err := kv.Get(t.Context(), account) + require.ErrorIs(t, err, jetstream.ErrKeyNotFound) + require.Equal(t, baseLimits(), svc.limitsFor(t.Context(), account)) +} + +// Clearing an account that never had a record must succeed rather than fail on +// the missing key. +func TestPutAccountQuotaEmptyOnAbsentKey(t *testing.T) { + svc, _ := newOverrideService(t) + require.NoError(t, svc.PutAccountQuota(t.Context(), "000000000009", Overrides{}, "operator")) +} + +// An explicit zero must survive the store and deny every request, which is the +// difference between a limit of zero and an absent override. +func TestPutAccountQuotaExplicitZero(t *testing.T) { + svc, _ := newOverrideService(t) + const account = "000000000002" + + require.NoError(t, svc.PutAccountQuota(t.Context(), account, Overrides{VPCs: ptr(0)}, "operator")) + + require.Equal(t, 0, svc.limitsFor(t.Context(), account).VPCs) + require.Error(t, exceeds(0, 1, svc.limitsFor(t.Context(), account).VPCs)) +} + +// A corrupt record is a read failure, not an absent one: it must fall back to +// the configured limits rather than to no limit at all. +func TestLimitsForCorruptRecordFailsClosed(t *testing.T) { + svc, kv := newOverrideService(t) + const account = "000000000002" + + _, err := kv.Put(t.Context(), account, []byte("{not json")) + require.NoError(t, err) + + require.Equal(t, baseLimits(), svc.limitsFor(t.Context(), account), + "a corrupt override must fall back to the configured limits") + + _, _, err = svc.GetAccountQuota(t.Context(), account) + require.Error(t, err, "an operator inspecting a quota must see the failure, not the fallback") +} + +// A gateway with no override bucket must refuse a write rather than silently +// accept one it cannot store. +func TestPutAccountQuotaWithoutBucket(t *testing.T) { + svc := New(baseLimits(), nil) + require.Error(t, svc.PutAccountQuota(t.Context(), "000000000002", Overrides{VCPUs: ptr(32)}, "operator")) +} diff --git a/spinifex/handlers/quota/override_test.go b/spinifex/handlers/quota/override_test.go new file mode 100644 index 000000000..77ce8ae0f --- /dev/null +++ b/spinifex/handlers/quota/override_test.go @@ -0,0 +1,145 @@ +package handlers_quota + +import ( + "encoding/json" + "testing" +) + +func ptr(v int) *int { return &v } + +// baseLimits is the sandbox baseline every account inherits when it has no +// override, mirroring the shipped [quota] block. +func baseLimits() Limits { + return Limits{ + Enabled: true, + VCPUs: 16, + VPCs: 4, + Subnets: 16, + EIPs: 4, + Volumes: 16, + VolumesGiB: 200, + RDSInstances: 2, + LoadBalancers: 2, + } +} + +// An override replaces only the dimensions it sets; every other dimension must +// come through from the configured baseline untouched. +func TestOverridesApplyIsSparse(t *testing.T) { + got := Overrides{VCPUs: ptr(32), RDSInstances: ptr(8)}.apply(baseLimits()) + + if got.VCPUs != 32 { + t.Errorf("VCPUs = %d, want 32", got.VCPUs) + } + if got.RDSInstances != 8 { + t.Errorf("RDSInstances = %d, want 8", got.RDSInstances) + } + want := baseLimits() + if got.VPCs != want.VPCs || got.EIPs != want.EIPs || got.Volumes != want.Volumes || + got.VolumesGiB != want.VolumesGiB || got.Subnets != want.Subnets || + got.LoadBalancers != want.LoadBalancers { + t.Errorf("unset dimensions changed: got %+v, want %+v", got, want) + } +} + +// An empty override must leave the baseline exactly as it is: this is the state +// of nearly every account and the reason no migration is needed. +func TestOverridesApplyEmptyInherits(t *testing.T) { + if got := (Overrides{}).apply(baseLimits()); got != baseLimits() { + t.Fatalf("empty override = %+v, want %+v", got, baseLimits()) + } +} + +// Enablement is a cluster-wide switch, so no override may turn quotas off for +// one account. +func TestOverridesCannotDisableEnforcement(t *testing.T) { + if got := (Overrides{VCPUs: ptr(0)}).apply(baseLimits()); !got.Enabled { + t.Fatal("override cleared Enabled, want it preserved") + } +} + +// Zero is a real limit that denies every request, so it must survive a JSON +// round trip rather than marshalling away and silently meaning "inherit". +func TestOverridesZeroSurvivesRoundTrip(t *testing.T) { + data, err := json.Marshal(Overrides{VCPUs: ptr(0)}) + if err != nil { + t.Fatalf("marshal: %v", err) + } + + var back Overrides + if err := json.Unmarshal(data, &back); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if back.VCPUs == nil { + t.Fatalf("explicit 0 was dropped by the round trip: %s", data) + } + if *back.VCPUs != 0 { + t.Fatalf("VCPUs = %d, want 0", *back.VCPUs) + } + if got := back.apply(baseLimits()); got.VCPUs != 0 { + t.Fatalf("applied VCPUs = %d, want 0", got.VCPUs) + } +} + +// An absent field must stay absent through a round trip, so inheritance is not +// silently converted into an explicit zero. +func TestOverridesAbsentStaysAbsent(t *testing.T) { + data, err := json.Marshal(Overrides{VPCs: ptr(8)}) + if err != nil { + t.Fatalf("marshal: %v", err) + } + + var back Overrides + if err := json.Unmarshal(data, &back); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if back.VCPUs != nil { + t.Fatalf("unset VCPUs decoded as %d, want nil", *back.VCPUs) + } + if back.apply(baseLimits()).VCPUs != baseLimits().VCPUs { + t.Fatal("unset VCPUs did not inherit the baseline") + } +} + +func TestOverridesEmpty(t *testing.T) { + tests := []struct { + name string + over Overrides + want bool + }{ + {"zero value", Overrides{}, true}, + {"provenance only", Overrides{UpdatedBy: "operator"}, true}, + {"one dimension", Overrides{EIPs: ptr(1)}, false}, + {"explicit zero", Overrides{VCPUs: ptr(0)}, false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := tt.over.Empty(); got != tt.want { + t.Fatalf("Empty() = %v, want %v", got, tt.want) + } + }) + } +} + +// A nil override bucket is what a quota-disabled gateway holds; resolution must +// fall through to the configured limits rather than panicking. +func TestLimitsForWithoutBucket(t *testing.T) { + svc := New(baseLimits(), nil) + if got := svc.limitsFor(t.Context(), "123456789012"); got != baseLimits() { + t.Fatalf("limitsFor = %+v, want the configured limits", got) + } +} + +// The Unlimited sentinel disables one dimension for one account, which zero +// cannot express because zero is a real limit that denies everything. +func TestExceedsUnlimited(t *testing.T) { + if err := exceeds(9000, 9000, Unlimited); err != nil { + t.Fatalf("exceeds with Unlimited = %v, want nil", err) + } + if err := exceeds(0, 1, 0); err == nil { + t.Fatal("a zero cap allowed a request, want it denied") + } + if err := exceeds(0, 0, 0); err != nil { + t.Fatalf("a zero cap denied a zero request: %v", err) + } +} diff --git a/spinifex/handlers/quota/quota.go b/spinifex/handlers/quota/quota.go index 43664eabe..ef0267980 100644 --- a/spinifex/handlers/quota/quota.go +++ b/spinifex/handlers/quota/quota.go @@ -25,6 +25,15 @@ const KVBucketAccountUsage = "spinifex-account-usage" // the two independent loops never share one mutex and block each other. const KVBucketQuotaReconcile = "spinifex-quota-reconcile" +// KVBucketAccountQuota holds the per-account limit overrides, keyed by +// accountID. It is separate from the accounts bucket so a quota write never +// contends with an identity write on the same revision. +const KVBucketAccountQuota = "spinifex-account-quota" + +// Unlimited disables one dimension for one account. Zero is a real limit that +// denies every request, so "no cap" needs a value of its own. +const Unlimited = -1 + // Limits mirrors the [quota] block in awsgw.toml. The zero value (Enabled // false) is a valid no-op, so gateways without a [quota] block are unaffected. // @@ -43,6 +52,12 @@ type Limits struct { EIPs int `toml:"eips"` VolumesGiB int `toml:"volumes_gib"` + // Volumes caps the number of EBS volumes, independently of VolumesGiB: a + // tenant can exhaust per-volume overhead long before it exhausts capacity. + Volumes int `toml:"volumes"` + RDSInstances int `toml:"rds_instances"` + LoadBalancers int `toml:"load_balancers"` + TokensPerMonthEnabled bool `toml:"tokens_per_month_enabled"` TokensPerMonth int64 `toml:"tokens_per_month"` RequestsPerMinuteEnabled bool `toml:"requests_per_minute_enabled"` @@ -69,6 +84,10 @@ type Service struct { // the requests-per-minute dimension. Always non-nil so CheckBedrockRPM // never needs a nil guard beyond Service itself. rpm *rpmLimiter + + // overrides holds the per-account limit overrides. Nil leaves every account + // on the configured limits, which is also what an absent key means. + overrides jetstream.KeyValue } // New constructs a quota Service from the configured limits and the gateway-owned @@ -78,6 +97,15 @@ func New(limits Limits, usage jetstream.KeyValue) *Service { return &Service{limits: limits, usage: usage, rpm: newRPMLimiter()} } +// SetOverrides attaches the per-account override bucket. Called once at wiring +// time; a nil Service or bucket leaves every account on the configured limits. +func (s *Service) SetOverrides(kv jetstream.KeyValue) { + if s == nil { + return + } + s.overrides = kv +} + // Exempt returns true for the global/system account and whenever quotas are // disabled, so those callers bypass every quota check. A nil Service (quotas // never built) is treated as disabled so callers need no nil guard. diff --git a/spinifex/handlers/quota/storage.go b/spinifex/handlers/quota/storage.go index ac2e4c3d4..698c7019e 100644 --- a/spinifex/handlers/quota/storage.go +++ b/spinifex/handlers/quota/storage.go @@ -13,15 +13,22 @@ import ( // VolumesGiB cap. Storage is live-counted, so a delete frees quota on the next // describe. Root/AMI volumes are created out-of-band by the launch path and // never appear in DescribeVolumes, so they are neither summed nor charged. +// The volume count is checked alongside the capacity sum: a tenant can exhaust +// per-volume overhead long before it exhausts GiB. func (s *Service) EnforceVolumeCreate(ctx context.Context, natsConn *nats.Conn, accountID string, requestedGiB int) error { if s.Exempt(accountID) { return nil } - total, _, err := s.volumeUsage(ctx, natsConn, accountID, "") + out, err := gateway_ec2_volume.DescribeVolumes(ctx, &ec2.DescribeVolumesInput{}, natsConn, accountID) if err != nil { return err } - return exceeds(total, requestedGiB, s.limits.VolumesGiB) + limits := s.limitsFor(ctx, accountID) + if err := exceeds(len(out.Volumes), 1, limits.Volumes); err != nil { + return err + } + total, _ := sumVolumeGiB(out.Volumes, "") + return exceeds(total, requestedGiB, limits.VolumesGiB) } // EnforceVolumeModify gates ModifyVolume on the account's live storage usage. @@ -36,7 +43,7 @@ func (s *Service) EnforceVolumeModify(ctx context.Context, natsConn *nats.Conn, if err != nil { return err } - return exceeds(total-oldGiB, newGiB, s.limits.VolumesGiB) + return exceeds(total-oldGiB, newGiB, s.limitsFor(ctx, accountID).VolumesGiB) } // volumeUsage sums the account's volume sizes in GiB via a single DescribeVolumes diff --git a/spinifex/handlers/quota/vcpu.go b/spinifex/handlers/quota/vcpu.go index 62b6a4d5d..a93cf9ed0 100644 --- a/spinifex/handlers/quota/vcpu.go +++ b/spinifex/handlers/quota/vcpu.go @@ -35,10 +35,7 @@ func (s *Service) CheckVCPU(ctx context.Context, accountID string, want int) err if err != nil { return err } - if current+want > s.limits.VCPUs { - return errors.New(awserrors.ErrorResourceLimitExceeded) - } - return nil + return exceeds(current, want, s.limitsFor(ctx, accountID).VCPUs) } // casVCPU applies next to accountID's counter under bounded JetStream CAS so diff --git a/spinifex/services/awsgw/awsgw.go b/spinifex/services/awsgw/awsgw.go index ba6c7e3ac..c72c7c16e 100644 --- a/spinifex/services/awsgw/awsgw.go +++ b/spinifex/services/awsgw/awsgw.go @@ -166,6 +166,32 @@ func openAccountUsageBucket(ctx context.Context, js jetstream.KeyValueManager, r return kv, nil } +// openAccountQuotaBucket opens (or idempotently creates) the per-account quota +// override bucket. History is 1: each key holds the current override set, and +// no earlier revision of a limit is worth keeping. +func openAccountQuotaBucket(ctx context.Context, js jetstream.KeyValueManager, replicas int) (jetstream.KeyValue, error) { + kv, err := js.KeyValue(ctx, handlers_quota.KVBucketAccountQuota) + if err == nil { + return kv, nil + } + if !errors.Is(err, jetstream.ErrBucketNotFound) { + return nil, fmt.Errorf("open account quota bucket: %w", err) + } + + kv, err = js.CreateKeyValue(ctx, jetstream.KeyValueConfig{ + Bucket: handlers_quota.KVBucketAccountQuota, + History: 1, + Replicas: max(replicas, 1), + }) + if errors.Is(err, jetstream.ErrBucketExists) { + return js.KeyValue(ctx, handlers_quota.KVBucketAccountQuota) + } + if err != nil { + return nil, fmt.Errorf("create account quota bucket: %w", err) + } + return kv, nil +} + func launchService(config *config.ClusterConfig) error { nodeConfig := config.Nodes[config.Node] @@ -480,14 +506,19 @@ func launchService(config *config.ClusterConfig) error { // Per-account service quotas. Only the enabled path opens the gateway-owned // usage KV bucket, leaving existing default-off gateways untouched; a disabled // config builds a no-op Service whose Exempt short-circuits every check. - var usageBucket jetstream.KeyValue + var usageBucket, quotaOverrides jetstream.KeyValue if quotaCfg.Enabled { usageBucket, err = openAccountUsageBucket(janitorCtx, js, len(config.Nodes)) if err != nil { return fmt.Errorf("init account usage bucket: %w", err) } + quotaOverrides, err = openAccountQuotaBucket(janitorCtx, js, len(config.Nodes)) + if err != nil { + return fmt.Errorf("init account quota bucket: %w", err) + } } gw.Quota = handlers_quota.New(quotaCfg, usageBucket) + gw.Quota.SetOverrides(quotaOverrides) // Bedrock token quota reads the stream-fed usage counters built above, // independent of whether the standing-infra dimensions are enabled.