fleetctl: Add restart and ssh-port functionality

Starts and stops a unit in one command.
Adds support to specify SSH port for SSH based commands:
* journal
* status
* ssh
* restart

refs coreos#760

fleetctl/fleetctl.go

@@ -74,6 +74,7 @@ var (
 		NoBlock       bool
 		BlockAttempts int
 		Fields        string
+		sshPort       int
 	}{}
 
 	// used to cache MachineStates
@@ -122,6 +123,7 @@ func init() {
 		cmdListUnitFiles,
 		cmdListUnits,
 		cmdLoadUnits,
+		cmdRestartUnit,
 		cmdSSH,
 		cmdStartUnit,
 		cmdStatusUnits,
@@ -716,3 +718,58 @@ func suToGlobal(su schema.Unit) bool {
 	}
 	return u.IsGlobal()
 }
+
+func waitForUnitsToRestart(units []schema.Unit, maxAttempts int, out io.Writer) chan error {
+	errchan := make(chan error)
+	var wg sync.WaitGroup
+	for _, unit := range units {
+		wg.Add(1)
+		go restartUnit(unit, maxAttempts, out, &wg, errchan)
+	}
+
+	go func() {
+		wg.Wait()
+		close(errchan)
+	}()
+
+	return errchan
+}
+
+func restartUnit(unit schema.Unit, maxAttempts int, out io.Writer, wg *sync.WaitGroup, errchan chan error) {
+	defer wg.Done()
+
+	log.V(1).Infof("Restarting unit: %s:%s", unit.Name, unit.MachineID)
+
+	sleep := 500 * time.Millisecond
+
+	if maxAttempts < 1 {
+		for {
+			if assertUnitRestart(unit, out) {
+				return
+			}
+			time.Sleep(sleep)
+		}
+	} else {
+		for attempt := 0; attempt < maxAttempts; attempt++ {
+			if assertUnitRestart(unit, out) {
+				return
+			}
+			time.Sleep(sleep)
+		}
+		errchan <- fmt.Errorf("timed out waiting for unit %s to restart", unit)
+	}
+}
+
+func assertUnitRestart(unit schema.Unit, out io.Writer) (ret bool) {
+	log.V(1).Infof("Running restart command on %s:%s", unit.Name, unit.MachineID)
+	command := fmt.Sprintf("sudo systemctl restart %s", unit.Name)
+	if exit := runCommand(command, unit.MachineID); exit == 0 {
+		log.V(1).Infof("Unit %s restarted", unit.Name)
+		msg := fmt.Sprintf("Unit %s restarted", unit.Name)
+		fmt.Fprintln(out, msg)
+		ret = true
+		return
+	} else {
+		return
+	}
+}

fleetctl/journal.go

