Allow the NTHManagedAsg tag to be configurable (#287)

* Make the tag to check for on a managed ASG configurable

* Fix the arguments print that occurs at program boot

* Fix broken tests from missing constant

* ManagedAsgTag needs a definition

* Add managedAsgTag support to the Helm chart

Author: blakestoddard

cmd/node-termination-handler.go

@@ -150,6 +150,7 @@ func main() {
 
 		sqsMonitor := sqsevent.SQSMonitor{
 			CheckIfManaged:   nthConfig.CheckASGTagBeforeDraining,
+			ManagedAsgTag:    nthConfig.ManagedAsgTag,
 			QueueURL:         nthConfig.QueueURL,
 			InterruptionChan: interruptionChan,
 			CancelChan:       cancelChan,

config/helm/aws-node-termination-handler/README.md

@@ -86,6 +86,8 @@ Parameter | Description | Default
 `enableSqsTerminationDraining` | If true, this turns on queue-processor mode which drains nodes when an SQS termination event is received| `false`
 `queueURL` | Listens for messages on the specified SQS queue URL | None
 `awsRegion` | If specified, use the AWS region for AWS API calls, else NTH will try to find the region through AWS_REGION env var, IMDS, or the specified queue URL | ``
+`checkASGTagBeforeDraining` | If true, check that the instance is tagged with "aws-node-termination-handler/managed" as the key before draining the node | `true`
+`managedAsgTag` | The tag to ensure is on a node if checkASGTagBeforeDraining is true | `aws-node-termination-handler/managed`
 
 ### AWS Node Termination Handler - IMDS Mode Configuration
 

config/helm/aws-node-termination-handler/templates/deployment.yaml

@@ -140,6 +140,8 @@ spec:
           {{- end }}
           - name: CHECK_ASG_TAG_BEFORE_DRAINING
             value: {{ .Values.checkASGTagBeforeDraining | quote }}
+          - name: MANAGED_ASG_TAG
+            value: {{ .Values.managedAsgTag | quote }}
           resources:
             {{- toYaml .Values.resources | nindent 12 }}
           {{- if .Values.enablePrometheusServer }}

config/helm/aws-node-termination-handler/values.yaml

@@ -45,6 +45,9 @@ queueURL: ""
 # checkASGTagBeforeDraining  If true, check that the instance is tagged with "aws-node-termination-handler/managed" as the key before draining the node
 checkASGTagBeforeDraining: true
 
+# managedAsgTag  The tag to ensure is on a node if checkASGTagBeforeDraining is true
+managedAsgTag: "aws-node-termination-handler/managed"
+
 # awsRegion If specified, use the AWS region for AWS API calls
 awsRegion: ""
 

pkg/config/config.go

@@ -58,6 +58,8 @@ const (
 	enableRebalanceMonitoringDefault        = false
 	checkASGTagBeforeDrainingConfigKey      = "CHECK_ASG_TAG_BEFORE_DRAINING"
 	checkASGTagBeforeDrainingDefault        = true
+	managedAsgTagConfigKey                  = "MANAGED_ASG_TAG"
+	managedAsgTagDefault                    = "aws-node-termination-handler/managed"
 	metadataTriesConfigKey                  = "METADATA_TRIES"
 	metadataTriesDefault                    = 3
 	cordonOnly                              = "CORDON_ONLY"
@@ -102,6 +104,7 @@ type Config struct {
 	EnableSQSTerminationDraining   bool
 	EnableRebalanceMonitoring      bool
 	CheckASGTagBeforeDraining      bool
+	ManagedAsgTag                  string
 	MetadataTries                  int
 	CordonOnly                     bool
 	TaintNode                      bool
@@ -147,6 +150,7 @@ func ParseCliArgs() (config Config, err error) {
 	flag.BoolVar(&config.EnableSQSTerminationDraining, "enable-sqs-termination-draining", getBoolEnv(enableSQSTerminationDrainingConfigKey, enableSQSTerminationDrainingDefault), "If true, drain nodes when an SQS termination event is received")
 	flag.BoolVar(&config.EnableRebalanceMonitoring, "enable-rebalance-monitoring", getBoolEnv(enableRebalanceMonitoringConfigKey, enableRebalanceMonitoringDefault), "If true, cordon nodes when the rebalance recommendation notice is received")
 	flag.BoolVar(&config.CheckASGTagBeforeDraining, "check-asg-tag-before-draining", getBoolEnv(checkASGTagBeforeDrainingConfigKey, checkASGTagBeforeDrainingDefault), "If true, check that the instance is tagged with \"aws-node-termination-handler/managed\" as the key before draining the node")
+	flag.StringVar(&config.ManagedAsgTag, "managed-asg-tag", getEnv(managedAsgTagConfigKey, managedAsgTagDefault), "Sets the tag to check for on instances that is propogated from the ASG before taking action, default to aws-node-termination-handler/managed")
 	flag.IntVar(&config.MetadataTries, "metadata-tries", getIntEnv(metadataTriesConfigKey, metadataTriesDefault), "The number of times to try requesting metadata. If you would like 2 retries, set metadata-tries to 3.")
 	flag.BoolVar(&config.CordonOnly, "cordon-only", getBoolEnv(cordonOnly, false), "If true, nodes will be cordoned but not drained when an interruption event occurs.")
 	flag.BoolVar(&config.TaintNode, "taint-node", getBoolEnv(taintNode, false), "If true, nodes will be tainted when an interruption event occurs.")
@@ -245,6 +249,7 @@ func (c Config) PrintJsonConfigArgs() {
 		Str("aws_endpoint", c.AWSEndpoint).
 		Str("queue_url", c.QueueURL).
 		Bool("check_asg_tag_before_draining", c.CheckASGTagBeforeDraining).
+		Str("ManagedAsgTag", c.ManagedAsgTag).
 		Msg("aws-node-termination-handler arguments")
 }
 
@@ -285,6 +290,7 @@ func (c Config) PrintHumanConfigArgs() {
 			"\taws-region: %s,\n"+
 			"\tqueue-url: %s,\n"+
 			"\tcheck-asg-tag-before-draining: %t,\n"+
+			"\tmanaged-asg-tag: %s,\n"+
 			"\taws-endpoint: %s,\n",
 		c.DryRun,
 		c.NodeName,
@@ -314,6 +320,7 @@ func (c Config) PrintHumanConfigArgs() {
 		c.AWSRegion,
 		c.QueueURL,
 		c.CheckASGTagBeforeDraining,
+		c.ManagedAsgTag,
 		c.AWSEndpoint,
 	)
 }

pkg/monitor/sqsevent/sqs-monitor.go

@@ -31,8 +31,6 @@ import (
 const (
 	// SQSTerminateKind is a const to define an SQS termination kind of interruption event
 	SQSTerminateKind = "SQS_TERMINATE"
-	// NTHManagedASG is the ASG tag key to determine if NTH is managing the ASG
-	NTHManagedASG = "aws-node-termination-handler/managed"
 )
 
 // SQSMonitor is a struct definition that knows how to process events from Amazon EventBridge
@@ -44,6 +42,7 @@ type SQSMonitor struct {
 	ASG              autoscalingiface.AutoScalingAPI
 	EC2              ec2iface.EC2API
 	CheckIfManaged   bool
+	ManagedAsgTag    string
 }
 
 // Kind denotes the kind of event that is processed
@@ -207,7 +206,7 @@ func (m SQSMonitor) isInstanceManaged(instanceID string) (bool, error) {
 	isManaged := false
 	err = m.ASG.DescribeTagsPages(&asgDescribeTagsInput, func(resp *autoscaling.DescribeTagsOutput, next bool) bool {
 		for _, tag := range resp.Tags {
-			if *tag.Key == NTHManagedASG {
+			if *tag.Key == m.ManagedAsgTag {
 				isManaged = true
 				// breaks paging loop
 				return false
@@ -220,7 +219,7 @@ func (m SQSMonitor) isInstanceManaged(instanceID string) (bool, error) {
 	if !isManaged {
 		log.Debug().
 			Str("instance_id", instanceID).
-			Msgf("The instance's Auto Scaling Group is not tagged as managed with tag key: %s", NTHManagedASG)
+			Msgf("The instance's Auto Scaling Group is not tagged as managed with tag key: %s", m.ManagedAsgTag)
 	}
 	return isManaged, err
 }

pkg/monitor/sqsevent/sqs-monitor_internal_test.go

@@ -48,11 +48,15 @@ func TestIsInstanceManaged(t *testing.T) {
 		},
 		DescribeTagsPagesResp: autoscaling.DescribeTagsOutput{
 			Tags: []*autoscaling.TagDescription{
-				{Key: aws.String(NTHManagedASG)},
+				{Key: aws.String("aws-node-termination-handler/managed")},
 			},
 		},
 	}
-	monitor := SQSMonitor{ASG: asgMock}
+	monitor := SQSMonitor{
+		ASG:            asgMock,
+		CheckIfManaged: true,
+		ManagedAsgTag:  "aws-node-termination-handler/managed",
+	}
 	isManaged, err := monitor.isInstanceManaged("i-0123456789")
 	h.Ok(t, err)
 	h.Equals(t, true, isManaged)

pkg/monitor/sqsevent/sqs-monitor_test.go

@@ -539,7 +539,7 @@ func mockIsManagedTrue(asg *h.MockedASG) h.MockedASG {
 	}
 	asg.DescribeTagsPagesResp = autoscaling.DescribeTagsOutput{
 		Tags: []*autoscaling.TagDescription{
-			{Key: aws.String(sqsevent.NTHManagedASG)},
+			{Key: aws.String("aws-node-termination-handler/managed")},
 		},
 	}
 	return *asg