Make helmfile respect signals send by kill command (not only Ctrl+C in terminal) (#750)

Fixes #746 

Signed-off-by: Dmitry Chepurovskiy <me@dm3ch.net>
Signed-off-by: yxxhero <aiopsclub@163.com>
Co-authored-by: yxxhero <aiopsclub@163.com>

pkg/helmexec/exec_test.go

@@ -31,7 +31,6 @@ func (mock *mockRunner) Execute(cmd string, args []string, env map[string]string
 	if len(mock.output) == 0 && strings.Join(args, " ") == "version --client --short" {
 		return []byte("v3.2.4+ge29ce2a"), nil
 	}
-
 	return mock.output, mock.err
 }
 

pkg/helmexec/runner.go

@@ -3,6 +3,7 @@ package helmexec
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -28,6 +29,7 @@ type ShellRunner struct {
 	Dir string
 
 	Logger *zap.SugaredLogger
+	Ctx    context.Context
 }
 
 // Execute a shell command
@@ -37,11 +39,11 @@ func (shell ShellRunner) Execute(cmd string, args []string, env map[string]strin
 	preparedCmd.Env = mergeEnv(os.Environ(), env)
 
 	if !enableLiveOutput {
-		return Output(preparedCmd, &logWriterGenerator{
+		return Output(shell.Ctx, preparedCmd, &logWriterGenerator{
 			log: shell.Logger,
 		})
 	} else {
-		return LiveOutput(preparedCmd, os.Stdout)
+		return LiveOutput(shell.Ctx, preparedCmd, os.Stdout)
 	}
 }
 
@@ -51,12 +53,12 @@ func (shell ShellRunner) ExecuteStdIn(cmd string, args []string, env map[string]
 	preparedCmd.Dir = shell.Dir
 	preparedCmd.Env = mergeEnv(os.Environ(), env)
 	preparedCmd.Stdin = stdin
-	return Output(preparedCmd, &logWriterGenerator{
+	return Output(shell.Ctx, preparedCmd, &logWriterGenerator{
 		log: shell.Logger,
 	})
 }
 
-func Output(c *exec.Cmd, logWriterGenerators ...*logWriterGenerator) ([]byte, error) {
+func Output(ctx context.Context, c *exec.Cmd, logWriterGenerators ...*logWriterGenerator) ([]byte, error) {
 	if c.Stdout != nil {
 		return nil, errors.New("exec: Stdout already set")
 	}
@@ -88,7 +90,17 @@ func Output(c *exec.Cmd, logWriterGenerators ...*logWriterGenerator) ([]byte, er
 	c.Stdout = io.MultiWriter(append([]io.Writer{&stdout, &combined}, logWriters...)...)
 	c.Stderr = io.MultiWriter(append([]io.Writer{&stderr, &combined}, logWriters...)...)
 
-	err := c.Run()
+	var err error
+	ch := make(chan error)
+	go func() {
+		ch <- c.Run()
+	}()
+	select {
+	case err = <-ch:
+	case <-ctx.Done():
+		_ = c.Process.Signal(os.Interrupt)
+		err = <-ch
+	}
 
 	if err != nil {
 		// TrimSpace is necessary, because otherwise helmfile prints the redundant new-lines after each error like:
@@ -111,7 +123,7 @@ func Output(c *exec.Cmd, logWriterGenerators ...*logWriterGenerator) ([]byte, er
 	return stdout.Bytes(), err
 }
 
-func LiveOutput(c *exec.Cmd, stdout io.Writer) ([]byte, error) {
+func LiveOutput(ctx context.Context, c *exec.Cmd, stdout io.Writer) ([]byte, error) {
 	reader, writer := io.Pipe()
 	scannerStopped := make(chan struct{})
 
@@ -126,8 +138,19 @@ func LiveOutput(c *exec.Cmd, stdout io.Writer) ([]byte, error) {
 	c.Stdout = writer
 	c.Stderr = writer
 	err := c.Start()
+
 	if err == nil {
-		err = c.Wait()
+		ch := make(chan error)
+		go func() {
+			ch <- c.Wait()
+		}()
+
+		select {
+		case err = <-ch:
+		case <-ctx.Done():
+			_ = c.Process.Signal(os.Interrupt)
+			err = <-ch
+		}
 		_ = writer.Close()
 		<-scannerStopped
 	}

pkg/helmexec/runner_test.go

@@ -2,6 +2,7 @@ package helmexec
 
 import (
 	"bytes"
+	"context"
 	"os/exec"
 	"reflect"
 	"strings"
@@ -32,6 +33,7 @@ func TestShellRunner_Execute(t *testing.T) {
 			var buffer bytes.Buffer
 			shell := ShellRunner{
 				Logger: NewLogger(&buffer, "debug"),
+				Ctx:    context.TODO(),
 			}
 			got, err := shell.Execute("echo", strings.Split("template", " "), map[string]string{}, tt.enableLiveOutput)
 
@@ -72,7 +74,7 @@ Usage:  helm template [NAME] [CHART] [flags]
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			w := &bytes.Buffer{}
-			got, err := LiveOutput(tt.cmd, w)
+			got, err := LiveOutput(context.Background(), tt.cmd, w)
 			if (err != nil) != tt.wantErr {
 				t.Errorf("LiveOutput() error = %v, wantErr %v", err, tt.wantErr)
 				return