@@ -12,7 +12,7 @@ var (
 	cmdJournal = &Command{
 		Name:    "journal",
 		Summary: "Print the journal of a unit in the cluster to stdout",
-		Usage:   "[--lines=N] [-f|--follow] <unit>",
+		Usage:   "[--lines=N] [ssh-port=N] [-f|--follow] <unit>",
 		Run:     runJournal,
 		Description: `Outputs the journal of a unit by connecting to the machine that the unit occupies.
 
@@ -30,6 +30,7 @@ func init() {
 	cmdJournal.Flags.IntVar(&flagLines, "lines", 10, "Number of recent log lines to return")
 	cmdJournal.Flags.BoolVar(&flagFollow, "follow", false, "Continuously print new entries as they are appended to the journal.")
 	cmdJournal.Flags.BoolVar(&flagFollow, "f", false, "Shorthand for --follow")
+	cmdJournal.Flags.IntVar(&sharedFlags.sshPort, "ssh-port", 22, "Use this SSH port to connect to host machine")
 }
 
 func runJournal(args []string) (exit int) {

fleetctl/restart.go

@@ -0,0 +1,79 @@
+package main
+
+import (
+	"os"
+	"time"
+
+	"github.com/coreos/fleet/schema"
+)
+
+var (
+	flagRolling    bool
+	cmdRestartUnit = &Command{
+		Name:    "restart",
+		Summary: "Instruct systemd to restart one or more units in the cluster.",
+		Usage:   "[--rolling] [--block-attempts=N] [--ssh-port=N] UNIT...",
+		Description: `Restart one or more units running in the cluster.
+
+Instructs systemd on the host machine to restart the unit, deferring to systemd
+completely for any custom start and stop directives (i.e. ExecStop or ExecStart
+options in the unit file).
+
+Restart a single unit:
+  fleetctl restart foo.service
+
+Restart an entire directory of units with glob matching, one at a time:
+  fleetctl restart --rolling myservice/*`,
+		Run: runRestartUnit,
+	}
+)
+
+func init() {
+	cmdRestartUnit.Flags.IntVar(&sharedFlags.BlockAttempts, "block-attempts", 0, "Run the restart command, performing up to N attempts before giving up. A value of 0 indicates no limit.")
+	cmdRestartUnit.Flags.BoolVar(&flagRolling, "rolling", false, "Restart each unit one at a time.")
+	cmdRestartUnit.Flags.IntVar(&sharedFlags.sshPort, "ssh-port", 22, "Use this SSH port to connect to host machine.")
+}
+
+func runRestartUnit(args []string) (exit int) {
+	units, err := findUnits(args)
+	if err != nil {
+		stderr("%v", err)
+		return 1
+	}
+
+	if !flagRolling {
+		errchan := waitForUnitsToRestart(units, sharedFlags.BlockAttempts, os.Stdout)
+		for err := range errchan {
+			stderr("Error waiting for units: %v", err)
+			exit = 1
+		}
+	} else {
+		for _, unit := range units {
+			rollingRestart(unit, sharedFlags.BlockAttempts)
+		}
+	}
+	return
+}
+
+func rollingRestart(unit schema.Unit, maxAttempts int) (ret bool) {
+	sleep := 500 * time.Millisecond
+	if maxAttempts < 1 {
+		for {
+			if assertUnitRestart(unit, out) {
+				ret = true
+				return
+			}
+			time.Sleep(sleep)
+		}
+		stderr("Error restarting unit %s", unit.Name)
+	} else {
+		for attempt := 0; attempt < maxAttempts; attempt++ {
+			if assertUnitRestart(unit, out) {
+				ret = true
+				return
+			}
+			time.Sleep(sleep)
+		}
+	}
+	return
+}

fleetctl/ssh.go

@@ -3,8 +3,10 @@ package main
 import (
 	"errors"
 	"fmt"
+	"net"
 	"os"
 	"os/exec"
+	"strconv"
 	"strings"
 	"syscall"
 
@@ -20,11 +22,11 @@ var (
 	cmdSSH                 = &Command{
 		Name:    "ssh",
 		Summary: "Open interactive shell on a machine in the cluster",
-		Usage:   "[-A|--forward-agent] [--machine|--unit] {MACHINE|UNIT}",
-		Description: `Open an interactive shell on a specific machine in the cluster or on the machine 
+		Usage:   "[-A|--forward-agent] [--ssh-port=N] [--machine|--unit] {MACHINE|UNIT}",
+		Description: `Open an interactive shell on a specific machine in the cluster or on the machine
 where the specified unit is located.
 
-fleetctl tries to detect whether your first argument is a machine or a unit. 
+fleetctl tries to detect whether your first argument is a machine or a unit.
 To skip this check use the --machine or --unit flags.
 
 Open a shell on a machine:
@@ -52,6 +54,7 @@ func init() {
 	cmdSSH.Flags.StringVar(&flagUnit, "unit", "", "Open SSH connection to machine running provided unit.")
 	cmdSSH.Flags.BoolVar(&flagSSHAgentForwarding, "forward-agent", false, "Forward local ssh-agent to target machine.")
 	cmdSSH.Flags.BoolVar(&flagSSHAgentForwarding, "A", false, "Shorthand for --forward-agent")
+	cmdSSH.Flags.IntVar(&sharedFlags.sshPort, "ssh-port", 22, "Use this SSH port to connect to host machine.")
 }
 
 func runSSH(args []string) (exit int) {
@@ -86,6 +89,8 @@ func runSSH(args []string) (exit int) {
 		return 1
 	}
 
+	addr = findSSHPort(addr)
+
 	args = pkg.TrimToDashes(args)
 
 	var sshClient *ssh.SSHForwardingClient
@@ -116,6 +121,15 @@ func runSSH(args []string) (exit int) {
 	return
 }
 
+func findSSHPort(addr string) string {
+	sshPort := sharedFlags.sshPort
+	if sshPort != 22 && !strings.Contains(addr, ":") {
+		return net.JoinHostPort(addr, strconv.Itoa(sshPort))
+	} else {
+		return addr
+	}
+}
+
 func globalMachineLookup(args []string) (string, error) {
 	if len(args) == 0 {
 		return "", errors.New("one machine or unit must be provided")
@@ -198,7 +212,8 @@ func runCommand(cmd string, machID string) (retcode int) {
 		if err != nil || ms == nil {
 			stderr("Error getting machine IP: %v", err)
 		} else {
-			err, retcode = runRemoteCommand(cmd, ms.PublicIP)
+			addr := findSSHPort(ms.PublicIP)
+			err, retcode = runRemoteCommand(cmd, addr)
 			if err != nil {
 				stderr("Error running remote command: %v", err)
 			}

fleetctl/status.go

@@ -10,7 +10,7 @@ import (
 var cmdStatusUnits = &Command{
 	Name:    "status",
 	Summary: "Output the status of one or more units in the cluster",
-	Usage:   "UNIT...",
+	Usage:   "[--ssh-port=N] UNIT...",
 	Description: `Output the status of one or more units currently running in the cluster.
 Supports glob matching of units in the current working directory or matches
 previously started units.
@@ -25,6 +25,10 @@ This command does not work with global units.`,
 	Run: runStatusUnits,
 }
 
+func init() {
+	cmdStatusUnits.Flags.IntVar(&sharedFlags.sshPort, "ssh-port", 22, "Use this SSH port to connect to host machine")
+}
+
 func runStatusUnits(args []string) (exit int) {
 	units, err := cAPI.Units()
 	if err != nil {