Listen to SIGINT and close scripts (#62)

Currently if a shuttle script is cancelled by sending a SIGINT signal the
process running in a script is not signalled. This leads to orphaned processes
and is especially an issue when launching web servers listening on a port that
can only by used once. Here the second run of the script will fail as the port
is already used.

This change introduces a context.Context down through the call chain and the
shell executor will listen for cancellations on this and signal the processes
accordingly. The context is cancelled when receiving SIGINT.

Module github.com/go-cmd/cmd is updated to the latest version which contains a
more stable shutdown API.

cmd/error.go

@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	stdcontext "context"
 	"errors"
 	"os"
 
@@ -16,6 +17,10 @@ func checkError(err error) {
 		uii.Errorln("shuttle failed\n%s", exitCode.Message)
 		os.Exit(exitCode.Code)
 	}
+	if errors.Is(err, stdcontext.Canceled) {
+		uii.Errorln("Operation cancelled")
+		os.Exit(2)
+	}
 	uii.Errorln("shuttle failed\nError: %s", err)
 	os.Exit(1)
 }

cmd/run.go

@@ -1,7 +1,9 @@
 package cmd
 
 import (
+	stdcontext "context"
 	"os"
+	"os/signal"
 
 	"github.com/lunarway/shuttle/pkg/executors"
 	"github.com/spf13/cobra"
@@ -16,8 +18,9 @@ var runCmd = &cobra.Command{
 		var commandName = args[0]
 		context, err := getProjectContext()
 		checkError(err)
-
-		err = executors.Execute(context, commandName, args[1:], validateArgs)
+		ctx, cancel := withSignal(stdcontext.Background())
+		defer cancel()
+		err = executors.Execute(ctx, context, commandName, args[1:], validateArgs)
 		checkError(err)
 	},
 }
@@ -44,3 +47,30 @@ func init() {
 	runCmd.Flags().BoolVar(&validateArgs, "validate", true, "Validate arguments against script definition in plan and exit with 1 on unknown or missing arguments")
 	rootCmd.AddCommand(runCmd)
 }
+
+// withSignal returns a copy of parent with a new Done channel. The returned
+// context's Done channel is closed when the returned cancel function is called,
+// if the parent context's Done channel is closed, if a SIGINT signal is
+// catched, whichever happens first.
+//
+// Canceling this context releases resources associated with it, so code should
+// call cancel as soon as the operations running in this Context complete.
+func withSignal(parent stdcontext.Context) (stdcontext.Context, func()) {
+	parent, cancel := stdcontext.WithCancel(parent)
+	c := make(chan os.Signal, 1)
+	signal.Notify(c, os.Interrupt)
+
+	go func() {
+		select {
+		case s := <-c:
+			uii.Infoln("Received %v signal...", s)
+			cancel()
+		case <-parent.Done():
+		}
+	}()
+
+	return parent, func() {
+		signal.Stop(c)
+		cancel()
+	}
+}

examples/moon-base/shuttle.yaml

@@ -6,3 +6,8 @@ vars:
     field:
       key: value
   run-as-root: false
+scripts:
+  serve:
+    description: Serve starts an nginx image as a long running process.
+    actions:
+    - shell: docker run --rm -i --name shuttle-nginx nginx

go.mod

