pkg/hostagent: Ensure calling HostAgent.close() before cancelling the context.

portfwd:
- Change `ClosableListeners.Close()` to return `error`
- Change `Forwarder.Close()` to return `error`
- Use `HostAgent.cleanUp()` instead of `defer` to calling `ClosableListeners.Close()`

Signed-off-by: Norio Nomura <[email protected]>

Author: norio-nomura

pkg/hostagent/hostagent.go

@@ -448,26 +448,19 @@ func (a *HostAgent) startRoutinesAndWait(ctx context.Context, errCh <-chan error
 		stRunning.Running = true
 		a.emitEvent(ctx, events.Event{Status: stRunning})
 	}()
-	for {
-		select {
-		case driverErr := <-errCh:
-			logrus.Infof("Driver stopped due to error: %q", driverErr)
-			cancelHA()
-			if closeErr := a.close(); closeErr != nil {
-				logrus.WithError(closeErr).Warn("an error during shutting down the host agent")
-			}
-			err := a.driver.Stop(ctx)
-			return err
-		case sig := <-a.signalCh:
-			logrus.Infof("Received %s, shutting down the host agent", osutil.SignalName(sig))
-			cancelHA()
-			if closeErr := a.close(); closeErr != nil {
-				logrus.WithError(closeErr).Warn("an error during shutting down the host agent")
-			}
-			err := a.driver.Stop(ctx)
-			return err
-		}
-	}
+	// wait for either the driver to stop or a signal to shut down
+	select {
+	case driverErr := <-errCh:
+		logrus.Infof("Driver stopped due to error: %q", driverErr)
+	case sig := <-a.signalCh:
+		logrus.Infof("Received %s, shutting down the host agent", osutil.SignalName(sig))
+	}
+	// close the host agent routines before cancelling the context
+	if closeErr := a.close(); closeErr != nil {
+		logrus.WithError(closeErr).Warn("an error during shutting down the host agent")
+	}
+	cancelHA()
+	return a.driver.Stop(ctx)
 }
 
 func (a *HostAgent) Info(_ context.Context) (*hostagentapi.Info, error) {
@@ -603,6 +596,7 @@ sudo chown -R "${USER}" /run/host-services`
 }
 
 // cleanUp registers a cleanup function to be called when the host agent is stopped.
+// The cleanup functions are called before the context is cancelled, in the reverse order of their registration.
 func (a *HostAgent) cleanUp(fn func() error) {
 	a.onCloseMu.Lock()
 	defer a.onCloseMu.Unlock()
@@ -674,6 +668,9 @@ func (a *HostAgent) watchGuestAgentEvents(ctx context.Context) {
 		}
 	}()
 
+	// ensure close before ctx is cancelled
+	a.cleanUp(a.grpcPortForwarder.Close)
+
 	for {
 		if a.client == nil || !isGuestAgentSocketAccessible(ctx, a.client) {
 			if a.driver.ForwardGuestAgent() {
@@ -806,7 +803,6 @@ func (a *HostAgent) processGuestAgentEvents(ctx context.Context, client *guestag
 			a.grpcPortForwarder.OnEvent(ctx, client, ev)
 		}
 	}
-	defer a.grpcPortForwarder.Close()
 
 	if err := client.Events(ctx, onEvent); err != nil {
 		if status.Code(err) == codes.Canceled {

pkg/portfwd/forward.go

@@ -34,8 +34,8 @@ func NewPortForwarder(rules []limatype.PortForward, ignoreTCP, ignoreUDP bool) *
 	}
 }
 
-func (fw *Forwarder) Close() {
-	fw.closableListeners.Close()
+func (fw *Forwarder) Close() error {
+	return fw.closableListeners.Close()
 }
 
 func (fw *Forwarder) OnEvent(ctx context.Context, client *guestagentclient.GuestAgentClient, ev *api.Event) {

pkg/portfwd/listener.go

@@ -38,19 +38,25 @@ func NewClosableListener() *ClosableListeners {
 	}
 }
 
-func (p *ClosableListeners) Close() {
+func (p *ClosableListeners) Close() error {
 	p.listenersRW.Lock()
 	defer p.listenersRW.Unlock()
+	var errs []error
 	for _, listener := range p.listeners {
-		listener.Close()
+		if err := listener.Close(); err != nil {
+			errs = append(errs, err)
+		}
 	}
 	clear(p.listeners)
 	p.udpListenersRW.Lock()
 	defer p.udpListenersRW.Unlock()
 	for _, listener := range p.udpListeners {
-		listener.Close()
+		if err := listener.Close(); err != nil {
+			errs = append(errs, err)
+		}
 	}
 	clear(p.udpListeners)
+	return errors.Join(errs...)
 }
 
 func (p *ClosableListeners) Forward(ctx context.Context, client *guestagentclient.GuestAgentClient,