feat(native): restart jetkvm_native automatically (#629)

Author: ym

native.go

@@ -8,6 +8,7 @@ import (
 	"io"
 	"net"
 	"os"
+	"os/exec"
 	"sync"
 	"time"
 
@@ -41,6 +42,11 @@ var ongoingRequests = make(map[int32]chan *CtrlResponse)
 
 var lock = &sync.Mutex{}
 
+var (
+	nativeCmd     *exec.Cmd
+	nativeCmdLock = &sync.Mutex{}
+)
+
 func CallCtrlAction(action string, params map[string]interface{}) (*CtrlResponse, error) {
 	lock.Lock()
 	defer lock.Unlock()
@@ -129,16 +135,26 @@ func StartNativeSocketServer(socketPath string, handleClient func(net.Conn), isC
 	scopedLogger.Info().Msg("server listening")
 
 	go func() {
-		conn, err := listener.Accept()
-		listener.Close()
-		if err != nil {
-			scopedLogger.Warn().Err(err).Msg("failed to accept socket")
-		}
-		if isCtrl {
-			close(ctrlClientConnected)
-			scopedLogger.Debug().Msg("first native ctrl socket client connected")
+		for {
+			conn, err := listener.Accept()
+
+			if err != nil {
+				scopedLogger.Warn().Err(err).Msg("failed to accept socket")
+				continue
+			}
+			if isCtrl {
+				// check if the channel is closed
+				select {
+				case <-ctrlClientConnected:
+					scopedLogger.Debug().Msg("ctrl client reconnected")
+				default:
+					close(ctrlClientConnected)
+					scopedLogger.Debug().Msg("first native ctrl socket client connected")
+				}
+			}
+
+			go handleClient(conn)
 		}
-		handleClient(conn)
 	}()
 
 	return listener
@@ -235,6 +251,51 @@ func handleVideoClient(conn net.Conn) {
 	}
 }
 
+func startNativeBinaryWithLock(binaryPath string) (*exec.Cmd, error) {
+	nativeCmdLock.Lock()
+	defer nativeCmdLock.Unlock()
+
+	cmd, err := startNativeBinary(binaryPath)
+	if err != nil {
+		return nil, err
+	}
+	nativeCmd = cmd
+	return cmd, nil
+}
+
+func restartNativeBinary(binaryPath string) error {
+	time.Sleep(10 * time.Second)
+	// restart the binary
+	nativeLogger.Info().Msg("restarting jetkvm_native binary")
+	cmd, err := startNativeBinary(binaryPath)
+	if err != nil {
+		nativeLogger.Warn().Err(err).Msg("failed to restart binary")
+	}
+	nativeCmd = cmd
+	return err
+}
+
+func superviseNativeBinary(binaryPath string) error {
+	nativeCmdLock.Lock()
+	defer nativeCmdLock.Unlock()
+
+	if nativeCmd == nil || nativeCmd.Process == nil {
+		return restartNativeBinary(binaryPath)
+	}
+
+	err := nativeCmd.Wait()
+
+	if err == nil {
+		nativeLogger.Info().Err(err).Msg("jetkvm_native binary exited with no error")
+	} else if exiterr, ok := err.(*exec.ExitError); ok {
+		nativeLogger.Warn().Int("exit_code", exiterr.ExitCode()).Msg("jetkvm_native binary exited with error")
+	} else {
+		nativeLogger.Warn().Err(err).Msg("jetkvm_native binary exited with unknown error")
+	}
+
+	return restartNativeBinary(binaryPath)
+}
+
 func ExtractAndRunNativeBin() error {
 	binaryPath := "/userdata/jetkvm/bin/jetkvm_native"
 	if err := ensureBinaryUpdated(binaryPath); err != nil {
@@ -246,12 +307,28 @@ func ExtractAndRunNativeBin() error {
 		return fmt.Errorf("failed to make binary executable: %w", err)
 	}
 	// Run the binary in the background
-	cmd, err := startNativeBinary(binaryPath)
+	cmd, err := startNativeBinaryWithLock(binaryPath)
 	if err != nil {
 		return fmt.Errorf("failed to start binary: %w", err)
 	}
 
-	//TODO: add auto restart
+	// check if the binary is still running every 10 seconds
+	go func() {
+		for {
+			select {
+			case <-appCtx.Done():
+				nativeLogger.Info().Msg("stopping native binary supervisor")
+				return
+			default:
+				err := superviseNativeBinary(binaryPath)
+				if err != nil {
+					nativeLogger.Warn().Err(err).Msg("failed to supervise native binary")
+					time.Sleep(1 * time.Second) // Add a short delay to prevent rapid successive calls
+				}
+			}
+		}
+	}()
+
 	go func() {
 		<-appCtx.Done()
 		nativeLogger.Info().Int("pid", cmd.Process.Pid).Msg("killing process")