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>

main.go

@@ -1,7 +1,6 @@
 package main
 
 import (
-	"fmt"
 	"os"
 	"os/signal"
 	"syscall"
@@ -15,19 +14,25 @@ import (
 func main() {
 	var sig os.Signal
 	sigs := make(chan os.Signal, 1)
+	errChan := make(chan error, 1)
+
 	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
 
 	go func() {
-		sig = <-sigs
+		globalConfig := new(config.GlobalOptions)
+		rootCmd, err := cmd.NewRootCmd(globalConfig)
+		if err != nil {
+			errChan <- err
+			return
+		}
+
+		errChan <- rootCmd.Execute()
 	}()
 
-	globalConfig := new(config.GlobalOptions)
-	rootCmd, err := cmd.NewRootCmd(globalConfig)
-	errors.HandleExitCoder(err)
-
-	if err := rootCmd.Execute(); err != nil {
+	select {
+	case sig = <-sigs:
 		if sig != nil {
-			fmt.Fprintln(os.Stderr, err)
+			app.Cancel()
 			app.CleanWaitGroup.Wait()
 
 			// See http://tldp.org/LDP/abs/html/exitcodes.html
@@ -38,6 +43,7 @@ func main() {
 				os.Exit(143)
 			}
 		}
+	case err := <-errChan:
 		errors.HandleExitCoder(err)
 	}
 }

pkg/app/app.go

@@ -2,6 +2,7 @@ package app
 
 import (
 	"bytes"
+	goContext "context"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -23,6 +24,7 @@ import (
 )
 
 var CleanWaitGroup sync.WaitGroup
+var Cancel goContext.CancelFunc
 
 // App is the main application object.
 type App struct {
@@ -50,6 +52,8 @@ type App struct {
 
 	helms      map[helmKey]helmexec.Interface
 	helmsMutex sync.Mutex
+
+	ctx goContext.Context
 }
 
 type HelmRelease struct {
@@ -63,6 +67,9 @@ type HelmRelease struct {
 }
 
 func New(conf ConfigProvider) *App {
+	ctx := goContext.Background()
+	ctx, Cancel = goContext.WithCancel(ctx)
+
 	return Init(&App{
 		OverrideKubeContext: conf.KubeContext(),
 		OverrideHelmBinary:  conf.HelmBinary(),
@@ -78,6 +85,7 @@ func New(conf ConfigProvider) *App {
 		ValuesFiles:         conf.StateValuesFiles(),
 		Set:                 conf.StateValuesSet(),
 		fs:                  filesystem.DefaultFileSystem(),
+		ctx:                 ctx,
 	})
 }
 
@@ -98,6 +106,7 @@ func Init(app *App) *App {
 func (a *App) Init(c InitConfigProvider) error {
 	runner := &helmexec.ShellRunner{
 		Logger: a.Logger,
+		Ctx:    a.ctx,
 	}
 	helmfileInit := NewHelmfileInit(a.OverrideHelmBinary, c, a.Logger, runner)
 	return helmfileInit.Initialize()
@@ -788,6 +797,7 @@ func (a *App) getHelm(st *state.HelmState) helmexec.Interface {
 	if _, ok := a.helms[key]; !ok {
 		a.helms[key] = helmexec.New(bin, helmexec.HelmExecOptions{EnableLiveOutput: a.EnableLiveOutput, DisableForceUpdate: a.DisableForceUpdate}, a.Logger, kubectx, &helmexec.ShellRunner{
 			Logger: a.Logger,
+			Ctx:    a.ctx,
 		})
 	}
 

pkg/event/bus.go

@@ -1,6 +1,7 @@
 package event
 
 import (
+	goContext "context"
 	"fmt"
 	"strings"
 
@@ -46,6 +47,9 @@ func (bus *Bus) Trigger(evt string, evtErr error, context map[string]interface{}
 		bus.Runner = helmexec.ShellRunner{
 			Dir:    bus.BasePath,
 			Logger: bus.Logger,
+			// It would be better to pass app.Ctx here, but it requires a lot of work.
+			// It seems that this code only for running hooks, which took not to long time as helm.
+			Ctx: goContext.TODO(),
 		}
 	}
 

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

pkg/tmpl/context_funcs.go

@@ -1,6 +1,7 @@
 package tmpl
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -174,7 +175,7 @@ func (c *Context) EnvExec(envs map[string]interface{}, command string, args []in
 	g.Go(func() error {
 		// We use CombinedOutput to produce helpful error messages
 		// See https://github.com/roboll/helmfile/issues/1158
-		bs, err := helmexec.Output(cmd)
+		bs, err := helmexec.Output(context.Background(), cmd)
 		if err != nil {
 			return err
 		}

test/e2e/template/helmfile/snapshot_test.go

@@ -75,6 +75,7 @@ func testHelmfileTemplateWithBuildCommand(t *testing.T, goccyGoYaml bool) {
 	logger := helmexec.NewLogger(os.Stderr, "info")
 	runner := &helmexec.ShellRunner{
 		Logger: logger,
+		Ctx:    context.TODO(),
 	}
 
 	c := fakeInit{}