Repository: coreos/fleet
Pull request #1433 (2016-02-16 12:19:30)

fleetctl: {load|unload|start|stop} and get unit code cleanups
This patches set aims to make the code more consistent and robust. It does not change the current behaviour just improves the code, consolidate it, more debug logs and code comments.

The main goal is to let stop, start, load and unload to share the same code base and probably avoid any corner case or bug that may related to one of these commands.

* Some logic can be easily shared later when we add stop to destroy command to solve  https://github.com/coreos/fleet/issues/1000  

* There were some TODOs in the code base that I cleaned up with the last two patches. These patches will make it easy for us to implement later https://github.com/coreos/fleet/pull/1295

tixxdz: PR updated, thanks
tixxdz: The functional tests pass. Thanks!
tixxdz: @jonboulle Investigating why semaphore fails, all comments were handled for this one and all functional tests pass here.... but semaphore fails for another error I guess...

Patch:
diff --git a/fleetctl/fleetctl.go b/fleetctl/fleetctl.go
index 71785606..69513e04 100644
--- a/fleetctl/fleetctl.go
+++ b/fleetctl/fleetctl.go
@@ -483,6 +483,51 @@ func getChecker() *ssh.HostKeyChecker {
 	return ssh.NewHostKeyChecker(keyFile)
 }
 
