Create WithGracePeriod,WithOnDeathSignalChildren options

Signed-off-by: Ed Warnicke <[email protected]>

.github/workflows/ci.yaml

@@ -1,6 +1,12 @@
 ---
 name: ci
-on: [pull_request, push]
+on:
+  push:
+    branches:
+      - master
+  pull_request:
+    branches:
+      - master
 jobs:
   build:
     name: build

.gitignore

@@ -15,3 +15,5 @@
 # vendor/
 
 .idea/
+afterterm
+!testcmds/afterterm

.golangci.yml

@@ -5,6 +5,8 @@ run:
   issues-exit-code: 1
   tests: true
 linters-settings:
+  funlen:
+    lines: 80
   gosec:
     settings:
       exclude: "G204"
@@ -22,8 +24,6 @@ linters-settings:
           - (github.com/sirupsen/logrus.FieldLogger).Fatalf
   golint:
     min-confidence: 0.8
-  goimports:
-    local-prefixes: github.com/networkservicemesh/cmd-forwarder-vppagent
   gocyclo:
     min-complexity: 15
   maligned:

README.md

@@ -6,7 +6,7 @@ Its main features are:
 2. It's Start() returns an errCh that will get zero or one errors, and be closed after the command has finished running
 3. You can use the WithXYZ pattern to do things like customize things like Stdout, Stdin, StdError, Dir, and Env variables
 
-Run Examples:
+## Run Examples:
 
 ```go
 if err := exechelper.Run("go list -m");err != nil {...}
@@ -20,7 +20,7 @@ errBuffer := bytes.NewBuffer([]bytes{})
 if err := exechelper.Run("go list -m",WithStdout(outputBuffer),WithStderr(errBuffer));err != nil {...}
 ```
 
