refactor(state): cleanup state polling and output

Author: Gabriel Monroy

client/create.go

@@ -2,7 +2,6 @@ package client
 
 import (
 	"fmt"
-	"os"
 	"strings"
 
 	"github.com/coreos/fleet/job"
@@ -31,9 +30,10 @@ func (c *FleetClient) Create(target string) (err error) {
 	if err != nil {
 		return err
 	}
-	errchan := waitForJobStates([]string{unitName}, testJobStateLoaded, 0, os.Stdout)
-	for err := range errchan {
-		return fmt.Errorf("error waiting for job %s: %v", unitName, err)
+	check := newStateCheck(testJobStateLoaded)
+	err = waitForJobStates([]string{unitName}, check)
+	if err != nil {
+		return err
 	}
 	return nil
 }

client/start.go

@@ -1,11 +1,6 @@
 package client
 
-import (
-	"fmt"
-	"os"
-
-	"github.com/coreos/fleet/job"
-)
+import "github.com/coreos/fleet/job"
 
 // Start launches target units and blocks until active
 func (c *FleetClient) Start(target string) (err error) {
@@ -20,9 +15,10 @@ func (c *FleetClient) Start(target string) (err error) {
 			return err
 		}
 	}
-	errchan := waitForJobStates(units, testUnitStateActive, 0, os.Stdout)
-	for err := range errchan {
-		return fmt.Errorf("error waiting for active: %v", err)
+	check := newStateCheck(testUnitStateActive)
+	err = waitForJobStates(units, check)
+	if err != nil {
+		return err
 	}
 	return nil
 }

client/state.go

@@ -2,8 +2,6 @@ package client
 
 import (
 	"fmt"
-	"io"
-	"os"
 	"sync"
 	"time"
 
@@ -13,99 +11,139 @@ import (
 type testJob func(j *job.Job) bool
 
 func testJobStateLoaded(j *job.Job) bool {
-	return j == nil || j.State == nil || *(j.State) != job.JobStateLoaded
+	if j == nil || j.State == nil {
+		return false
+	}
+	return *(j.State) == job.JobStateLoaded
 }
 
 func testJobStateLaunched(j *job.Job) bool {
-	return j == nil || j.State == nil || *(j.State) != job.JobStateLaunched
+	if j == nil || j.State == nil {
+		return false
+	}
+	return *(j.State) == job.JobStateLaunched
 }
 
 func testJobStateInactive(j *job.Job) bool {
-	return j == nil || j.State == nil || *(j.State) != job.JobStateInactive
+	if j == nil || j.State == nil {
+		return false
+	}
+	return *(j.State) == job.JobStateInactive
 }
 
 func testUnitStateActive(j *job.Job) bool {
-	return j == nil || j.UnitState == nil || j.UnitState.ActiveState != "active"
+	if j == nil || j.UnitState == nil {
+		return false
+	}
+	return j.UnitState.ActiveState == "active"
 
 }
 
-// TODO: refactor to separate presentation to io.Writer from status polling
+// stateCheck defines how to monitor a job state
+type stateCheck struct {
+	test      testJob
+	statechan chan *jobState
+	errchan   chan error
+}
 
-// waitForJobStates polls each of the indicated jobs until each of their
-// states is equal to that which the caller indicates, or until the
-// polling operation times out. waitForJobStates will retry forever, or
-// up to maxAttempts times before timing out if maxAttempts is greater
-// than zero. Returned is an error channel used to communicate when
-// timeouts occur. The returned error channel will be closed after all
-// polling operation is complete.
-func waitForJobStates(jobs []string, test testJob, maxAttempts int, out io.Writer) chan error {
+// newStateCheck returns a StateCheck struct with new channels for monitoring
+func newStateCheck(test testJob) *stateCheck {
+	statechan := make(chan *jobState)
 	errchan := make(chan error)
+	return &stateCheck{test, statechan, errchan}
+}
+
+// waitForJobStates polls each of the indicated jobs until each of their
+// states is equal to that which the caller indicates via stateCheck test
+func waitForJobStates(jobs []string, check *stateCheck) error {
 	var wg sync.WaitGroup
+
+	// check each job with the stateCheck
 	for _, name := range jobs {
 		wg.Add(1)
-		go checkJobState(name, test, maxAttempts, out, &wg, errchan)
+		go checkJobState(name, check, &wg)
 	}
+
+	// wait for all jobs to complete
 	go func() {
 		wg.Wait()
-		close(errchan)
+		close(check.errchan)
 	}()
-	return errchan
-}
 
-func checkJobState(jobName string, test testJob, maxAttempts int, out io.Writer, wg *sync.WaitGroup, errchan chan error) {
-	defer wg.Done()
-	sleep := 100 * time.Millisecond
-	if maxAttempts < 1 {
-		for {
-			if assertJobState(jobName, test, out) {
-				return
+	// print output while jobs are transitioning
+	defer fmt.Printf("\n")
+	for {
+		select {
+		// read from state channel
+		case state := <-check.statechan:
+			// return on closed channel
+			if state == nil {
+				return nil
 			}
-			time.Sleep(sleep)
-		}
-	} else {
-		for attempt := 0; attempt < maxAttempts; attempt++ {
-			if assertJobState(jobName, test, out) {
-				return
+			// otherwise print output
+			if state.sub == "" || state.sub == "dead" {
+				fmt.Printf("\033[0;33m%v:\033[0m %v, %v                              \r",
+					state.name, state.loaded, state.active)
+			} else {
+				fmt.Printf("\033[0;33m%v:\033[0m %v, %v (%v)                        \r",
+					state.name, state.loaded, state.active, state.sub)
 			}
-			time.Sleep(sleep)
+		// read from error channel
+		case err := <-check.errchan:
+			return err
 		}
-		errchan <- fmt.Errorf("timed out waiting for job %s to report state", jobName)
+		time.Sleep(200 * time.Millisecond)
 	}
 }
 
-func assertJobState(name string, test testJob, out io.Writer) (ret bool) {
-	j, err := cAPI.Job(name)
-	if err != nil {
-		fmt.Fprintf(os.Stderr, "Error retrieving Job(%s) from Registry: %v", name, err)
-		return
-	}
-	if test(j) {
-		if j.State != nil && j.UnitState != nil && j.UnitState.ActiveState != "" && j.UnitState.SubState != "" {
-			fmt.Fprintf(out, "\033[0;33m%v:\033[0m %v, %v (%v)           \r", j.Name, *(j.State), j.UnitState.ActiveState, j.UnitState.SubState)
+func checkJobState(jobName string, check *stateCheck, wg *sync.WaitGroup) {
+	defer wg.Done()
+	sleep := 100 * time.Millisecond
+	for {
+		if assertJobState(jobName, check) {
+			return
 		}
-		return
+		time.Sleep(sleep)
 	}
-	ret = true
+}
+
+type jobState struct {
+	name   string
+	loaded string
+	active string
+	sub    string
+}
 
-	var msg string
-	if j.State != nil && j.UnitState != nil && j.UnitState.ActiveState != "" && j.UnitState.SubState != "" {
-		msg = fmt.Sprintf("\033[0;33m%v:\033[0m %v, %v (%v)", name, *(j.State), j.UnitState.ActiveState, j.UnitState.SubState)
-	} else {
-		msg = fmt.Sprintf("\033[0;33m%v:\033[0m %v", name, *(j.State))
+func newJobState(name string, j *job.Job) *jobState {
+	var (
+		loaded string
+		active string
+		sub    string
+	)
+	if j.State != nil {
+		loaded = fmt.Sprintf("%v", *(j.State))
 	}
+	if j.UnitState != nil {
+		active = j.UnitState.ActiveState
+		sub = j.UnitState.SubState
+	}
+	return &jobState{name, loaded, active, sub}
+}
 
-	machines, err := cAPI.Machines()
+func assertJobState(name string, check *stateCheck) bool {
+	j, err := cAPI.Job(name)
 	if err != nil {
-		fmt.Fprintf(os.Stderr, "Failed retrieving list of Machines from Registry: %v", err)
-	}
-	for _, ms := range machines {
-		if ms.ID != j.TargetMachineID {
-			continue
-		}
-		msg = fmt.Sprintf("%s on %s", msg, machineFullLegend(ms, false))
-		break
+		check.errchan <- fmt.Errorf("Error retrieving Job(%s) from Registry: %v", name, err)
+		return false
 	}
 
-	fmt.Fprintln(out, msg)
-	return
+	// send current state to the output channel
+	check.statechan <- newJobState(name, j)
+
+	// if test function
+	if check.test(j) {
+		close(check.statechan)
+		return true
+	}
+	return false
 }

client/stop.go

@@ -1,11 +1,6 @@
 package client
 
-import (
-	"fmt"
-	"os"
-
-	"github.com/coreos/fleet/job"
-)
+import "github.com/coreos/fleet/job"
 
 // Stop sets target units to inactive and blocks until complete
 func (c *FleetClient) Stop(target string) (err error) {
@@ -20,9 +15,10 @@ func (c *FleetClient) Stop(target string) (err error) {
 			return err
 		}
 	}
-	errchan := waitForJobStates(units, testJobStateInactive, 0, os.Stdout)
-	for err := range errchan {
-		return fmt.Errorf("error waiting for inactive: %v", err)
+	check := newStateCheck(testJobStateInactive)
+	err = waitForJobStates(units, check)
+	if err != nil {
+		return err
 	}
 	return nil
 }