+// getUnitFile attempts to get a UnitFile configuration
+// It takes a unit file name as a parameter and tries first to lookup
+// the unit from the local disk. If it fails, it checks if the provided
+// file name may reference an instance of a template unit, if so, it
+// tries to get the template configuration either from the registry or
+// the local disk.
+// It returns a UnitFile configuration or nil; and any error ecountered
+func getUnitFile(file string) (*unit.UnitFile, error) {
+	var uf *unit.UnitFile
+	name := unitNameMangle(file)
+
+	log.Debugf("Looking for Unit(%s) or its corresponding template", name)
+
+	// Assume that the file references a local unit file on disk and
+	// attempt to load it, if it exists
+	if _, err := os.Stat(file); !os.IsNotExist(err) {
+		uf, err = getUnitFromFile(file)
+		if err != nil {
+			return nil, fmt.Errorf("failed getting Unit(%s) from file: %v", file, err)
+		}
+	} else {
+		// Otherwise (if the unit file does not exist), check if the
+		// name appears to be an instance of a template unit
+		info := unit.NewUnitNameInfo(name)
+		if info == nil {
+			return nil, fmt.Errorf("error extracting information from unit name %s", name)
+		} else if !info.IsInstance() {
+			return nil, fmt.Errorf("unable to find Unit(%s) in Registry or on filesystem", name)
+		}
+
+		// If it is an instance check for a corresponding template
+		// unit in the Registry or disk.
+		// If we found a template unit, later we create a
+		// near-identical instance unit in the Registry - same
+		// unit file as the template, but different name
+		uf, err = getUnitFileFromTemplate(info, file)
+		if err != nil {
+			return nil, fmt.Errorf("failed getting Unit(%s) from template: %v", file, err)
+		}
+	}
+
+	log.Debugf("Found Unit(%s)", name)
+	return uf, nil
+}
+
 // getUnitFromFile attempts to load a Unit from a given filename
 // It returns the Unit or nil, and any error encountered
 func getUnitFromFile(file string) (*unit.UnitFile, error) {
@@ -497,6 +542,39 @@ func getUnitFromFile(file string) (*unit.UnitFile, error) {
 	return unit.NewUnitFile(string(out))
 }
 
+// getUnitFileFromTemplate attempts to get a Unit from a template unit that
+// is either in the registry or on the file system
+// It takes two arguments, the template information and the unit file name
+// It returns the Unit or nil; and any error encountered
+func getUnitFileFromTemplate(uni *unit.UnitNameInfo, fileName string) (*unit.UnitFile, error) {
+	var uf *unit.UnitFile
+
+	tmpl, err := cAPI.Unit(uni.Template)
+	if err != nil {
+		return nil, fmt.Errorf("error retrieving template Unit(%s) from Registry: %v", uni.Template, err)
+	}
+
+	if tmpl != nil {
+		warnOnDifferentLocalUnit(fileName, tmpl)
+		uf = schema.MapSchemaUnitOptionsToUnitFile(tmpl.Options)
+		log.Debugf("Template Unit(%s) found in registry", uni.Template)
+	} else {
+		// Finally, if we could not find a template unit in the Registry,
+		// check the local disk for one instead
+		filePath := path.Join(path.Dir(fileName), uni.Template)
+		if _, err := os.Stat(filePath); os.IsNotExist(err) {
+			return nil, fmt.Errorf("unable to find template Unit(%s) in Registry or on filesystem", uni.Template)
+		}
+
+		uf, err = getUnitFromFile(filePath)
+		if err != nil {
+			return nil, fmt.Errorf("unable to load template Unit(%s) from file: %v", uni.Template, err)
+		}
+	}
+
+	return uf, nil
+}
+
 func getTunnelFlag() string {
 	tun := globalFlags.Tunnel
 	if tun != "" && !strings.Contains(tun, ":") {
@@ -599,8 +677,6 @@ func lazyCreateUnits(args []string) error {
 	errchan := make(chan error)
 	var wg sync.WaitGroup
 	for _, arg := range args {
-		// TODO(jonboulle): this loop is getting too unwieldy; factor it out
-
 		arg = maybeAppendDefaultUnitType(arg)
 		name := unitNameMangle(arg)
 
@@ -615,45 +691,12 @@ func lazyCreateUnits(args []string) error {
 			continue
 		}
 
-		var uf *unit.UnitFile
-		// Failing that, assume the name references a local unit file on disk, and attempt to load that, if it exists
-		// TODO(mischief): consolidate these two near-identical codepaths
-		if _, err := os.Stat(arg); !os.IsNotExist(err) {
-			uf, err = getUnitFromFile(arg)
-			if err != nil {
-				return fmt.Errorf("failed getting Unit(%s) from file: %v", arg, err)
-			}
-		} else {
-			// Otherwise (if the unit file does not exist), check if the name appears to be an instance unit,
-			// and if so, check for a corresponding template unit in the Registry
-			uni := unit.NewUnitNameInfo(name)
-			if uni == nil {
-				return fmt.Errorf("error extracting information from unit name %s", name)
-			} else if !uni.IsInstance() {
-				return fmt.Errorf("unable to find Unit(%s) in Registry or on filesystem", name)
-			}
-			tmpl, err := cAPI.Unit(uni.Template)
-			if err != nil {
-				return fmt.Errorf("error retrieving template Unit(%s) from Registry: %v", uni.Template, err)
-			}
-
-			// Finally, if we could not find a template unit in the Registry, check the local disk for one instead
-			if tmpl == nil {
-				file := path.Join(path.Dir(arg), uni.Template)
-				if _, err := os.Stat(file); os.IsNotExist(err) {
-					return fmt.Errorf("unable to find Unit(%s) or template Unit(%s) in Registry or on filesystem", name, uni.Template)
-				}
-				uf, err = getUnitFromFile(file)
-				if err != nil {
-					return fmt.Errorf("failed getting template Unit(%s) from file: %v", uni.Template, err)
-				}
-			} else {
-				warnOnDifferentLocalUnit(arg, tmpl)
-				uf = schema.MapSchemaUnitOptionsToUnitFile(tmpl.Options)
-			}
-
-			// If we found a template unit, create a near-identical instance unit in
-			// the Registry - same unit file as the template, but different name
+		// Assume that the name references a local unit file on
+		// disk or if it is an instance unit and if so get its
+		// corresponding unit
+		uf, err := getUnitFile(arg)
+		if err != nil {
+			return err
 		}
 
 		_, err = createUnit(name, uf)
@@ -746,6 +789,54 @@ func setTargetStateOfUnits(units []string, state job.JobState) ([]*schema.Unit,
 	return triggered, nil
 }
 
+// getBlockAttempts gets the correct value of how many attempts to try
+// before giving up on an operation.
+// It returns a negative value which means do not block, if zero is
+// returned then it means try forever, and if a positive value is
+// returned then try up to that value
+func getBlockAttempts() int {
+	// By default we wait forever
+	var attempts int = 0
+
+	if sharedFlags.BlockAttempts > 0 {
+		attempts = sharedFlags.BlockAttempts
+	}
+
+	if sharedFlags.NoBlock {
+		attempts = -1
+	}
+
+	return attempts
+}
+
+// tryWaitForUnitStates tries to wait for units to reach the desired state.
+// It takes 5 arguments, the units to wait for, the desired state, the
+// desired JobState, how many attempts before timing out and a writer
+// interface.
+// tryWaitForUnitStates polls each of the indicated units until they
+// reach the desired state. If maxAttempts is negative, then it will not
+// wait, it will assume that all units reached their desired state.
+// If maxAttempts is zero tryWaitForUnitStates will retry forever, and
+// if it is greater than zero, it will retry up to the indicated value.
+// It returns 0 on success or 1 on errors.
+func tryWaitForUnitStates(units []string, state string, js job.JobState, maxAttempts int, out io.Writer) (ret int) {
+	// We do not wait just assume we reached the desired state
+	if maxAttempts <= -1 {
+		for _, name := range units {
+			stdout("Triggered unit %s %s", name, state)
+		}
+		return
+	}
+
+	errchan := waitForUnitStates(units, js, maxAttempts, out)
+	for err := range errchan {
+		stderr("Error waiting for units: %v", err)
+		ret = 1
+	}
+
+	return
+}
+
 // waitForUnitStates polls each of the indicated units until each of their
 // states is equal to that which the caller indicates, or until the
 // polling operation times out. waitForUnitStates will retry forever, or
diff --git a/fleetctl/load.go b/fleetctl/load.go
index 497f8605..0f7b0923 100644
--- a/fleetctl/load.go
+++ b/fleetctl/load.go
@@ -66,17 +66,7 @@ func runLoadUnits(args []string) (exit int) {
 		}
 	}
 
-	if !sharedFlags.NoBlock {
-		errchan := waitForUnitStates(loading, job.JobStateLoaded, sharedFlags.BlockAttempts, os.Stdout)
-		for err := range errchan {
-			stderr("Error waiting for units: %v", err)
-			exit = 1
-		}
-	} else {
-		for _, name := range loading {
-			stdout("Triggered unit %s load", name)
-		}
-	}
+	exit = tryWaitForUnitStates(loading, "load", job.JobStateLoaded, getBlockAttempts(), os.Stdout)
 
 	return
 }
diff --git a/fleetctl/start.go b/fleetctl/start.go
index 77ea0d34..2ded2953 100644
--- a/fleetctl/start.go
+++ b/fleetctl/start.go
@@ -74,17 +74,7 @@ func runStartUnit(args []string) (exit int) {
 		}
 	}
 
-	if !sharedFlags.NoBlock {
-		errchan := waitForUnitStates(starting, job.JobStateLaunched, sharedFlags.BlockAttempts, os.Stdout)
-		for err := range errchan {
-			stderr("Error waiting for units: %v", err)
-			exit = 1
-		}
-	} else {
-		for _, name := range starting {
-			stdout("Triggered unit %s start", name)
-		}
-	}
+	exit = tryWaitForUnitStates(starting, "start", job.JobStateLaunched, getBlockAttempts(), os.Stdout)
 
 	return
 }
diff --git a/fleetctl/stop.go b/fleetctl/stop.go
index 262946bf..941b33d7 100644
--- a/fleetctl/stop.go
+++ b/fleetctl/stop.go
@@ -79,17 +79,7 @@ func runStopUnit(args []string) (exit int) {
 		}
 	}
 
-	if !sharedFlags.NoBlock {
-		errchan := waitForUnitStates(stopping, job.JobStateLoaded, sharedFlags.BlockAttempts, os.Stdout)
-		for err := range errchan {
-			stderr("Error waiting for units: %v", err)
-			exit = 1
-		}
-	} else {
-		for _, name := range stopping {
-			stdout("Triggered unit %s stop", name)
-		}
-	}
+	exit = tryWaitForUnitStates(stopping, "stop", job.JobStateLoaded, getBlockAttempts(), os.Stdout)
 
 	return
 }
