Forcefully terminate all processes with a configured shutdown timeout

src/app/process.go

@@ -73,6 +73,8 @@ type Process struct {
 	stdin               io.WriteCloser
 	passProvided        bool
 	isTuiEnabled        bool
+	canceller           context.CancelFunc
+	waitDone            chan struct{}
 }
 
 func NewProcess(opts ...ProcOpts) *Process {
@@ -85,6 +87,7 @@ func NewProcess(opts ...ProcOpts) *Process {
 		started:       false,
 		done:          false,
 		procStateChan: make(chan string, 1),
+		waitDone:      make(chan struct{}, 1),
 	}
 
 	for _, opt := range opts {
@@ -113,6 +116,10 @@ func (p *Process) run() int {
 
 	p.onProcessStart()
 	for {
+		if len(p.waitDone) > 0 {
+			<-p.waitDone
+		}
+
 		err := p.setStateAndRun(p.getStartingStateName(), p.getProcessStarter())
 		if err != nil {
 			log.Error().Err(err).Msgf(`Failed to run command ["%v"] for process %s`, strings.Join(p.getCommand(), `" "`), p.getName())
@@ -137,6 +144,8 @@ func (p *Process) run() int {
 		//TODO Fix this
 		time.Sleep(50 * time.Millisecond)
 		_ = p.command.Wait()
+		p.waitDone <- struct{}{}
+
 		p.Lock()
 		p.setExitCode(p.command.ExitCode())
 		p.Unlock()
@@ -202,18 +211,50 @@ func (p *Process) getProcessStarter() func() error {
 }
 
 func (p *Process) getCommander() command.Commander {
+	ctx, canceller := context.WithCancel(context.Background())
+	p.canceller = canceller
+	onCancel := func() error {
+		err := p.command.Stop(
+			p.procConf.ShutDownParams.Signal,
+			p.procConf.ShutDownParams.ParentOnly,
+		)
+
+		if err != nil {
+			return err
+		}
+
+		timeoutInt := p.procConf.ShutDownParams.ShutDownTimeout
+		if timeoutInt != UndefinedShutdownTimeoutSec {
+			timeout := time.Duration(timeoutInt) * time.Second
+			select {
+			case <-p.waitDone:
+				break
+			case <-time.After(timeout):
+				return p.command.Stop(
+					int(syscall.SIGKILL),
+					p.procConf.ShutDownParams.ParentOnly,
+				)
+			}
+		}
+
+		return nil
+	}
+
 	if p.procConf.IsTty && !p.isMain {
-		return command.BuildPtyCommand(
+		return command.BuildPtyCommandContext(
+			ctx,
+			onCancel,
 			p.procConf.Executable,
 			p.mergeExtraArgs(),
 		)
 	} else {
-		return command.BuildCommand(
+		return command.BuildCommandContext(
+			ctx,
+			onCancel,
 			p.procConf.Executable,
 			p.mergeExtraArgs(),
 		)
 	}
-
 }
 
 func (p *Process) mergeExtraArgs() []string {
@@ -364,8 +405,8 @@ func (p *Process) stopProcess(cancelReadinessFuncs bool) error {
 	if isStringDefined(p.procConf.ShutDownParams.ShutDownCommand) {
 		return p.doConfiguredStop(p.procConf.ShutDownParams)
 	}
-
-	return p.command.Stop(p.procConf.ShutDownParams.Signal, p.procConf.ShutDownParams.ParentOnly)
+	p.canceller()
+	return nil
 }
 
 func (p *Process) doConfiguredStop(params types.ShutDownParams) error {

src/app/system_unix_test.go

@@ -0,0 +1,69 @@
+//go:build !windows
+
+package app
+
+import (
+	"github.com/f1bonacc1/process-compose/src/command"
+	"github.com/f1bonacc1/process-compose/src/types"
+	"syscall"
+	"testing"
+	"time"
+)
+
+func assertProcessStatus(t *testing.T, runner *ProjectRunner, procName string, wantStatus string) {
+	t.Helper()
+	state, err := runner.GetProcessState(procName)
+	if err != nil {
+		t.Fatalf("%s", err)
+	}
+	if state.Status != wantStatus {
+		t.Fatalf("process %s status want %s got %s", procName, wantStatus, state.Status)
+	}
+}
+
+func TestSystem_TestProcShutDownWithConfiguredTimeOut(t *testing.T) {
+	ignoresSigTerm := "IgnoresSIGTERM"
+	shell := command.DefaultShellConfig()
+	timeout := 5
+
+	project := &types.Project{
+		Processes: map[string]types.ProcessConfig{
+			ignoresSigTerm: {
+				Name:        ignoresSigTerm,
+				ReplicaName: ignoresSigTerm,
+				Executable:  shell.ShellCommand,
+				Args:        []string{shell.ShellArgument, "trap '' SIGTERM && sleep 60"},
+				ShutDownParams: types.ShutDownParams{
+					ShutDownTimeout: timeout,
+					Signal:          int(syscall.SIGTERM),
+				},
+			},
+		},
+		ShellConfig: shell,
+	}
+	runner, err := NewProjectRunner(&ProjectOpts{project: project})
+	if err != nil {
+		t.Fatalf("%s", err)
+	}
+	go runner.Run()
+
+	time.Sleep(100 * time.Millisecond)
+	assertProcessStatus(t, runner, ignoresSigTerm, types.ProcessStateRunning)
+
+	// If the test fails, cleanup after ourselves
+	proc := runner.getRunningProcess(ignoresSigTerm)
+	defer proc.command.Stop(int(syscall.SIGKILL), true)
+
+	err = runner.StopProcess(ignoresSigTerm)
+	if err != nil {
+		t.Fatalf("%s", err)
+	}
+
+	for i := 0; i < timeout-1; i++ {
+		time.Sleep(time.Second)
+		assertProcessStatus(t, runner, ignoresSigTerm, types.ProcessStateTerminating)
+	}
+
+	time.Sleep(2 * time.Second)
+	assertProcessStatus(t, runner, ignoresSigTerm, types.ProcessStateCompleted)
+}

src/command/command.go

@@ -8,21 +8,15 @@ import (
 	"runtime"
 )
 
-func BuildCommand(cmd string, args []string) *CmdWrapper {
-	return &CmdWrapper{
-		cmd: exec.Command(cmd, args...),
-	}
+func BuildCommandContext(ctx context.Context, onCancel func() error, cmd string, args []string) *CmdWrapper {
+	cmdCtx := exec.CommandContext(ctx, cmd, args...)
+	cmdCtx.Cancel = onCancel
+	return &CmdWrapper{cmd: cmdCtx}
 }
 
-func BuildPtyCommand(cmd string, args []string) *CmdWrapperPty {
+func BuildPtyCommandContext(ctx context.Context, onCancel func() error, cmd string, args []string) *CmdWrapperPty {
 	return &CmdWrapperPty{
-		CmdWrapper: BuildCommand(cmd, args),
-	}
-}
-
-func BuildCommandContext(ctx context.Context, shellCmd string) *CmdWrapper {
-	return &CmdWrapper{
-		cmd: exec.CommandContext(ctx, getRunnerShell(), getRunnerArg(), shellCmd),
+		CmdWrapper: BuildCommandContext(ctx, onCancel, cmd, args),
 	}
 }
 

src/health/exec_checker.go

@@ -16,7 +16,11 @@ func (c *execChecker) Status() (interface{}, error) {
 	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(c.timeout)*time.Second)
 	defer cancel()
 
-	cmd := command.BuildCommandContext(ctx, c.command)
+	cmd := command.BuildCommandShellArgContext(
+		ctx,
+		*command.DefaultShellConfig(),
+		c.command,
+	)
 	cmd.SetDir(c.workingDir)
 
 	if err := cmd.Run(); err != nil {