Add instance update and reactivation

Adds functionality to perform Bottlerocket version upgrades and
reactivates drained instances after successful upgrade.

[Issue: #9]

updater/aws.go

@@ -62,20 +62,20 @@ func containsAttribute(attrs []*ecs.Attribute, searchString string) bool {
 
 // checkAndDrain drains eligible container instances. Container instances are eligible if all running
 // tasks were started by a service, or if there are no running tasks.
-func (u *updater) checkAndDrain(containerInstance string) error {
+func (u *updater) checkAndDrain(ec2ID string, containerInstance string) (string, error) {
 	if !eligible(u.ecs, &containerInstance, &u.cluster) {
-		return nil
+		return "", nil
 	}
 
 	err := changeInstanceState(u.ecs, aws.String(containerInstance), &u.cluster, "DRAINING")
 	if err != nil {
 		err2 := changeInstanceState(u.ecs, aws.String(containerInstance), &u.cluster, "ACTIVE")
 		if err2 != nil {
-			return fmt.Errorf("failed to undrain: %w", err2.Error())
+			return "", fmt.Errorf("failed to undrain: %w", err2)
 		}
-		return fmt.Errorf("failed to drain: %w", err)
+		return "", fmt.Errorf("failed to drain: %w", err)
 	}
-	return nil
+	return ec2ID, nil
 }
 
 func eligible(ecsClient *ecs.ECS, containerInstance, cluster *string) bool {
@@ -185,7 +185,7 @@ func (u *updater) sendCommand(instanceIDs []string, ssmCommand string) (string,
 	return commandID, nil
 }
 
-func (u *updater) checkCommandOutput(commandID string, instanceIDs []string) ([]string, error) {
+func (u *updater) checkUpdateStatus(commandID string, instanceIDs []string) ([]string, error) {
 	updateCandidates := make([]string, 0)
 	for _, v := range instanceIDs {
 		resp, err := u.ssm.GetCommandInvocation(&ssm.GetCommandInvocationInput{
@@ -218,3 +218,53 @@ func (u *updater) checkCommandOutput(commandID string, instanceIDs []string) ([]
 	}
 	return updateCandidates, nil
 }
+
+func (u *updater) checkUpdateResult(commandID string, instanceID string) error {
+	resp, err := u.ssm.GetCommandInvocation(&ssm.GetCommandInvocationInput{
+		CommandId:  aws.String(commandID),
+		InstanceId: aws.String(instanceID),
+	})
+	if err != nil {
+		return fmt.Errorf("failed to retrieve command invocation output: %#v", err)
+	}
+	log.Printf("Response: %s", *resp.Status)
+	if *resp.Status == "Success" {
+		log.Printf("Updated %s succeessfully", instanceID)
+	} else {
+		log.Printf("Instance failed to update")
+	}
+	return nil
+}
+
+func (u *updater) getVersion(commandID string, instanceIDs []string) (map[string]string, error) {
+	updatedVersions := make(map[string]string)
+	for _, EC2ID := range instanceIDs {
+		resp, err := u.ssm.GetCommandInvocation(&ssm.GetCommandInvocationInput{
+			CommandId:  aws.String(commandID),
+			InstanceId: aws.String(EC2ID),
+		})
+		if err != nil {
+			return nil, fmt.Errorf("failed to retrieve command invocation output: %#v", err)
+		}
+
+		type Version struct {
+			Version string `json:"version"`
+		}
+
+		type Image struct {
+			Image Version `json:"image"`
+		}
+
+		type Partition struct {
+			ActivePartition Image `json:"active_partition"`
+		}
+
+		var version Partition
+		err = json.Unmarshal([]byte(*resp.StandardOutputContent), &version)
+		if err != nil {
+			log.Printf("failed to unmarshal command invocation output: %#v", err)
+		}
+		updatedVersions[EC2ID] = version.ActivePartition.Image.Version
+	}
+	return updatedVersions, nil
+}

updater/main.go

@@ -78,7 +78,7 @@ func _main() error {
 		return err
 	}
 
-	instancesToUpdate, err := up.checkCommandOutput(commandID, instances)
+	instancesToUpdate, err := up.checkUpdateStatus(commandID, instances)
 	if err != nil {
 		return err
 	}
@@ -90,11 +90,43 @@ func _main() error {
 	fmt.Println("Instances ready for update: ", instancesToUpdate)
 
 	for ec2ID, containerARN := range bottlerocketInstanceMap {
-		err := up.checkAndDrain(containerARN)
+		checkedEC2ID, err := up.checkAndDrain(ec2ID, containerARN)
 		if err != nil {
 			return err
 		}
-		log.Printf("Instance %s drained", ec2ID)
+		log.Printf("Instance %s drained", checkedEC2ID)
+
+		if checkedEC2ID == "" {
+			return nil
+		}
+
+		checkedEC2IDAsSlice := []string{checkedEC2ID}
+		if len(checkedEC2IDAsSlice) != 0 {
+			updateCommandID, err := up.sendCommand(checkedEC2IDAsSlice, "apiclient update apply -r")
+			if err != nil {
+				return err
+			}
+
+			err = up.checkUpdateResult(updateCommandID, checkedEC2ID)
+			if err != nil {
+				return err
+			}
+
+			err = changeInstanceState(up.ecs, aws.String(containerARN), flagCluster, "ACTIVE")
+			if err != nil {
+				return err
+			}
+
+			updateStatus, err := up.sendCommand(checkedEC2IDAsSlice, "apiclient update check")
+			if err != nil {
+				return err
+			}
+
+			updatedVersions, err := up.getVersion(updateStatus, checkedEC2IDAsSlice)
+			for _ec2ID, version := range updatedVersions {
+				log.Printf("Instance %s updated to Bottlerocket: %s", _ec2ID, version)
+			}
+		}
 	}
 	return nil
 }