refactor(plugin): use context for health check lifecycle

Replace the stop channel with context.Context cancellation and
sync.WaitGroup for health check goroutine lifecycle management:

- Initialize health channel and healthCtx before Serve() to eliminate
  race where kubelet could call ListAndWatch before fields are set
- Use sync.WaitGroup to ensure health goroutine completes before
  channel close in Stop()
- Add nil guards in Stop() to handle partial Start() failure safely
- Add structured error handling for health check completion (success,
  canceled, error cases)
- Add debug logging in ListAndWatch for context cancellation and
  channel closure
- Fix updateResponseForMPS receiver from value to pointer

Signed-off-by: Carlos Eduardo Arango Gutierrez <carangog@redhat.com>
Co-authored-by: Evan Lezar <elezar@nvidia.com>

Author: ArangoGutierrez

internal/plugin/server.go

@@ -25,6 +25,7 @@ import (
 	"path"
 	"path/filepath"
 	"strings"
+	"sync"
 	"time"
 
 	cdiapi "tags.cncf.io/container-device-interface/pkg/cdi"
@@ -62,7 +63,12 @@ type nvidiaDevicePlugin struct {
 	socket string
 	server *grpc.Server
 	health chan *rm.Device
-	stop   chan interface{}
+
+	// healthCtx and healthCancel control the health check goroutine lifecycle.
+	healthCtx    context.Context
+	healthCancel context.CancelFunc
+	// healthWg is used to wait for the health check goroutine to complete during cleanup.
+	healthWg sync.WaitGroup
 
 	imexChannels imex.Channels
 
@@ -90,11 +96,6 @@ func (o *options) devicePluginForResource(ctx context.Context, resourceManager r
 		mps: mpsOptions,
 
 		socket: getPluginSocketPath(resourceManager.Resource()),
-		// These will be reinitialized every
-		// time the plugin server is restarted.
-		server: nil,
-		health: nil,
-		stop:   nil,
 	}
 	return &plugin, nil
 }
@@ -106,19 +107,6 @@ func getPluginSocketPath(resource spec.ResourceName) string {
 	return filepath.Join(pluginapi.DevicePluginPath, pluginName) + ".sock"
 }
 
-func (plugin *nvidiaDevicePlugin) initialize() {
-	plugin.server = grpc.NewServer([]grpc.ServerOption{}...)
-	plugin.health = make(chan *rm.Device)
-	plugin.stop = make(chan interface{})
-}
-
-func (plugin *nvidiaDevicePlugin) cleanup() {
-	close(plugin.stop)
-	plugin.server = nil
-	plugin.health = nil
-	plugin.stop = nil
-}
-
 // Devices returns the full set of devices associated with the plugin.
 func (plugin *nvidiaDevicePlugin) Devices() rm.Devices {
 	return plugin.rm.Devices()
@@ -127,16 +115,16 @@ func (plugin *nvidiaDevicePlugin) Devices() rm.Devices {
 // Start starts the gRPC server, registers the device plugin with the Kubelet,
 // and starts the device healthchecks.
 func (plugin *nvidiaDevicePlugin) Start(kubeletSocket string) error {
-	plugin.initialize()
-
 	if err := plugin.mps.waitForDaemon(); err != nil {
 		return fmt.Errorf("error waiting for MPS daemon: %w", err)
 	}
 
+	plugin.health = make(chan *rm.Device)
+	plugin.healthCtx, plugin.healthCancel = context.WithCancel(plugin.ctx)
+
 	err := plugin.Serve()
 	if err != nil {
 		klog.Errorf("Could not start device plugin for '%s': %s", plugin.rm.Resource(), err)
-		plugin.cleanup()
 		return err
 	}
 	klog.Infof("Starting to serve '%s' on %s", plugin.rm.Resource(), plugin.socket)
@@ -148,10 +136,17 @@ func (plugin *nvidiaDevicePlugin) Start(kubeletSocket string) error {
 	}
 	klog.Infof("Registered device plugin for '%s' with Kubelet", plugin.rm.Resource())
 
+	plugin.healthWg.Add(1)
 	go func() {
+		defer plugin.healthWg.Done()
 		// TODO: add MPS health check
-		err := plugin.rm.CheckHealth(plugin.stop, plugin.health)
-		if err != nil {
+		err := plugin.rm.CheckHealth(plugin.healthCtx, plugin.health)
+		switch {
+		case err == nil:
+			klog.Infof("Health check completed successfully for '%s'", plugin.rm.Resource())
+		case errors.Is(err, context.Canceled):
+			klog.V(4).Infof("Health check canceled for '%s' (plugin shutdown)", plugin.rm.Resource())
+		default:
 			klog.Errorf("Failed to start health check: %v; continuing with health checks disabled", err)
 		}
 	}()
@@ -164,12 +159,21 @@ func (plugin *nvidiaDevicePlugin) Stop() error {
 	if plugin == nil || plugin.server == nil {
 		return nil
 	}
+	// Stop health checks if they were started.
+	if plugin.healthCancel != nil {
+		plugin.healthCancel()
+		plugin.healthWg.Wait()
+	}
+	if plugin.health != nil {
+		close(plugin.health)
+	}
+
 	klog.Infof("Stopping to serve '%s' on %s", plugin.rm.Resource(), plugin.socket)
 	plugin.server.Stop()
+	plugin.server = nil
 	if err := os.Remove(plugin.socket); err != nil && !os.IsNotExist(err) {
 		return err
 	}
-	plugin.cleanup()
 	return nil
 }
 
@@ -181,6 +185,7 @@ func (plugin *nvidiaDevicePlugin) Serve() error {
 		return err
 	}
 
+	plugin.server = grpc.NewServer([]grpc.ServerOption{}...)
 	pluginapi.RegisterDevicePluginServer(plugin.server, plugin)
 
 	go func() {
@@ -271,9 +276,14 @@ func (plugin *nvidiaDevicePlugin) ListAndWatch(e *pluginapi.Empty, s pluginapi.D
 
 	for {
 		select {
-		case <-plugin.stop:
+		case <-plugin.healthCtx.Done():
+			klog.V(4).Infof("Stopping health checks for '%s'", plugin.rm.Resource())
 			return nil
-		case d := <-plugin.health:
+		case d, ok := <-plugin.health:
+			if !ok {
+				klog.V(4).Infof("Health channel closed for '%s'", plugin.rm.Resource())
+				return nil
+			}
 			// FIXME: there is no way to recover from the Unhealthy state.
 			d.Health = pluginapi.Unhealthy
 			klog.Infof("'%s' device marked unhealthy: %s", plugin.rm.Resource(), d.ID)
@@ -368,7 +378,7 @@ func (plugin *nvidiaDevicePlugin) getAllocateResponse(requestIds []string) (*plu
 // updateResponseForMPS ensures that the ContainerAllocate response contains the information required to use MPS.
 // This includes per-resource pipe and log directories as well as a global daemon-specific shm
 // and assumes that an MPS control daemon has already been started.
-func (plugin nvidiaDevicePlugin) updateResponseForMPS(response *pluginapi.ContainerAllocateResponse) {
+func (plugin *nvidiaDevicePlugin) updateResponseForMPS(response *pluginapi.ContainerAllocateResponse) {
 	plugin.mps.updateReponse(response)
 }