From 7abd809d7b49980a1a2b9db9b4c9eaa0a27983e6 Mon Sep 17 00:00:00 2001
From: Aniruddh Agarwal <me@anrddh.me>
Date: Mon, 19 Aug 2024 20:58:35 +0530
Subject: [PATCH] feat: SIGKILL procs without kill command afer timeout

We start all procs as CommandContexts instead of Commands, with a
WithCancel context attached. A context cancellation handler which tries
to stop the application gracefully, waits the specified amount of time,
and subsequently SIGKILLs if the process hasn't stopped already is
attached to the CommandContext.

The CommandContext type internally has some synchronization to try to
minimize the chances of races between the context cancellation handler,
and normal termination of applications.
---
 src/app/process.go     | 51 +++++++++++++++++++++++++++++++++++++-----
 src/command/command.go | 12 +++++-----
 2 files changed, 52 insertions(+), 11 deletions(-)

diff --git a/src/app/process.go b/src/app/process.go
index c3e46e4..8f96ad5 100644
--- a/src/app/process.go
+++ b/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 {
diff --git a/src/command/command.go b/src/command/command.go
index 75555ff..c7ccb83 100644
--- a/src/command/command.go
+++ b/src/command/command.go
@@ -8,15 +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),
+		CmdWrapper: BuildCommandContext(ctx, onCancel, cmd, args),
 	}
 }