@@ -8,7 +8,7 @@ require (
 	github.com/aokoli/goutils v1.0.1 // indirect
 	github.com/cli/safeexec v1.0.0
 	github.com/davecgh/go-spew v1.1.1 // indirect
-	github.com/go-cmd/cmd v1.2.0
+	github.com/go-cmd/cmd v1.3.0
 	github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510
 	github.com/google/uuid v1.0.0 // indirect
 	github.com/huandu/xstrings v1.2.0 // indirect

go.sum

@@ -8,10 +8,10 @@ github.com/cli/safeexec v1.0.0 h1:0VngyaIyqACHdcMNWfo6+KdUYnqEr2Sg+bSP1pdF+dI=
 github.com/cli/safeexec v1.0.0/go.mod h1:Z/D4tTN8Vs5gXYHDCbaM1S/anmEDnJb1iW0+EJ5zx3Q=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
-github.com/go-cmd/cmd v1.2.0 h1:Aohz0ZG0nQbvT4z55Mh+fdegX48GSAXL3cSsbYxRfvI=
-github.com/go-cmd/cmd v1.2.0/go.mod h1:XgKkd0L6sv9WcYV0FS8RfG1RJCSTVHTsLeAD2pTgHt0=
-github.com/go-test/deep v1.0.5 h1:AKODKU3pDH1RzZzm6YZu77YWtEAq6uh1rLIAQlay2qc=
-github.com/go-test/deep v1.0.5/go.mod h1:QV8Hv/iy04NyLBxAdO9njL0iVPN1S4d/A3NVv1V36o8=
+github.com/go-cmd/cmd v1.3.0 h1:Wet2eYkLouFqyiG+x6P6l8CICRywhRD6sjMNalTSvbs=
+github.com/go-cmd/cmd v1.3.0/go.mod h1:l/X/csRuYRDqiQIz9PPJBn4xDrdxgBXeLE9x1BeFU6M=
+github.com/go-test/deep v1.0.6 h1:UHSEyLZUwX9Qoi99vVwvewiMC8mM2bf7XEM2nqvzEn8=
+github.com/go-test/deep v1.0.6/go.mod h1:QV8Hv/iy04NyLBxAdO9njL0iVPN1S4d/A3NVv1V36o8=
 github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 h1:El6M4kTTCOh6aBiKaUGG7oYTSPP8MxqL4YI3kZKwcP4=
 github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510/go.mod h1:pupxD2MaaD3pAXIBCelhxNneeOaAeabZDe5s4K6zSpQ=
 github.com/google/uuid v1.0.0 h1:b4Gk+7WdP/d3HZH8EJsZpvV7EtDOgaZLtnaNGIu1adA=

pkg/executors/docker.go

@@ -1,6 +1,7 @@
 package executors
 
 import (
+	"context"
 	"io"
 	"log"
 	"os"
@@ -11,7 +12,7 @@ import (
 )
 
 // Build builds the docker image from a shuttle plan
-func executeDocker(context ActionExecutionContext) error {
+func executeDocker(ctx context.Context, context ActionExecutionContext) error {
 	dockerFilePath := path.Join(context.ScriptContext.Project.LocalPlanPath, context.Action.Dockerfile)
 	projectPath := context.ScriptContext.Project.ProjectPath
 	execCmd := exec.Command("docker", "build", "-f", dockerFilePath, projectPath)

pkg/executors/executor.go

@@ -1,6 +1,7 @@
 package executors
 
 import (
+	"context"
 	"fmt"
 	"sort"
 	"strings"
@@ -25,7 +26,7 @@ type ActionExecutionContext struct {
 }
 
 // Execute is the command executor for the plan files
-func Execute(p config.ShuttleProjectContext, command string, args []string, validateArgs bool) error {
+func Execute(ctx context.Context, p config.ShuttleProjectContext, command string, args []string, validateArgs bool) error {
 	script, ok := p.Scripts[command]
 	if !ok {
 		return errors.NewExitCode(2, "Script '%s' not found", command)
@@ -49,7 +50,7 @@ func Execute(p config.ShuttleProjectContext, command string, args []string, vali
 			Action:        action,
 			ActionIndex:   actionIndex,
 		}
-		err := executeAction(actionContext)
+		err := executeAction(ctx, actionContext)
 		if err != nil {
 			return err
 		}
@@ -162,18 +163,18 @@ func expectedArgumentsHelp(command string, args []config.ShuttleScriptArgs) stri
 	return s.String()
 }
 
-func executeAction(context ActionExecutionContext) error {
+func executeAction(ctx context.Context, context ActionExecutionContext) error {
 	for _, executor := range executors {
 		if executor.match(context.Action) {
-			return executor.execute(context)
+			return executor.execute(ctx, context)
 		}
 	}
 
 	panic(fmt.Sprintf("Could not find an executor for %v.actions[%v]!", context.ScriptContext.ScriptName, context.ActionIndex))
 }
 
 type actionMatchFunc = func(config.ShuttleAction) bool
-type actionExecutionFunc = func(ActionExecutionContext) error
+type actionExecutionFunc = func(context.Context, ActionExecutionContext) error
 
 type actionExecutor struct {
 	match   actionMatchFunc

pkg/executors/executor_test.go

@@ -1,8 +1,12 @@
 package executors
 
 import (
+	"context"
+	"fmt"
 	"testing"
+	"time"
 
+	"github.com/go-cmd/cmd"
 	"github.com/lunarway/shuttle/pkg/config"
 	"github.com/stretchr/testify/assert"
 )
@@ -140,3 +144,58 @@ func TestSortValidationErrors(t *testing.T) {
 		})
 	}
 }
+
+// TestExecute_contextCancellation tests that scripts are closed when the
+// context is cancelled.
+func TestExecute_contextCancellation(t *testing.T) {
+	imageName := fmt.Sprintf("shuttle-test-execute-cancellation-%d", time.Now().UnixNano())
+	t.Logf("Starting image %s", imageName)
+	projectContext := config.ShuttleProjectContext{
+		Scripts: map[string]config.ShuttlePlanScript{
+			"serve": {
+				Description: "",
+				Actions: []config.ShuttleAction{
+					{
+						Shell: fmt.Sprintf("docker run --rm -i --name %s nginx", imageName),
+					},
+				},
+				Args: []config.ShuttleScriptArgs{},
+			},
+		},
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	go func() {
+		// let the container start before stopping it
+		time.Sleep(500 * time.Millisecond)
+		cancel()
+	}()
+
+	err := Execute(ctx, projectContext, "serve", nil, true)
+	assert.EqualError(t, err, context.Canceled.Error())
+
+	// sadly we need to give the docker some time before "docker ps" shows the
+	// containers
+	time.Sleep(500 * time.Millisecond)
+	images := runningDockerImages(t, imageName)
+	assert.Len(t, images, 0, "expected no images to be running")
+}
+
+func runningDockerImages(t *testing.T, imageName string) []string {
+	t.Helper()
+	cmd := cmd.NewCmd("docker", "ps", "-a", "--format", "{{ .Names }}")
+	status := <-cmd.Start()
+	t.Logf("docker ps: stderr: %v", status.Stderr)
+
+	t.Logf("Docker containers")
+	for _, container := range status.Stdout {
+		t.Logf("- %s", container)
+		if container == imageName {
+			t.Errorf("Container '%s still exists in docker but shouldn't", container)
+		}
+	}
+	if status.Exit != 0 {
+		t.Fatal("Failed to check running docker images")
+	}
+	return nil
+}

pkg/executors/shell.go

@@ -1,6 +1,7 @@
 package executors
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -12,7 +13,7 @@ import (
 )
 
 // Build builds the docker image from a shuttle plan
-func executeShell(context ActionExecutionContext) error {
+func executeShell(ctx context.Context, context ActionExecutionContext) error {
 	//log.Printf("Exec: %s", context.Action.Shell)
 	//cmdAndArgs := strings.Split(s.Shell, " ")
 	//cmd := cmdAndArgs[0]
@@ -61,13 +62,29 @@ func executeShell(context ActionExecutionContext) error {
 
 	// Run and wait for Cmd to return, discard Status
 	context.ScriptContext.Project.UI.Titleln("shell: %s", context.Action.Shell)
-	status := <-execCmd.Start()
-	<-doneChan
 
-	if status.Exit > 0 {
-		return errors.NewExitCode(4, "Failed executing script `%s`: shell script `%s`\nExit code: %v", context.ScriptContext.ScriptName, context.Action.Shell, status.Exit)
+	// stop cmd if context is cancelled
+	go func() {
+		select {
+		case <-ctx.Done():
+			err := execCmd.Stop()
+			if err != nil {
+				context.ScriptContext.Project.UI.Errorln("Failed to stop script '%s': %v", context.Action.Shell, err)
+			}
+		case <-doneChan:
+		}
+	}()
+
+	select {
+	case status := <-execCmd.Start():
+		<-doneChan
+		if status.Exit > 0 {
+			return errors.NewExitCode(4, "Failed executing script `%s`: shell script `%s`\nExit code: %v", context.ScriptContext.ScriptName, context.Action.Shell, status.Exit)
+		}
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
 	}
-	return nil
 }
 
 func init() {