fleetctl: optimize stop/unload commands

Use the Unit API instead of Units API to speed up lookups on environments with a large number of distinct services. Fixes coreos#1343
octoblu · Sep 4, 2015 · 5a8d47b · 5a8d47b
1 parent cd23d91
commit 5a8d47b
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 57 deletions.
diff --git a/fleetctl/fleetctl.go b/fleetctl/fleetctl.go
@@ -528,31 +528,6 @@ func machineFullLegend(ms machine.MachineState, full bool) string {
 	return legend
 }
 
-func findUnits(args []string) (sus []schema.Unit, err error) {
-	units, err := cAPI.Units()
-	if err != nil {
-		return nil, err
-	}
-
-	uMap := make(map[string]*schema.Unit, len(units))
-	for _, u := range units {
-		u := u
-		uMap[u.Name] = u
-	}
-
-	filtered := make([]schema.Unit, 0)
-	for _, v := range args {
-		v = unitNameMangle(v)
-		u, ok := uMap[v]
-		if !ok {
-			continue
-		}
-		filtered = append(filtered, *u)
-	}
-
-	return filtered, nil
-}
-
 func createUnit(name string, uf *unit.UnitFile) (*schema.Unit, error) {
 	if uf == nil {
 		return nil, fmt.Errorf("nil unit provided")

diff --git a/fleetctl/stop.go b/fleetctl/stop.go
@@ -52,30 +52,37 @@ func init() {
 }
 
 func runStopUnit(args []string) (exit int) {
-	units, err := findUnits(args)
-	if err != nil {
-		stderr("%v", err)
-		return 1
-	}
-
 	stopping := make([]string, 0)
-	for _, u := range units {
-		if !suToGlobal(u) {
-			if job.JobState(u.CurrentState) == job.JobStateInactive {
-				stderr("Unable to stop unit %s in state %s", u.Name, job.JobStateInactive)
+
+	for _, arg := range args {
+		name := unitNameMangle(arg)
+		unit, err := cAPI.Unit(name)
+		if err != nil {
+			stderr("Error retrieving unit: %v", err)
+			return 1
+		}
+
+		if unit == nil {
+			stderr("Unit %s does not exist.", name)
+			continue
+		}
+
+		if !suToGlobal(*unit) {
+			if job.JobState(unit.CurrentState) == job.JobStateInactive {
+				stderr("Unable to stop unit %s in state %s", unit.Name, job.JobStateInactive)
 				return 1
-			} else if job.JobState(u.CurrentState) == job.JobStateLoaded {
-				log.Debugf("Unit(%s) already %s, skipping.", u.Name, job.JobStateLoaded)
+			} else if job.JobState(unit.CurrentState) == job.JobStateLoaded {
+				log.Debugf("Unit(%s) already %s, skipping.", unit.Name, job.JobStateLoaded)
 				continue
 			}
 		}
 
-		log.Debugf("Setting target state of Unit(%s) to %s", u.Name, job.JobStateLoaded)
-		cAPI.SetUnitTargetState(u.Name, string(job.JobStateLoaded))
-		if suToGlobal(u) {
-			stdout("Triggered global unit %s stop", u.Name)
+		log.Debugf("Setting target state of Unit(%s) to %s", unit.Name, job.JobStateLoaded)
+		cAPI.SetUnitTargetState(unit.Name, string(job.JobStateLoaded))
+		if suToGlobal(*unit) {
+			stdout("Triggered global unit %s stop", unit.Name)
 		} else {
-			stopping = append(stopping, u.Name)
+			stopping = append(stopping, unit.Name)
 		}
 	}
 

diff --git a/fleetctl/unload.go b/fleetctl/unload.go
@@ -36,27 +36,34 @@ func init() {
 }
 
 func runUnloadUnit(args []string) (exit int) {
-	units, err := findUnits(args)
-	if err != nil {
-		stderr("%v", err)
-		return 1
-	}
-
 	wait := make([]string, 0)
-	for _, s := range units {
-		if !suToGlobal(s) {
-			if job.JobState(s.CurrentState) == job.JobStateInactive {
-				log.Debugf("Target state of Unit(%s) already %s, skipping.", s.Name, job.JobStateInactive)
+
+	for _, arg := range args {
+		name := unitNameMangle(arg)
+		unit, err := cAPI.Unit(name)
+		if err != nil {
+			stderr("Error retrieving unit: %v", err)
+			return 1
+		}
+
+		if unit == nil {
+			stderr("Unit %s does not exist.", name)
+			continue
+		}
+
+		if !suToGlobal(*unit) {
+			if job.JobState(unit.CurrentState) == job.JobStateInactive {
+				log.Debugf("Target state of Unit(%s) already %s, skipping.", unit.Name, job.JobStateInactive)
 				continue
 			}
 		}
 
-		log.Debugf("Setting target state of Unit(%s) to %s", s.Name, job.JobStateInactive)
-		cAPI.SetUnitTargetState(s.Name, string(job.JobStateInactive))
-		if suToGlobal(s) {
-			stdout("Triggered global unit %s unload", s.Name)
+		log.Debugf("Setting target state of Unit(%s) to %s", unit.Name, job.JobStateInactive)
+		cAPI.SetUnitTargetState(unit.Name, string(job.JobStateInactive))
+		if suToGlobal(*unit) {
+			stdout("Triggered global unit %s unload", unit.Name)
 		} else {
-			wait = append(wait, s.Name)
+			wait = append(wait, unit.Name)
 		}
 	}