Add tags to Auto Scaling Group and Fleet request

Author: haugenj

docs/deployment/aws/README.md

@@ -19,6 +19,7 @@ Escalator requires the following IAM policy to be able to properly integrate wit
       "Effect": "Allow",
       "Action": [
         "autoscaling:AttachInstances",
+        "autoscaling:CreateOrUpdateTags",
         "autoscaling:DescribeAutoScalingGroups",
         "autoscaling:SetDesiredCapacity",
         "autoscaling:TerminateInstanceInAutoScalingGroup",

pkg/cloudprovider/aws/aws.go

@@ -26,6 +26,10 @@ const (
 	LifecycleSpot = "spot"
 	// The AttachInstances API only supports adding 20 instances at a time
 	batchSize = 20
+	// tagKey is the key for the tag applied to the ASG
+	tagKey = "k8s.io/atlassian-escalator/enabled"
+	// tagValue is the value for the tag applied to the ASG
+	tagValue = "true"
 )
 
 func instanceToProviderID(instance *autoscaling.Instance) string {
@@ -92,6 +96,34 @@ func (c *CloudProvider) RegisterNodeGroups(groups ...cloudprovider.NodeGroupConf
 			continue
 		}
 
+		// Search the asg for tagKey, then add it if it's not present
+		hasTag := false
+		tags := group.Tags
+		for _, tag := range tags {
+			if *tag.Key == tagKey {
+				hasTag = true
+				break
+			}
+		}
+		if !hasTag {
+			tagInput := &autoscaling.CreateOrUpdateTagsInput{
+				Tags: []*autoscaling.Tag{
+					{
+						Key:               awsapi.String(tagKey),
+						PropagateAtLaunch: awsapi.Bool(true),
+						ResourceId:        awsapi.String(id),
+						ResourceType:      awsapi.String("auto-scaling-group"),
+						Value:             awsapi.String(tagValue),
+					},
+				},
+			}
+			log.WithField("asg", id).Infof("creating auto scaling tag")
+			_, err := c.service.CreateOrUpdateTags(tagInput)
+			if err != nil {
+				log.Errorf("failed to create auto scaling tag for ASG %v", id)
+			}
+		}
+
 		c.nodeGroups[id] = NewNodeGroup(configs[id], group, c)
 	}
 
@@ -485,6 +517,17 @@ func createFleetInput(n NodeGroup, addCount int64) (*ec2.CreateFleetInput, error
 				Overrides: launchTemplateOverrides,
 			},
 		},
+		TagSpecifications: []*ec2.TagSpecification{
+			{
+				ResourceType: awsapi.String(ec2.ResourceTypeFleet),
+				Tags: []*ec2.Tag{
+					{
+						Key:   awsapi.String(tagKey),
+						Value: awsapi.String(tagValue),
+					},
+				},
+			},
+		},
 	}
 
 	if lifecycle == LifecycleOnDemand {

pkg/cloudprovider/aws/cloud_provider_test.go

@@ -91,10 +91,12 @@ func TestCloudProvider_GetNodeGroup(t *testing.T) {
 
 func TestCloudProvider_RegisterNodeGroups(t *testing.T) {
 	tests := []struct {
-		name       string
-		nodeGroups map[string]bool
-		response   *autoscaling.DescribeAutoScalingGroupsOutput
-		err        error
+		name        string
+		nodeGroups  map[string]bool
+		response    *autoscaling.DescribeAutoScalingGroupsOutput
+		err         error
+		tagResponse *autoscaling.CreateOrUpdateTagsOutput
+		tagErr      error
 	}{
 		{
 			"register node group that does not exist",
@@ -103,6 +105,8 @@ func TestCloudProvider_RegisterNodeGroups(t *testing.T) {
 			},
 			&autoscaling.DescribeAutoScalingGroupsOutput{},
 			nil,
+			&autoscaling.CreateOrUpdateTagsOutput{},
+			nil,
 		},
 		{
 			"register node groups that exist",
@@ -121,6 +125,8 @@ func TestCloudProvider_RegisterNodeGroups(t *testing.T) {
 				},
 			},
 			nil,
+			&autoscaling.CreateOrUpdateTagsOutput{},
+			nil,
 		},
 		{
 			"register node groups, some don't exist",
@@ -136,12 +142,32 @@ func TestCloudProvider_RegisterNodeGroups(t *testing.T) {
 				},
 			},
 			nil,
+			&autoscaling.CreateOrUpdateTagsOutput{},
+			nil,
 		},
 		{
 			"register no node groups",
 			map[string]bool{},
 			&autoscaling.DescribeAutoScalingGroupsOutput{},
 			fmt.Errorf("no groups"),
+			&autoscaling.CreateOrUpdateTagsOutput{},
+			nil,
+		},
+		{
+			"register existing node group with error from CreateOrUpdateTags",
+			map[string]bool{
+				"1": true,
+			},
+			&autoscaling.DescribeAutoScalingGroupsOutput{
+				AutoScalingGroups: []*autoscaling.Group{
+					{
+						AutoScalingGroupName: aws.String("1"),
+					},
+				},
+			},
+			nil,
+			&autoscaling.CreateOrUpdateTagsOutput{},
+			fmt.Errorf("unauthorized operation"),
 		},
 	}
 
@@ -151,6 +177,8 @@ func TestCloudProvider_RegisterNodeGroups(t *testing.T) {
 			service := test.MockAutoscalingService{
 				DescribeAutoScalingGroupsOutput: tt.response,
 				DescribeAutoScalingGroupsErr:    tt.err,
+				CreateOrUpdateTagsOutput:        tt.tagResponse,
+				CreateOrUpdateTagsErr:           tt.tagErr,
 			}
 
 			var ids []string

pkg/test/aws.go

@@ -16,6 +16,9 @@ type MockAutoscalingService struct {
 	AttachInstanceOutput *autoscaling.AttachInstancesOutput
 	AttachInstanceErr    error
 
+	CreateOrUpdateTagsOutput *autoscaling.CreateOrUpdateTagsOutput
+	CreateOrUpdateTagsErr    error
+
 	DescribeAutoScalingGroupsOutput *autoscaling.DescribeAutoScalingGroupsOutput
 	DescribeAutoScalingGroupsErr    error
 
@@ -31,6 +34,11 @@ func (m MockAutoscalingService) AttachInstances(*autoscaling.AttachInstancesInpu
 	return m.AttachInstanceOutput, m.AttachInstanceErr
 }
 
+// CreateOrUpdateTags mock implementation for MockAutoscalingService
+func (m MockAutoscalingService) CreateOrUpdateTags(*autoscaling.CreateOrUpdateTagsInput) (*autoscaling.CreateOrUpdateTagsOutput, error) {
+	return m.CreateOrUpdateTagsOutput, m.CreateOrUpdateTagsErr
+}
+
 // DescribeAutoScalingGroups mock implementation for MockAutoscalingService
 func (m MockAutoscalingService) DescribeAutoScalingGroups(*autoscaling.DescribeAutoScalingGroupsInput) (*autoscaling.DescribeAutoScalingGroupsOutput, error) {
 	return m.DescribeAutoScalingGroupsOutput, m.DescribeAutoScalingGroupsErr