Merge pull request #3859 from unsuman/fix/ext-driver-shutdown-issue

Fix: Server Shutdown Bug

Author: AkihiroSuda

pkg/driver/external/client/methods.go

@@ -93,9 +93,7 @@ func (d *DriverClient) Start(ctx context.Context) (chan error, error) {
 func (d *DriverClient) Stop(ctx context.Context) error {
 	d.logger.Debug("Stopping driver instance")
 
-	connCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
-	defer cancel()
-	_, err := d.DriverSvc.Stop(connCtx, &emptypb.Empty{})
+	_, err := d.DriverSvc.Stop(ctx, &emptypb.Empty{})
 	if err != nil {
 		d.logger.Errorf("Failed to stop driver instance: %v", err)
 		return err

pkg/driver/external/server/server.go

@@ -14,6 +14,7 @@ import (
 	"os/exec"
 	"os/signal"
 	"path/filepath"
+	"sync"
 	"syscall"
 	"time"
 
@@ -35,6 +36,20 @@ type DriverServer struct {
 	logger *logrus.Logger
 }
 
+type listenerTracker struct {
+	net.Listener
+	connected chan struct{}
+	once      sync.Once
+}
+
+func (t *listenerTracker) Accept() (net.Conn, error) {
+	c, err := t.Listener.Accept()
+	if err == nil {
+		t.once.Do(func() { close(t.connected) })
+	}
+	return c, err
+}
+
 func Serve(driver driver.Driver) {
 	logger := logrus.New()
 	logger.SetLevel(logrus.DebugLevel)
@@ -57,6 +72,11 @@ func Serve(driver driver.Driver) {
 	}
 	defer listener.Close()
 
+	tListener := &listenerTracker{
+		Listener:  listener,
+		connected: make(chan struct{}),
+	}
+
 	output := map[string]string{"socketPath": socketPath}
 	if err := json.NewEncoder(os.Stdout).Encode(output); err != nil {
 		logger.Fatalf("Failed to encode socket path as JSON: %v", err)
@@ -86,21 +106,26 @@ func Serve(driver driver.Driver) {
 	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
 
 	shutdownCh := make(chan struct{})
+	var closeOnce sync.Once
+	closeShutdown := func() { closeOnce.Do(func() { close(shutdownCh) }) }
 
 	go func() {
 		<-sigs
 		logger.Info("Received shutdown signal, stopping server...")
-		close(shutdownCh)
+		closeShutdown()
 	}()
 
 	go func() {
 		timer := time.NewTimer(60 * time.Second)
 		defer timer.Stop()
 
 		select {
+		case <-tListener.connected:
+			logger.Debug("Client connected; disabling 60s startup shutdown")
+			return
 		case <-timer.C:
 			logger.Info("No client connected within 60 seconds, shutting down server...")
-			close(shutdownCh)
+			closeShutdown()
 		case <-shutdownCh:
 			return
 		}
@@ -109,7 +134,7 @@ func Serve(driver driver.Driver) {
 	go func() {
 		logger.Infof("Starting external driver server for %s", driver.Info().DriverName)
 		logger.Infof("Server starting on Unix socket: %s", socketPath)
-		if err := server.Serve(listener); err != nil {
+		if err := server.Serve(tListener); err != nil {
 			if errors.Is(err, grpc.ErrServerStopped) {
 				logger.Errorf("Server stopped: %v", err)
 			} else {

pkg/driver/qemu/qemu_driver.go

@@ -96,8 +96,8 @@ func (l *LimaQemuDriver) CreateDisk(ctx context.Context) error {
 	return EnsureDisk(ctx, qCfg)
 }
 
-func (l *LimaQemuDriver) Start(ctx context.Context) (chan error, error) {
-	ctx, cancel := context.WithCancel(ctx)
+func (l *LimaQemuDriver) Start(_ context.Context) (chan error, error) {
+	ctx, cancel := context.WithCancel(context.Background())
 	defer func() {
 		if l.qCmd == nil {
 			cancel()
@@ -220,8 +220,9 @@ func (l *LimaQemuDriver) Start(ctx context.Context) (chan error, error) {
 		return nil, err
 	}
 	l.qCmd = qCmd
-	l.qWaitCh = make(chan error)
+	l.qWaitCh = make(chan error, 1)
 	go func() {
+		defer close(l.qWaitCh)
 		l.qWaitCh <- qCmd.Wait()
 	}()
 	l.vhostCmds = vhostCmds
@@ -380,20 +381,32 @@ func (l *LimaQemuDriver) shutdownQEMU(ctx context.Context, timeout time.Duration
 		logrus.WithError(err).Warnf("failed to send system_powerdown command via the QMP socket %q, forcibly killing QEMU", qmpSockPath)
 		return l.killQEMU(ctx, timeout, qCmd, qWaitCh)
 	}
-	deadline := time.After(timeout)
+	timeoutCtx, timeoutCancel := context.WithTimeout(context.Background(), timeout)
+	defer timeoutCancel()
+
 	select {
-	case qWaitErr := <-qWaitCh:
+	case qWaitErr, ok := <-qWaitCh:
+		if !ok {
+			logrus.Info("QEMU wait channel was closed")
+			_ = l.removeVNCFiles()
+			return l.killVhosts()
+		}
 		entry := logrus.NewEntry(logrus.StandardLogger())
 		if qWaitErr != nil {
 			entry = entry.WithError(qWaitErr)
 		}
 		entry.Info("QEMU has exited")
 		_ = l.removeVNCFiles()
 		return errors.Join(qWaitErr, l.killVhosts())
-	case <-deadline:
+	case <-timeoutCtx.Done():
+		if qCmd.ProcessState != nil {
+			logrus.Info("QEMU has already exited")
+			_ = l.removeVNCFiles()
+			return l.killVhosts()
+		}
+		logrus.Warnf("QEMU did not exit in %v, forcibly killing QEMU", timeout)
+		return l.killQEMU(ctx, timeout, qCmd, qWaitCh)
 	}
-	logrus.Warnf("QEMU did not exit in %v, forcibly killing QEMU", timeout)
-	return l.killQEMU(ctx, timeout, qCmd, qWaitCh)
 }
 
 func (l *LimaQemuDriver) killQEMU(_ context.Context, _ time.Duration, qCmd *exec.Cmd, qWaitCh <-chan error) error {