Improves tests for setting and restoring desired count.

README.md

@@ -42,6 +42,8 @@ The update methodology is simple:
 
 ASG Roller will check both launch configurations, comparing names of the launch configuration used, and launch templates, comparing ID or Name, and version.
 
+ASG Roller will store the original desired value of the ASG as a tag on the ASG, with the key `aws-asg-roller/OriginalDesired`. This helps maintain state in the situation where the process terminates.
+
 ## App Awareness
 In addition to the above, ASG Roller is able to insert app-specific logic at two distinct points:
 
@@ -89,18 +91,13 @@ These permissions are as follows:
   - "autoscaling:TerminateInstanceInAutoScalingGroup"
   - "autoscaling:UpdateAutoScalingGroup"
   - "autoscaling:DescribeTags"
+  - "autoscaling:CreateOrUpdateTags"
   - "autoscaling:DescribeLaunchConfigurations"
   - "ec2:DescribeLaunchTemplates"
   - "ec2:DescribeInstances"
   Resource: "*"
 ```
 
-If the `ROLLER_ORIGINAL_DESIRED_ON_TAG` tag option is enabled, the following permission is also required:
-
-```
-autoscaling:CreateOrUpdateTags
-```
-
 These permissions can be set either via running ASG Roller on an AWS node that has the correct role, or via API keys to a user that has the correct roles/permissions.
 
 * If the AWS environment variables `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY`are set, it will use those
@@ -230,7 +227,6 @@ ASG Roller takes its configuration via environment variables. All environment va
 * `ROLLER_DELETE_LOCAL_DATA`: If set to `false` (default), will not reclaim a node until there are no pods with [emptyDir](https://kubernetes.io/docs/concepts/storage/volumes/#emptydir) running on the node; if set to `true`, will continue to terminate the pod and delete the local data before reclaiming the node. The default is `false` to maintain backward compatibility. 
 * `ROLLER_CHECK_DELAY`: Time, in seconds, between checks of ASG status.
 * `ROLLER_CAN_INCREASE_MAX`: If set to `true`, will increase the ASG maximum size to accommodate the increase in desired count. If set to `false`, will instead error when desired is higher than max.
-* `ROLLER_ORIGINAL_DESIRED_ON_TAG`: If set to `true`, will store the original desired value of the ASG as a tag on the ASG, with the key `aws-asg-roller/OriginalDesired`. This helps maintain state in the situation where the process terminates.
 * `ROLLER_VERBOSE`: If set to `true`, will increase verbosity of logs.
 * `KUBECONFIG`: Path to kubernetes config file for authenticating to the kubernetes cluster. Required only if `ROLLER_KUBERNETES` is `true` and we are not operating in a kubernetes cluster.
 

aws.go

@@ -23,7 +23,6 @@ func setAsgDesired(svc autoscalingiface.AutoScalingAPI, asg *autoscaling.Group,
 			return fmt.Errorf("unable to increase ASG %s desired size to %d as greater than max size %d", *asg.AutoScalingGroupName, count, *asg.MaxSize)
 		}
 	}
-
 	if verbose {
 		log.Printf("increasing ASG %s desired count to %d", *asg.AutoScalingGroupName, count)
 	}

aws_internal_test.go

@@ -153,6 +153,19 @@ func (m *mockAsgSvc) UpdateAutoScalingGroup(in *autoscaling.UpdateAutoScalingGro
 	ret := &autoscaling.UpdateAutoScalingGroupOutput{}
 	return ret, m.err
 }
+func (m *mockAsgSvc) DescribeTags(in *autoscaling.DescribeTagsInput) (*autoscaling.DescribeTagsOutput, error) {
+	m.counter.add("DescribeTags", in)
+	ret := &autoscaling.DescribeTagsOutput{
+		// value of "auto-scaling-group" tag is the ASG name
+		Tags: m.groups[*in.Filters[0].Values[0]].Tags,
+	}
+	return ret, m.err
+}
+func (m *mockAsgSvc) CreateOrUpdateTags(in *autoscaling.CreateOrUpdateTagsInput) (*autoscaling.CreateOrUpdateTagsOutput, error) {
+	m.counter.add("CreateOrUpdateTags", in)
+	ret := &autoscaling.CreateOrUpdateTagsOutput{}
+	return ret, m.err
+}
 
 func TestAwsGetHostnames(t *testing.T) {
 	tests := []struct {

original_desired.go

@@ -3,7 +3,6 @@ package main
 import (
 	"fmt"
 	"log"
-	"os"
 	"strconv"
 
 	"github.com/aws/aws-sdk-go/aws"
@@ -13,36 +12,28 @@ import (
 
 const asgTagNameOriginalDesired = "aws-asg-roller/OriginalDesired"
 
-var (
-	storeOriginalDesiredOnTag = os.Getenv("ROLLER_ORIGINAL_DESIRED_ON_TAG") == "true"
-)
-
 // Populates the original desired values for each ASG, based on the current 'desired' value if unkonwn.
 // The original desired value is recorded as a tag on the respective ASG. Subsequent runs attempt to
 // read the value of the tag to preserve state in the case of the process terminating.
 func populateOriginalDesired(originalDesired map[string]int64, asgs []*autoscaling.Group, asgSvc autoscalingiface.AutoScalingAPI) error {
 	for _, asg := range asgs {
 		asgName := *asg.AutoScalingGroupName
-		if storeOriginalDesiredOnTag {
-			tagOriginalDesired, err := getOriginalDesiredTag(asgSvc, asgName)
-			if err != nil {
-				return err
-			}
-			if tagOriginalDesired >= 0 {
-				originalDesired[asgName] = tagOriginalDesired
-				continue
-			}
+		tagOriginalDesired, err := getOriginalDesiredTag(asgSvc, asgName)
+		if err != nil {
+			return err
+		}
+		if tagOriginalDesired >= 0 {
+			originalDesired[asgName] = tagOriginalDesired
+			continue
 		}
 		// guess based on the current value
 		originalDesired[asgName] = *asg.DesiredCapacity
 		if verbose {
 			log.Printf("guessed desired value of %d from current desired on ASG: %s", *asg.DesiredCapacity, asgName)
 		}
-		if storeOriginalDesiredOnTag {
-			err := setOriginalDesiredTag(asgSvc, asgName, asg)
-			if err != nil {
-				return err
-			}
+		err = setOriginalDesiredTag(asgSvc, asgName, asg)
+		if err != nil {
+			return err
 		}
 	}
 	return nil

roller.go

@@ -77,7 +77,9 @@ func adjust(asgList []string, ec2Svc ec2iface.EC2API, asgSvc autoscalingiface.Au
 			log.Printf("[%s] error calculating adjustment - skipping: %v\n", *asg.AutoScalingGroupName, err)
 			continue
 		}
-		newDesired[*asg.AutoScalingGroupName] = newDesiredA
+		if newDesiredA != *asg.DesiredCapacity {
+			newDesired[*asg.AutoScalingGroupName] = newDesiredA
+		}
 		if terminateID != "" {
 			log.Printf("[%s] scheduled termination: %s", *asg.AutoScalingGroupName, terminateID)
 			newTerminate[*asg.AutoScalingGroupName] = terminateID
@@ -130,13 +132,14 @@ func calculateAdjustment(asg *autoscaling.Group, ec2Svc ec2iface.EC2API, hostnam
 
 	// Possibilities:
 	// 1- we have some old ones, but have not started updates yet: set the desired, increment and loop
-	// 2- we have no old ones, but have started updates: we must be at end, so finish
+	// 2- we have no old ones: we must be at end or have no work to do, so finish
 	// 3- we have some old ones, but have started updates: run the updates
 	if len(oldInstances) == 0 {
-		if desired != originalDesired {
-			// we are done; return to desired to original value
-			return originalDesired, "", nil
+		// we are done
+		if verbose && desired != originalDesired {
+			log.Printf("[%s] returning desired to original value %d", *asg.AutoScalingGroupName, originalDesired)
 		}
+		return originalDesired, "", nil
 	}
 	if originalDesired == desired {
 		// we have not started updates; raise the desired count

roller_internal_test.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"fmt"
+	"strconv"
 	"strings"
 	"testing"
 
@@ -39,7 +40,6 @@ func TestCalculateAdjustment(t *testing.T) {
 					- state of each new node
 			outputs
 			 - new desired number
-			 - new original desired number
 			 - node id to terminated (if any)
 			 - errors (if any)
 	*/
@@ -65,33 +65,32 @@ func TestCalculateAdjustment(t *testing.T) {
 		originalDesired       int64
 		readiness             readiness
 		targetDesired         int64
-		targetOriginalDesired int64
 		targetTerminate       string
 		err                   error
 	}{
 		// 1 old, 2 new healthy, 0 new unhealthy, should terminate old
-		{[]string{"1"}, []string{"2", "3"}, []string{}, 3, 2, nil, 3, 2, "1", nil},
+		{[]string{"1"}, []string{"2", "3"}, []string{}, 3, 2, nil, 3, "1", nil},
 		// 0 old, 2 new healthy, 0 new unhealthy, should indicate end of process
-		{[]string{}, []string{"2", "3"}, []string{}, 3, 2, nil, 2, 0, "", nil},
+		{[]string{}, []string{"2", "3"}, []string{}, 2, 2, nil, 2, "", nil},
 		// 2 old, 0 new healthy, 0 new unhealthy, should indicate start of process
-		{[]string{"1", "2"}, []string{}, []string{}, 2, 0, nil, 3, 2, "", nil},
+		{[]string{"1", "2"}, []string{}, []string{}, 2, 2, nil, 3, "", nil},
 		// 2 old, 0 new healthy, 0 new unhealthy, started, should not do anything until new healthy one
-		{[]string{"1", "2"}, []string{}, []string{}, 3, 2, nil, 3, 2, "", nil},
+		{[]string{"1", "2"}, []string{}, []string{}, 3, 2, nil, 3, "", nil},
 		// 2 old, 1 new healthy, 0 new unhealthy, remove an old one
-		{[]string{"1", "2"}, []string{"3"}, []string{}, 3, 2, nil, 3, 2, "1", nil},
+		{[]string{"1", "2"}, []string{"3"}, []string{}, 3, 2, nil, 3, "1", nil},
 		// 2 old, 0 new healthy, 1 new unhealthy, started, should not do anything until new one is healthy
-		{[]string{"1", "2"}, []string{}, []string{"3"}, 3, 2, nil, 3, 2, "", nil},
+		{[]string{"1", "2"}, []string{}, []string{"3"}, 3, 2, nil, 3, "", nil},
 
 		// 2 old, 1 new healthy, 0 new unhealthy, 1 new unready, should not change anything
-		{[]string{"1", "2"}, []string{"3"}, []string{}, 3, 2, unreadyCountHandler, 3, 2, "", nil},
+		{[]string{"1", "2"}, []string{"3"}, []string{}, 3, 2, unreadyCountHandler, 3, "", nil},
 		// 2 old, 1 new healthy, 0 new unhealthy, 0 new unready, 1 error: should not change anything
-		{[]string{"1", "2"}, []string{"3"}, []string{}, 3, 2, unreadyErrorHandler, 3, 2, "", fmt.Errorf("Error")},
+		{[]string{"1", "2"}, []string{"3"}, []string{}, 3, 2, unreadyErrorHandler, 3, "", fmt.Errorf("Error")},
 		// 2 old, 1 new healthy, 0 new unhealthy, 0 unready, remove an old one
-		{[]string{"1", "2"}, []string{"3"}, []string{}, 3, 2, readyHandler, 3, 2, "1", nil},
+		{[]string{"1", "2"}, []string{"3"}, []string{}, 3, 2, readyHandler, 3, "1", nil},
 		// 2 old, 1 new healthy, 0 new unhealthy, 0 new unready, 1 error: should not change anything
-		{[]string{"1", "2"}, []string{"3"}, []string{}, 3, 2, terminateErrorHandler, 3, 2, "", fmt.Errorf("Unexpected error")},
+		{[]string{"1", "2"}, []string{"3"}, []string{}, 3, 2, terminateErrorHandler, 3, "", fmt.Errorf("Unexpected error")},
 		// 2 old, 1 new healthy, 0 new unhealthy, 0 unready, successful terminate: remove an old one
-		{[]string{"1", "2"}, []string{"3"}, []string{}, 3, 2, terminateHandler, 3, 2, "1", nil},
+		{[]string{"1", "2"}, []string{"3"}, []string{}, 3, 2, terminateHandler, 3, "1", nil},
 	}
 	hostnameMap := map[string]string{}
 	for i := 0; i < 20; i++ {
@@ -155,45 +154,41 @@ func TestCalculateAdjustment(t *testing.T) {
 
 func TestAdjust(t *testing.T) {
 	tests := []struct {
-		desc                    string
-		asgs                    []string
-		handler                 readiness
-		err                     error
-		oldIds                  map[string][]string
-		newIds                  map[string][]string
-		asgOriginalDesired      map[string]int64
-		originalDesired         map[string]int64
-		newOriginalDesired      map[string]int64
-		newDesired              map[string]int64
-		expectedOriginalDesired map[string]int64
-		max                     map[string]int64
-		terminate               []string
-		canIncreaseMax          bool
+		desc              string
+		asgs              []string
+		handler           readiness
+		err               error
+		oldIds            map[string][]string
+		newIds            map[string][]string
+		asgCurrentDesired map[string]int64
+		originalDesired   map[string]int64
+		newDesired        map[string]int64
+		max               map[string]int64
+		terminate         []string
+		canIncreaseMax    bool
 	}{
 		{
-			"2 asgs adjust in progress",
+			"2 asgs adjust first run",
 			[]string{"myasg", "anotherasg"},
 			nil,
 			nil,
 			map[string][]string{
-				"myasg":      {"1"},
+				"myasg":      {"1", "2"},
 				"anotherasg": {},
 			},
 			map[string][]string{
-				"myasg":      {"2", "3"},
+				"myasg":      {},
 				"anotherasg": {"8", "9", "10"},
 			},
-			map[string]int64{"myasg": 2, "anotherasg": 10},
-			map[string]int64{"myasg": 2, "anotherasg": 10},
-			map[string]int64{"myasg": 2, "anotherasg": 0},
+			map[string]int64{"myasg": 2, "anotherasg": 3},
 			map[string]int64{"myasg": 2},
-			map[string]int64{"myasg": 2, "anotherasg": 0},
-			map[string]int64{"myasg": 3, "anotherasg": 11},
-			[]string{"1"},
+			map[string]int64{"myasg": 3},
+			map[string]int64{"myasg": 3, "anotherasg": 4},
+			[]string{},
 			false,
 		},
 		{
-			"2 asgs adjust first run",
+			"2 asgs adjust in progress",
 			[]string{"myasg", "anotherasg"},
 			nil,
 			nil,
@@ -205,13 +200,11 @@ func TestAdjust(t *testing.T) {
 				"myasg":      {"2", "3"},
 				"anotherasg": {"8", "9", "10"},
 			},
-			map[string]int64{"myasg": 2},
+			map[string]int64{"myasg": 3, "anotherasg": 3},
+			map[string]int64{"myasg": 2, "anotherasg": 3},
 			map[string]int64{},
-			map[string]int64{"myasg": 2},
-			map[string]int64{"myasg": 3},
-			map[string]int64{"myasg": 2},
-			map[string]int64{"myasg": 3},
-			[]string{},
+			map[string]int64{"myasg": 3, "anotherasg": 4},
+			[]string{"1"},
 			false,
 		},
 		{
@@ -228,10 +221,8 @@ func TestAdjust(t *testing.T) {
 				"anotherasg": {"8", "9", "10"},
 			},
 			map[string]int64{"myasg": 2},
-			map[string]int64{},
 			map[string]int64{"myasg": 2},
 			map[string]int64{},
-			map[string]int64{"myasg": 2},
 			map[string]int64{"myasg": 3},
 			[]string{},
 			false,
@@ -240,7 +231,7 @@ func TestAdjust(t *testing.T) {
 			"2 asgs adjust increase max fail",
 			[]string{"myasg", "anotherasg"},
 			nil,
-			fmt.Errorf("Error setting desired to 3 for ASG myasg: unable to increase ASG myasg desired size to 3 as greater than max size 2"),
+			fmt.Errorf("[myasg] error setting desired to 3: unable to increase ASG myasg desired size to 3 as greater than max size 2"),
 			map[string][]string{
 				"myasg":      {"1"},
 				"anotherasg": {},
@@ -250,11 +241,9 @@ func TestAdjust(t *testing.T) {
 				"anotherasg": {"8", "9", "10"},
 			},
 			map[string]int64{"myasg": 2},
-			map[string]int64{},
 			map[string]int64{"myasg": 2},
 			map[string]int64{},
 			map[string]int64{"myasg": 2},
-			map[string]int64{"myasg": 2},
 			[]string{},
 			false,
 		},
@@ -272,11 +261,9 @@ func TestAdjust(t *testing.T) {
 				"anotherasg": {"8", "9", "10"},
 			},
 			map[string]int64{"myasg": 2},
-			map[string]int64{},
 			map[string]int64{"myasg": 2},
 			map[string]int64{"myasg": 3},
 			map[string]int64{"myasg": 2},
-			map[string]int64{"myasg": 2},
 			[]string{},
 			true,
 		},
@@ -290,7 +277,7 @@ func TestAdjust(t *testing.T) {
 				lcName := "lconfig"
 				oldLcName := fmt.Sprintf("old%s", lcName)
 				myHealthy := healthy
-				desired := tt.asgOriginalDesired[name]
+				desired := tt.asgCurrentDesired[name]
 				max := tt.max[name]
 				instances := make([]*autoscaling.Instance, 0)
 				for _, id := range tt.oldIds[name] {
@@ -310,13 +297,25 @@ func TestAdjust(t *testing.T) {
 					})
 				}
 				// construct the Group we will pass
-				validGroups[n] = &autoscaling.Group{
+				validGroup := &autoscaling.Group{
 					AutoScalingGroupName:    &name,
 					DesiredCapacity:         &desired,
 					Instances:               instances,
 					LaunchConfigurationName: &lcName,
 					MaxSize:                 &max,
 				}
+				if originalDesired, ok := tt.originalDesired[name]; ok {
+					validGroup.Tags = []*autoscaling.TagDescription{
+						{
+							Key:               aws.String(asgTagNameOriginalDesired),
+							PropagateAtLaunch: aws.Bool(false),
+							ResourceId:        &name,
+							ResourceType:      aws.String("auto-scaling-group"),
+							Value:             aws.String(strconv.FormatInt(originalDesired, 10)),
+						},
+					}
+				}
+				validGroups[n] = validGroup
 			}
 			asgSvc := &mockAsgSvc{
 				groups: validGroups,
@@ -343,10 +342,6 @@ func TestAdjust(t *testing.T) {
 				t.Errorf("%d: mismatched errors, actual then expected", i)
 				t.Logf("%v", err)
 				t.Logf("%v", tt.err)
-			case !testStringInt64MapEq(tt.newOriginalDesired, tt.expectedOriginalDesired):
-				t.Errorf("%d: Mismatched desired, actual then expected", i)
-				t.Logf("%v", tt.originalDesired)
-				t.Logf("%v", tt.newOriginalDesired)
 			}
 
 			// check each svc with its correct calls