Graceful scale-in: send SIGTERM only to healthy Linux instances

Author: scadu

scaler/asg.go

@@ -25,12 +25,13 @@ const (
 )
 
 type AutoscaleGroupDetails struct {
-	Pending      int64
-	DesiredCount int64
-	MinSize      int64
-	MaxSize      int64
-	InstanceIDs  []string // Instance IDs in the ASG
-	ActualCount  int64    // Actual number of running instances
+	Pending            int64
+	DesiredCount       int64
+	MinSize            int64
+	MaxSize            int64
+	InstanceIDs        []string // All instance IDs in the ASG
+	HealthyInstanceIDs []string // Instance IDs that are healthy (for graceful termination)
+	ActualCount        int64    // Actual number of running instances
 }
 
 type ASGDriver struct {
@@ -84,27 +85,22 @@ func (a *ASGDriver) getASGPlatform(ctx context.Context, instances []ec2Types.Ins
 	return "linux"
 }
 
-// getCheckCommand returns the appropriate check command for the platform
+// getCheckCommand returns the appropriate check command for the platform.
+// It checks if the buildkite-agent service is running.
 func (a *ASGDriver) getCheckCommand(platform string) string {
 	if platform == "windows" {
 		return `
 $AgentStatus = nssm status buildkite-agent 2>&1
 if ($AgentStatus -match "SERVICE_RUNNING") {
     Write-Output "RUNNING"
 } else {
-    Write-Output "NOT_RUNNING"
+    Write-Output "NOT_RUNNING: $AgentStatus"
 }
 `
 	}
 
 	// Default to Linux
 	return `#!/bin/bash
-# Linux check command
-if [ -f /tmp/buildkite-agent-termination-marker ]; then
-  echo "MARKER_EXISTS: Instance is already marked for termination"
-  exit 0
-fi
-
 ACTIVE_STATE=$(systemctl show buildkite-agent -p ActiveState | cut -d= -f2)
 case "$ACTIVE_STATE" in
   "active"|"activating") echo "RUNNING" ;;
@@ -271,19 +267,26 @@ func (a *ASGDriver) Describe(ctx context.Context) (AutoscaleGroupDetails, error)
 	}
 
 	instanceIDs := make([]string, 0, len(asg.Instances))
+	healthyInstanceIDs := make([]string, 0, len(asg.Instances))
 	for _, instance := range asg.Instances {
 		if instance.InstanceId != nil {
 			instanceIDs = append(instanceIDs, *instance.InstanceId)
+			// Only include instances that are healthy and in-service for graceful termination
+			if instance.HealthStatus != nil && *instance.HealthStatus == "Healthy" &&
+				instance.LifecycleState == types.LifecycleStateInService {
+				healthyInstanceIDs = append(healthyInstanceIDs, *instance.InstanceId)
+			}
 		}
 	}
 
 	details := AutoscaleGroupDetails{
-		Pending:      pending,
-		DesiredCount: int64(*result.AutoScalingGroups[0].DesiredCapacity),
-		MinSize:      int64(*result.AutoScalingGroups[0].MinSize),
-		MaxSize:      int64(*result.AutoScalingGroups[0].MaxSize),
-		InstanceIDs:  instanceIDs,
-		ActualCount:  running,
+		Pending:            pending,
+		DesiredCount:       int64(*result.AutoScalingGroups[0].DesiredCapacity),
+		MinSize:            int64(*result.AutoScalingGroups[0].MinSize),
+		MaxSize:            int64(*result.AutoScalingGroups[0].MaxSize),
+		InstanceIDs:        instanceIDs,
+		HealthyInstanceIDs: healthyInstanceIDs,
+		ActualCount:        running,
 	}
 
 	log.Printf("↳ Got pending=%d, desired=%d, actual=%d, min=%d, max=%d (took %v)",
@@ -370,70 +373,47 @@ type dryRunASG struct {
 
 func (a *ASGDriver) SendSIGTERMToAgents(ctx context.Context, instanceID string) error {
 	ssmClient := ssm.NewFromConfig(a.Cfg)
+	ec2Client := ec2.NewFromConfig(a.Cfg)
 
-	// Wait for SSM agent to be ready before sending command
-	if err := a.waitForSSMReady(ctx, instanceID, 30*time.Second); err != nil {
-		log.Printf("SSM agent not ready on instance %s, cannot send SIGTERM: %v", instanceID, err)
-		return err
-	}
+	// Detect platform for EC2 instance
+	platform := "linux"
+	documentName := "AWS-RunShellScript"
+	stopCommand := "sudo systemctl stop buildkite-agent.service"
 
-	// First check if this instance already has a termination marker.
-	// This makes SIGTERM idempotent: even if we're called multiple times
-	// across different Lambda invocations, we only send one SIGTERM.
-	checkMarkerCommand := `#!/bin/bash
-if [ -f /tmp/buildkite-agent-termination-marker ]; then
-  echo "MARKER_EXISTS"
-  cat /tmp/buildkite-agent-termination-marker
-  exit 0
-else
-  echo "NO_MARKER"
-  exit 1
-fi
-`
-	checkOutput, err := ssmClient.SendCommand(ctx, &ssm.SendCommandInput{
-		InstanceIds:  []string{instanceID},
-		DocumentName: aws.String("AWS-RunShellScript"),
-		Parameters:   map[string][]string{"commands": {checkMarkerCommand}},
-		Comment:      aws.String("Check if instance already has termination marker"),
+	descResp, descErr := ec2Client.DescribeInstances(ctx, &ec2.DescribeInstancesInput{
+		InstanceIds: []string{instanceID},
 	})
-	if err != nil {
-		log.Printf("[Elastic CI Mode] Warning: Could not check termination marker on %s: %v", instanceID, err)
-	} else {
-		time.Sleep(2 * time.Second)
-
-		checkResult, getErr := ssmClient.GetCommandInvocation(ctx, &ssm.GetCommandInvocationInput{
-			CommandId:  checkOutput.Command.CommandId,
-			InstanceId: aws.String(instanceID),
-		})
-		if getErr == nil && checkResult.Status == ssmTypes.CommandInvocationStatusSuccess &&
-			checkResult.StandardOutputContent != nil && strings.Contains(*checkResult.StandardOutputContent, "MARKER_EXISTS") {
-			log.Printf("[Elastic CI Mode] Instance %s already received SIGTERM (marker exists), skipping", instanceID)
-			return nil
+	if descErr != nil {
+		log.Printf("[Elastic CI Mode] Warning: Failed to detect platform for %s, defaulting to Linux: %v", instanceID, descErr)
+	} else if len(descResp.Reservations) > 0 && len(descResp.Reservations[0].Instances) > 0 {
+		instance := descResp.Reservations[0].Instances[0]
+		if strings.EqualFold(string(instance.Platform), "windows") {
+			platform = "windows"
+			documentName = "AWS-RunPowerShellScript"
+			stopCommand = "nssm stop buildkite-agent"
 		}
 	}
 
-	// Command that creates a marker file (for idempotency) and then stops the agent gracefully.
-	// The marker file is checked by CleanupDanglingInstances to avoid treating this instance
-	// as "dangling" while it's still in the process of shutting down.
-	command := `#!/bin/bash
-set -euo pipefail
-echo "Termination requested at $(date)" > /tmp/buildkite-agent-termination-marker
-sudo systemctl stop buildkite-agent.service || sudo /opt/buildkite-agent/bin/buildkite-agent stop --signal SIGTERM
-`
-	log.Printf("[Elastic CI Mode] Sending SIGTERM to instance %s", instanceID)
+	// Wait for SSM agent to be ready before sending command
+	if err := a.waitForSSMReady(ctx, instanceID, 30*time.Second); err != nil {
+		log.Printf("[Elastic CI Mode] SSM agent not ready on %s, cannot send stop command: %v", instanceID, err)
+		return err
+	}
+
+	log.Printf("[Elastic CI Mode] Sending stop command to %s instance %s", platform, instanceID)
 
-	_, err = ssmClient.SendCommand(ctx, &ssm.SendCommandInput{
+	_, err := ssmClient.SendCommand(ctx, &ssm.SendCommandInput{
 		InstanceIds:  []string{instanceID},
-		DocumentName: aws.String("AWS-RunShellScript"),
-		Parameters:   map[string][]string{"commands": {command}},
+		DocumentName: aws.String(documentName),
+		Parameters:   map[string][]string{"commands": {stopCommand}},
 		Comment:      aws.String("Gracefully stop Buildkite agent"),
 	})
 
 	if err != nil {
-		log.Printf("[Elastic CI Mode] Error sending SIGTERM to instance %s: %v", instanceID, err)
+		log.Printf("[Elastic CI Mode] Error sending stop command to %s: %v", instanceID, err)
 		return err
 	}
-	log.Printf("[Elastic CI Mode] Successfully sent SIGTERM command to instance %s", instanceID)
+	log.Printf("[Elastic CI Mode] Successfully sent stop command to %s", instanceID)
 	return nil
 }
 

scaler/asg_test.go

@@ -90,11 +90,9 @@ func TestGetCheckCommand(t *testing.T) {
 				"ActiveState",
 				"RUNNING",
 				"NOT_RUNNING",
-				"MARKER_EXISTS",
 			},
 			expectedNotContains: []string{
-				"PowerShell",
-				"Get-Service",
+				"nssm",
 			},
 		},
 		{

scaler/scaler.go

@@ -388,11 +388,13 @@ func (s *Scaler) scaleIn(ctx context.Context, desired int64, current AutoscaleGr
 		log.Printf("[Elastic CI Mode] Using graceful termination for %d instances", instancesToTerminate)
 
 		// Determine instances to terminate by sorting by launch time (oldest first)
+		// Only consider healthy instances for graceful termination - unhealthy ones are already
+		// being terminated by the ASG (e.g., instances marked unhealthy by CleanupDanglingInstances)
 		maxToTerminate := instancesToTerminate
 
 		instancesForTermination := make([]string, 0, maxToTerminate)
 
-		if len(current.InstanceIDs) > 0 {
+		if len(current.HealthyInstanceIDs) > 0 {
 			// Define a struct to hold instance info for sorting
 			type instanceInfo struct {
 				ID         string
@@ -401,15 +403,15 @@ func (s *Scaler) scaleIn(ctx context.Context, desired int64, current AutoscaleGr
 
 			ec2Svc := ec2.NewFromConfig(s.cfg)
 
-			instances := make([]instanceInfo, 0, len(current.InstanceIDs))
+			instances := make([]instanceInfo, 0, len(current.HealthyInstanceIDs))
 			describeResult, err := ec2Svc.DescribeInstances(ctx, &ec2.DescribeInstancesInput{
-				InstanceIds: current.InstanceIDs,
+				InstanceIds: current.HealthyInstanceIDs,
 			})
 
 			if err != nil {
 				log.Printf("[Elastic CI Mode] Warning: Could not get instance launch times: %v", err)
 				// Fall back to unsorted if we can't get launch times
-				instancesForTermination = current.InstanceIDs
+				instancesForTermination = current.HealthyInstanceIDs
 				if int64(len(instancesForTermination)) > maxToTerminate {
 					instancesForTermination = instancesForTermination[:maxToTerminate]
 				}