Changing HostMaintenance to return a map[string]bool where true indicates success, false indicates failure to transition to the desired state.

monitors.go

@@ -153,22 +153,20 @@ func (m *Monitor) Instances(key *aurora.JobKey, instances int32, interval int, t
 	return false, nil
 }
 
-// Monitor host status until all hosts match the status provided. May return an error along with a non nil map which contains
-// the hosts that did not transition to the desired modes(s).
-func (m *Monitor) HostMaintenance(hosts []string, modes []aurora.MaintenanceMode, sleepTime, steps int) (map[string]struct{}, error) {
+// Monitor host status until all hosts match the status provided. Returns a map where the value is true if the host
+// is in one of the desired mode(s) or false if it is not.
+func (m *Monitor) HostMaintenance(hosts []string, modes []aurora.MaintenanceMode, sleepTime, steps int) (map[string]bool, error) {
 
 	//  Transform modes to monitor for into a set for easy lookup
 	desiredMode := make(map[aurora.MaintenanceMode]struct{})
-	for _,mode := range modes {
+	for _, mode := range modes {
 		desiredMode[mode] = struct{}{}
 	}
 
-	// Turn slice into a host set to eliminate duplicates. Delete hosts that have entered the desired mode from
-	// observed list. We are done when the number of observed hosts reaches zero.
-	// This avoids having to go through and check the list one by one each cycle.
-	observedHosts := make(map[string]struct{})
-	for _,host := range hosts {
-		observedHosts[host] = struct{}{}
+	// Turn slice into a host set to eliminate duplicates.
+	observedHosts := make(map[string]bool)
+	for _, host := range hosts {
+		observedHosts[host] = false
 	}
 
 	for step := 0; step < steps; step++ {
@@ -176,20 +174,24 @@ func (m *Monitor) HostMaintenance(hosts []string, modes []aurora.MaintenanceMode
 		_, result, err := m.Client.MaintenanceStatus(hosts...)
 		if err != nil {
 			// Error is either a payload error or a severe connection error
-			return observedHosts, errors.Wrap(err,"client error")
+			return observedHosts, errors.Wrap(err, "client error")
 		}
 
 		for status := range result.GetStatuses() {
-			if  _, ok := desiredMode[status.GetMode()]; ok {
-				fmt.Printf("host %s entered %s state\n", status.GetHost(), status.GetMode())
-				delete(observedHosts, status.GetHost())
-			}
-		}
+			if _, ok := desiredMode[status.GetMode()]; ok {
+				observedHosts[status.GetHost()] = true
 
-		if len(observedHosts) == 0{
-			return observedHosts, nil
-		} else {
-			fmt.Printf("%d host(s) not in desired state\n", len(observedHosts))
+				transitionedHosts := 0
+				for _, val := range observedHosts {
+					if val {
+						transitionedHosts++
+					}
+				}
+
+				if len(observedHosts) == transitionedHosts {
+					return observedHosts, nil
+				}
+			}
 		}
 
 		time.Sleep(time.Duration(sleepTime) * time.Second)