-Start Examples
+## Start Examples
 
 ```go
 ctx,cancel := context.WithCancel(context.Background())
@@ -34,3 +34,34 @@ errCh := exechelper.Start(startContext,"spire-server run",WithContext(ctx))
 
 Similarly, exechelper.Output(...), exechelper.CombinedOutput(...) are provided.
 
+## Multiplexing Stdout/Stderr
+
+WithStdout and WithStderr can be used multiple times, with each provided io.Writer provided getting a copy of the stdout/stderr.
+As many WithStdout or WithStderr options may be used on the same Run/Start/Output/CombinedOutput as you wish
+
+Example:
+
+```go
+outputBuffer := bytes.NewBuffer([]byte{})
+errBuffer := bytes.NewBuffer([]bytes{})
+if err := exechelper.Run("go list -m",WithStdout(outputBuffer),WithStderr(errBuffer),WithStdout(os.Stdout),WithStderr(os.Stderr));err != nil {...}
+// stdout of "go list -m" is written to both outputBuffer and os.Stdout.  stderr is written to both errBuffer and os.Stdout
+```
+
+## Terminating with SIGTERM and a grace period
+
+By default, canceling the context on a Start will result in SIGKILL being sent to the child process (this is how exec.Cmd works).
+It is often useful to send SIGTERM first, and allow a grace period before sending SIGKILL to allow processes the opportunity
+to clean themselves up.
+
+Example:
+```go
+ctx,cancel := context.WithCancel(context.Background())
+errCh := exechelper.Start(startContext,"spire-server run",WithContext(ctx),WithGracePeriod(30*time.Second))
+cancel()
+<-errCh
+// Calling cancel will send SIGTERM to the spire-server and wait up to 30 seconds for it to exit before sending SIGKILL
+// errCh will receive any errors from the exit of spire-server immediately after spire-server exiting, whether that
+// is immediately, any time during the 30 second grace period, or after the SIGKILL is sent
+```
+

exechelper.go

@@ -21,8 +21,11 @@ import (
 	"bytes"
 	"context"
 	"os/exec"
+	"syscall"
+	"time"
 
 	"github.com/google/shlex"
+	"github.com/pkg/errors"
 )
 
 // Run - Creates a exec.Cmd using cmdStr.  Runs exec.Cmd.Run and returns the resulting error
@@ -34,22 +37,104 @@ func Run(cmdStr string, options ...*Option) error {
 func Start(cmdStr string, options ...*Option) <-chan error {
 	errCh := make(chan error, 1)
 
+	// Extract context from options
+	optionCtx := extractContextFromOptions(options)
+
+	// Extract graceperiod from options
+	graceperiod, err := extractGracePeriodFromOptions(optionCtx, options)
+	if err != nil {
+		errCh <- err
+		close(errCh)
+		return errCh
+	}
+
+	// By default, the context passed to StartContext (ie, cmdCtx) is the same as the context we got from the options
+	// (ie, optionsCtx)
+	cmdCtx := optionCtx
+
+	// But if we have a graceperiod, we need a separate cmdCtx and cmdCancel so we can insert our SIGTERM
+	// between the optionsCtx.Done() and time.After(graceperiod) before *actually* canceling the cmdCtx
+	// and thus sending SIGKILL to the cmd
+	var cmdCancel context.CancelFunc
+	if graceperiod != 0 {
+		cmdCtx, cmdCancel = context.WithCancel(context.Background())
+	}
+
+	cmd, err := constructCommand(cmdCtx, cmdStr, options)
+	if err != nil {
+		errCh <- err
+		close(errCh)
+		if cmdCancel != nil {
+			cmdCancel()
+		}
+		return errCh
+	}
+
+	// Start the *exec.Cmd
+	if err = cmd.Start(); err != nil {
+		errCh <- err
+		close(errCh)
+		if cmdCancel != nil {
+			cmdCancel()
+		}
+		return errCh
+	}
+
+	// By default, the error channel we send any error from the wait to (waitErrCh) is the one we return (errCh)
+	waitErrCh := errCh
+
+	// But if we have a graceperiod and a cmdCancel, we need a distinct waitErrCh from the one we return,
+	// so that we can select on waitErrCh after sending SIGTERM and then forward any errors to errCh
+	if cmdCancel != nil && graceperiod > 0 {
+		waitErrCh = make(chan error, len(errCh))
+	}
+
+	// Collect the wait
+	go func(waitErrCh chan error) {
+		if err := cmd.Wait(); err != nil {
+			waitErrCh <- err
+		}
+		close(waitErrCh)
+	}(waitErrCh)
+
+	// Handle SIGTERM and graceperiod
+	if cmdCancel != nil && graceperiod > 0 {
+		go handleGracePeriod(optionCtx, cmd, cmdCancel, graceperiod, waitErrCh, errCh)
+	}
+
+	return errCh
+}
+
+func extractGracePeriodFromOptions(ctx context.Context, options []*Option) (time.Duration, error) {
+	var graceperiod time.Duration
+	for _, option := range options {
+		if option.GracePeriod != 0 {
+			graceperiod = option.GracePeriod
+			if ctx == nil {
+				return 0, errors.New("graceperiod cannot be set without WithContext option")
+			}
+		}
+	}
+	return graceperiod, nil
+}
+
+func extractContextFromOptions(options []*Option) context.Context {
 	// Set the context
-	var ctx context.Context
+	var optionCtx context.Context
 	for _, option := range options {
 		if option.Context != nil {
-			ctx = option.Context
+			optionCtx = option.Context
 		}
 	}
+	return optionCtx
+}
 
+func constructCommand(ctx context.Context, cmdStr string, options []*Option) (*exec.Cmd, error) {
 	// Construct the command args
 	args, err := shlex.Split(cmdStr)
 	if err != nil {
-		errCh <- err
-		close(errCh)
-		return errCh
+		return nil, err
 	}
-
 	// Create the *exec.Cmd
 	var cmd *exec.Cmd
 	switch ctx {
@@ -63,30 +148,36 @@ func Start(cmdStr string, options ...*Option) <-chan error {
 	for _, option := range options {
 		// Apply the CmdOptions
 		if option.CmdOption != nil {
-			if err = option.CmdOption(cmd); err != nil {
-				errCh <- err
-				close(errCh)
-				return errCh
+			if err := option.CmdOption(cmd); err != nil {
+				return nil, err
 			}
 		}
 	}
+	return cmd, nil
+}
 
-	// Start the *exec.Cmd
-	if err = cmd.Start(); err != nil {
-		errCh <- err
-		close(errCh)
-		return errCh
-	}
+func handleGracePeriod(optionCtx context.Context, cmd *exec.Cmd, cmdCancel context.CancelFunc, graceperiod time.Duration, waitErrCh <-chan error, errCh chan<- error) {
+	// Wait for the optionCtx to be done
+	<-optionCtx.Done()
 
-	// Collect the wait
-	go func(chan error) {
-		if err := cmd.Wait(); err != nil {
-			errCh <- err
-		}
-		close(errCh)
-	}(errCh)
+	// Send SIGTERM
+	_ = cmd.Process.Signal(syscall.SIGTERM)
 
-	return errCh
+	// Wait for either the waitErrCh to be closed or have an error (ie, cmd exited) or graceperiod
+	// either way
+	select {
+	case <-waitErrCh:
+	case <-time.After(graceperiod):
+	}
+	// Cancel the cmdCtx passed to exec.StartContext
+	cmdCancel()
+
+	// Move all errors from waitErrCh to errCh
+	for err := range waitErrCh {
+		errCh <- err
+	}
+	// Close errCh
+	close(errCh)
 }
 
 // Output - Creates a exec.Cmd using cmdStr.  Runs exec.Cmd.Output and returns the resulting output as []byte and error

exechelper_test.go

@@ -12,10 +12,13 @@ import (
 	"os/exec"
 	"path"
 	"strings"
+	"sync"
+	"syscall"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 
 	"github.com/edwarnicke/exechelper"
 )
@@ -68,6 +71,9 @@ func TestStartWithContext(t *testing.T) {
 	select {
 	case err := <-errCh:
 		assert.IsType(t, &exec.ExitError{}, err) // Because we canceled we will get an exec.ExitError{}
+		exitErr := err.(*exec.ExitError)
+		status := exitErr.ProcessState.Sys().(syscall.WaitStatus)
+		assert.Equal(t, status.Signal(), syscall.SIGKILL)
 		assert.Empty(t, errCh)
 	case <-time.After(time.Second):
 		assert.Fail(t, "Failed to cancel context")
@@ -168,3 +174,76 @@ func TestWithEnvKV(t *testing.T) {
 	_, err := exechelper.Output("printenv", exechelper.WithEnvKV(key1))
 	assert.Error(t, err)
 }
+
+func TestWithGracePeriodWithContext(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	graceperiod := 200 * time.Millisecond
+	err := exechelper.Run("go build ./testcmds/afterterm")
+	require.NoError(t, err)
+	errCh := exechelper.Start(
+		fmt.Sprintf("./afterterm %s", graceperiod-50*time.Millisecond),
+		exechelper.WithContext(ctx),
+		exechelper.WithGracePeriod(graceperiod),
+		exechelper.WithStdout(os.Stdout),
+	)
+	time.Sleep(100 * time.Millisecond)
+	cancelTime := time.Now()
+	cancel()
+	ok := true
+	for ok {
+		select {
+		case err, ok = <-errCh:
+			require.NoError(t, err)
+		case <-time.After(graceperiod + 50*time.Millisecond):
+			require.Failf(t, "", "failed to stop within graceperiod(%s): %s", graceperiod+50*time.Millisecond, time.Since(cancelTime))
+			ok = false
+		}
+	}
+}
+
+func TestWithGracePeriodExceeded(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	graceperiod := 200 * time.Millisecond
+	err := exechelper.Run("go build ./testcmds/afterterm")
+	require.NoError(t, err)
+	errCh := exechelper.Start(
+		fmt.Sprintf("./afterterm %s", 2*graceperiod),
+		exechelper.WithContext(ctx),
+		exechelper.WithGracePeriod(graceperiod),
+		exechelper.WithStdout(os.Stdout),
+	)
+	time.Sleep(100 * time.Millisecond)
+	cancelTime := time.Now()
+	cancel()
+	ok := true
+	errOnce := sync.Once{}
+	for ok {
+		select {
+		case err, ok = <-errCh:
+			errOnce.Do(func() {
+				assert.True(t, ok, "at least one real err should be returned on errCh")
+			})
+			if !ok {
+				break
+			}
+			require.IsType(t, &exec.ExitError{}, err) // Because graceperiod is exceeded, we get an ExitError for killed
+			exitErr := err.(*exec.ExitError)
+			status := exitErr.ProcessState.Sys().(syscall.WaitStatus)
+			assert.Equal(t, status.Signal(), syscall.SIGKILL)
+			assert.Empty(t, errCh)
+		case <-time.After(graceperiod + 100*time.Millisecond):
+			require.Failf(t, "", "failed to stop within graceperiod(%s): %s", graceperiod, time.Since(cancelTime))
+			ok = false
+		}
+	}
+}
+
+func TestWithGracePeriodWithoutContext(t *testing.T) {
+	graceperiod := 1 * time.Second
+	errCh := exechelper.Start(
+		"sleep 600",
+		exechelper.WithGracePeriod(graceperiod),
+		exechelper.WithStdout(os.Stdout),
+	)
+	require.Error(t, <-errCh)
+}

options.go

@@ -22,6 +22,7 @@ import (
 	"os"
 	"os/exec"
 	"strings"
+	"time"
 
 	"github.com/pkg/errors"
 )
@@ -33,6 +34,8 @@ type CmdFunc func(cmd *exec.Cmd) error
 type Option struct {
 	// Context - context (if any) for running the exec.Cmd
 	Context context.Context
+	// SIGTERM grace period
+	GracePeriod time.Duration
 	// CmdFunc to be applied to the exec.Cmd
 	CmdOption CmdFunc
 }
@@ -155,3 +158,9 @@ func WithEnvMap(envMap map[string]string) *Option {
 	}
 	return WithEnvKV(envs...)
 }
+
+// WithGracePeriod - will send a SIGTERM when ctx.Done() and wait up to gracePeriod before
+// SIGKILLing the process.
+func WithGracePeriod(gracePeriod time.Duration) *Option {
+	return &Option{GracePeriod: gracePeriod}
+}