diff --git a/autoscaling.go b/autoscaling.go index de954b9..53846ad 100644 --- 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 index 1c1e039..a5f5887 100644 --- 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 index b068341..8a8de63 100644 --- 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.