diff --git a/fleetctl/unload.go b/fleetctl/unload.go
index 6758f282..48eb833f 100644
--- a/fleetctl/unload.go
+++ b/fleetctl/unload.go
@@ -60,17 +60,7 @@ func runUnloadUnit(args []string) (exit int) {
 		}
 	}
 
-	if !sharedFlags.NoBlock {
-		errchan := waitForUnitStates(wait, job.JobStateInactive, sharedFlags.BlockAttempts, os.Stdout)
-		for err := range errchan {
-			stderr("Error waiting for units: %v", err)
-			exit = 1
-		}
-	} else {
-		for _, name := range wait {
-			stdout("Triggered unit %s unload", name)
-		}
-	}
+	exit = tryWaitForUnitStates(wait, "unload", job.JobStateInactive, getBlockAttempts(), os.Stdout)
 
 	return
 }


Test patch:
diff --git a/fleetctl/fleetctl_test.go b/fleetctl/fleetctl_test.go
index 0d50de40..64329280 100644
--- a/fleetctl/fleetctl_test.go
+++ b/fleetctl/fleetctl_test.go
@@ -181,6 +181,37 @@ func TestUnitNameMangle(t *testing.T) {
 	}
 }
 
+func TestGetBlockAttempts(t *testing.T) {
+	oldNoBlock := sharedFlags.NoBlock
+	oldBlockAttempts := sharedFlags.BlockAttempts
+
+	defer func() {
+		sharedFlags.NoBlock = oldNoBlock
+		sharedFlags.BlockAttempts = oldBlockAttempts
+	}()
+
+	var blocktests = []struct {
+		noBlock       bool
+		blockAttempts int
+		expected      int
+	}{
+		{true, 0, -1},
+		{true, -1, -1},
+		{true, 9999, -1},
+		{false, 0, 0},
+		{false, -1, 0},
+		{false, 9999, 9999},
+	}
+
+	for _, tt := range blocktests {
+		sharedFlags.NoBlock = tt.noBlock
+		sharedFlags.BlockAttempts = tt.blockAttempts
+		if n := getBlockAttempts(); n != tt.expected {
+			t.Errorf("got %d, want %d", n, tt.expected)
+		}
+	}
+}
+
 func newUnitFile(t *testing.T, contents string) *unit.UnitFile {
 	uf, err := unit.NewUnitFile(contents)
 	if err != nil {