Handle EC2 throttling and capacity errors when launching runners (#19)

Co-authored-by: Henry Skiba <henry.skiba@frgrisk.com>

Author: luqmanbazran

.gitignore

@@ -24,6 +24,9 @@ go.work
 
 bootstrap
 
+# Local `go build` output (named after the module directory)
+github-runner-autoscaler
+
 .aws-sam
 
 # Local SAM configuration (not committed)

main.go

@@ -20,6 +20,7 @@ import (
 	"github.com/aws/aws-lambda-go/events"
 	"github.com/aws/aws-lambda-go/lambda"
 	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/aws/retry"
 	"github.com/aws/aws-sdk-go-v2/config"
 	"github.com/aws/aws-sdk-go-v2/service/ec2"
 	"github.com/aws/aws-sdk-go-v2/service/ec2/types"
@@ -38,16 +39,129 @@ type RunnerConfiguration struct {
 	KeyName        string   `json:"key"`
 }
 
-// retryableLaunchErrors are EC2 error codes for which launching the runner in
-// the next configured subnet (potentially a different AZ) may succeed.
-var retryableLaunchErrors = []string{
+// launchCycleErrorCodes are EC2 error codes for which launching the runner in
+// the next configured subnet (potentially a different AZ) or the next candidate
+// instance type may succeed. launchInstance handles these itself by cycling to
+// the next subnet/type, so excludeLaunchCycleErrors also marks them
+// non-retryable on the client (see there for why).
+var launchCycleErrorCodes = []string{
 	"InsufficientFreeAddressesInSubnet",
 	"InsufficientInstanceCapacity",
 	"InvalidSubnetID.NotFound",
 	"Unsupported",
 }
 
-func handler(request events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
+// isLaunchCycleError reports whether err is one launchInstance handles by moving
+// on to the next subnet/instance type, rather than a fatal error or a throttle
+// the SDK's retryer already backed off and retried.
+func isLaunchCycleError(err error) bool {
+	apiErr, ok := errors.AsType[smithy.APIError](err)
+
+	return ok && slices.Contains(launchCycleErrorCodes, apiErr.ErrorCode())
+}
+
+// excludeLaunchCycleErrors keeps the SDK retryer from retrying the errors that
+// launchInstance handles by cycling to the next subnet/instance type.
+// InsufficientInstanceCapacity is an HTTP 500, so the default retryer would
+// otherwise burn its own attempts (with backoff) retrying the same subnet before
+// launchInstance ever gets to try the next one. Returning FalseTernary here
+// short-circuits that so the error surfaces immediately; throttling and other
+// transient errors fall through to UnknownTernary and keep default retry.
+var excludeLaunchCycleErrors = retry.IsErrorRetryableFunc(func(err error) aws.Ternary {
+	if isLaunchCycleError(err) {
+		return aws.FalseTernary
+	}
+
+	return aws.UnknownTernary
+})
+
+// newEC2Retryer returns the SDK's standard retryer (throttle + transient backoff
+// left at defaults) with cycle errors excluded so launchInstance can move to the
+// next subnet/type without waiting on same-call retries.
+func newEC2Retryer() *retry.Standard {
+	return retry.NewStandard(func(o *retry.StandardOptions) {
+		o.Retryables = append(
+			[]retry.IsErrorRetryable{excludeLaunchCycleErrors},
+			retry.DefaultRetryables...,
+		)
+	})
+}
+
+// launchInstance attempts to launch a runner across the candidate instance types
+// and subnets. It tries each instance type in turn, sweeping every subnet; a
+// capacity or subnet error (see retryableLaunchErrors) moves straight on to the
+// next subnet, then the next type, so a type that is out of capacity in every AZ
+// falls through to the next candidate. Throttling and other transient errors are
+// backed off and retried by the EC2 client's own retryer before they reach here,
+// so any error that is not a cycle error is treated as fatal. It returns the
+// launched instance ID, or an error if no type/subnet combination succeeds.
+func launchInstance(
+	ctx context.Context,
+	svc *ec2.Client,
+	runInput *ec2.RunInstancesInput,
+	instanceTypes []types.InstanceType,
+	subnets []string,
+) (string, error) {
+	var lastErr error
+
+	for _, instanceType := range instanceTypes {
+		runInput.InstanceType = instanceType
+
+		for _, subnet := range subnets {
+			runInput.NetworkInterfaces[0].SubnetId = aws.String(subnet)
+
+			output, err := svc.RunInstances(ctx, runInput)
+			if err != nil {
+				if !isLaunchCycleError(err) {
+					slog.Error("failed to run instances", "error", err.Error())
+
+					return "", err
+				}
+
+				slog.Warn(
+					"capacity/subnet error, trying next",
+					"instanceType", instanceType,
+					"subnet", subnet,
+					"error", err.Error(),
+				)
+
+				lastErr = err
+
+				continue
+			}
+
+			if len(output.Instances) == 0 || output.Instances[0].InstanceId == nil {
+				slog.Warn(
+					"no instance created in subnet, trying next",
+					"instanceType", instanceType,
+					"subnet", subnet,
+				)
+
+				lastErr = errors.New("run instances returned no instance id")
+
+				continue
+			}
+
+			return aws.ToString(output.Instances[0].InstanceId), nil
+		}
+
+		slog.Warn(
+			"all subnets failed for instance type, trying next type",
+			"instanceType", instanceType,
+		)
+	}
+
+	if lastErr == nil {
+		lastErr = errors.New("failed to launch instance in any subnet")
+	}
+
+	return "", fmt.Errorf("failed to launch instance: %w", lastErr)
+}
+
+func handler(
+	ctx context.Context,
+	request events.APIGatewayProxyRequest,
+) (events.APIGatewayProxyResponse, error) {
 	var githubEventHeader string
 
 	for k, v := range request.MultiValueHeaders {
@@ -103,19 +217,29 @@ func handler(request events.APIGatewayProxyRequest) (events.APIGatewayProxyRespo
 
 		region := cmp.Or(os.Getenv("AWS_DEFAULT_REGION"), os.Getenv("AWS_REGION"))
 
-		instanceType := types.InstanceTypeC7aLarge
-		instanceTypes := instanceType.Values()
+		validInstanceTypes := types.InstanceTypeC7aLarge.Values()
+
+		// Candidate instance types to try, in the order the labels appear on the
+		// job. Trying several lets the launch fall back when a type is out of
+		// capacity (InsufficientInstanceCapacity) in every configured subnet.
+		var instanceTypes []types.InstanceType
 
 		for _, label := range event.GetWorkflowJob().Labels {
 			if _, ok := runnerConfig[label]; ok {
 				region = label
 			}
 
-			if slices.Contains(instanceTypes, types.InstanceType(label)) {
-				instanceType = types.InstanceType(label)
+			candidate := types.InstanceType(label)
+			if slices.Contains(validInstanceTypes, candidate) &&
+				!slices.Contains(instanceTypes, candidate) {
+				instanceTypes = append(instanceTypes, candidate)
 			}
 		}
 
+		if len(instanceTypes) == 0 {
+			instanceTypes = []types.InstanceType{types.InstanceTypeC7aLarge}
+		}
+
 		regionCfg, ok := runnerConfig[region]
 		if !ok {
 			return events.APIGatewayProxyResponse{StatusCode: http.StatusInternalServerError},
@@ -127,14 +251,16 @@ func handler(request events.APIGatewayProxyRequest) (events.APIGatewayProxyRespo
 				fmt.Errorf("no subnets configured for region %s", region)
 		}
 
-		cfg, err := config.LoadDefaultConfig(context.TODO(), config.WithRegion(region))
+		cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(region))
 		if err != nil {
 			return events.APIGatewayProxyResponse{StatusCode: http.StatusInternalServerError}, err
 		}
 
 		slog.Info("creating runner in region", "region", region)
 
-		svc := ec2.NewFromConfig(cfg)
+		svc := ec2.NewFromConfig(cfg, func(o *ec2.Options) {
+			o.Retryer = newEC2Retryer()
+		})
 		sm := secretsmanager.NewFromConfig(cfg)
 
 		secretName := os.Getenv("GITHUB_PAT_SECRET_NAME")
@@ -149,7 +275,7 @@ func handler(request events.APIGatewayProxyRequest) (events.APIGatewayProxyRespo
 		}
 
 		secretOut, err := sm.GetSecretValue(
-			context.TODO(),
+			ctx,
 			&secretsmanager.GetSecretValueInput{SecretId: aws.String(secretName)},
 		)
 		if err != nil {
@@ -194,7 +320,7 @@ func handler(request events.APIGatewayProxyRequest) (events.APIGatewayProxyRespo
 			return events.APIGatewayProxyResponse{StatusCode: http.StatusOK}, nil
 		}
 
-		slog.Info("creating instance", "instanceType", instanceType)
+		slog.Info("creating instance", "instanceTypes", instanceTypes)
 
 		tpl, err := template.New("userdata").Parse(userData)
 		if err != nil {
@@ -219,7 +345,8 @@ func handler(request events.APIGatewayProxyRequest) (events.APIGatewayProxyRespo
 			EbsOptimized:                      aws.Bool(true),
 			ImageId:                           aws.String(regionCfg.ImageID),
 			InstanceInitiatedShutdownBehavior: types.ShutdownBehaviorTerminate,
-			InstanceType:                      instanceType,
+			// InstanceType is set per-attempt by launchInstance so it can fall
+			// back across the candidate instanceTypes on capacity errors.
 			IamInstanceProfile: &types.IamInstanceProfileSpecification{
 				Arn: aws.String(instanceProfileArn),
 			},
@@ -247,67 +374,22 @@ func handler(request events.APIGatewayProxyRequest) (events.APIGatewayProxyRespo
 			UserData: aws.String(base64.StdEncoding.EncodeToString([]byte(finalUserData))),
 		}
 
-		var lastErr error
-
-		for _, subnet := range regionCfg.SubnetID {
-			runInput.NetworkInterfaces[0].SubnetId = aws.String(subnet)
-
-			output, err := svc.RunInstances(context.TODO(), runInput)
-			if err != nil {
-				apiErr, ok := errors.AsType[smithy.APIError](err)
-				if ok && slices.Contains(retryableLaunchErrors, apiErr.ErrorCode()) {
-					slog.Warn(
-						"retrying in next subnet",
-						"subnet",
-						subnet,
-						"reason",
-						apiErr.ErrorCode(),
-					)
-
-					lastErr = err
-
-					continue
-				}
-
-				slog.Error("failed to run instances", "error", err.Error())
-
-				return events.APIGatewayProxyResponse{
-					Body:       err.Error(),
-					StatusCode: http.StatusInternalServerError,
-				}, err
-			}
-
-			if len(output.Instances) == 0 || output.Instances[0].InstanceId == nil {
-				slog.Warn(
-					"no instance created in subnet, trying next",
-					"subnet",
-					subnet,
-				)
-
-				lastErr = errors.New("run instances returned no instance id")
-
-				continue
-			}
-
-			instanceID := aws.ToString(output.Instances[0].InstanceId)
-			slog.Info("instance created", "instanceID", instanceID)
+		instanceID, err := launchInstance(ctx, svc, runInput, instanceTypes, regionCfg.SubnetID)
+		if err != nil {
+			slog.Error("failed to launch instance", "error", err.Error())
 
 			return events.APIGatewayProxyResponse{
-				Body:       instanceID,
-				StatusCode: http.StatusOK,
-			}, nil
-		}
-
-		if lastErr == nil {
-			lastErr = errors.New("failed to launch instance in any subnet")
+				Body:       err.Error(),
+				StatusCode: http.StatusInternalServerError,
+			}, err
 		}
 
-		slog.Error("failed to launch instance in any subnet", "error", lastErr.Error())
+		slog.Info("instance created", "instanceID", instanceID)
 
 		return events.APIGatewayProxyResponse{
-			Body:       lastErr.Error(),
-			StatusCode: http.StatusInternalServerError,
-		}, lastErr
+			Body:       instanceID,
+			StatusCode: http.StatusOK,
+		}, nil
 
 	default:
 		err = fmt.Errorf("unknown event type %T", event)