From aa5be82834563bfaa650a4d6e7535c19b873cac3 Mon Sep 17 00:00:00 2001 From: Dmitry Chepurovskiy Date: Sat, 29 Apr 2023 09:25:29 +0300 Subject: [PATCH] Make helmfile respect signals send by kill command (not only Ctrl+C in terminal) (#750) Fixes #746 Signed-off-by: Dmitry Chepurovskiy Signed-off-by: yxxhero Co-authored-by: yxxhero --- main.go | 22 +++++++----- pkg/app/app.go | 10 ++++++ pkg/event/bus.go | 4 +++ pkg/helmexec/exec_test.go | 1 - pkg/helmexec/runner.go | 37 +++++++++++++++++---- pkg/helmexec/runner_test.go | 4 ++- pkg/tmpl/context_funcs.go | 3 +- test/e2e/template/helmfile/snapshot_test.go | 1 + 8 files changed, 64 insertions(+), 18 deletions(-) diff --git a/main.go b/main.go index 7b66573f..0197027a 100644 --- a/main.go +++ b/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) } } diff --git a/pkg/app/app.go b/pkg/app/app.go index d4c10127..a3b22b10 100644 --- a/pkg/app/app.go +++ b/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, }) } diff --git a/pkg/event/bus.go b/pkg/event/bus.go index 92b69bff..883acb4b 100644 --- a/pkg/event/bus.go +++ b/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(), } } diff --git a/pkg/helmexec/exec_test.go b/pkg/helmexec/exec_test.go index 04a55b07..7d66f1a4 100644 --- a/pkg/helmexec/exec_test.go +++ b/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 } diff --git a/pkg/helmexec/runner.go b/pkg/helmexec/runner.go index 9aca1cef..08be29c4 100644 --- a/pkg/helmexec/runner.go +++ b/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 } diff --git a/pkg/helmexec/runner_test.go b/pkg/helmexec/runner_test.go index d174c893..220b0325 100644 --- a/pkg/helmexec/runner_test.go +++ b/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 diff --git a/pkg/tmpl/context_funcs.go b/pkg/tmpl/context_funcs.go index f63adc48..c4e5adc5 100644 --- a/pkg/tmpl/context_funcs.go +++ b/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 } diff --git a/test/e2e/template/helmfile/snapshot_test.go b/test/e2e/template/helmfile/snapshot_test.go index be449c94..dd1c5fa5 100644 --- a/test/e2e/template/helmfile/snapshot_test.go +++ b/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{}