fix: handle detached not terminated instances (#18)

Fix the issue where a detached but not terminated instance could lead to incorrect scaling decisions.

Author: ilya-hontarau

internal/auto_scaler_test.go

@@ -4,8 +4,10 @@ import (
 	"bytes"
 	"context"
 	"testing"
+	"time"
 
 	"github.com/aws/aws-sdk-go-v2/service/autoscaling/types"
+	ec2types "github.com/aws/aws-sdk-go-v2/service/ec2/types"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/exp/slog"
@@ -68,7 +70,7 @@ func TestAutoScalerScalingUp(t *testing.T) {
 		MaxSize:              ptr(int32(3)),
 		DesiredCapacity:      ptr(int32(2)),
 		Instances: []types.Instance{
-			{},
+			{InstanceId: ptr("instance")},
 		},
 	}, nil)
 	ctrl.On("ScaleUpASG", mock.Anything, int32(2)).Return(nil)
@@ -107,8 +109,8 @@ func TestAutoScalerScalingDown(t *testing.T) {
 		MaxSize:              ptr(int32(3)),
 		DesiredCapacity:      ptr(int32(2)),
 		Instances: []types.Instance{
-			{},
-			{},
+			{InstanceId: ptr("instance")},
+			{InstanceId: ptr("instance2")},
 		},
 	}, nil)
 	ctrl.On("DrainWorker", mock.Anything, "1").Return(true, nil)
@@ -117,6 +119,54 @@ func TestAutoScalerScalingDown(t *testing.T) {
 	require.NoError(t, err)
 }
 
+func TestAutoScalerDetachedNotTerminatedInstances(t *testing.T) {
+	var buf bytes.Buffer
+	h := slog.NewTextHandler(&buf, nil)
+
+	cfg := internal.RuntimeConfig{}
+
+	ctrl := new(MockController)
+	defer ctrl.AssertExpectations(t)
+
+	scaler := internal.NewAutoScaler(ctrl, slog.New(h))
+
+	ctrl.On("GetWorkerPool", mock.Anything).Return(&internal.WorkerPool{
+		Workers: []internal.Worker{
+			{
+				ID:       "1",
+				Metadata: `{"asg_id": "group", "instance_id": "instance"}`,
+			},
+			{
+				ID:       "2",
+				Drained:  true,
+				Metadata: `{"asg_id": "group", "instance_id": "detached"}`,
+			},
+		},
+		PendingRuns: 2,
+	}, nil)
+	ctrl.On("GetAutoscalingGroup", mock.Anything).Return(&types.AutoScalingGroup{
+		AutoScalingGroupName: ptr("group"),
+		MinSize:              ptr(int32(1)),
+		MaxSize:              ptr(int32(3)),
+		DesiredCapacity:      ptr(int32(2)),
+		Instances: []types.Instance{
+			{InstanceId: ptr("instance")},
+		},
+	}, nil)
+	ctrl.On("KillInstance", mock.Anything, "detached").Return(nil)
+	output := []ec2types.Instance{{
+		InstanceId: ptr("detached"),
+		LaunchTime: nullable(time.Now().Add(-time.Hour)),
+	}}
+	ctrl.On(
+		"DescribeInstances",
+		mock.Anything,
+		[]string{"detached"},
+	).Return(output, nil)
+	err := scaler.Scale(context.Background(), cfg)
+	require.NoError(t, err)
+}
+
 func ptr[T any](v T) *T {
 	return &v
 }

internal/state.go

@@ -85,14 +85,37 @@ func (s *State) IdleWorkers() []Worker {
 
 // StrayInstances returns a list of instance IDs that don't have a corresponding
 // worker in the worker pool.
-func (s *State) StrayInstances() (out []string) {
+func (s *State) StrayInstances() []string {
+	var res []string
 	for instanceID := range s.inServiceInstanceIDs {
 		if _, ok := s.workersByInstanceID[instanceID]; !ok {
-			out = append(out, string(instanceID))
+			res = append(res, string(instanceID))
 		}
 	}
 
-	return
+	res = append(res, s.detachedNotTerminatedInstances()...)
+
+	return res
+}
+
+func (s *State) detachedNotTerminatedInstances() []string {
+	instanceIDs := make(map[InstanceID]struct{})
+	for _, instance := range s.ASG.Instances {
+		instanceIDs[InstanceID(*instance.InstanceId)] = struct{}{}
+	}
+
+	var res []string
+	for instanceID, worker := range s.workersByInstanceID {
+		if !worker.Drained {
+			continue
+		}
+		if _, ok := instanceIDs[instanceID]; ok {
+			continue
+		}
+
+		res = append(res, string(instanceID))
+	}
+	return res
 }
 
 func (s *State) Decide(maxCreate, maxKill int) Decision {

internal/state_test.go

@@ -7,10 +7,52 @@ import (
 	"github.com/aws/aws-sdk-go-v2/service/autoscaling/types"
 	"github.com/franela/goblin"
 	. "github.com/onsi/gomega"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 
 	"github.com/spacelift-io/awsautoscalr/internal"
 )
 
+func TestState_StrayInstances(t *testing.T) {
+	const asgName = "asg-name"
+	const instanceID = "instance-id"
+	const failedToTerminateInstanceID = "instance-id2"
+	asg := &types.AutoScalingGroup{
+		AutoScalingGroupName: nullable(asgName),
+		MinSize:              nullable(int32(1)),
+		MaxSize:              nullable(int32(5)),
+		DesiredCapacity:      nullable(int32(3)),
+		Instances: []types.Instance{
+			{
+				InstanceId: nullable(instanceID),
+			},
+		},
+	}
+	workerPool := &internal.WorkerPool{
+		Workers: []internal.Worker{
+			{
+				Metadata: mustJSON(map[string]any{
+					"asg_id":      asgName,
+					"instance_id": instanceID,
+				}),
+			},
+			{
+				Drained: true,
+				Metadata: mustJSON(map[string]any{
+					"asg_id":      asgName,
+					"instance_id": failedToTerminateInstanceID,
+				}),
+			},
+		},
+	}
+
+	state, err := internal.NewState(workerPool, asg)
+	require.NoError(t, err)
+
+	strayInstances := state.StrayInstances()
+	assert.Equal(t, []string{failedToTerminateInstanceID}, strayInstances)
+}
+
 func TestState(t *testing.T) {
 	g := goblin.Goblin(t)
 	RegisterFailHandler(func(m string, _ ...int) { g.Fail(m) })