Merge pull request #97 from stevenmatthewt/master

Support configuring intervals for spot instance termination and autoscaling heartbeats
buildkite · Oct 18, 2022 · f67af3d · f67af3d
2 parents d354903 + 5b6f16d
commit f67af3d
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 33 deletions.
diff --git a/autoscaling.go b/autoscaling.go
@@ -12,6 +12,7 @@ import (
 )
 
 // AutoscalingClient for testing purposes
+//
 //go:generate mockgen -destination=mocks/mock_autoscaling_client.go -package=mocks github.com/buildkite/lifecycled AutoscalingClient
 type AutoscalingClient autoscalingiface.AutoScalingAPI
 
@@ -34,21 +35,23 @@ type Message struct {
 }
 
 // NewAutoscalingListener ...
-func NewAutoscalingListener(instanceID string, queue *Queue, autoscaling AutoscalingClient) *AutoscalingListener {
+func NewAutoscalingListener(instanceID string, queue *Queue, autoscaling AutoscalingClient, heartbeatInterval time.Duration) *AutoscalingListener {
 	return &AutoscalingListener{
-		listenerType: "autoscaling",
-		instanceID:   instanceID,
-		queue:        queue,
-		autoscaling:  autoscaling,
+		listenerType:      "autoscaling",
+		instanceID:        instanceID,
+		queue:             queue,
+		autoscaling:       autoscaling,
+		heartbeatInterval: heartbeatInterval,
 	}
 }
 
 // AutoscalingListener ...
 type AutoscalingListener struct {
-	listenerType string
-	instanceID   string
-	queue        *Queue
-	autoscaling  AutoscalingClient
+	listenerType      string
+	instanceID        string
+	queue             *Queue
+	autoscaling       AutoscalingClient
+	heartbeatInterval time.Duration
 }
 
 // Type returns a string describing the listener type.
@@ -126,9 +129,10 @@ func (l *AutoscalingListener) Start(ctx context.Context, notices chan<- Terminat
 				}
 
 				notices <- &autoscalingTerminationNotice{
-					noticeType:  l.Type(),
-					message:     &msg,
-					autoscaling: l.autoscaling,
+					noticeType:        l.Type(),
+					message:           &msg,
+					autoscaling:       l.autoscaling,
+					heartbeatInterval: l.heartbeatInterval,
 				}
 				return nil
 			}
@@ -137,9 +141,10 @@ func (l *AutoscalingListener) Start(ctx context.Context, notices chan<- Terminat
 }
 
 type autoscalingTerminationNotice struct {
-	noticeType  string
-	message     *Message
-	autoscaling AutoscalingClient
+	noticeType        string
+	message           *Message
+	autoscaling       AutoscalingClient
+	heartbeatInterval time.Duration
 }
 
 func (n *autoscalingTerminationNotice) Type() string {
@@ -162,7 +167,7 @@ func (n *autoscalingTerminationNotice) Handle(ctx context.Context, handler Handl
 		}
 	}()
 
-	ticker := time.NewTicker(10 * time.Second)
+	ticker := time.NewTicker(n.heartbeatInterval)
 	defer ticker.Stop()
 
 	go func() {

diff --git a/cmd/lifecycled/main.go b/cmd/lifecycled/main.go
@@ -29,14 +29,16 @@ func main() {
 	app.DefaultEnvars()
 
 	var (
-		instanceID          string
-		snsTopic            string
-		disableSpotListener bool
-		handler             *os.File
-		jsonLogging         bool
-		debugLogging        bool
-		cloudwatchGroup     string
-		cloudwatchStream    string
+		instanceID                   string
+		snsTopic                     string
+		disableSpotListener          bool
+		handler                      *os.File
+		jsonLogging                  bool
+		debugLogging                 bool
+		cloudwatchGroup              string
+		cloudwatchStream             string
+		spotListenerInterval         time.Duration
+		autoscalingHeartbeatInterval time.Duration
 	)
 
 	app.Flag("instance-id", "The instance id to listen for events for").
@@ -64,6 +66,14 @@ func main() {
 	app.Flag("debug", "Show debugging info").
 		BoolVar(&debugLogging)
 
+	app.Flag("spot-listener-interval", "Interval to check for spot instance termination notices").
+		Default("5s").
+		DurationVar(&spotListenerInterval)
+
+	app.Flag("autoscaling-heartbeat-interval", "Interval to send AWS Lifecycle Heartbeat Actions").
+		Default("10s").
+		DurationVar(&autoscalingHeartbeatInterval)
+
 	app.Action(func(c *kingpin.ParseContext) error {
 		logger := logrus.New()
 		if jsonLogging {
@@ -148,10 +158,11 @@ func main() {
 
 		handler := lifecycled.NewFileHandler(handler)
 		daemon := lifecycled.New(&lifecycled.Config{
-			InstanceID:           instanceID,
-			SNSTopic:             snsTopic,
-			SpotListener:         !disableSpotListener,
-			SpotListenerInterval: 5 * time.Second,
+			InstanceID:                   instanceID,
+			SNSTopic:                     snsTopic,
+			SpotListener:                 !disableSpotListener,
+			SpotListenerInterval:         spotListenerInterval,
+			AutoscalingHeartbeatInterval: autoscalingHeartbeatInterval,
 		}, sess, logger)
 
 		notice, err := daemon.Start(ctx)

diff --git a/daemon.go b/daemon.go
@@ -52,17 +52,18 @@ func NewDaemon(
 			sqsClient,
 			snsClient,
 		)
-		daemon.AddListener(NewAutoscalingListener(config.InstanceID, queue, asgClient))
+		daemon.AddListener(NewAutoscalingListener(config.InstanceID, queue, asgClient, config.AutoscalingHeartbeatInterval))
 	}
 	return daemon
 }
 
 // Config for the Lifecycled Daemon.
 type Config struct {
-	InstanceID           string
-	SNSTopic             string
-	SpotListener         bool
-	SpotListenerInterval time.Duration
+	InstanceID                   string
+	SNSTopic                     string
+	SpotListener                 bool
+	SpotListenerInterval         time.Duration
+	AutoscalingHeartbeatInterval time.Duration
 }
 
 // Daemon is what orchestrates the listening and execution of the handler on a termination notice.