diff --git a/cmd/spinifex/cmd/templates/awsgw.toml b/cmd/spinifex/cmd/templates/awsgw.toml index fcf96f7d1..04b9e16a8 100644 --- a/cmd/spinifex/cmd/templates/awsgw.toml +++ b/cmd/spinifex/cmd/templates/awsgw.toml @@ -1,5 +1,5 @@ # AWS Gateway configuration file in TOML format. -version = "2" +version = "3" region = "{{.Region}}" # Enable verbose debugging @@ -8,8 +8,8 @@ debug = false # API request throttling (per-account, per-action token bucket) [ratelimit] enabled = false -rate = 20 # Default sustained requests per second per account per action -burst = 100 # Default maximum burst capacity +rate = 20 # Default sustained requests per second per account per action +burst = 100 # Default maximum burst capacity [ratelimit.action.RunInstances] rate = 2 @@ -26,3 +26,12 @@ burst = 40 [ratelimit.action.StopInstances] rate = 2 burst = 40 + +# Per-account service quotas (default off) +[quota] +enabled = false +vcpus = 8 +vpcs = 8 +subnets = 16 +eips = 4 +volumes_gib = 100 diff --git a/spinifex/admin/admin.go b/spinifex/admin/admin.go index 46abfefdf..4194e5443 100644 --- a/spinifex/admin/admin.go +++ b/spinifex/admin/admin.go @@ -372,6 +372,7 @@ func SetServiceOwnership() { // so /etc/spinifex stays at 0750 (no group-write needed). for path, mode := range map[string]os.FileMode{ "/etc/spinifex/spinifex.toml": 0640, + "/etc/spinifex/awsgw/awsgw.toml": 0640, "/etc/spinifex/master.key": 0640, "/etc/spinifex/viperblock/encryption.key": 0640, "/etc/spinifex/server.pem": 0644, diff --git a/spinifex/gateway/ec2.go b/spinifex/gateway/ec2.go index f1593b7ae..da49c7f2e 100644 --- a/spinifex/gateway/ec2.go +++ b/spinifex/gateway/ec2.go @@ -7,6 +7,7 @@ import ( "net/http" "strings" + "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/service/ec2" "github.com/mulgadc/spinifex/spinifex/awsec2query" "github.com/mulgadc/spinifex/spinifex/awserrors" @@ -26,6 +27,7 @@ import ( gateway_ec2_volume "github.com/mulgadc/spinifex/spinifex/gateway/ec2/volume" gateway_ec2_vpc "github.com/mulgadc/spinifex/spinifex/gateway/ec2/vpc" gateway_ec2_zone "github.com/mulgadc/spinifex/spinifex/gateway/ec2/zone" + handlers_quota "github.com/mulgadc/spinifex/spinifex/handlers/quota" "github.com/mulgadc/spinifex/spinifex/utils" ) @@ -95,7 +97,19 @@ var ec2Actions = map[string]EC2Handler{ passRoleCheck := func(roleARN string) error { return gw.checkPolicyResource(r, "iam", "PassRole", roleARN) } - return gateway_ec2_instance.RunInstances(input, gw.NATSConn, gw.IAMService, accountID, passRoleCheck, gw.ExpectedNodes) + launchQuotaCheck := func() error { + return gw.Quota.EnforceLaunch(accountID, aws.StringValue(input.InstanceType), int(aws.Int64Value(input.MaxCount))) + } + reservation, err := gateway_ec2_instance.RunInstances(input, gw.NATSConn, gw.IAMService, accountID, passRoleCheck, launchQuotaCheck, gw.ExpectedNodes) + if err != nil { + return nil, err + } + // Charge the actual launched vCPUs; a counter write failure is drift for + // reconcile to correct, so it must not fail the already-successful launch. + if err := gw.Quota.ChargeLaunch(accountID, &reservation); err != nil { + slog.Warn("RunInstances: vcpu quota charge failed, reconcile will correct", "account", accountID, "err", err) + } + return reservation, nil }), "AssociateIamInstanceProfile": ec2HandlerWithReq(func(input *ec2.AssociateIamInstanceProfileInput, gw *GatewayConfig, accountID string, r *http.Request) (any, error) { passRoleCheck := func(roleARN string) error { @@ -137,7 +151,25 @@ var ec2Actions = map[string]EC2Handler{ return gateway_ec2_instance.GetConsoleOutput(input, gw.NATSConn, accountID) }), "ModifyInstanceAttribute": ec2Handler(func(input *ec2.ModifyInstanceAttributeInput, gw *GatewayConfig, accountID string) (any, error) { - return gateway_ec2_instance.ModifyInstanceAttribute(input, gw.NATSConn, accountID) + var delta int + if input.InstanceType != nil { + resolve := handlers_quota.NATSInstanceTypeResolver(gw.NATSConn, func() int { return gw.ExpectedNodes }) + d, err := gw.Quota.EnforceRetype(resolve, accountID, aws.StringValue(input.InstanceId), aws.StringValue(input.InstanceType.Value)) + if err != nil { + return nil, err + } + delta = d + } + out, err := gateway_ec2_instance.ModifyInstanceAttribute(input, gw.NATSConn, accountID) + if err != nil { + return nil, err + } + // Charge the retype's vCPU growth; a counter write failure is drift for + // reconcile to correct, so it must not fail the applied retype. + if err := gw.Quota.AddVCPU(accountID, delta); err != nil { + slog.Warn("ModifyInstanceAttribute: vcpu quota charge failed, reconcile will correct", "account", accountID, "err", err) + } + return out, nil }), "DescribeInstanceAttribute": ec2Handler(func(input *ec2.DescribeInstanceAttributeInput, gw *GatewayConfig, accountID string) (any, error) { return gateway_ec2_instance.DescribeInstanceAttribute(input, gw.NATSConn, gw.DiscoverActiveNodes(), accountID) @@ -197,9 +229,15 @@ var ec2Actions = map[string]EC2Handler{ return gateway_ec2_volume.DescribeVolumes(input, gw.NATSConn, accountID) }), "ModifyVolume": ec2Handler(func(input *ec2.ModifyVolumeInput, gw *GatewayConfig, accountID string) (any, error) { + if err := gw.Quota.EnforceVolumeModify(gw.NATSConn, accountID, aws.StringValue(input.VolumeId), int(aws.Int64Value(input.Size))); err != nil { + return nil, err + } return gateway_ec2_volume.ModifyVolume(input, gw.NATSConn, accountID) }), "CreateVolume": ec2Handler(func(input *ec2.CreateVolumeInput, gw *GatewayConfig, accountID string) (any, error) { + if err := gw.Quota.EnforceVolumeCreate(gw.NATSConn, accountID, int(aws.Int64Value(input.Size))); err != nil { + return nil, err + } return gateway_ec2_volume.CreateVolume(input, gw.NATSConn, accountID) }), "DeleteVolume": ec2Handler(func(input *ec2.DeleteVolumeInput, gw *GatewayConfig, accountID string) (any, error) { @@ -302,6 +340,9 @@ var ec2Actions = map[string]EC2Handler{ return gateway_ec2_capacityreservation.CancelCapacityReservation(input, gw.NATSConn, gw.ExpectedNodes, accountID) }), "CreateVpc": ec2Handler(func(input *ec2.CreateVpcInput, gw *GatewayConfig, accountID string) (any, error) { + if err := gw.Quota.EnforceVPCs(gw.NATSConn, accountID, 1); err != nil { + return nil, err + } return gateway_ec2_vpc.CreateVpc(input, gw.NATSConn, accountID) }), "DeleteVpc": ec2Handler(func(input *ec2.DeleteVpcInput, gw *GatewayConfig, accountID string) (any, error) { @@ -317,6 +358,9 @@ var ec2Actions = map[string]EC2Handler{ return gateway_ec2_vpc.DescribeVpcAttribute(input, gw.NATSConn, accountID) }), "CreateSubnet": ec2Handler(func(input *ec2.CreateSubnetInput, gw *GatewayConfig, accountID string) (any, error) { + if err := gw.Quota.EnforceSubnets(gw.NATSConn, accountID, 1); err != nil { + return nil, err + } return gateway_ec2_vpc.CreateSubnet(input, gw.NATSConn, accountID) }), "DeleteSubnet": ec2Handler(func(input *ec2.DeleteSubnetInput, gw *GatewayConfig, accountID string) (any, error) { @@ -398,6 +442,9 @@ var ec2Actions = map[string]EC2Handler{ return gateway_ec2_vpc.RevokeSecurityGroupEgress(input, gw.NATSConn, accountID) }), "AllocateAddress": ec2Handler(func(input *ec2.AllocateAddressInput, gw *GatewayConfig, accountID string) (any, error) { + if err := gw.Quota.EnforceEIPs(gw.NATSConn, accountID, 1); err != nil { + return nil, err + } return gateway_ec2_eip.AllocateAddress(input, gw.NATSConn, accountID) }), "ReleaseAddress": ec2Handler(func(input *ec2.ReleaseAddressInput, gw *GatewayConfig, accountID string) (any, error) { diff --git a/spinifex/gateway/ec2/instance/DescribeInstanceStatus.go b/spinifex/gateway/ec2/instance/DescribeInstanceStatus.go index 19ed2ac1a..671490362 100644 --- a/spinifex/gateway/ec2/instance/DescribeInstanceStatus.go +++ b/spinifex/gateway/ec2/instance/DescribeInstanceStatus.go @@ -77,7 +77,7 @@ func queryStoppedInstancesForStatus(natsConn *nats.Conn, input *ec2.DescribeInst slog.Error("DescribeInstanceStatus: failed to marshal stopped query", "err", err) return nil } - reservations := queryInstanceBucket(natsConn, "ec2.DescribeStoppedInstances", jsonData, accountID) + reservations, _ := queryInstanceBucket(natsConn, "ec2.DescribeStoppedInstances", jsonData, accountID) var statuses []*ec2.InstanceStatus for _, res := range reservations { if res == nil { diff --git a/spinifex/gateway/ec2/instance/DescribeInstances.go b/spinifex/gateway/ec2/instance/DescribeInstances.go index e7e09e8af..d96e17951 100644 --- a/spinifex/gateway/ec2/instance/DescribeInstances.go +++ b/spinifex/gateway/ec2/instance/DescribeInstances.go @@ -16,17 +16,49 @@ import ( ) // DescribeInstances fans out to all nodes via NATS and aggregates the results. +// It is lenient: partial results from a slow or unreachable node are returned +// without error, which is correct for user-facing describes. Callers that must +// not act on a partial view (the quota reconcile) use DescribeInstancesForReconcile. func DescribeInstances(input *ec2.DescribeInstancesInput, natsConn *nats.Conn, expectedNodes int, accountID string) (*ec2.DescribeInstancesOutput, error) { + reservations, _, firstClient4xx, err := gatherInstances(input, natsConn, expectedNodes, accountID) + if err != nil { + return nil, err + } + // Propagate a deterministic 4xx only when nothing was collected (fan-out + KV). + if firstClient4xx != "" && len(reservations) == 0 { + return nil, errors.New(firstClient4xx) + } + 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. +func DescribeInstancesForReconcile(input *ec2.DescribeInstancesInput, natsConn *nats.Conn, expectedNodes int, accountID string) (reservations []*ec2.Reservation, complete bool, err error) { + reservations, complete, _, err = gatherInstances(input, natsConn, expectedNodes, accountID) + return reservations, complete, err +} + +// gatherInstances runs the running-instance fan-out plus the stopped/terminated +// KV bucket queries and aggregates every reservation. complete reports whether +// the sweep saw a success frame from all expectedNodes without timing out and +// both bucket queries succeeded — the precondition reconcile needs before it may +// lower a counter. firstClient4xx carries the first deterministic 4xx for the +// lenient caller to surface when nothing was collected. +func gatherInstances(input *ec2.DescribeInstancesInput, natsConn *nats.Conn, expectedNodes int, accountID string) (reservations []*ec2.Reservation, complete bool, firstClient4xx string, err error) { jsonData, err := json.Marshal(input) if err != nil { slog.Error("DescribeInstances: Failed to marshal input", "err", err) - return nil, fmt.Errorf("failed to marshal input: %w", err) + return nil, false, "", fmt.Errorf("failed to marshal input: %w", err) } frames, sum, err := utils.Gather(natsConn, "ec2.DescribeInstances", jsonData, utils.GatherOpts{Timeout: 3 * time.Second, ExpectedNodes: expectedNodes, AccountID: accountID}) if err != nil { - return nil, err + return nil, false, "", err } var allReservations []*ec2.Reservation @@ -37,34 +69,32 @@ func DescribeInstances(input *ec2.DescribeInstancesInput, natsConn *nats.Conn, e } } + // The fan-out is complete only when every expected node answered with a + // success frame and the deadline was not hit; an error frame or a missing + // node leaves the view partial. + fanoutComplete := expectedNodes > 0 && !sum.TimedOut && sum.Successes >= expectedNodes + // Both topics use queue groups (single responder each); query in parallel. var kvMu sync.Mutex var kvWg sync.WaitGroup + bucketsOK := true for _, topic := range []string{"ec2.DescribeStoppedInstances", "ec2.DescribeTerminatedInstances"} { kvWg.Add(1) go func(topic string) { defer kvWg.Done() - reservations := queryInstanceBucket(natsConn, topic, jsonData, accountID) - if len(reservations) > 0 { - kvMu.Lock() - allReservations = append(allReservations, reservations...) - kvMu.Unlock() + reservations, ok := queryInstanceBucket(natsConn, topic, jsonData, accountID) + kvMu.Lock() + defer kvMu.Unlock() + if !ok { + bucketsOK = false } + allReservations = append(allReservations, reservations...) }(topic) } kvWg.Wait() - // Propagate a deterministic 4xx only when nothing was collected (fan-out + KV). - if sum.FirstClient4xx != "" && len(allReservations) == 0 { - return nil, errors.New(sum.FirstClient4xx) - } - - output := &ec2.DescribeInstancesOutput{ - Reservations: allReservations, - } - slog.Info("DescribeInstances: Aggregated response", "total_reservations", len(allReservations)) - return output, nil + return allReservations, fanoutComplete && bucketsOK, sum.FirstClient4xx, nil } // EnrichInstanceProfileIDs resolves IamInstanceProfile.Id for every instance @@ -107,27 +137,30 @@ func EnrichInstanceProfileIDs(out *ec2.DescribeInstancesOutput, iamSvc handlers_ } } -// queryInstanceBucket queries a single describe topic and returns its reservations. -func queryInstanceBucket(natsConn *nats.Conn, topic string, jsonData []byte, accountID string) []*ec2.Reservation { +// queryInstanceBucket queries a single describe topic and returns its +// reservations. ok is false when the query failed (request error or error +// payload), so a reconcile caller can treat the sweep as incomplete rather than +// silently dropping the bucket's instances. +func queryInstanceBucket(natsConn *nats.Conn, topic string, jsonData []byte, accountID string) (reservations []*ec2.Reservation, ok bool) { reqMsg := nats.NewMsg(topic) reqMsg.Data = jsonData reqMsg.Header.Set(utils.AccountIDHeader, accountID) msg, err := natsConn.RequestMsg(reqMsg, 3*time.Second) if err != nil { slog.Warn("DescribeInstances: Failed to query instance bucket", "topic", topic, "err", err) - return nil + return nil, false } if responseError, parseErr := utils.ValidateErrorPayload(msg.Data); parseErr != nil { slog.Warn("DescribeInstances: Instance bucket query returned error", "topic", topic, "code", responseError.Code) - return nil + return nil, false } var output ec2.DescribeInstancesOutput if err := json.Unmarshal(msg.Data, &output); err != nil { slog.Error("DescribeInstances: Failed to unmarshal instance bucket response", "topic", topic, "err", err) - return nil + return nil, false } if len(output.Reservations) > 0 { slog.Info("DescribeInstances: Collected reservations from bucket", "topic", topic, "count", len(output.Reservations)) } - return output.Reservations + return output.Reservations, true } diff --git a/spinifex/gateway/ec2/instance/RebootInstances.go b/spinifex/gateway/ec2/instance/RebootInstances.go index 951a876b6..2b76957b2 100644 --- a/spinifex/gateway/ec2/instance/RebootInstances.go +++ b/spinifex/gateway/ec2/instance/RebootInstances.go @@ -69,7 +69,7 @@ func RebootInstances(input *ec2.RebootInstancesInput, natsConn *nats.Conn, accou if err != nil { return nil, fmt.Errorf("marshal describe input: %w", err) } - if reservations := queryInstanceBucket(natsConn, "ec2.DescribeStoppedInstances", describeData, accountID); len(reservations) > 0 { + if reservations, _ := queryInstanceBucket(natsConn, "ec2.DescribeStoppedInstances", describeData, accountID); len(reservations) > 0 { return nil, errors.New(awserrors.ErrorIncorrectInstanceState) } diff --git a/spinifex/gateway/ec2/instance/RunInstances.go b/spinifex/gateway/ec2/instance/RunInstances.go index 64fcf4b5a..0bd1fa06e 100644 --- a/spinifex/gateway/ec2/instance/RunInstances.go +++ b/spinifex/gateway/ec2/instance/RunInstances.go @@ -19,6 +19,10 @@ import ( // Defined as a callback to avoid an import cycle with the gateway package. type PassRoleChecker func(roleARN string) error +// LaunchQuotaChecker enforces any gateway-side quota that must run after input +// validation and PassRole authorization, but before the launch is dispatched. +type LaunchQuotaChecker func() error + type RunInstancesResponse struct { Reservation *ec2.Reservation `locationName:"RunInstancesResponse"` } @@ -72,14 +76,16 @@ func ValidateRunInstancesInput(input *ec2.RunInstancesInput) (err error) { // When ClientToken is set, wraps the launch in idempotency via the KV store. // iamSvc may be nil only when no IamInstanceProfile is supplied. // passRoleCheck may be nil to skip PassRole enforcement. -func RunInstances(input *ec2.RunInstancesInput, natsConn *nats.Conn, iamSvc handlers_iam.IAMService, accountID string, passRoleCheck PassRoleChecker, expectedNodes int) (reservation ec2.Reservation, err error) { +// launchQuotaCheck may be nil to skip quota enforcement; when present it runs +// after validation and PassRole authorization, but before any launch dispatch. +func RunInstances(input *ec2.RunInstancesInput, natsConn *nats.Conn, iamSvc handlers_iam.IAMService, accountID string, passRoleCheck PassRoleChecker, launchQuotaCheck LaunchQuotaChecker, expectedNodes int) (reservation ec2.Reservation, err error) { if err = ValidateRunInstancesInput(input); err != nil { return reservation, err } token := aws.StringValue(input.ClientToken) if token == "" { - return runInstancesInner(input, natsConn, iamSvc, accountID, passRoleCheck, expectedNodes) + return runInstancesInner(input, natsConn, iamSvc, accountID, passRoleCheck, launchQuotaCheck, expectedNodes) } // ClientToken set: dedup concurrent/retried launches; store failure is fatal @@ -92,17 +98,22 @@ func RunInstances(input *ec2.RunInstancesInput, natsConn *nats.Conn, iamSvc hand // Hash before any mutation so the same token+params always matches. paramHash := clientTokenParamHash(input) return runInstancesWithClientToken(store, accountID, token, paramHash, func() (ec2.Reservation, error) { - return runInstancesInner(input, natsConn, iamSvc, accountID, passRoleCheck, expectedNodes) + return runInstancesInner(input, natsConn, iamSvc, accountID, passRoleCheck, launchQuotaCheck, expectedNodes) }) } // runInstancesInner performs the actual launch: profile resolution, placement // routing, and capacity-aware distribution. Wrapped by RunInstances for idempotency. -func runInstancesInner(input *ec2.RunInstancesInput, natsConn *nats.Conn, iamSvc handlers_iam.IAMService, accountID string, passRoleCheck PassRoleChecker, expectedNodes int) (reservation ec2.Reservation, err error) { +func runInstancesInner(input *ec2.RunInstancesInput, natsConn *nats.Conn, iamSvc handlers_iam.IAMService, accountID string, passRoleCheck PassRoleChecker, launchQuotaCheck LaunchQuotaChecker, expectedNodes int) (reservation ec2.Reservation, err error) { resolvedProfile, err := resolveAndAuthorizeInstanceProfile(input, iamSvc, accountID, passRoleCheck) if err != nil { return reservation, err } + if launchQuotaCheck != nil { + if err := launchQuotaCheck(); err != nil { + return reservation, err + } + } // Targeted capacity-reservation launch routes straight to the owning node's // cr-subject. The gateway does input-only checks (malformed id, the PG+CR diff --git a/spinifex/gateway/ec2/instance/RunInstancesIntoReservation_test.go b/spinifex/gateway/ec2/instance/RunInstancesIntoReservation_test.go index 3ec6e6992..b8441d8db 100644 --- a/spinifex/gateway/ec2/instance/RunInstancesIntoReservation_test.go +++ b/spinifex/gateway/ec2/instance/RunInstancesIntoReservation_test.go @@ -39,7 +39,7 @@ func TestCapacityReservationTargetID(t *testing.T) { // A malformed reservation id is rejected at the gateway before any NATS call. func TestRunInstancesInner_TargetedMalformedID(t *testing.T) { - _, err := runInstancesInner(crTargetedInput("bogus-id"), nil, nil, "123456789012", nil, 1) + _, err := runInstancesInner(crTargetedInput("bogus-id"), nil, nil, "123456789012", nil, nil, 1) require.Error(t, err) assert.Equal(t, awserrors.ErrorInvalidCapacityReservationIdMalformed, err.Error()) } @@ -48,7 +48,7 @@ func TestRunInstancesInner_TargetedMalformedID(t *testing.T) { func TestRunInstancesInner_TargetedWithPlacementGroup(t *testing.T) { input := crTargetedInput(testCRID) input.Placement = &ec2.Placement{GroupName: aws.String("pg-cluster")} - _, err := runInstancesInner(input, nil, nil, "123456789012", nil, 1) + _, err := runInstancesInner(input, nil, nil, "123456789012", nil, nil, 1) require.Error(t, err) assert.Equal(t, awserrors.ErrorInvalidParameterValue, err.Error()) } @@ -117,7 +117,7 @@ func TestRunInstancesInner_TargetedRoutesToCRSubject(t *testing.T) { defer func() { _ = sub.Unsubscribe() }() time.Sleep(50 * time.Millisecond) - reservation, err := runInstancesInner(crTargetedInput(testCRID), nc, nil, "123456789012", nil, 1) + reservation, err := runInstancesInner(crTargetedInput(testCRID), nc, nil, "123456789012", nil, nil, 1) require.NoError(t, err) require.Len(t, reservation.Instances, 1) assert.Equal(t, "i-cr1", aws.StringValue(reservation.Instances[0].InstanceId)) diff --git a/spinifex/gateway/ec2/instance/RunInstances_test.go b/spinifex/gateway/ec2/instance/RunInstances_test.go index ee339b322..5379e058c 100644 --- a/spinifex/gateway/ec2/instance/RunInstances_test.go +++ b/spinifex/gateway/ec2/instance/RunInstances_test.go @@ -210,7 +210,7 @@ func TestParseRunInstances(t *testing.T) { for _, test := range tests { t.Run(test.name, func(t *testing.T) { // For validation tests, we can pass nil conn since validation happens before NATS call - response, err := RunInstances(test.input, nil, nil, "123456789012", nil, 1) + response, err := RunInstances(test.input, nil, nil, "123456789012", nil, nil, 1) // Use assert to check if the error is as expected assert.Equal(t, test.want, err) diff --git a/spinifex/gateway/ec2/instance/placement_test.go b/spinifex/gateway/ec2/instance/placement_test.go index dc181268d..aa187ec63 100644 --- a/spinifex/gateway/ec2/instance/placement_test.go +++ b/spinifex/gateway/ec2/instance/placement_test.go @@ -631,7 +631,7 @@ func TestRunInstances_SingleInstanceDistributes(t *testing.T) { MaxCount: aws.Int64(1), } - reservation, err := RunInstances(input, nc, nil, "test-account", nil, 1) + reservation, err := RunInstances(input, nc, nil, "test-account", nil, nil, 1) require.NoError(t, err) assert.Len(t, reservation.Instances, 1) assert.Equal(t, "i-single", aws.StringValue(reservation.Instances[0].InstanceId)) @@ -1136,7 +1136,7 @@ func TestRunInstances_ClusterPlacementGroupRouting(t *testing.T) { }, } - reservation, err := RunInstances(input, nc, nil, "test-account", nil, 1) + reservation, err := RunInstances(input, nc, nil, "test-account", nil, nil, 1) require.NoError(t, err) assert.Len(t, reservation.Instances, 2) } @@ -1184,7 +1184,7 @@ func TestRunInstances_MultiInstanceUsesDistribution(t *testing.T) { MaxCount: aws.Int64(2), } - reservation, err := RunInstances(input, nc, nil, "test-account", nil, 1) + reservation, err := RunInstances(input, nc, nil, "test-account", nil, nil, 1) require.NoError(t, err) assert.Len(t, reservation.Instances, 2) assert.True(t, statusQueried, "multi-instance launch should query node status") diff --git a/spinifex/gateway/gateway.go b/spinifex/gateway/gateway.go index 6cb7256bc..a0b7e3c8a 100644 --- a/spinifex/gateway/gateway.go +++ b/spinifex/gateway/gateway.go @@ -23,6 +23,7 @@ import ( gateway_ecrauth "github.com/mulgadc/spinifex/spinifex/gateway/ecrauth" "github.com/mulgadc/spinifex/spinifex/gateway/policy" handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam" + handlers_quota "github.com/mulgadc/spinifex/spinifex/handlers/quota" handlers_sts "github.com/mulgadc/spinifex/spinifex/handlers/sts" "github.com/mulgadc/spinifex/spinifex/types" "github.com/mulgadc/spinifex/spinifex/utils" @@ -91,8 +92,12 @@ type GatewayConfig struct { STSService handlers_sts.STSService RateLimiter *AuthRateLimiter // Per-IP auth failure rate limiter Throttler *ratelimit.Throttler // Per-account+action API request throttler - Version string // Build-time version string (set from cmd.Version) - Commit string // Build-time commit hash (set from cmd.Commit) + // Quota enforces per-account service quotas. Built unconditionally; a disabled + // config yields a no-op Service whose Exempt always returns true. Nil only in + // unit tests of unrelated routes, where no handler reaches the quota checks. + Quota *handlers_quota.Service + Version string // Build-time version string (set from cmd.Version) + Commit string // Build-time commit hash (set from cmd.Commit) // ECRRegistry serves the OCI Distribution v2 (/v2/*) surface. Nil falls back // to the 501 stub (e.g. in unit tests of unrelated routes). ECRRegistry *gateway_ecr.Registry diff --git a/spinifex/handlers/quota/live.go b/spinifex/handlers/quota/live.go new file mode 100644 index 000000000..9ca957ac7 --- /dev/null +++ b/spinifex/handlers/quota/live.go @@ -0,0 +1,59 @@ +package handlers_quota + +import ( + "errors" + + "github.com/aws/aws-sdk-go/service/ec2" + "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" + "github.com/nats-io/nats.go" +) + +// exceeds rejects with ResourceLimitExceeded when an account already holding +// count of a live-counted dimension would pass its cap by adding want more. It is +// 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 count+want > limit { + return errors.New(awserrors.ErrorResourceLimitExceeded) + } + return nil +} + +// EnforceVPCs gates CreateVpc on the account's live DescribeVpcs count. +func (s *Service) EnforceVPCs(natsConn *nats.Conn, accountID string, want int) error { + if s.Exempt(accountID) { + return nil + } + out, err := gateway_ec2_vpc.DescribeVpcs(&ec2.DescribeVpcsInput{}, natsConn, accountID) + if err != nil { + return err + } + return exceeds(len(out.Vpcs), want, s.limits.VPCs) +} + +// EnforceSubnets gates CreateSubnet on the account's live DescribeSubnets count. +// Subnets are capped per-account in aggregate, not per-VPC. +func (s *Service) EnforceSubnets(natsConn *nats.Conn, accountID string, want int) error { + if s.Exempt(accountID) { + return nil + } + out, err := gateway_ec2_vpc.DescribeSubnets(&ec2.DescribeSubnetsInput{}, natsConn, accountID) + if err != nil { + return err + } + return exceeds(len(out.Subnets), want, s.limits.Subnets) +} + +// EnforceEIPs gates AllocateAddress on the account's live DescribeAddresses count. +func (s *Service) EnforceEIPs(natsConn *nats.Conn, accountID string, want int) error { + if s.Exempt(accountID) { + return nil + } + out, err := gateway_ec2_eip.DescribeAddresses(&ec2.DescribeAddressesInput{}, natsConn, accountID) + if err != nil { + return err + } + return exceeds(len(out.Addresses), want, s.limits.EIPs) +} diff --git a/spinifex/handlers/quota/live_test.go b/spinifex/handlers/quota/live_test.go new file mode 100644 index 000000000..f1688d573 --- /dev/null +++ b/spinifex/handlers/quota/live_test.go @@ -0,0 +1,67 @@ +package handlers_quota + +import ( + "testing" + + "github.com/mulgadc/spinifex/spinifex/awserrors" + "github.com/mulgadc/spinifex/spinifex/utils" +) + +// liveLimits is a representative enabled tier with distinct caps per dimension. +var liveLimits = Limits{Enabled: true, VPCs: 8, Subnets: 16, EIPs: 2, VolumesGiB: 100} + +// TestExceeds covers the shared comparison every live dimension runs after its +// Describe* count: under the cap passes, at or over it rejects, and a single +// oversized want is rejected on its own. +func TestExceeds(t *testing.T) { + tests := []struct { + name string + count, want, limit int + wantErr string + }{ + {"under limit", 7, 1, 8, ""}, + {"at limit rejects", 8, 1, 8, awserrors.ErrorResourceLimitExceeded}, + {"empty fills to limit", 0, 8, 8, ""}, + {"want overshoots limit", 0, 9, 8, awserrors.ErrorResourceLimitExceeded}, + {"count plus want overshoots", 6, 3, 8, awserrors.ErrorResourceLimitExceeded}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := exceeds(tt.count, tt.want, tt.limit) + switch { + case tt.wantErr == "" && err != nil: + t.Fatalf("exceeds(%d, %d, %d) = %v, want nil", tt.count, tt.want, tt.limit, err) + case tt.wantErr != "" && (err == nil || err.Error() != tt.wantErr): + t.Fatalf("exceeds(%d, %d, %d) = %v, want %q", tt.count, tt.want, tt.limit, err, tt.wantErr) + } + }) + } +} + +// Exempt callers must short-circuit before the Describe* round trip, so the +// per-dimension methods return nil even with a nil NATS connection (which would +// otherwise panic on describe). This covers both the disabled-config and the +// system-account exemptions for all three live dimensions. +func TestEnforceLiveExemptShortCircuits(t *testing.T) { + const normalAccount = "123456789012" + disabled := New(Limits{Enabled: false}, nil) + enabled := New(liveLimits, nil) + cases := []struct { + name string + fn func() error + }{ + {"vpc disabled", func() error { return disabled.EnforceVPCs(nil, normalAccount, 1) }}, + {"subnet disabled", func() error { return disabled.EnforceSubnets(nil, normalAccount, 1) }}, + {"eip disabled", func() error { return disabled.EnforceEIPs(nil, normalAccount, 1) }}, + {"vpc system account", func() error { return enabled.EnforceVPCs(nil, utils.GlobalAccountID, 1) }}, + {"subnet system account", func() error { return enabled.EnforceSubnets(nil, utils.GlobalAccountID, 1) }}, + {"eip system account", func() error { return enabled.EnforceEIPs(nil, utils.GlobalAccountID, 1) }}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if err := tc.fn(); err != nil { + t.Fatalf("exempt path returned %v, want nil", err) + } + }) + } +} diff --git a/spinifex/handlers/quota/quota.go b/spinifex/handlers/quota/quota.go new file mode 100644 index 000000000..c281133ec --- /dev/null +++ b/spinifex/handlers/quota/quota.go @@ -0,0 +1,65 @@ +// Package handlers_quota enforces per-account service quotas in the AWS gateway. +// It caps how much standing infrastructure a single account can hold (vCPUs, +// VPCs, subnets, Elastic IPs, EBS storage). Limits are a single config-tunable +// tier; the system account and disabled configs bypass every check. +package handlers_quota + +import ( + "time" + + "github.com/mulgadc/spinifex/spinifex/utils" + "github.com/nats-io/nats.go" +) + +// ReconcileInterval is how often the gateway recomputes vCPU counters. It is +// fixed rather than configurable to keep the [quota] block simple. +const ReconcileInterval = 30 * time.Second + +// KVBucketAccountUsage is the gateway-owned KV bucket holding one integer vCPU +// counter per account, keyed by accountID. It is the only counter-backed +// dimension; the live-counted dimensions need no stored state. +const KVBucketAccountUsage = "spinifex-account-usage" + +// KVBucketQuotaReconcile is the leader-lock bucket for the gateway quota +// reconcile. It is deliberately distinct from vpcd's network-reconcile lock so +// the two independent loops never share one mutex and block each other. +const KVBucketQuotaReconcile = "spinifex-quota-reconcile" + +// 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. +type Limits struct { + Enabled bool `toml:"enabled"` + VCPUs int `toml:"vcpus"` + VPCs int `toml:"vpcs"` + Subnets int `toml:"subnets"` + EIPs int `toml:"eips"` + VolumesGiB int `toml:"volumes_gib"` +} + +// Service enforces per-account quotas for one gateway. It holds the configured +// limits and the per-account vCPU counter bucket; the live-counted dimensions +// derive their usage from account-filtered Describe* calls and need no state. +type Service struct { + limits Limits + // usage holds the per-account vCPU counters (key {accountID}). Nil when + // quotas are disabled, in which case Exempt short-circuits every check + // before the counter is touched. + usage nats.KeyValue +} + +// New constructs a quota Service from the configured limits and the gateway-owned +// account-usage KV bucket. usage may be nil when quotas are disabled; Exempt then +// short-circuits every check before the counter is read. +func New(limits Limits, usage nats.KeyValue) *Service { + return &Service{limits: limits, usage: usage} +} + +// 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. +func (s *Service) Exempt(accountID string) bool { + if s == nil || !s.limits.Enabled { + return true + } + return accountID == utils.GlobalAccountID +} diff --git a/spinifex/handlers/quota/quota_test.go b/spinifex/handlers/quota/quota_test.go new file mode 100644 index 000000000..b4072aa67 --- /dev/null +++ b/spinifex/handlers/quota/quota_test.go @@ -0,0 +1,89 @@ +package handlers_quota + +import ( + "testing" + + "github.com/mulgadc/spinifex/spinifex/utils" + toml "github.com/pelletier/go-toml/v2" +) + +// quotaTOML mirrors how awsgw.toml nests the limits under a [quota] table. +type quotaTOML struct { + Quota Limits `toml:"quota"` +} + +func TestLimitsTOMLRoundTrip(t *testing.T) { + const src = ` +[quota] +enabled = true +vcpus = 8 +vpcs = 8 +subnets = 16 +eips = 2 +volumes_gib = 100 +` + want := Limits{ + Enabled: true, + VCPUs: 8, + VPCs: 8, + Subnets: 16, + EIPs: 2, + VolumesGiB: 100, + } + + var parsed quotaTOML + if err := toml.Unmarshal([]byte(src), &parsed); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if parsed.Quota != want { + t.Fatalf("parsed = %+v, want %+v", parsed.Quota, want) + } + + // Marshal then unmarshal again must reproduce the same struct. + encoded, err := toml.Marshal(parsed) + if err != nil { + t.Fatalf("marshal: %v", err) + } + var back quotaTOML + if err := toml.Unmarshal(encoded, &back); err != nil { + t.Fatalf("re-unmarshal: %v", err) + } + if back.Quota != want { + t.Fatalf("round-trip = %+v, want %+v", back.Quota, want) + } +} + +// A config without a [quota] block must decode to the zero value, which is a +// disabled no-op so existing gateways are unaffected. +func TestLimitsAbsentIsZeroValue(t *testing.T) { + var parsed quotaTOML + if err := toml.Unmarshal([]byte("region = \"us-west-1\"\n"), &parsed); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if parsed.Quota != (Limits{}) { + t.Fatalf("absent [quota] = %+v, want zero value", parsed.Quota) + } +} + +func TestExempt(t *testing.T) { + const normalAccount = "123456789012" + tests := []struct { + name string + limits Limits + accountID string + want bool + }{ + {"system account when enabled", Limits{Enabled: true}, utils.GlobalAccountID, true}, + {"normal account when enabled", Limits{Enabled: true}, normalAccount, false}, + {"normal account when disabled", Limits{Enabled: false}, normalAccount, true}, + {"system account when disabled", Limits{Enabled: false}, utils.GlobalAccountID, true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + s := New(tt.limits, nil) + if got := s.Exempt(tt.accountID); got != tt.want { + t.Errorf("Exempt(%q) = %v, want %v", tt.accountID, got, tt.want) + } + }) + } +} diff --git a/spinifex/handlers/quota/reconcile.go b/spinifex/handlers/quota/reconcile.go new file mode 100644 index 000000000..8d233b10e --- /dev/null +++ b/spinifex/handlers/quota/reconcile.go @@ -0,0 +1,119 @@ +package handlers_quota + +import ( + "context" + "fmt" + "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. +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( + &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 { + if s == nil || !s.limits.Enabled { + return nil + } + ids, err := accounts() + if err != nil { + return fmt.Errorf("quota reconcile: list accounts: %w", err) + } + var firstErr error + for _, accountID := range ids { + 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(accountID, sumReservationVCPUs(reservations), complete); err != nil { + slog.Warn("quota reconcile: counter overwrite failed", "account", accountID, "err", err) + if firstErr == nil { + firstErr = err + } + } + } + return firstErr +} + +// 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. +func sumReservationVCPUs(reservations []*ec2.Reservation) int { + total := 0 + for _, res := range reservations { + if res == nil { + continue + } + for _, inst := range res.Instances { + if inst == nil || inst.InstanceType == nil || isTerminalState(inst.State) { + continue + } + if vcpus, ok := instancetypes.DefaultVCPUs(*inst.InstanceType); ok { + total += vcpus + } + } + } + return total +} + +// isTerminalState reports whether an instance 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 on the next pass. +func isTerminalState(state *ec2.InstanceState) bool { + if state == nil || state.Name == nil { + return false + } + switch *state.Name { + case ec2.InstanceStateNameShuttingDown, ec2.InstanceStateNameTerminated: + return true + default: + return false + } +} diff --git a/spinifex/handlers/quota/reconcile_test.go b/spinifex/handlers/quota/reconcile_test.go new file mode 100644 index 000000000..a1447d4a9 --- /dev/null +++ b/spinifex/handlers/quota/reconcile_test.go @@ -0,0 +1,210 @@ +package handlers_quota + +import ( + "context" + "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) + } + return byAccount[accountID], true, nil + } +} + +func accountList(ids ...string) AccountLister { + return func() ([]string, error) { return ids, nil } +} + +func TestReconcileCorrectsOverCount(t *testing.T) { + s := newVCPUService(t, Limits{Enabled: true, VCPUs: 100}) + + // Counter drifted high (e.g. a stale increment); the account truly holds one + // m5.xlarge (4) plus one t3.micro (2) = 6. + if err := s.AddVCPU(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 { + 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(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) { + s := newVCPUService(t, Limits{Enabled: true, VCPUs: 100}) + + if err := s.AddVCPU(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 { + 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. +func TestReconcileSkipsSystemAccount(t *testing.T) { + s := newVCPUService(t, Limits{Enabled: true, VCPUs: 100}) + + // Park a non-zero value directly under the system key; AddVCPU would no-op it. + if _, err := s.usage.PutString(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 { + 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) { + s := newVCPUService(t, Limits{Enabled: true, VCPUs: 100}) + + const bad, good = "111111111111", "222222222222" + if err := s.AddVCPU(bad, 8); err != nil { + t.Fatalf("seed bad: %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 + } + + err := s.Reconcile(context.Background(), accountList(bad, good), lister) + 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) +} + +// 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. +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(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 { + 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 { + t.Fatalf("Reconcile raise: %v", err) + } + assertCounter(t, s, testAccount, 12) // raised to the observed 3 x 4 = 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. +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) + } + var nilService *Service + if err := nilService.Reconcile(context.Background(), nil, nil); err != nil { + t.Fatalf("nil Reconcile = %v, want nil", err) + } +} + +func assertCounter(t *testing.T, s *Service, accountID string, want int) { + t.Helper() + got, _, err := s.readVCPU(accountID) + if err != nil { + t.Fatalf("readVCPU(%s): %v", accountID, err) + } + if got != want { + t.Fatalf("counter[%s] = %d, want %d", accountID, got, want) + } +} diff --git a/spinifex/handlers/quota/storage.go b/spinifex/handlers/quota/storage.go new file mode 100644 index 000000000..376ba3bbc --- /dev/null +++ b/spinifex/handlers/quota/storage.go @@ -0,0 +1,66 @@ +package handlers_quota + +import ( + "github.com/aws/aws-sdk-go/service/ec2" + gateway_ec2_volume "github.com/mulgadc/spinifex/spinifex/gateway/ec2/volume" + "github.com/nats-io/nats.go" +) + +// EnforceVolumeCreate gates CreateVolume on the account's live storage usage: +// the sum of its existing volume sizes plus requestedGiB must stay within the +// 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. +func (s *Service) EnforceVolumeCreate(natsConn *nats.Conn, accountID string, requestedGiB int) error { + if s.Exempt(accountID) { + return nil + } + total, _, err := s.volumeUsage(natsConn, accountID, "") + if err != nil { + return err + } + return exceeds(total, requestedGiB, s.limits.VolumesGiB) +} + +// EnforceVolumeModify gates ModifyVolume on the account's live storage usage. +// The resized volume's current size is already in the live sum, so the growth +// charged is newGiB - oldGiB; AWS only grows volumes, so this delta is >= 0 and +// a shrink can never falsely trip the cap. +func (s *Service) EnforceVolumeModify(natsConn *nats.Conn, accountID, volumeID string, newGiB int) error { + if s.Exempt(accountID) { + return nil + } + total, oldGiB, err := s.volumeUsage(natsConn, accountID, volumeID) + if err != nil { + return err + } + return exceeds(total-oldGiB, newGiB, s.limits.VolumesGiB) +} + +// volumeUsage sums the account's volume sizes in GiB via a single DescribeVolumes +// sweep and, when volumeID is non-empty, also returns that volume's current size +// (0 if absent). One pass serves both the create and modify checks. +func (s *Service) volumeUsage(natsConn *nats.Conn, accountID, volumeID string) (total, target int, err error) { + out, err := gateway_ec2_volume.DescribeVolumes(&ec2.DescribeVolumesInput{}, natsConn, accountID) + if err != nil { + return 0, 0, err + } + total, target = sumVolumeGiB(out.Volumes, volumeID) + return total, target, nil +} + +// sumVolumeGiB totals the sizes of volumes in GiB and, when volumeID is non-empty, +// returns that volume's size as target. Unsized volumes contribute nothing. +func sumVolumeGiB(volumes []*ec2.Volume, volumeID string) (total, target int) { + for _, v := range volumes { + if v == nil || v.Size == nil { + continue + } + size := int(*v.Size) + total += size + if volumeID != "" && v.VolumeId != nil && *v.VolumeId == volumeID { + target = size + } + } + return total, target +} diff --git a/spinifex/handlers/quota/storage_test.go b/spinifex/handlers/quota/storage_test.go new file mode 100644 index 000000000..b412ec5d1 --- /dev/null +++ b/spinifex/handlers/quota/storage_test.go @@ -0,0 +1,111 @@ +package handlers_quota + +import ( + "testing" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/service/ec2" + "github.com/mulgadc/spinifex/spinifex/awserrors" + "github.com/mulgadc/spinifex/spinifex/utils" +) + +func vol(id string, size int64) *ec2.Volume { + return &ec2.Volume{VolumeId: aws.String(id), Size: aws.Int64(size)} +} + +// TestSumVolumeGiB covers the size-sum helper: totals are summed, the named +// volume is reported back, and nil/unsized entries contribute nothing. Root/AMI +// volumes never appear in DescribeVolumes, so an empty list models a fresh +// account and confirms the helper charges only data volumes it is given. +func TestSumVolumeGiB(t *testing.T) { + tests := []struct { + name string + volumes []*ec2.Volume + volumeID string + wantTotal int + wantTarget int + }{ + {"empty", nil, "", 0, 0}, + {"sum without target", []*ec2.Volume{vol("vol-a", 30), vol("vol-b", 30)}, "", 60, 0}, + {"sum with target", []*ec2.Volume{vol("vol-a", 30), vol("vol-b", 30)}, "vol-a", 60, 30}, + {"nil and unsized skipped", []*ec2.Volume{vol("vol-a", 40), nil, {VolumeId: aws.String("vol-b")}}, "vol-a", 40, 40}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + total, target := sumVolumeGiB(tt.volumes, tt.volumeID) + if total != tt.wantTotal || target != tt.wantTarget { + t.Fatalf("sumVolumeGiB() = (%d, %d), want (%d, %d)", total, target, tt.wantTotal, tt.wantTarget) + } + }) + } +} + +// TestEnforceStorage drives the full per-volume pipeline the wiring methods run +// after DescribeVolumes: sum the live volumes, then compare. CreateVolume charges +// total + requested; ModifyVolume charges total - oldSize + newSize so the +// resized volume's existing size is not double-counted. A snapshot restore is a +// normal sized create, so it is gated identically. +func TestEnforceStorage(t *testing.T) { + const limit = 100 // liveLimits.VolumesGiB + existing := []*ec2.Volume{vol("vol-a", 30), vol("vol-b", 30)} // 60 GiB in use + + tests := []struct { + name string + op string // "create" or "modify" + volumeID string // modify target + want int // create: requested; modify: new size + wantErr string + }{ + {"create under limit", "create", "", 30, ""}, + {"create to limit", "create", "", 40, ""}, + {"create over limit", "create", "", 41, awserrors.ErrorResourceLimitExceeded}, + {"snapshot restore past limit", "create", "", 50, awserrors.ErrorResourceLimitExceeded}, + {"resize to limit", "modify", "vol-a", 70, ""}, + {"resize past limit", "modify", "vol-a", 80, awserrors.ErrorResourceLimitExceeded}, + {"resize shrink never trips", "modify", "vol-a", 1, ""}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + total, target := sumVolumeGiB(existing, tt.volumeID) + var err error + switch tt.op { + case "create": + err = exceeds(total, tt.want, limit) + case "modify": + err = exceeds(total-target, tt.want, limit) + } + switch { + case tt.wantErr == "" && err != nil: + t.Fatalf("%s want nil, got %v (limit %d)", tt.name, err, limit) + case tt.wantErr != "" && (err == nil || err.Error() != tt.wantErr): + t.Fatalf("%s = %v, want %q", tt.name, err, tt.wantErr) + } + }) + } +} + +// Exempt callers must short-circuit before the DescribeVolumes round trip, so the +// gates return nil even with a nil NATS connection (which would otherwise panic +// on describe). Covers both the disabled-config and system-account exemptions for +// create and modify. +func TestEnforceVolumeExemptShortCircuits(t *testing.T) { + const normalAccount = "123456789012" + disabled := New(Limits{Enabled: false}, nil) + enabled := New(liveLimits, nil) + cases := []struct { + name string + fn func() error + }{ + {"create disabled", func() error { return disabled.EnforceVolumeCreate(nil, normalAccount, 1) }}, + {"modify disabled", func() error { return disabled.EnforceVolumeModify(nil, normalAccount, "vol-a", 1) }}, + {"create system account", func() error { return enabled.EnforceVolumeCreate(nil, utils.GlobalAccountID, 1) }}, + {"modify system account", func() error { return enabled.EnforceVolumeModify(nil, utils.GlobalAccountID, "vol-a", 1) }}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if err := tc.fn(); err != nil { + t.Fatalf("exempt path returned %v, want nil", err) + } + }) + } +} diff --git a/spinifex/handlers/quota/vcpu.go b/spinifex/handlers/quota/vcpu.go new file mode 100644 index 000000000..235eff84b --- /dev/null +++ b/spinifex/handlers/quota/vcpu.go @@ -0,0 +1,266 @@ +package handlers_quota + +import ( + "encoding/json" + "errors" + "fmt" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/service/ec2" + "github.com/mulgadc/spinifex/spinifex/awserrors" + gateway_ec2_instance "github.com/mulgadc/spinifex/spinifex/gateway/ec2/instance" + "github.com/mulgadc/spinifex/spinifex/instancetypes" + "github.com/nats-io/nats.go" +) + +// vcpuCASRetries bounds AddVCPU's retry on a revision conflict. Each retry +// implies another writer committed, so a single grow can conflict at most as +// many times as there are concurrent grows on the account; the bound sits well +// above any realistic in-flight launch burst and trips only as a circuit +// breaker, never on genuine contention. +const vcpuCASRetries = 100 + +// CheckVCPU rejects with ResourceLimitExceeded when charging want more vCPUs to +// accountID would exceed the configured cap. It only reads the counter; the +// caller increments via AddVCPU after the grow succeeds. Two concurrent checks +// can both pass before either increments (the documented soft-cap window); the +// per-node physical gate still backstops real overcommit. +func (s *Service) CheckVCPU(accountID string, want int) error { + if s.Exempt(accountID) { + return nil + } + current, _, err := s.readVCPU(accountID) + if err != nil { + return err + } + if current+want > s.limits.VCPUs { + return errors.New(awserrors.ErrorResourceLimitExceeded) + } + return nil +} + +// casVCPU applies next to accountID's counter under bounded JetStream CAS so +// concurrent writers cannot lose updates. next maps the current value to the new +// value and a skip flag; skip leaves the counter untouched and never creates a +// key. A revision conflict is retried; exhausting the bound is a hard error. +func (s *Service) casVCPU(accountID string, next func(current int) (value int, skip bool)) error { + for range vcpuCASRetries { + current, revision, err := s.readVCPU(accountID) + if err != nil { + return err + } + value, skip := next(current) + if skip { + return nil + } + data, err := json.Marshal(value) + if err != nil { + return err + } + if revision == 0 { + _, err = s.usage.Create(accountID, data) + } else { + _, err = s.usage.Update(accountID, data, revision) + } + if err == nil { + return nil + } + if !isCASConflict(err) { + return err + } + } + return fmt.Errorf("vcpu counter CAS exhausted for %s after %d attempts", accountID, vcpuCASRetries) +} + +// AddVCPU adds delta vCPUs to accountID's counter under CAS so concurrent grows +// cannot lose updates. A non-positive delta is a no-op: shrinks are never charged +// and are left to reconcile to lower the counter. +func (s *Service) AddVCPU(accountID string, delta int) error { + if s.Exempt(accountID) || delta <= 0 { + return nil + } + return s.casVCPU(accountID, func(current int) (int, bool) { + return current + delta, false + }) +} + +// setVCPU overwrites accountID's counter to value, the only operation that lowers +// a counter. A counter already equal to value is left untouched, so a steady-state +// pass writes nothing and an account with no usage and no key is never created. +// Reconcile is the sole caller and runs under the leader lock. +func (s *Service) setVCPU(accountID string, value int) error { + return s.casVCPU(accountID, func(current int) (int, bool) { + return value, current == value + }) +} + +// 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(accountID string, value int, complete bool) error { + if complete { + return s.setVCPU(accountID, value) + } + current, _, err := s.readVCPU(accountID) + if err != nil { + return err + } + if value <= current { + return nil + } + return s.setVCPU(accountID, value) +} + +// isCASConflict reports whether err is a lost optimistic-concurrency race: a +// Create on an existing key or an Update against a stale revision. Both map to +// JetStream's wrong-last-sequence code and are retryable; any other error is a +// genuine failure the caller must surface. +func isCASConflict(err error) bool { + if errors.Is(err, nats.ErrKeyExists) { + return true + } + var apiErr *nats.APIError + if errors.As(err, &apiErr) { + return apiErr.ErrorCode == nats.JSErrCodeStreamWrongLastSequence + } + return false +} + +// EnforceLaunch is the check-before gate for RunInstances: it rejects when the +// worst-case charge of maxCount instances of instanceType would push accountID +// over its vCPU cap. The actual launched count is charged afterwards via +// ChargeLaunch. An unknown instance type contributes nothing and is left for the +// daemon to reject as InvalidInstanceType; the Exempt short-circuit lives in +// CheckVCPU. +func (s *Service) EnforceLaunch(accountID, instanceType string, maxCount int) error { + perType, ok := instancetypes.DefaultVCPUs(instanceType) + if !ok { + return nil + } + return s.CheckVCPU(accountID, perType*maxCount) +} + +// ChargeLaunch is the increment-after for a successful RunInstances: it adds the +// vCPUs actually launched, summed from the returned reservation, to accountID's +// counter. The daemon may launch fewer than maxCount, so the charge is the real +// reservation rather than the checked worst case. The caller treats a write +// failure as drift for reconcile to correct, never failing the live launch. +func (s *Service) ChargeLaunch(accountID string, reservation *ec2.Reservation) error { + return s.AddVCPU(accountID, sumReservationVCPUs([]*ec2.Reservation{reservation})) +} + +// InstanceTypeResolver returns the current instance type of instanceID owned by +// accountID. ok is false when the instance does not exist for the account, in +// which case the retype gate charges nothing and leaves the modify for the daemon +// to reject. It mirrors reconcile's InstanceLister so unit tests can inject a +// static type without a NATS round trip. +type InstanceTypeResolver func(accountID, instanceID string) (instanceType string, ok bool, err error) + +// NATSInstanceTypeResolver builds the production resolver: an account-filtered, +// single-instance describe returning the live type of a running or stopped +// instance. It uses the strict reconcile variant so an incomplete sweep does not +// masquerade as "instance absent" and wave a retype past the cap (see +// confirmInstanceType). expectedNodes is the configured node total, re-evaluated +// per call so a config change is reflected without a restart. +func NATSInstanceTypeResolver(natsConn *nats.Conn, expectedNodes func() int) InstanceTypeResolver { + return func(accountID, instanceID string) (string, bool, error) { + reservations, complete, err := gateway_ec2_instance.DescribeInstancesForReconcile( + &ec2.DescribeInstancesInput{InstanceIds: []*string{aws.String(instanceID)}}, + natsConn, expectedNodes(), accountID) + if err != nil { + return "", false, err + } + return confirmInstanceType(reservations, instanceID, complete) + } +} + +// confirmInstanceType resolves instanceID's type from a describe sweep. A genuine +// absence on a complete sweep returns ok false with no error, so the retype gate +// charges nothing and the daemon rejects the modify. An instance that cannot be +// found on an incomplete sweep (a node or bucket missed) fails closed with +// ErrorServerInternal: a partial view must not be read as absent and let a retype +// skip its cap check. The client retries once the sweep is clean. +func confirmInstanceType(reservations []*ec2.Reservation, instanceID string, complete bool) (string, bool, error) { + instanceType, ok := instanceTypeFromReservations(reservations, instanceID) + if !ok && !complete { + return "", false, errors.New(awserrors.ErrorServerInternal) + } + return instanceType, ok, nil +} + +// EnforceRetype is the check-before gate for a ModifyInstanceAttribute that +// changes InstanceType: it rejects when retyping accountID's instance to newType +// would push its vCPU counter over the cap. It returns the vCPU delta the caller +// charges via AddVCPU once the daemon applies the retype. A retype-down or +// same-size change yields a non-positive delta that charges nothing and is left +// to reconcile. Exempt accounts, an unknown newType (the daemon rejects it), and +// an instance the resolver cannot find (the daemon rejects the modify) all return +// 0 with no check. +func (s *Service) EnforceRetype(resolve InstanceTypeResolver, accountID, instanceID, newType string) (int, error) { + if s.Exempt(accountID) { + return 0, nil + } + newVCPUs, ok := instancetypes.DefaultVCPUs(newType) + if !ok { + return 0, nil + } + oldType, ok, err := resolve(accountID, instanceID) + if err != nil { + return 0, err + } + if !ok { + return 0, nil + } + oldVCPUs, ok := instancetypes.DefaultVCPUs(oldType) + if !ok { + return 0, nil + } + delta := newVCPUs - oldVCPUs + if delta <= 0 { + return delta, nil + } + if err := s.CheckVCPU(accountID, delta); err != nil { + return 0, err + } + return delta, nil +} + +// instanceTypeFromReservations finds instanceID among reservations and returns +// its type. ok is false when the instance is absent, terminal, or untyped, so the +// retype gate charges nothing and defers to the daemon. +func instanceTypeFromReservations(reservations []*ec2.Reservation, instanceID string) (string, bool) { + for _, res := range reservations { + if res == nil { + continue + } + for _, inst := range res.Instances { + if inst == nil || inst.InstanceType == nil || isTerminalState(inst.State) { + continue + } + if aws.StringValue(inst.InstanceId) == instanceID { + return *inst.InstanceType, true + } + } + } + return "", false +} + +// readVCPU returns accountID's current reserved vCPU count and the KV revision +// it was read at. A missing key is the zero counter at revision 0, which AddVCPU +// treats as a create. +func (s *Service) readVCPU(accountID string) (count int, revision uint64, err error) { + entry, err := s.usage.Get(accountID) + if err != nil { + if errors.Is(err, nats.ErrKeyNotFound) { + return 0, 0, nil + } + return 0, 0, err + } + if err := json.Unmarshal(entry.Value(), &count); err != nil { + return 0, 0, fmt.Errorf("unmarshal vcpu counter for %s: %w", accountID, err) + } + return count, entry.Revision(), nil +} diff --git a/spinifex/handlers/quota/vcpu_launch_test.go b/spinifex/handlers/quota/vcpu_launch_test.go new file mode 100644 index 000000000..d302ab8a7 --- /dev/null +++ b/spinifex/handlers/quota/vcpu_launch_test.go @@ -0,0 +1,143 @@ +package handlers_quota + +import ( + "testing" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/service/ec2" + "github.com/mulgadc/spinifex/spinifex/awserrors" + "github.com/mulgadc/spinifex/spinifex/utils" +) + +// reservationOf builds a reservation of n instances of one type, modelling what +// the daemon returns from a launch so ChargeLaunch sums the actual vCPUs. +func reservationOf(instanceType string, n int) *ec2.Reservation { + res := &ec2.Reservation{} + for range n { + res.Instances = append(res.Instances, &ec2.Instance{InstanceType: aws.String(instanceType)}) + } + return res +} + +// TestEnforceLaunch covers the check-before gate: the charge tested is the +// worst case maxCount * the type's vCPUs, an at-limit launch passes, one over +// rejects, headroom is respected, and an unknown type is left for the daemon to +// reject (charged nothing here). Each case uses its own account so the embedded +// JetStream is started once for the whole table. +func TestEnforceLaunch(t *testing.T) { + s := newVCPUService(t, Limits{Enabled: true, VCPUs: 8}) // t3.micro = 2 vCPU + + tests := []struct { + name string + account string + instanceType string + maxCount int + reserve int // pre-charged vCPUs + wantErr string + }{ + {"empty to limit", "100000000001", "t3.micro", 4, 0, ""}, + {"empty over limit", "100000000002", "t3.micro", 5, 0, awserrors.ErrorResourceLimitExceeded}, + {"headroom respected", "100000000003", "t3.micro", 1, 6, ""}, + {"headroom exceeded", "100000000004", "t3.micro", 2, 6, awserrors.ErrorResourceLimitExceeded}, + {"unknown type passes", "100000000005", "made.up", 1000, 0, ""}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if tt.reserve > 0 { + if err := s.AddVCPU(tt.account, tt.reserve); err != nil { + t.Fatalf("AddVCPU(%d): %v", tt.reserve, err) + } + } + err := s.EnforceLaunch(tt.account, tt.instanceType, tt.maxCount) + switch { + case tt.wantErr == "" && err != nil: + t.Fatalf("EnforceLaunch(%s, %d) = %v, want nil", tt.instanceType, tt.maxCount, err) + case tt.wantErr != "" && (err == nil || err.Error() != tt.wantErr): + t.Fatalf("EnforceLaunch(%s, %d) = %v, want %q", tt.instanceType, tt.maxCount, err, tt.wantErr) + } + }) + } +} + +// TestChargeLaunchActual charges the vCPUs the daemon actually launched, not the +// checked worst case: an EnforceLaunch sized for maxCount=4 passes, but when the +// daemon returns only 2 instances the counter rises by their vCPUs alone. +func TestChargeLaunchActual(t *testing.T) { + s := newVCPUService(t, Limits{Enabled: true, VCPUs: 8}) + + if err := s.EnforceLaunch(testAccount, "t3.micro", 4); err != nil { + t.Fatalf("EnforceLaunch = %v, want nil", err) + } + // Daemon launched only 2; charge the actual 4 vCPUs, not the checked 8. + if err := s.ChargeLaunch(testAccount, reservationOf("t3.micro", 2)); err != nil { + t.Fatalf("ChargeLaunch = %v", err) + } + got, _, err := s.readVCPU(testAccount) + if err != nil { + t.Fatalf("readVCPU: %v", err) + } + if got != 4 { + t.Fatalf("counter = %d, want 4 (actual launched, not the checked worst case)", got) + } +} + +// TestLaunchSoftCapBounded documents the soft-cap window: two launches read the +// counter before either charges, so both pass the check and the counter then +// overshoots the cap — but only by the in-flight launches. CAS loses no update, +// so the final counter is exactly their sum, never unbounded. +func TestLaunchSoftCapBounded(t *testing.T) { + s := newVCPUService(t, Limits{Enabled: true, VCPUs: 8}) + + // Pre-charge to 6/8 so each 2-vCPU launch individually fits the headroom. + if err := s.AddVCPU(testAccount, 6); err != nil { + t.Fatalf("AddVCPU(6): %v", err) + } + + // Two launches enter the gate before either charges: both read 6/8 and pass. + if err := s.EnforceLaunch(testAccount, "t3.micro", 1); err != nil { + t.Fatalf("first EnforceLaunch = %v, want nil in soft-cap window", err) + } + if err := s.EnforceLaunch(testAccount, "t3.micro", 1); err != nil { + t.Fatalf("second EnforceLaunch = %v, want nil in soft-cap window", err) + } + + // Both charges then land under CAS without losing an update. + for range 2 { + if err := s.ChargeLaunch(testAccount, reservationOf("t3.micro", 1)); err != nil { + t.Fatalf("ChargeLaunch = %v", err) + } + } + + got, _, err := s.readVCPU(testAccount) + if err != nil { + t.Fatalf("readVCPU: %v", err) + } + // 6 + 2*2 = 10: the cap of 8 is overshot by exactly one launch and no more. + if got != 10 { + t.Fatalf("counter = %d, want 10 (soft-cap overshoot must stay bounded)", got) + } +} + +// Exempt callers must short-circuit before the counter is read or written, so +// both gates return nil even with a nil KV handle. Covers the disabled-config +// and system-account exemptions for the check and the charge. +func TestLaunchExemptShortCircuits(t *testing.T) { + disabled := New(Limits{Enabled: false}, nil) + enabled := New(Limits{Enabled: true, VCPUs: 8}, nil) + cases := []struct { + name string + fn func() error + }{ + {"enforce disabled", func() error { return disabled.EnforceLaunch(testAccount, "t3.micro", 1000) }}, + {"charge disabled", func() error { return disabled.ChargeLaunch(testAccount, reservationOf("t3.micro", 1000)) }}, + {"enforce system account", func() error { return enabled.EnforceLaunch(utils.GlobalAccountID, "t3.micro", 1000) }}, + {"charge system account", func() error { return enabled.ChargeLaunch(utils.GlobalAccountID, reservationOf("t3.micro", 1000)) }}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if err := tc.fn(); err != nil { + t.Fatalf("exempt path returned %v, want nil", err) + } + }) + } +} diff --git a/spinifex/handlers/quota/vcpu_retype_test.go b/spinifex/handlers/quota/vcpu_retype_test.go new file mode 100644 index 000000000..518871f00 --- /dev/null +++ b/spinifex/handlers/quota/vcpu_retype_test.go @@ -0,0 +1,228 @@ +package handlers_quota + +import ( + "errors" + "testing" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/service/ec2" + "github.com/mulgadc/spinifex/spinifex/awserrors" + "github.com/mulgadc/spinifex/spinifex/utils" +) + +// 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 { + return func(string, string) (string, bool, error) { + return instanceType, ok, err + } +} + +// TestEnforceRetype covers the check-before gate. The account counter already +// holds the instance's old vCPUs, so the gate charges only the grow delta: +// retype-up within headroom and at the limit pass and return the delta, one over +// rejects, retype-down and same-size return a non-positive delta left to +// reconcile, and an unknown type on either side or a missing instance is deferred +// to the daemon. Each case uses its own account so the embedded JetStream is +// started once for the whole table. +func TestEnforceRetype(t *testing.T) { + s := newVCPUService(t, Limits{Enabled: true, VCPUs: 8}) // t3.micro=2, m5.xlarge=4 + + tests := []struct { + name string + account string + oldType string + newType string + resolveOK bool + reserve int // pre-charged vCPUs (includes the instance's old vCPUs) + wantDelta int + wantErr string + }{ + {"retype up within headroom", "100000000001", "t3.micro", "m5.xlarge", true, 2, 2, ""}, + {"retype up to limit", "100000000002", "t3.micro", "m5.xlarge", true, 6, 2, ""}, + {"retype up over limit", "100000000003", "t3.micro", "m5.xlarge", true, 7, 0, awserrors.ErrorResourceLimitExceeded}, + {"retype down left to reconcile", "100000000004", "m5.xlarge", "t3.micro", true, 4, -2, ""}, + {"same size no charge", "100000000005", "t3.micro", "c5.large", true, 2, 0, ""}, + {"unknown new type left to daemon", "100000000006", "t3.micro", "made.up", true, 2, 0, ""}, + {"unknown old type left to daemon", "100000000007", "made.up", "t3.micro", true, 2, 0, ""}, + {"instance not found", "100000000008", "t3.micro", "m5.xlarge", false, 2, 0, ""}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if tt.reserve > 0 { + if err := s.AddVCPU(tt.account, tt.reserve); err != nil { + t.Fatalf("AddVCPU(%d): %v", tt.reserve, err) + } + } + delta, err := s.EnforceRetype(staticResolver(tt.oldType, tt.resolveOK, nil), tt.account, "i-abc", tt.newType) + switch { + case tt.wantErr == "" && err != nil: + t.Fatalf("EnforceRetype = %v, want nil", err) + case tt.wantErr != "" && (err == nil || err.Error() != tt.wantErr): + t.Fatalf("EnforceRetype err = %v, want %q", err, tt.wantErr) + } + if delta != tt.wantDelta { + t.Fatalf("EnforceRetype delta = %d, want %d", delta, tt.wantDelta) + } + }) + } +} + +// TestRetypeChargesDelta mirrors the handler's two-phase flow: the gate returns +// the grow delta, then AddVCPU charges it so the counter reflects the new type. +// A retype from t3.micro (2) to m5.xlarge (4) on a 2/8 account lands at 4/8. +func TestRetypeChargesDelta(t *testing.T) { + s := newVCPUService(t, Limits{Enabled: true, VCPUs: 8}) + if err := s.AddVCPU(testAccount, 2); err != nil { // the t3.micro is already counted + t.Fatalf("seed AddVCPU: %v", err) + } + delta, err := s.EnforceRetype(staticResolver("t3.micro", true, nil), testAccount, "i-a", "m5.xlarge") + if err != nil { + t.Fatalf("EnforceRetype: %v", err) + } + if err := s.AddVCPU(testAccount, delta); err != nil { + t.Fatalf("AddVCPU(%d): %v", delta, err) + } + assertCounter(t, s, testAccount, 4) +} + +// TestRetypeDownDoesNotCharge confirms a shrink leaves the counter untouched: the +// gate returns a negative delta and AddVCPU no-ops it, so reconcile remains the +// only path that lowers the counter. +func TestRetypeDownDoesNotCharge(t *testing.T) { + s := newVCPUService(t, Limits{Enabled: true, VCPUs: 8}) + if err := s.AddVCPU(testAccount, 4); err != nil { // the m5.xlarge is already counted + t.Fatalf("seed AddVCPU: %v", err) + } + delta, err := s.EnforceRetype(staticResolver("m5.xlarge", true, nil), testAccount, "i-a", "t3.micro") + if err != nil { + t.Fatalf("EnforceRetype: %v", err) + } + if delta > 0 { + t.Fatalf("delta = %d, want <= 0 (shrink left to reconcile)", delta) + } + if err := s.AddVCPU(testAccount, delta); err != nil { + t.Fatalf("AddVCPU(%d): %v", delta, err) + } + assertCounter(t, s, testAccount, 4) // unchanged: reconcile lowers, not the retype +} + +// TestEnforceRetypeResolverError surfaces a describe failure to the caller so the +// modify is rejected rather than charged on a guessed type. +func TestEnforceRetypeResolverError(t *testing.T) { + s := newVCPUService(t, Limits{Enabled: true, VCPUs: 8}) + sentinel := errors.New("describe boom") + delta, err := s.EnforceRetype(staticResolver("t3.micro", false, sentinel), testAccount, "i-a", "m5.xlarge") + if !errors.Is(err, sentinel) || delta != 0 { + t.Fatalf("EnforceRetype = (%d, %v), want (0, %v)", delta, err, sentinel) + } +} + +// Exempt callers (disabled config, system account) short-circuit before the +// resolver is consulted, so the gate returns a zero delta and the describe round +// trip is skipped entirely. +func TestEnforceRetypeExemptShortCircuits(t *testing.T) { + failResolver := func(string, string) (string, bool, error) { + t.Fatal("resolver called for exempt account") + return "", false, nil + } + disabled := New(Limits{Enabled: false}, nil) + enabled := New(Limits{Enabled: true, VCPUs: 8}, nil) + cases := []struct { + name string + fn func() (int, error) + }{ + {"disabled", func() (int, error) { return disabled.EnforceRetype(failResolver, testAccount, "i-a", "m5.xlarge") }}, + {"system account", func() (int, error) { + return enabled.EnforceRetype(failResolver, utils.GlobalAccountID, "i-a", "m5.xlarge") + }}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + delta, err := tc.fn() + if err != nil || delta != 0 { + t.Fatalf("exempt EnforceRetype = (%d, %v), want (0, nil)", delta, err) + } + }) + } +} + +// TestInstanceTypeFromReservations covers the resolver's extraction helper: it +// returns a live instance's type by ID, skips terminal instances, and reports +// absent for an unknown ID. +func TestInstanceTypeFromReservations(t *testing.T) { + withID := func(id, instanceType, state string) *ec2.Instance { + inst := instance(instanceType, state) + inst.InstanceId = aws.String(id) + return inst + } + reservations := []*ec2.Reservation{ + reservation( + withID("i-term", "m5.xlarge", ec2.InstanceStateNameTerminated), + withID("i-stopped", "t3.micro", ec2.InstanceStateNameStopped), + ), + } + tests := []struct { + name string + instanceID string + wantType string + wantOK bool + }{ + {"found stopped", "i-stopped", "t3.micro", true}, + {"terminal skipped", "i-term", "", false}, + {"absent", "i-missing", "", false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + gotType, gotOK := instanceTypeFromReservations(reservations, tt.instanceID) + if gotType != tt.wantType || gotOK != tt.wantOK { + t.Fatalf("instanceTypeFromReservations(%q) = (%q, %v), want (%q, %v)", + tt.instanceID, gotType, gotOK, tt.wantType, tt.wantOK) + } + }) + } +} + +// TestConfirmInstanceType locks in the partial-view fix the production resolver +// relies on. A found instance resolves regardless of sweep completeness. A +// genuine absence on a complete sweep returns ok false with no error, so the gate +// charges nothing and the daemon rejects the modify; but an absence on an +// incomplete sweep fails closed with ErrorServerInternal, so a node or bucket that +// missed the sweep cannot be read as "absent" and wave a retype past the cap. +func TestConfirmInstanceType(t *testing.T) { + withID := func(id, instanceType, state string) *ec2.Instance { + inst := instance(instanceType, state) + inst.InstanceId = aws.String(id) + return inst + } + reservations := []*ec2.Reservation{ + reservation(withID("i-stopped", "t3.micro", ec2.InstanceStateNameStopped)), + } + tests := []struct { + name string + instanceID string + complete bool + wantType string + wantOK bool + wantErr string + }{ + {"found on complete sweep", "i-stopped", true, "t3.micro", true, ""}, + {"found on incomplete sweep", "i-stopped", false, "t3.micro", true, ""}, + {"absent on complete sweep defers to daemon", "i-missing", true, "", false, ""}, + {"absent on incomplete sweep fails closed", "i-missing", false, "", false, awserrors.ErrorServerInternal}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + gotType, gotOK, err := confirmInstanceType(reservations, tt.instanceID, tt.complete) + if gotType != tt.wantType || gotOK != tt.wantOK { + t.Fatalf("confirmInstanceType() = (%q, %v), want (%q, %v)", gotType, gotOK, tt.wantType, tt.wantOK) + } + switch { + case tt.wantErr == "" && err != nil: + t.Fatalf("confirmInstanceType() err = %v, want nil", err) + case tt.wantErr != "" && (err == nil || err.Error() != tt.wantErr): + t.Fatalf("confirmInstanceType() err = %v, want %q", err, tt.wantErr) + } + }) + } +} diff --git a/spinifex/handlers/quota/vcpu_test.go b/spinifex/handlers/quota/vcpu_test.go new file mode 100644 index 000000000..a72912d5b --- /dev/null +++ b/spinifex/handlers/quota/vcpu_test.go @@ -0,0 +1,119 @@ +package handlers_quota + +import ( + "sync" + "testing" + + "github.com/mulgadc/spinifex/spinifex/awserrors" + "github.com/mulgadc/spinifex/spinifex/testutil" + "github.com/mulgadc/spinifex/spinifex/utils" + "github.com/nats-io/nats.go" +) + +const testAccount = "123456789012" + +// newVCPUService starts an embedded JetStream, creates the account-usage bucket +// the way the gateway does, and returns a quota Service bound to it. +func newVCPUService(t *testing.T, limits Limits) *Service { + t.Helper() + _, _, js := testutil.StartTestJetStream(t) + kv, err := js.CreateKeyValue(&nats.KeyValueConfig{ + Bucket: KVBucketAccountUsage, + History: 1, + }) + if err != nil { + t.Fatalf("create usage bucket: %v", err) + } + return New(limits, kv) +} + +func TestCheckVCPUBoundaries(t *testing.T) { + s := newVCPUService(t, Limits{Enabled: true, VCPUs: 8}) + + // Empty counter: at-limit passes, over-limit rejects. + if err := s.CheckVCPU(testAccount, 8); err != nil { + t.Fatalf("CheckVCPU(8) on empty = %v, want nil", err) + } + if err := s.CheckVCPU(testAccount, 9); err == nil || err.Error() != awserrors.ErrorResourceLimitExceeded { + t.Fatalf("CheckVCPU(9) on empty = %v, want %q", err, awserrors.ErrorResourceLimitExceeded) + } + + // Reserve 6, then the remaining headroom is 2. + if err := s.AddVCPU(testAccount, 6); err != nil { + t.Fatalf("AddVCPU(6) = %v", err) + } + if err := s.CheckVCPU(testAccount, 2); err != nil { + t.Fatalf("CheckVCPU(2) at 6/8 = %v, want nil", err) + } + if err := s.CheckVCPU(testAccount, 3); err == nil || err.Error() != awserrors.ErrorResourceLimitExceeded { + t.Fatalf("CheckVCPU(3) at 6/8 = %v, want %q", err, awserrors.ErrorResourceLimitExceeded) + } +} + +// Exempt callers (disabled config, system account) must short-circuit before +// touching the counter, so the checks return nil even with a nil KV handle. +func TestVCPUExemptShortCircuits(t *testing.T) { + disabled := New(Limits{Enabled: false}, nil) + enabled := New(Limits{Enabled: true, VCPUs: 8}, nil) + cases := []struct { + name string + fn func() error + }{ + {"check disabled", func() error { return disabled.CheckVCPU(testAccount, 1000) }}, + {"add disabled", func() error { return disabled.AddVCPU(testAccount, 1000) }}, + {"check system account", func() error { return enabled.CheckVCPU(utils.GlobalAccountID, 1000) }}, + {"add system account", func() error { return enabled.AddVCPU(utils.GlobalAccountID, 1000) }}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if err := tc.fn(); err != nil { + t.Fatalf("exempt path returned %v, want nil", err) + } + }) + } +} + +func TestAddVCPUAccumulatesAndIgnoresShrinks(t *testing.T) { + s := newVCPUService(t, Limits{Enabled: true, VCPUs: 100}) + + for _, delta := range []int{4, 2, 0, -3} { + if err := s.AddVCPU(testAccount, delta); err != nil { + t.Fatalf("AddVCPU(%d) = %v", delta, err) + } + } + // 4 + 2 land; the 0 and -3 are no-ops left to reconcile. + got, _, err := s.readVCPU(testAccount) + if err != nil { + t.Fatalf("readVCPU: %v", err) + } + if got != 6 { + t.Fatalf("counter = %d, want 6", got) + } +} + +// Concurrent grows under CAS must not lose updates: the final counter equals the +// sum of every increment. +func TestAddVCPUConcurrentNoLostUpdates(t *testing.T) { + s := newVCPUService(t, Limits{Enabled: true, VCPUs: 1000}) + + const goroutines, perGoroutine = 10, 5 + var wg sync.WaitGroup + for range goroutines { + wg.Go(func() { + for range perGoroutine { + if err := s.AddVCPU(testAccount, 1); err != nil { + t.Errorf("AddVCPU: %v", err) + } + } + }) + } + wg.Wait() + + got, _, err := s.readVCPU(testAccount) + if err != nil { + t.Fatalf("readVCPU: %v", err) + } + if want := goroutines * perGoroutine; got != want { + t.Fatalf("counter = %d, want %d (lost updates under CAS)", got, want) + } +} diff --git a/spinifex/instancetypes/generate.go b/spinifex/instancetypes/generate.go index 02686aa8d..ce8cad80e 100644 --- a/spinifex/instancetypes/generate.go +++ b/spinifex/instancetypes/generate.go @@ -32,6 +32,29 @@ func SpecForSystemType(name string) (vcpu int, memGB float64, ok bool) { true } +// defaultVCPUsByInstanceType maps every instance type name to its default vCPU +// count. vCPUs depend only on the size suffix, not the CPU generation or +// architecture, so this static map covers every family the cluster can run, +// including ones a given host cannot itself detect. +var defaultVCPUsByInstanceType = func() map[string]int { + m := make(map[string]int) + for _, def := range instanceFamilyDefs { + for _, size := range def.sizes { + m[def.name+"."+size.suffix] = size.vcpus + } + } + return m +}() + +// DefaultVCPUs returns the default vCPU count for an instance type name (e.g. +// "c5.large"); ok is false for an unknown type. The result is independent of +// host CPU generation, so a gateway can size any account's instances even for +// families its own host cannot run. +func DefaultVCPUs(instanceType string) (vcpus int, ok bool) { + v, ok := defaultVCPUsByInstanceType[instanceType] + return v, ok +} + // generateForGeneration creates instance types for the given CPU generation. // Cross-vendor siblings are included on x86_64 so a mixed Intel+AMD cluster // can serve either vendor family. diff --git a/spinifex/instancetypes/generate_test.go b/spinifex/instancetypes/generate_test.go index fe30deee2..ee9a51983 100644 --- a/spinifex/instancetypes/generate_test.go +++ b/spinifex/instancetypes/generate_test.go @@ -31,6 +31,42 @@ func countFamily(types map[string]*ec2.InstanceTypeInfo, prefix string) int { return count } +func TestDefaultVCPUs(t *testing.T) { + cases := []struct { + name string + vcpus int + ok bool + }{ + {"t3.micro", 2, true}, // burstable + {"c5.large", 2, true}, // compute optimized + {"m5.4xlarge", 16, true}, // general purpose + {"r5.2xlarge", 8, true}, // memory optimized + {"g5.xlarge", 4, true}, // GPU family, host-arch independent + {"sys.micro", 1, true}, // internal system type + {"unknown.size", 0, false}, // not a real type + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got, ok := DefaultVCPUs(tc.name) + assert.Equal(t, tc.ok, ok) + assert.Equal(t, tc.vcpus, got) + }) + } +} + +func TestDefaultVCPUs_IndependentOfHostGeneration(t *testing.T) { + // Every family/size in the catalog resolves regardless of which CPU the + // host detects, so reconcile can size instances launched on other nodes. + for _, def := range instanceFamilyDefs { + for _, size := range def.sizes { + name := def.name + "." + size.suffix + got, ok := DefaultVCPUs(name) + require.True(t, ok, "%s must resolve", name) + assert.Equal(t, size.vcpus, got, "%s vCPUs", name) + } + } +} + func TestIsSystemType(t *testing.T) { assert.True(t, IsSystemType("sys.micro")) assert.True(t, IsSystemType("sys.small")) diff --git a/spinifex/network/reconcile/drift.go b/spinifex/network/reconcile/drift.go index 26a2cbab3..57b28010e 100644 --- a/spinifex/network/reconcile/drift.go +++ b/spinifex/network/reconcile/drift.go @@ -36,7 +36,7 @@ func DriftLoop(ctx context.Context, rec Reconciler, nc *nats.Conn, localAZ, hold // runDriftCycle is one tick body, split out so tests can drive it directly. func runDriftCycle(ctx context.Context, rec Reconciler, nc *nats.Conn, js nats.JetStreamContext, localAZ, holder string) { - release, elected := AcquireLeader(nc, holder) + release, elected := AcquireLeader(nc, KVBucketVPCDReconcile, holder) if !elected { return } diff --git a/spinifex/network/reconcile/lock.go b/spinifex/network/reconcile/lock.go index 1e06c02d5..c9f0d75e5 100644 --- a/spinifex/network/reconcile/lock.go +++ b/spinifex/network/reconcile/lock.go @@ -10,7 +10,7 @@ import ( // Single CAS-elected leader key; TTL bounds crash-recovery. const ( - reconcileLeaderBucket = "spinifex-vpcd-reconcile" + KVBucketVPCDReconcile = "spinifex-vpcd-reconcile" reconcileLeaderKey = "leader" reconcileLeaderTTL = 60 * time.Second ) @@ -22,9 +22,11 @@ var ( leaderRetryStep = 1 * time.Second ) -// AcquireLeader returns (release, true) exactly once across all vpcds, or -// (nil, false) for losers and JetStream-unreachable cases. -func AcquireLeader(nc *nats.Conn, holder string) (func(), bool) { +// AcquireLeader elects one leader on the named lock bucket. Independent +// reconcile loops pass distinct buckets so they never share a single mutex: the +// gateway quota reconcile must not block vpcd's network reconcile, and vice +// versa. +func AcquireLeader(nc *nats.Conn, bucket, holder string) (func(), bool) { js, _ := nc.JetStream() var ( @@ -35,10 +37,10 @@ func AcquireLeader(nc *nats.Conn, holder string) (func(), bool) { for { // Get-or-create: CreateKeyValue returns "stream name already in use" if // the bucket exists; attach first, create only when genuinely absent. - kv, err = js.KeyValue(reconcileLeaderBucket) + kv, err = js.KeyValue(bucket) if errors.Is(err, nats.ErrBucketNotFound) { kv, err = js.CreateKeyValue(&nats.KeyValueConfig{ - Bucket: reconcileLeaderBucket, + Bucket: bucket, History: 1, TTL: reconcileLeaderTTL, }) @@ -48,22 +50,22 @@ func AcquireLeader(nc *nats.Conn, holder string) (func(), bool) { } if time.Now().After(deadline) { slog.Error("reconcile/lock: JetStream KV unreachable after retry, skipping reconcile", - "holder", holder, "waited", leaderRetryFor, "err", err) + "holder", holder, "bucket", bucket, "waited", leaderRetryFor, "err", err) return nil, false } - slog.Debug("reconcile/lock: JetStream KV not ready, retrying", "holder", holder, "err", err) + slog.Debug("reconcile/lock: JetStream KV not ready, retrying", "holder", holder, "bucket", bucket, "err", err) time.Sleep(leaderRetryStep) } if _, err := kv.Create(reconcileLeaderKey, []byte(holder)); err != nil { - slog.Info("reconcile/lock: another vpcd is leader, skipping reconcile", "holder", holder, "err", err) + slog.Info("reconcile/lock: another holder is leader, skipping reconcile", "holder", holder, "bucket", bucket, "err", err) return nil, false } - slog.Info("reconcile/lock: elected", "holder", holder) + slog.Info("reconcile/lock: elected", "holder", holder, "bucket", bucket) return func() { if err := kv.Delete(reconcileLeaderKey); err != nil { - slog.Warn("reconcile/lock: failed to release lock (TTL will reap)", "holder", holder, "err", err) + slog.Warn("reconcile/lock: failed to release lock (TTL will reap)", "holder", holder, "bucket", bucket, "err", err) } }, true } diff --git a/spinifex/network/reconcile/lock_test.go b/spinifex/network/reconcile/lock_test.go index 52715b63e..699f28094 100644 --- a/spinifex/network/reconcile/lock_test.go +++ b/spinifex/network/reconcile/lock_test.go @@ -14,19 +14,19 @@ import ( func TestAcquireLeader_GetOrCreateAcrossCalls(t *testing.T) { _, nc, _ := testutil.StartTestJetStream(t) - release, ok := AcquireLeader(nc, "node-1") + release, ok := AcquireLeader(nc, KVBucketVPCDReconcile, "node-1") require.True(t, ok, "first acquire must win and create the bucket") require.NotNil(t, release) // Bucket now exists; the second acquire must reach the Create-key contention // path (and lose) rather than dead-ending on CreateKeyValue. - loserRelease, ok := AcquireLeader(nc, "node-2") + loserRelease, ok := AcquireLeader(nc, KVBucketVPCDReconcile, "node-2") assert.False(t, ok, "second acquire must lose while the lock is held") assert.Nil(t, loserRelease) // Releasing frees the key so a subsequent acquire can take over. release() - release2, ok := AcquireLeader(nc, "node-3") + release2, ok := AcquireLeader(nc, KVBucketVPCDReconcile, "node-3") require.True(t, ok, "acquire after release must win") require.NotNil(t, release2) release2() diff --git a/spinifex/services/awsgw/awsgw.go b/spinifex/services/awsgw/awsgw.go index a09f15ac9..fbd90a8c5 100644 --- a/spinifex/services/awsgw/awsgw.go +++ b/spinifex/services/awsgw/awsgw.go @@ -3,6 +3,7 @@ package awsgw import ( "context" "crypto/tls" + "errors" "fmt" "log/slog" "net" @@ -21,7 +22,9 @@ import ( gateway_ecrauth "github.com/mulgadc/spinifex/spinifex/gateway/ecrauth" "github.com/mulgadc/spinifex/spinifex/handlers/ecr" handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam" + handlers_quota "github.com/mulgadc/spinifex/spinifex/handlers/quota" handlers_sts "github.com/mulgadc/spinifex/spinifex/handlers/sts" + "github.com/mulgadc/spinifex/spinifex/network/reconcile" "github.com/mulgadc/spinifex/spinifex/objectstore" "github.com/mulgadc/spinifex/spinifex/utils" "github.com/nats-io/nats.go" @@ -88,22 +91,48 @@ func (svc *Service) Reload() (err error) { } // awsgwTOML is the top-level structure of awsgw.toml used to extract the -// ratelimit section. Other fields are parsed elsewhere (e.g. region, debug). +// ratelimit and quota sections. Other fields are parsed elsewhere (e.g. region, +// debug). type awsgwTOML struct { - Ratelimit ratelimit.Config `toml:"ratelimit"` + Ratelimit ratelimit.Config `toml:"ratelimit"` + Quota handlers_quota.Limits `toml:"quota"` } -// loadThrottleConfig parses the [ratelimit] section from the awsgw TOML config. -func loadThrottleConfig(path string) (ratelimit.Config, error) { +// loadAWSGWConfig reads and parses awsgw.toml once, returning the [ratelimit] and +// [quota] sections together. Both sections default to their zero value (a +// disabled no-op) when absent, so a config without either block stays valid. +func loadAWSGWConfig(path string) (awsgwTOML, error) { data, err := os.ReadFile(path) if err != nil { - return ratelimit.Config{}, fmt.Errorf("read awsgw config %s: %w", path, err) + return awsgwTOML{}, fmt.Errorf("read awsgw config %s: %w", path, err) } var cfg awsgwTOML if err := toml.Unmarshal(data, &cfg); err != nil { - return ratelimit.Config{}, fmt.Errorf("parse awsgw config %s: %w", path, err) + return awsgwTOML{}, fmt.Errorf("parse awsgw config %s: %w", path, err) } - return cfg.Ratelimit, nil + return cfg, nil +} + +// openAccountUsageBucket opens (or idempotently creates) the gateway-owned +// per-account vCPU usage bucket. History is 1: each account key holds a single +// CAS-updated integer counter. It attaches first and creates only when the bucket +// is genuinely absent, so a transient create error is not masked by the fallback. +func openAccountUsageBucket(js nats.JetStreamContext, replicas int) (nats.KeyValue, error) { + if replicas < 1 { + replicas = 1 + } + kv, err := js.KeyValue(handlers_quota.KVBucketAccountUsage) + if errors.Is(err, nats.ErrBucketNotFound) { + kv, err = js.CreateKeyValue(&nats.KeyValueConfig{ + Bucket: handlers_quota.KVBucketAccountUsage, + History: 1, + Replicas: replicas, + }) + } + if err != nil { + return nil, fmt.Errorf("open account usage bucket: %w", err) + } + return kv, nil } func launchService(config *config.ClusterConfig) error { @@ -194,12 +223,15 @@ func launchService(config *config.ClusterConfig) error { return fmt.Errorf("load bootstrap from %s: %w", bootstrapPath, err) } - // Load API throttle config from awsgw.toml [ratelimit] section. + // Load the awsgw config once: [ratelimit] throttling and [quota] per-account + // service quotas. A load error leaves both at their zero value (disabled). awsgwTomlPath := filepath.Join(nodeConfig.BaseDir, "config", "awsgw", "awsgw.toml") - throttleCfg, err := loadThrottleConfig(awsgwTomlPath) + awsgwCfg, err := loadAWSGWConfig(awsgwTomlPath) if err != nil { - slog.Warn("Failed to load throttle config, throttling disabled", "err", err) + slog.Warn("Failed to load awsgw config, throttling and quotas disabled", "err", err) } + throttleCfg := awsgwCfg.Ratelimit + quotaCfg := awsgwCfg.Quota // OCI Distribution v2 registry: blob/manifest bytes stream straight to // predastore from the gateway; repo/tag/manifest metadata and in-progress @@ -285,6 +317,30 @@ func launchService(config *config.ClusterConfig) error { defer gw.Throttler.Stop() } + // 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 nats.KeyValue + if quotaCfg.Enabled { + usageBucket, err = openAccountUsageBucket(js, max(len(config.Nodes), 1)) + if err != nil { + return fmt.Errorf("init account usage bucket: %w", err) + } + } + gw.Quota = handlers_quota.New(quotaCfg, usageBucket) + + // Leader-locked vCPU reconcile: the only path that lowers the counter, + // recomputing it from the running-plus-stopped sweep so out-of-band + // 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) + } + handler := gw.SetupRoutes() // Load TLS certificate @@ -344,6 +400,39 @@ func initIAMService(natsConn *nats.Conn, masterKey []byte, clusterSize int) (*ha } } +// 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) { + holder, _ := os.Hostname() + list := handlers_quota.NATSInstanceLister(natsConn, expectedNodes) + + runPass := func() { + release, elected := reconcile.AcquireLeader(natsConn, handlers_quota.KVBucketQuotaReconcile, holder) + if !elected { + return + } + defer release() + if err := quota.Reconcile(ctx, accounts, list); err != nil { + slog.Warn("quota reconcile pass failed", "err", err) + } + } + + runPass() + ticker := time.NewTicker(handlers_quota.ReconcileInterval) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + runPass() + } + } +} + // accountLister is the slice of IAMService the lifecycle sweeper needs. type accountLister interface { ListAccounts() ([]*handlers_iam.Account, error) diff --git a/spinifex/services/awsgw/awsgw_test.go b/spinifex/services/awsgw/awsgw_test.go index 6a04d1541..e51a60f19 100644 --- a/spinifex/services/awsgw/awsgw_test.go +++ b/spinifex/services/awsgw/awsgw_test.go @@ -27,8 +27,9 @@ burst = 40 ` require.NoError(t, os.WriteFile(path, []byte(content), 0o644)) - cfg, err := loadThrottleConfig(path) + parsed, err := loadAWSGWConfig(path) require.NoError(t, err) + cfg := parsed.Ratelimit assert.True(t, cfg.Enabled) assert.Equal(t, 20, cfg.Rate) assert.Equal(t, 100, cfg.Burst) @@ -48,9 +49,9 @@ burst = 100 ` require.NoError(t, os.WriteFile(path, []byte(content), 0o644)) - cfg, err := loadThrottleConfig(path) + parsed, err := loadAWSGWConfig(path) require.NoError(t, err) - assert.False(t, cfg.Enabled) + assert.False(t, parsed.Ratelimit.Enabled) } func TestLoadThrottleConfig_NoSection(t *testing.T) { @@ -63,14 +64,75 @@ debug = false ` require.NoError(t, os.WriteFile(path, []byte(content), 0o644)) - cfg, err := loadThrottleConfig(path) + parsed, err := loadAWSGWConfig(path) require.NoError(t, err) // Missing section → zero-value config (disabled, rate=0, burst=0). - assert.False(t, cfg.Enabled) - assert.Equal(t, 0, cfg.Rate) + assert.False(t, parsed.Ratelimit.Enabled) + assert.Equal(t, 0, parsed.Ratelimit.Rate) } -func TestLoadThrottleConfig_MissingFile(t *testing.T) { - _, err := loadThrottleConfig("/nonexistent/awsgw.toml") +func TestLoadAWSGWConfig_MissingFile(t *testing.T) { + _, err := loadAWSGWConfig("/nonexistent/awsgw.toml") assert.Error(t, err) } + +func TestLoadQuotaConfig_Enabled(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "awsgw.toml") + content := ` +version = "3" +region = "us-east-1" + +[quota] +enabled = true +vcpus = 8 +vpcs = 8 +subnets = 16 +eips = 2 +volumes_gib = 100 +` + require.NoError(t, os.WriteFile(path, []byte(content), 0o644)) + + parsed, err := loadAWSGWConfig(path) + require.NoError(t, err) + cfg := parsed.Quota + assert.True(t, cfg.Enabled) + assert.Equal(t, 8, cfg.VCPUs) + assert.Equal(t, 8, cfg.VPCs) + assert.Equal(t, 16, cfg.Subnets) + assert.Equal(t, 2, cfg.EIPs) + assert.Equal(t, 100, cfg.VolumesGiB) +} + +func TestLoadQuotaConfig_Disabled(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "awsgw.toml") + content := ` +version = "3" +[quota] +enabled = false +vcpus = 8 +` + require.NoError(t, os.WriteFile(path, []byte(content), 0o644)) + + parsed, err := loadAWSGWConfig(path) + require.NoError(t, err) + assert.False(t, parsed.Quota.Enabled) +} + +func TestLoadQuotaConfig_NoSection(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "awsgw.toml") + content := ` +version = "3" +region = "us-east-1" +debug = false +` + require.NoError(t, os.WriteFile(path, []byte(content), 0o644)) + + parsed, err := loadAWSGWConfig(path) + require.NoError(t, err) + // Missing section → zero-value Limits, a disabled no-op. + assert.False(t, parsed.Quota.Enabled) + assert.Equal(t, 0, parsed.Quota.VCPUs) +} diff --git a/spinifex/services/awsgw/quota_reconcile_test.go b/spinifex/services/awsgw/quota_reconcile_test.go new file mode 100644 index 000000000..9935ccda0 --- /dev/null +++ b/spinifex/services/awsgw/quota_reconcile_test.go @@ -0,0 +1,82 @@ +package awsgw + +import ( + "context" + "encoding/json" + "strconv" + "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/testutil" + "github.com/nats-io/nats.go" + "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() }) +} + +// 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) { + _, nc, js := testutil.StartTestJetStream(t) + bucket, err := openAccountUsageBucket(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{}) + + 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 }) + 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") + + cancel() + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("runQuotaReconcile did not return after ctx cancel") + } +} + +// 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 nats.KeyValue, account string) int { + t.Helper() + entry, err := bucket.Get(account) + if err != nil { + return -1 + } + n, err := strconv.Atoi(string(entry.Value())) + require.NoError(t, err) + return n +} diff --git a/spinifex/vpcd/vpcd.go b/spinifex/vpcd/vpcd.go index 85e25dedc..963491ec8 100644 --- a/spinifex/vpcd/vpcd.go +++ b/spinifex/vpcd/vpcd.go @@ -531,7 +531,7 @@ func launchService(cfg *Config) error { // Elect one vpcd for startup reconcile; without this, concurrent Get-then-Create on Logical_Router // produces duplicate rows that ovn-nbctl rejects. Runtime events still fan out via queue groups. holder, _ := os.Hostname() - releaseLeader, isLeader := reconcile.AcquireLeader(nc, holder) + releaseLeader, isLeader := reconcile.AcquireLeader(nc, reconcile.KVBucketVPCDReconcile, holder) subscriber, err := subscribers.New(subscribers.Config{ Topology: topoMgr,