feat: add a delay to scale in (#28)

Author: Apollorion

README.md

@@ -25,6 +25,7 @@ Two additional environment variables are optional, but very useful if you're run
 
 - `AUTOSCALING_MAX_KILL` (defaults to 1) - the maximum number of instances the utility is allowed to terminate in a single run;
 - `AUTOSCALING_MAX_CREATE` (defaults to 1) - the maximum number of instances the utility is allowed to create in a single run;
+- `AUTOSCALING_SCALE_DOWN_DELAY` (defaults to 0) - the number of minutes a worker must be registered to Spacelift before its eligible to be scaled in.
 
 ## Important note on concurrency
 

internal/auto_scaler.go

@@ -45,7 +45,7 @@ func (s AutoScaler) Scale(ctx context.Context, cfg RuntimeConfig) error {
 		return fmt.Errorf("could not get autoscaling group: %w", err)
 	}
 
-	state, err := NewState(workerPool, asg)
+	state, err := NewState(workerPool, asg, cfg)
 	if err != nil {
 		return fmt.Errorf("could not create state: %w", err)
 	}
@@ -109,10 +109,10 @@ func (s AutoScaler) Scale(ctx context.Context, cfg RuntimeConfig) error {
 	// If we got this far, we're scaling down.
 	logger.With("instances", decision.ScalingSize).Info("scaling down ASG")
 
-	idleWorkers := state.IdleWorkers()
+	scalableWorkers := state.ScalableWorkers()
 
 	for i := 0; i < decision.ScalingSize; i++ {
-		worker := idleWorkers[i]
+		worker := scalableWorkers[i]
 
 		_, instanceID, _ := worker.InstanceIdentity()
 

internal/runtime_config.go

@@ -6,8 +6,9 @@ type RuntimeConfig struct {
 	SpaceliftAPIEndpoint   string `env:"SPACELIFT_API_KEY_ENDPOINT,notEmpty"`
 	SpaceliftWorkerPoolID  string `env:"SPACELIFT_WORKER_POOL_ID,notEmpty"`
 
-	AutoscalingGroupARN  string `env:"AUTOSCALING_GROUP_ARN,notEmpty"`
-	AutoscalingRegion    string `env:"AUTOSCALING_REGION,notEmpty"`
-	AutoscalingMaxKill   int    `env:"AUTOSCALING_MAX_KILL" envDefault:"1"`
-	AutoscalingMaxCreate int    `env:"AUTOSCALING_MAX_CREATE" envDefault:"1"`
+	AutoscalingScaleDownDelay int    `env:"AUTOSCALING_SCALE_DOWN_DELAY" envDefault:"0"`
+	AutoscalingGroupARN       string `env:"AUTOSCALING_GROUP_ARN,notEmpty"`
+	AutoscalingRegion         string `env:"AUTOSCALING_REGION,notEmpty"`
+	AutoscalingMaxKill        int    `env:"AUTOSCALING_MAX_KILL" envDefault:"1"`
+	AutoscalingMaxCreate      int    `env:"AUTOSCALING_MAX_CREATE" envDefault:"1"`
 }

internal/state.go

@@ -2,8 +2,8 @@ package internal
 
 import (
 	"fmt"
-
 	"github.com/aws/aws-sdk-go-v2/service/autoscaling/types"
+	"time"
 )
 
 // State represents the state of the world, as far as the autoscaler is
@@ -15,9 +15,10 @@ type State struct {
 
 	inServiceInstanceIDs map[InstanceID]struct{}
 	workersByInstanceID  map[InstanceID]Worker
+	cfg                  RuntimeConfig
 }
 
-func NewState(workerPool *WorkerPool, asg *types.AutoScalingGroup) (*State, error) {
+func NewState(workerPool *WorkerPool, asg *types.AutoScalingGroup, cfg RuntimeConfig) (*State, error) {
 	workersByInstanceID := make(map[InstanceID]Worker)
 	inServiceInstanceIDs := make(map[InstanceID]struct{})
 
@@ -65,18 +66,28 @@ func NewState(workerPool *WorkerPool, asg *types.AutoScalingGroup) (*State, erro
 		ASG:                  asg,
 		inServiceInstanceIDs: inServiceInstanceIDs,
 		workersByInstanceID:  workersByInstanceID,
+		cfg:                  cfg,
 	}, nil
 }
 
-// IdleWorkers returns a list of workers that are not currently busy.
-func (s *State) IdleWorkers() []Worker {
+// ScalableWorkers returns a list of workers that are not currently busy.
+func (s *State) ScalableWorkers() []Worker {
 	var out []Worker
-
 	for _, worker := range s.WorkerPool.Workers {
 		if worker.Busy {
 			continue
 		}
 
+		// Even though the worker might be idle, we will give it some time to pick up more work
+		// if the customer wants
+		if s.cfg.AutoscalingScaleDownDelay != 0 {
+			workerCreationTime := time.Unix(int64(worker.CreatedAt), 0)
+			minimumAliveTime := workerCreationTime.Add(time.Duration(s.cfg.AutoscalingScaleDownDelay) * time.Minute)
+			if !time.Now().After(minimumAliveTime) {
+				continue
+			}
+		}
+
 		out = append(out, worker)
 	}
 
@@ -126,9 +137,9 @@ func (s *State) Decide(maxCreate, maxKill int) Decision {
 		}
 	}
 
-	idle := s.IdleWorkers()
+	scalable := s.ScalableWorkers()
 
-	difference := int(s.WorkerPool.PendingRuns) - len(idle)
+	difference := int(s.WorkerPool.PendingRuns) - len(scalable)
 
 	if difference > 0 {
 		return s.determineScaleUp(difference, maxCreate)

internal/state_test.go

@@ -2,13 +2,13 @@ package internal_test
 
 import (
 	"encoding/json"
-	"testing"
-
 	"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"
+	"testing"
+	"time"
 
 	"github.com/spacelift-io/awsautoscalr/internal"
 )
@@ -17,6 +17,7 @@ func TestState_StrayInstances(t *testing.T) {
 	const asgName = "asg-name"
 	const instanceID = "instance-id"
 	const failedToTerminateInstanceID = "instance-id2"
+	cfg := internal.RuntimeConfig{}
 	asg := &types.AutoScalingGroup{
 		AutoScalingGroupName: nullable(asgName),
 		MinSize:              nullable(int32(1)),
@@ -46,7 +47,7 @@ func TestState_StrayInstances(t *testing.T) {
 		},
 	}
 
-	state, err := internal.NewState(workerPool, asg)
+	state, err := internal.NewState(workerPool, asg, cfg)
 	require.NoError(t, err)
 
 	strayInstances := state.StrayInstances()
@@ -60,6 +61,7 @@ func TestState(t *testing.T) {
 	g.Describe("State", func() {
 		var asg *types.AutoScalingGroup
 		var workerPool *internal.WorkerPool
+		var cfg internal.RuntimeConfig
 
 		var sut *internal.State
 
@@ -76,9 +78,10 @@ func TestState(t *testing.T) {
 					DesiredCapacity:      nullable(int32(3)),
 				}
 				workerPool = &internal.WorkerPool{}
+				cfg = internal.RuntimeConfig{}
 			})
 
-			g.JustBeforeEach(func() { sut, err = internal.NewState(workerPool, asg) })
+			g.JustBeforeEach(func() { sut, err = internal.NewState(workerPool, asg, cfg) })
 
 			g.Describe("when the ASG is invalid", func() {
 				g.Describe("when the name is not set", func() {
@@ -210,10 +213,10 @@ func TestState(t *testing.T) {
 			})
 		})
 
-		g.Describe("IdleWorkers", func() {
-			var idleWorkers []internal.Worker
+		g.Describe("ScaleableWorkers", func() {
+			var scalableWorkers []internal.Worker
 
-			g.JustBeforeEach(func() { idleWorkers = sut.IdleWorkers() })
+			g.JustBeforeEach(func() { scalableWorkers = sut.ScalableWorkers() })
 
 			g.BeforeEach(func() {
 				workerPool.Workers = []internal.Worker{
@@ -223,8 +226,8 @@ func TestState(t *testing.T) {
 			})
 
 			g.It("should return the idle workers", func() {
-				Expect(idleWorkers).To(HaveLen(1))
-				Expect(idleWorkers[0].ID).To(Equal("idle"))
+				Expect(scalableWorkers).To(HaveLen(1))
+				Expect(scalableWorkers[0].ID).To(Equal("idle"))
 			})
 		})
 
@@ -243,10 +246,20 @@ func TestState(t *testing.T) {
 				}
 				workerPool = &internal.WorkerPool{}
 
-				sut = &internal.State{
-					WorkerPool: workerPool,
-					ASG:        asg,
+				asg = &types.AutoScalingGroup{
+					AutoScalingGroupName: nullable("asg-name"),
+					MinSize:              nullable(int32(1)),
+					MaxSize:              nullable(int32(2)),
+					DesiredCapacity:      nullable(int32(2)),
+				}
+				workerPool = &internal.WorkerPool{}
+				cfg = internal.RuntimeConfig{
+					AutoscalingScaleDownDelay: 50,
 				}
+
+				var err error
+				sut, err = internal.NewState(workerPool, asg, cfg)
+				Expect(err).NotTo(HaveOccurred())
 			})
 
 			g.JustBeforeEach(func() {
@@ -401,6 +414,24 @@ func TestState(t *testing.T) {
 									Expect(decision.Comments).To(Equal([]string{"removing idle workers"}))
 								})
 							})
+
+							g.Describe("when waiting some time for idle workers", func() {
+								g.BeforeEach(func() {
+									workerPool.Workers = []internal.Worker{
+										{ID: "busy", Busy: true},
+										{ID: "idle", Busy: false, CreatedAt: int32(time.Now().Unix())},
+									}
+								})
+
+								g.It("should not scale down because the worker is busy and not enough time has passed", func() {
+									Expect(decision.ScalingDirection).To(Equal(internal.ScalingDirectionNone))
+									Expect(decision.ScalingSize).To(BeZero())
+									Expect(decision.Comments).To(Equal([]string{
+										"autoscaling group exactly at the right size",
+									}))
+								})
+							})
+
 						})
 					})
 				})