fix(kube-run): handle SIGINT/SIGTERM properly (#6893)

Signed-off-by: Alexandr Zaytsev <alexandr.zaytsev@flant.com>

Author: nervgh

cmd/werf/kube_run/kube_run.go

@@ -696,9 +696,6 @@ func copyFromPod(ctx context.Context, namespace, pod, container string, copyFrom
 	}
 
 	if err := cmd.Run(); err != nil {
-		if errors.Is(ctx.Err(), context.Canceled) {
-			graceful.Terminate(ctx, err, werfExec.ExitCode(err))
-		}
 		logboek.Context(ctx).Warn().LogF("Error copying %q from pod %s/s: %s\n", copyFrom.Src, namespace, pod, err)
 	}
 }
@@ -749,15 +746,33 @@ func stopContainer(ctx context.Context, namespace, pod, container string, extraA
 	}
 
 	if err := cmd.Run(); err != nil {
-		if errors.Is(ctx.Err(), context.Canceled) {
-			graceful.Terminate(ctx, err, werfExec.ExitCode(err))
-		}
 		logboek.Context(ctx).Warn().LogF("Error stopping service container %s/%s/%s for copying files: %s\n", namespace, pod, container, err)
 	}
 }
 
+func signalContainer(ctx context.Context, namespace, pod, container string, extraArgs []string) error {
+	ctx = context.WithoutCancel(ctx)
+
+	logboek.Context(ctx).LogF("Signal container %q in pod ...\n", container)
+
+	args := []string{
+		"exec", pod, "-q", "--pod-running-timeout", "5s", "-c", container,
+	}
+
+	args = append(args, extraArgs...)
+	args = append(args, "--", "pkill", "-P", "0")
+
+	cmd := werfExec.PrepareGracefulCancellation(util.ExecKubectlCmdContext(ctx, args...))
+
+	if err := cmd.Run(); err != nil {
+		return fmt.Errorf("signal container error: %w", err)
+	}
+
+	return nil
+}
+
 func execCommandInPod(ctx context.Context, namespace, pod, container string, command, extraArgs []string) error {
-	logboek.Context(ctx).LogF("Execing into pod ...\n")
+	logboek.Context(ctx).LogF("Executing into pod ...\n")
 
 	args := []string{
 		"exec", pod, "-q", "--pod-running-timeout", "5h", "-c", container,
@@ -777,7 +792,9 @@ func execCommandInPod(ctx context.Context, namespace, pod, container string, com
 	args = append(args, command...)
 
 	cmd := werfExec.PrepareGracefulCancellation(util.ExecKubectlCmdContext(ctx, args...))
-	cmd.Stdin = os.Stdin
+	cmd.Cancel = func() error {
+		return signalContainer(ctx, namespace, pod, container, extraArgs)
+	}
 
 	if *commonCmdData.DryRun {
 		fmt.Println(cmd.String())

test/e2e/kube-run/_fixtures/simple/state0/app.sh

@@ -0,0 +1,23 @@
+#!/bin/bash
+
+set -eu
+
+is_terminated=0
+
+log () {
+  echo "[$(date +"%T")] $1"
+}
+
+signal_handler () {
+  log "Signal handled"
+  is_terminated=1
+}
+
+trap signal_handler SIGINT SIGTERM EXIT
+
+while [ ${is_terminated} -eq 0 ]; do
+  log "Looping ..."
+  sleep 1
+done
+
+log "Script completed"

test/e2e/kube-run/_fixtures/simple/state0/werf.yaml

@@ -4,3 +4,6 @@ configVersion: 1
 ---
 image: main
 from: ubuntu:22.04
+git:
+  - add: /app.sh
+    to: /opt/app.sh

test/e2e/kube-run/simple_test.go

@@ -3,6 +3,7 @@ package e2e_kube_run_test
 import (
 	"os"
 	"strings"
+	"time"
 
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
@@ -25,6 +26,12 @@ var _ = Describe("Simple kube-run", Label("e2e", "kube-run", "simple"), func() {
 
 			By("state0: execute kube-run")
 			combinedOut := werfProject.KubeRun(ctx, kubeRunOpts)
+			Expect(combinedOut).To(ContainSubstring("Creating namespace"))
+			Expect(combinedOut).To(ContainSubstring("Running pod"))
+			Expect(combinedOut).To(ContainSubstring("Executing into pod"))
+			Expect(combinedOut).To(ContainSubstring("Stopping container"))
+			Expect(combinedOut).To(ContainSubstring("Cleaning up pod"))
+
 			expectOutFn(combinedOut)
 		},
 		Entry(
@@ -77,6 +84,25 @@ var _ = Describe("Simple kube-run", Label("e2e", "kube-run", "simple"), func() {
 				Expect(out).To(ContainSubstring("ID=ubuntu"))
 			},
 		),
+		Entry(
+			"should cancel long running process",
+			&werf.KubeRunOptions{
+				Command: []string{"/opt/app.sh"},
+				Image:   "main",
+				CommonOptions: werf.CommonOptions{
+					ShouldFail:            true,
+					ExtraArgs:             []string{}, // be able to work without "-it" options
+					CancelOnOutput:        "Looping ...",
+					CancelOnOutputTimeout: time.Minute,
+				},
+			},
+			func(out string) {
+				Expect(out).To(ContainSubstring("Signal container"))
+				Expect(out).To(ContainSubstring("Signal handled"))   // from script
+				Expect(out).To(ContainSubstring("Script completed")) // from script
+			},
+			SpecTimeout(time.Minute*3),
+		),
 	)
 })
 

test/pkg/utils/command.go

@@ -1,18 +1,22 @@
 package utils
 
 import (
+	"bufio"
 	"bytes"
 	"context"
 	"encoding/base64"
 	"fmt"
+	"io"
 	"os"
 	"os/exec"
 	"strings"
+	"time"
 
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
 
-	"github.com/werf/werf/v2/pkg/logging"
+	"github.com/werf/werf/v2/pkg/util/option"
+	werfExec "github.com/werf/werf/v2/pkg/werf/exec"
 )
 
 func RunCommand(ctx context.Context, dir, command string, args ...string) ([]byte, error) {
@@ -33,10 +37,14 @@ type RunCommandOptions struct {
 	ToStdin       string
 	ShouldSucceed bool
 	NoStderr      bool
+
+	CancelOnOutput        string
+	CancelOnOutputTimeout time.Duration
 }
 
 func RunCommandWithOptions(ctx context.Context, dir, command string, args []string, options RunCommandOptions) ([]byte, error) {
-	cmd := exec.CommandContext(logging.WithLogger(ctx), command, args...)
+	cmd := exec.CommandContext(ctx, command, args...)
+	cmd = werfExec.PrepareGracefulCancellation(cmd)
 
 	if dir != "" {
 		cmd.Dir = dir
@@ -45,29 +53,65 @@ func RunCommandWithOptions(ctx context.Context, dir, command string, args []stri
 	cmd.Env = append(os.Environ(), options.ExtraEnv...)
 
 	if options.ToStdin != "" {
-		var b bytes.Buffer
-		b.Write([]byte(options.ToStdin))
-		cmd.Stdin = &b
+		cmd.Stdin = bytes.NewReader([]byte(options.ToStdin))
 	}
 
-	var res []byte
-	var err error
+	stdout, err := cmd.StdoutPipe()
+	Expect(err).To(Succeed())
+
+	stderr, err := cmd.StderrPipe()
+	Expect(err).To(Succeed())
+
+	var outputReader io.Reader
+
 	if options.NoStderr {
-		res, err = cmd.Output()
+		outputReader = stdout
 	} else {
-		res, err = cmd.CombinedOutput()
+		outputReader = io.MultiReader(stdout, stderr)
 	}
 
-	_, _ = GinkgoWriter.Write(res)
+	res := &bytes.Buffer{}
+
+	Expect(cmd.Start()).To(Succeed())
+
+	if options.CancelOnOutput != "" {
+		copyReader := io.TeeReader(outputReader, res)
+		waitForOutput(copyReader, options.CancelOnOutput, options.CancelOnOutputTimeout)
+		Expect(cmd.Cancel()).To(Succeed())
+	}
+
+	_, err = io.Copy(res, outputReader)
+	Expect(err).To(Succeed())
+
+	err = cmd.Wait()
+
+	_, _ = GinkgoWriter.Write(res.Bytes())
 
 	if options.ShouldSucceed {
 		errorDesc := fmt.Sprintf("%[2]s %[3]s (dir: %[1]s)", dir, command, strings.Join(args, " "))
 		Expect(err).ShouldNot(HaveOccurred(), errorDesc)
 	}
 
-	return res, err
+	return res.Bytes(), err
 }
 
 func ShelloutPack(command string) string {
 	return fmt.Sprintf("eval $(echo %s | base64 -d)", base64.StdEncoding.EncodeToString([]byte(command)))
 }
+
+// waitForOutput waits for output and exits early or exits by timeout
+func waitForOutput(reader io.Reader, output string, timeout time.Duration) {
+	scanner := bufio.NewScanner(reader)
+	tmr := time.NewTimer(option.ValueOrDefault(timeout, time.Minute))
+
+	for scanner.Scan() {
+		select {
+		case <-tmr.C:
+			return
+		default:
+			if strings.Contains(scanner.Text(), output) {
+				return
+			}
+		}
+	}
+}

test/pkg/werf/project.go

@@ -9,6 +9,7 @@ import (
 	"path/filepath"
 	"strings"
 	"sync"
+	"time"
 
 	. "github.com/onsi/gomega"
 
@@ -36,6 +37,9 @@ type Project struct {
 type CommonOptions struct {
 	ShouldFail bool
 	ExtraArgs  []string
+
+	CancelOnOutput        string
+	CancelOnOutputTimeout time.Duration
 }
 
 type BuildOptions struct {
@@ -87,6 +91,9 @@ type KubeCtlOptions struct {
 type runCommandOptions struct {
 	ShouldFail bool
 	Args       []string
+
+	CancelOnOutput        string
+	CancelOnOutputTimeout time.Duration
 }
 
 func (p *Project) Build(ctx context.Context, opts *BuildOptions) (combinedOut string) {
@@ -217,7 +224,12 @@ func (p *Project) KubeRun(ctx context.Context, opts *KubeRunOptions) string {
 		args = append(args, opts.Command...)
 	}
 
-	return p.runCommand(ctx, runCommandOptions{Args: args, ShouldFail: opts.ShouldFail})
+	return p.runCommand(ctx, runCommandOptions{
+		Args:                  args,
+		ShouldFail:            opts.ShouldFail,
+		CancelOnOutput:        opts.CancelOnOutput,
+		CancelOnOutputTimeout: opts.CancelOnOutputTimeout,
+	})
 }
 
 func (p *Project) KubeCtl(ctx context.Context, opts *KubeCtlOptions) string {
@@ -287,13 +299,11 @@ func (p *Project) CreateRegistryPullSecretFromDockerConfig(ctx context.Context)
 }
 
 func (p *Project) runCommand(ctx context.Context, opts runCommandOptions) string {
-	outb, err := iutils.RunCommand(ctx, p.GitRepoPath, p.WerfBinPath, opts.Args...)
-	if opts.ShouldFail {
-		Expect(err).To(HaveOccurred())
-	} else {
-		Expect(err).NotTo(HaveOccurred())
-	}
-
+	outb, _ := iutils.RunCommandWithOptions(ctx, p.GitRepoPath, p.WerfBinPath, opts.Args, iutils.RunCommandOptions{
+		ShouldSucceed:         !opts.ShouldFail,
+		CancelOnOutput:        opts.CancelOnOutput,
+		CancelOnOutputTimeout: opts.CancelOnOutputTimeout,
+	})
 	return string(outb)
 }