Merge pull request #241 from QuentinPerez/fix239

Fix239
scaleway · Nov 27, 2015 · be9b58c · be9b58c
2 parents 7281312 + 83765ab
commit be9b58c
Show file tree

Hide file tree

Showing 3 changed files with 67 additions and 28 deletions.
diff --git a/README.md b/README.md
@@ -677,13 +677,14 @@ Options:
   -a, --attach=false    Attach to serial console
   --bootscript=""       Assign a bootscript
   -d, --detach=false    Run server in background and print server ID
-  -e, --env=""          Provide metadata tags passed to initrd (i.e., boot=resue INITRD_DEBUG=1)
+  -e, --env=""          Provide metadata tags passed to initrd (i.e., boot=rescue INITRD_DEBUG=1)
   -g, --gateway=""      Use a SSH gateway
   -h, --help=false      Print usage
   --ip-address=""       Assign an IP
   --name=""             Assign a name
   --rm=false            Automatically remove the server when it exits
   --show-boot=false     Allows to show the boot
+  -T, --timeout=0       Set timeout value to seconds
   --tmp-ssh-key=false   Access your server without uploading your SSH key to your account
   -v, --volume=""       Attach additional volume (i.e., 50G)
 
@@ -1142,6 +1143,7 @@ $ scw inspect myserver | jq '.[0].public_ip.address'
 
 ### master (unreleased)
 
