Merge pull request #193 from haugenj/master

Jacobious52 · web-flow · commit ee864bf14ee6 · 2020-07-14T16:49:01.000+10:00
Add tags to Auto Scaling Group and Fleet request
diff --git a/cmd/main.go b/cmd/main.go
@@ -77,6 +77,7 @@ func setupCloudProvider(nodegroups []controller.NodeGroupOptions) cloudprovider.
 				FleetInstanceReadyTimeout: n.AWS.FleetInstanceReadyTimeoutDuration(),
 				Lifecycle:                 n.AWS.Lifecycle,
 				InstanceTypeOverrides:     n.AWS.InstanceTypeOverrides,
+				ResourceTagging:           n.AWS.ResourceTagging,
 			},
 		})
 	}
diff --git a/docs/configuration/nodegroup.md b/docs/configuration/nodegroup.md
@@ -32,6 +32,7 @@ node_groups:
         launch_template_id: "1"
         lifecycle: on-demand
         instance_type_overrides: ["t2.large", "t3.large"]
+        resource_tagging: false
 ```
 
 ## Options
@@ -242,3 +243,9 @@ Dependent on Launch Template ID being specified.
 An optional list of instance types to override the instance type within the launch template. Providing multiple instance
 types here increases the likelihood of a Spot request being successful. If omitted the instance type to request will
 be taken from the launch template.
+
+### `aws.resource_tagging`
+
+Tag ASG and Fleet Request resources used by Escalator with the metatdata key-value pair
+`k8s.io/atlassian-escalator/enabled`:`true`. Tagging doesn't alter the functionality of Escalator. Read more about
+tagging your AWS resources [here](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/Using_Tags.html).
diff --git a/docs/deployment/aws/README.md b/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",
diff --git a/pkg/cloudprovider/aws/aws.go b/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 ASGs and Fleet requests
+	tagKey = "k8s.io/atlassian-escalator/enabled"
+	// tagValue is the value for the tag applied to ASGs and Fleet requests
+	tagValue = "true"
 )
 
 func instanceToProviderID(instance *autoscaling.Instance) string {
@@ -92,6 +96,8 @@ func (c *CloudProvider) RegisterNodeGroups(groups ...cloudprovider.NodeGroupConf
 			continue
 		}
 
+		addASGTags(configs[id], group, c)
+
 		c.nodeGroups[id] = NewNodeGroup(configs[id], group, c)
 	}
 
@@ -499,6 +505,20 @@ func createFleetInput(n NodeGroup, addCount int64) (*ec2.CreateFleetInput, error
 		}
 	}
 
+	if n.config.AWSConfig.ResourceTagging {
+		fleetInput.TagSpecifications = []*ec2.TagSpecification{
+			{
+				ResourceType: awsapi.String(ec2.ResourceTypeFleet),
+				Tags: []*ec2.Tag{
+					{
+						Key:   awsapi.String(tagKey),
+						Value: awsapi.String(tagValue),
+					},
+				},
+			},
+		}
+	}
+
 	return fleetInput, nil
 }
 
@@ -547,3 +567,37 @@ func createTemplateOverrides(n NodeGroup) ([]*ec2.FleetLaunchTemplateOverridesRe
 
 	return launchTemplateOverrides, nil
 }
+
+// addASGTags will search an ASG for the tagKey and add the tag if it's not found
+func addASGTags(config *cloudprovider.NodeGroupConfig, asg *autoscaling.Group, provider *CloudProvider) {
+	if !config.AWSConfig.ResourceTagging {
+		return
+	}
+
+	tags := asg.Tags
+	for _, tag := range tags {
+		if *tag.Key == tagKey {
+			return
+		}
+	}
+
+	id := awsapi.StringValue(asg.AutoScalingGroupName)
+
+	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 := provider.service.CreateOrUpdateTags(tagInput)
+	if err != nil {
+		log.Errorf("failed to create auto scaling tag for ASG %v", id)
+	}
+}
diff --git a/pkg/cloudprovider/aws/aws_test.go b/pkg/cloudprovider/aws/aws_test.go
@@ -28,6 +28,7 @@ func setupAWSMocks() {
 		MaxSize:              aws.Int64(int64(25)),
 		DesiredCapacity:      aws.Int64(int64(1)),
 		VPCZoneIdentifier:    aws.String("subnetID-1,subnetID-2"),
+		Tags:                 []*autoscaling.TagDescription{},
 	}
 
 	mockAWSConfig = cloudprovider.AWSNodeGroupConfig{
@@ -36,6 +37,7 @@ func setupAWSMocks() {
 		FleetInstanceReadyTimeout: tickerTimeout,
 		Lifecycle:                 LifecycleOnDemand,
 		InstanceTypeOverrides:     []string{"instance-1", "instance-2"},
+		ResourceTagging:           false,
 	}
 
 	mockNodeGroup = NodeGroup{
@@ -128,6 +130,29 @@ func TestCreateFleetInput(t *testing.T) {
 	}
 }
 
+func TestCreateFleetInput_WithResourceTagging(t *testing.T) {
+	setupAWSMocks()
+	autoScalingGroups := []*autoscaling.Group{&mockASG}
+	nodeGroups := map[string]*NodeGroup{mockNodeGroup.id: &mockNodeGroup}
+	addCount := int64(2)
+
+	awsCloudProvider, _ := newMockCloudProviderUsingInjection(
+		nodeGroups,
+		&test.MockAutoscalingService{
+			DescribeAutoScalingGroupsOutput: &autoscaling.DescribeAutoScalingGroupsOutput{
+				AutoScalingGroups: autoScalingGroups,
+			},
+		},
+		&test.MockEc2Service{},
+	)
+	mockNodeGroup.provider = awsCloudProvider
+	mockAWSConfig.ResourceTagging = true
+	mockNodeGroupConfig.AWSConfig = mockAWSConfig
+
+	_, err := createFleetInput(mockNodeGroup, addCount)
+	assert.Nil(t, err, "Expected no error from createFleetInput")
+}
+
 func TestCreateTemplateOverrides_FailedCall(t *testing.T) {
 	setupAWSMocks()
 	expectedError := errors.New("call failed")
@@ -234,3 +259,66 @@ func TestCreateTemplateOverrides_NoInstanceTypeOverrides_Success(t *testing.T) {
 	_, err := createTemplateOverrides(mockNodeGroup)
 	assert.Nil(t, err, "Expected no error from createTemplateOverrides")
 }
+func TestAddASGTags_ResourceTaggingFalse(t *testing.T) {
+	setupAWSMocks()
+	mockNodeGroupConfig.AWSConfig.ResourceTagging = false
+	awsCloudProvider, _ := newMockCloudProviderUsingInjection(
+		nil,
+		&test.MockAutoscalingService{},
+		&test.MockEc2Service{},
+	)
+	addASGTags(&mockNodeGroupConfig, &mockASG, awsCloudProvider)
+}
+
+func TestAddASGTags_ResourceTaggingTrue(t *testing.T) {
+	setupAWSMocks()
+	mockNodeGroupConfig.AWSConfig.ResourceTagging = true
+
+	// Mock service call
+	awsCloudProvider, _ := newMockCloudProviderUsingInjection(
+		nil,
+		&test.MockAutoscalingService{
+			CreateOrUpdateTagsOutput: &autoscaling.CreateOrUpdateTagsOutput{},
+		},
+		&test.MockEc2Service{},
+	)
+	addASGTags(&mockNodeGroupConfig, &mockASG, awsCloudProvider)
+}
+
+func TestAddASGTags_ASGAlreadyTagged(t *testing.T) {
+	setupAWSMocks()
+	mockNodeGroupConfig.AWSConfig.ResourceTagging = true
+
+	// Mock existing tags
+	key := tagKey
+	asgTag := autoscaling.TagDescription{
+		Key: &key,
+	}
+	mockASG.Tags = append(mockASG.Tags, &asgTag)
+
+	// Mock service call
+	awsCloudProvider, _ := newMockCloudProviderUsingInjection(
+		nil,
+		&test.MockAutoscalingService{
+			CreateOrUpdateTagsOutput: &autoscaling.CreateOrUpdateTagsOutput{},
+		},
+		&test.MockEc2Service{},
+	)
+	addASGTags(&mockNodeGroupConfig, &mockASG, awsCloudProvider)
+}
+
+func TestAddASGTags_WithErrorResponse(t *testing.T) {
+	setupAWSMocks()
+	mockNodeGroupConfig.AWSConfig.ResourceTagging = true
+
+	// Mock service call and error
+	awsCloudProvider, _ := newMockCloudProviderUsingInjection(
+		nil,
+		&test.MockAutoscalingService{
+			CreateOrUpdateTagsOutput: &autoscaling.CreateOrUpdateTagsOutput{},
+			CreateOrUpdateTagsErr:    errors.New("unauthorized"),
+		},
+		&test.MockEc2Service{},
+	)
+	addASGTags(&mockNodeGroupConfig, &mockASG, awsCloudProvider)
+}
diff --git a/pkg/cloudprovider/aws/cloud_provider_test.go b/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
diff --git a/pkg/cloudprovider/interface.go b/pkg/cloudprovider/interface.go
@@ -117,4 +117,5 @@ type AWSNodeGroupConfig struct {
 	FleetInstanceReadyTimeout time.Duration
 	Lifecycle                 string
 	InstanceTypeOverrides     []string
+	ResourceTagging           bool
 }
diff --git a/pkg/controller/node_group.go b/pkg/controller/node_group.go
@@ -59,6 +59,7 @@ type AWSNodeGroupOptions struct {
 	FleetInstanceReadyTimeout string   `json:"fleet_instance_ready_timeout,omitempty" yaml:"fleet_instance_ready_timeout,omitempty"`
 	Lifecycle                 string   `json:"lifecycle,omitempty" yaml:"lifecycle,omitempty"`
 	InstanceTypeOverrides     []string `json:"instance_type_overrides,omitempty" yaml:"instance_type_overrides,omitempty"`
+	ResourceTagging           bool     `json:"resource_tagging,omitempty" yaml:"resource_tagging,omitempty"`
 
 	// Private variables for storing the parsed duration from the string
 	fleetInstanceReadyTimeout time.Duration
diff --git a/pkg/test/aws.go b/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

Original file line number	Diff line number	Diff line change
`@@ -77,6 +77,7 @@ func setupCloudProvider(nodegroups []controller.NodeGroupOptions) cloudprovider.`
`77`	`77`	`FleetInstanceReadyTimeout: n.AWS.FleetInstanceReadyTimeoutDuration(),`
`78`	`78`	`Lifecycle: n.AWS.Lifecycle,`
`79`	`79`	`InstanceTypeOverrides: n.AWS.InstanceTypeOverrides,`
	`80`	`+ ResourceTagging: n.AWS.ResourceTagging,`
`80`	`81`	`},`
`81`	`82`	`})`
`82`	`83`	`}`
Original file line number	Diff line number	Diff line change
`@@ -117,4 +117,5 @@ type AWSNodeGroupConfig struct {`
`117`	`117`	`FleetInstanceReadyTimeout time.Duration`
`118`	`118`	`Lifecycle string`
`119`	`119`	`InstanceTypeOverrides []string`
	`120`	`+ ResourceTagging bool`
`120`	`121`	`}`