connection handling / block controller handling (#326)

pkg/shellexec/conninterface.go

@@ -2,7 +2,9 @@ package shellexec
 
 import (
 	"io"
+	"os"
 	"os/exec"
+	"time"
 
 	"github.com/creack/pty"
 	"golang.org/x/crypto/ssh"
@@ -10,6 +12,7 @@ import (
 
 type ConnInterface interface {
 	Kill()
+	KillGraceful(time.Duration)
 	Wait() error
 	Start() error
 	StdinPipe() (io.WriteCloser, error)
@@ -32,6 +35,22 @@ func (cw CmdWrap) Wait() error {
 	return cw.Cmd.Wait()
 }
 
+func (cw CmdWrap) KillGraceful(timeout time.Duration) {
+	if cw.Cmd.Process == nil {
+		return
+	}
+	if cw.Cmd.ProcessState != nil && cw.Cmd.ProcessState.Exited() {
+		return
+	}
+	cw.Cmd.Process.Signal(os.Interrupt)
+	go func() {
+		time.Sleep(timeout)
+		if cw.Cmd.ProcessState == nil || !cw.Cmd.ProcessState.Exited() {
+			cw.Cmd.Process.Kill() // force kill if it is already not exited
+		}
+	}()
+}
+
 func (cw CmdWrap) Start() error {
 	defer func() {
 		for _, extraFile := range cw.Cmd.ExtraFiles {
@@ -75,6 +94,10 @@ func (sw SessionWrap) Kill() {
 	sw.Session.Close()
 }
 
+func (sw SessionWrap) KillGraceful(timeout time.Duration) {
+	sw.Kill()
+}
+
 func (sw SessionWrap) Wait() error {
 	return sw.Session.Wait()
 }

pkg/shellexec/shellexec.go

@@ -15,16 +15,19 @@ import (
 	"strings"
 	"sync"
 	"syscall"
+	"time"
 
 	"github.com/creack/pty"
 	"github.com/wavetermdev/thenextwave/pkg/remote"
+	"github.com/wavetermdev/thenextwave/pkg/remote/conncontroller"
 	"github.com/wavetermdev/thenextwave/pkg/util/shellutil"
 	"github.com/wavetermdev/thenextwave/pkg/wavebase"
 	"github.com/wavetermdev/thenextwave/pkg/waveobj"
 	"github.com/wavetermdev/thenextwave/pkg/wshutil"
-	"golang.org/x/crypto/ssh"
 )
 
+const DefaultGracefulKillWait = 400 * time.Millisecond
+
 type CommandOptsType struct {
 	Interactive bool              `json:"interactive,omitempty"`
 	Login       bool              `json:"login,omitempty"`
@@ -33,6 +36,7 @@ type CommandOptsType struct {
 }
 
 type ShellProc struct {
+	ConnName  string
 	Cmd       ConnInterface
 	CloseOnce *sync.Once
 	DoneCh    chan any // closed after proc.Wait() returns
@@ -40,7 +44,7 @@ type ShellProc struct {
 }
 
 func (sp *ShellProc) Close() {
-	sp.Cmd.Kill()
+	sp.Cmd.KillGraceful(DefaultGracefulKillWait)
 	go func() {
 		waitErr := sp.Cmd.Wait()
 		sp.SetWaitErrorAndSignalDone(waitErr)
@@ -134,7 +138,8 @@ func (pp *PipePty) WriteString(s string) (n int, err error) {
 	return pp.Write([]byte(s))
 }
 
-func StartRemoteShellProc(termSize waveobj.TermSize, cmdStr string, cmdOpts CommandOptsType, client *ssh.Client) (*ShellProc, error) {
+func StartRemoteShellProc(termSize waveobj.TermSize, cmdStr string, cmdOpts CommandOptsType, conn *conncontroller.SSHConn) (*ShellProc, error) {
+	client := conn.GetClient()
 	shellPath, err := remote.DetectShell(client)
 	if err != nil {
 		return nil, err
@@ -244,7 +249,7 @@ func StartRemoteShellProc(termSize waveobj.TermSize, cmdStr string, cmdOpts Comm
 		pipePty.Close()
 		return nil, err
 	}
-	return &ShellProc{Cmd: sessionWrap, CloseOnce: &sync.Once{}, DoneCh: make(chan any)}, nil
+	return &ShellProc{Cmd: sessionWrap, ConnName: conn.GetName(), CloseOnce: &sync.Once{}, DoneCh: make(chan any)}, nil
 }
 
 func isZshShell(shellPath string) bool {