+* Support of `scw run --timeout=X` ([#239](https://github.com/scaleway/scaleway-cli/issues/239))
 * Check the "stopped" state for `scw run | exec -w`([#229](https://github.com/scaleway/scaleway-cli/issues/229))
 
 View full [commits list](https://github.com/scaleway/scaleway-cli/compare/v1.6.0...master)

diff --git a/pkg/cli/cmd_run.go b/pkg/cli/cmd_run.go
@@ -33,9 +33,10 @@ var cmdRun = &Command{
 func init() {
 	cmdRun.Flag.StringVar(&runCreateName, []string{"-name"}, "", "Assign a name")
 	cmdRun.Flag.StringVar(&runCreateBootscript, []string{"-bootscript"}, "", "Assign a bootscript")
-	cmdRun.Flag.StringVar(&runCreateEnv, []string{"e", "-env"}, "", "Provide metadata tags passed to initrd (i.e., boot=resue INITRD_DEBUG=1)")
+	cmdRun.Flag.StringVar(&runCreateEnv, []string{"e", "-env"}, "", "Provide metadata tags passed to initrd (i.e., boot=rescue INITRD_DEBUG=1)")
 	cmdRun.Flag.StringVar(&runCreateVolume, []string{"v", "-volume"}, "", "Attach additional volume (i.e., 50G)")
 	cmdRun.Flag.BoolVar(&runHelpFlag, []string{"h", "-help"}, false, "Print usage")
+	cmdRun.Flag.Int64Var(&runTimeout, []string{"T", "-timeout"}, 0, "Set timeout value to seconds")
 	cmdRun.Flag.StringVar(&runIPAddress, []string{"-ip-address"}, "", "Assign an IP")
 	cmdRun.Flag.BoolVar(&runAttachFlag, []string{"a", "-attach"}, false, "Attach to serial console")
 	cmdRun.Flag.BoolVar(&runDetachFlag, []string{"d", "-detach"}, false, "Run server in background and print server ID")
@@ -59,6 +60,7 @@ var runDetachFlag bool         // -d, --detach flag
 var runGateway string          // -g, --gateway flag
 var runTmpSSHKey bool          // --tmp-ssh-key flag
 var runShowBoot bool           // --show-boot flag
+var runTimeout int64           // --timeout flag
 
 func runRun(cmd *Command, rawArgs []string) error {
 	if runHelpFlag {
@@ -101,6 +103,7 @@ func runRun(cmd *Command, rawArgs []string) error {
 		TmpSSHKey:  runTmpSSHKey,
 		ShowBoot:   runShowBoot,
 		IP:         runIPAddress,
+		Timeout:    runTimeout,
 		// FIXME: DynamicIPRequired
 		// FIXME: Timeout
 	}

diff --git a/pkg/commands/run.go b/pkg/commands/run.go
@@ -10,6 +10,7 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+	"time"
 
 	"github.com/scaleway/scaleway-cli/pkg/api"
 	"github.com/scaleway/scaleway-cli/pkg/config"
@@ -32,6 +33,7 @@ type RunArgs struct {
 	AutoRemove bool
 	TmpSSHKey  bool
 	ShowBoot   bool
+	Timeout    int64
 	// DynamicIPRequired
 	// Timeout
 }
@@ -116,34 +118,60 @@ func Run(ctx CommandContext, args RunArgs) error {
 	// Sync cache on disk
 	ctx.API.Sync()
 
+	closeTimeout := make(chan struct{})
+	timeoutExit := make(chan struct{})
+
+	if args.Timeout > 0 {
+		go func() {
+			select {
+			case <-time.After(time.Duration(args.Timeout) * time.Second):
+				close(timeoutExit)
+			case <-closeTimeout:
+				break
+			}
+		}()
+	}
 	if args.ShowBoot {
 		// Attach to server serial
 		logrus.Info("Attaching to server console ...")
 		gottycli, done, err := utils.AttachToSerial(serverID, ctx.API.Token)
 		if err != nil {
+			close(closeTimeout)
 			return fmt.Errorf("cannot attach to server serial: %v", err)
 		}
 		utils.Quiet(true)
 		notif, gateway, err := waitSSHConnection(ctx, args, serverID)
 		if err != nil {
+			close(closeTimeout)
+			gottycli.ExitLoop()
+			<-done
 			return err
 		}
-		sshConnection := <-notif
-		gottycli.ExitLoop()
-		<-done
-		utils.Quiet(false)
-		if sshConnection.err != nil {
-			return sshConnection.err
-		}
-		server := sshConnection.server
-		logrus.Info("Connecting to server ...")
-		if err = utils.SSHExec(server.PublicAddress.IP, server.PrivateIP, []string{}, false, gateway); err != nil {
-			return fmt.Errorf("Connection to server failed: %v", err)
+		select {
+		case <-timeoutExit:
+			gottycli.ExitLoop()
+			<-done
+			utils.Quiet(false)
+			return fmt.Errorf("Operation timed out")
+		case sshConnection := <-notif:
+			close(closeTimeout)
+			gottycli.ExitLoop()
+			<-done
+			utils.Quiet(false)
+			if sshConnection.err != nil {
+				return sshConnection.err
+			}
+			server := sshConnection.server
+			logrus.Info("Connecting to server ...")
+			if err = utils.SSHExec(server.PublicAddress.IP, server.PrivateIP, []string{}, false, gateway); err != nil {
+				return fmt.Errorf("Connection to server failed: %v", err)
+			}
 		}
 	} else if args.Attach {
 		// Attach to server serial
 		logrus.Info("Attaching to server console ...")
 		gottycli, done, err := utils.AttachToSerial(serverID, ctx.API.Token)
+		close(closeTimeout)
 		if err != nil {
 			return fmt.Errorf("cannot attach to server serial: %v", err)
 		}
@@ -152,25 +180,31 @@ func Run(ctx CommandContext, args RunArgs) error {
 	} else {
 		notif, gateway, err := waitSSHConnection(ctx, args, serverID)
 		if err != nil {
+			close(closeTimeout)
 			return err
 		}
-		sshConnection := <-notif
-		if sshConnection.err != nil {
-			return sshConnection.err
-		}
-		server := sshConnection.server
-		// exec -w SERVER COMMAND ARGS...
-		if len(args.Command) < 1 {
-			logrus.Info("Connecting to server ...")
-			if err = utils.SSHExec(server.PublicAddress.IP, server.PrivateIP, []string{}, false, gateway); err != nil {
-				return fmt.Errorf("Connection to server failed: %v", err)
+		select {
+		case <-timeoutExit:
+			return fmt.Errorf("Operation timed out")
+		case sshConnection := <-notif:
+			close(closeTimeout)
+			if sshConnection.err != nil {
+				return sshConnection.err
 			}
-		} else {
-			logrus.Infof("Executing command: %s ...", args.Command)
-			if err = utils.SSHExec(server.PublicAddress.IP, server.PrivateIP, args.Command, false, gateway); err != nil {
-				return fmt.Errorf("command execution failed: %v", err)
+			server := sshConnection.server
+			// exec -w SERVER COMMAND ARGS...
+			if len(args.Command) < 1 {
+				logrus.Info("Connecting to server ...")
+				if err = utils.SSHExec(server.PublicAddress.IP, server.PrivateIP, []string{}, false, gateway); err != nil {
+					return fmt.Errorf("Connection to server failed: %v", err)
+				}
+			} else {
+				logrus.Infof("Executing command: %s ...", args.Command)
+				if err = utils.SSHExec(server.PublicAddress.IP, server.PrivateIP, args.Command, false, gateway); err != nil {
+					return fmt.Errorf("command execution failed: %v", err)
+				}
+				logrus.Info("Command successfuly executed")
 			}
-			logrus.Info("Command successfuly executed")
 		}
 	}
 	return nil