pass a single context throughout the device-plugin method call stack (#1284)

This change follows the Go best practices. With a single ctx reference,
we allow for the proper propagation of cancellations and graceful terminations
across all goroutines of the device-plugin application.

Signed-off-by: Tariq Ibrahim <[email protected]>

Author: tariq1890

cmd/nvidia-device-plugin/main.go

@@ -353,7 +353,7 @@ func startPlugins(c *cli.Context, o *options) ([]plugin.Interface, bool, error)
 
 	// Get the set of plugins.
 	klog.Info("Retrieving plugins.")
-	plugins, err := GetPlugins(infolib, nvmllib, devicelib, config)
+	plugins, err := GetPlugins(c.Context, infolib, nvmllib, devicelib, config)
 	if err != nil {
 		return nil, false, fmt.Errorf("error getting plugins: %v", err)
 	}

cmd/nvidia-device-plugin/plugin-manager.go

@@ -17,6 +17,7 @@
 package main
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/NVIDIA/go-nvlib/pkg/nvlib/device"
@@ -30,7 +31,7 @@ import (
 )
 
 // GetPlugins returns a set of plugins for the specified configuration.
-func GetPlugins(infolib info.Interface, nvmllib nvml.Interface, devicelib device.Interface, config *spec.Config) ([]plugin.Interface, error) {
+func GetPlugins(ctx context.Context, infolib info.Interface, nvmllib nvml.Interface, devicelib device.Interface, config *spec.Config) ([]plugin.Interface, error) {
 	// TODO: We could consider passing this as an argument since it should already be used to construct nvmllib.
 	driverRoot := root(*config.Flags.Plugin.ContainerDriverRoot)
 
@@ -61,7 +62,7 @@ func GetPlugins(infolib info.Interface, nvmllib nvml.Interface, devicelib device
 		return nil, fmt.Errorf("unable to create cdi handler: %v", err)
 	}
 
-	plugins, err := plugin.New(infolib, nvmllib, devicelib,
+	plugins, err := plugin.New(ctx, infolib, nvmllib, devicelib,
 		plugin.WithCDIHandler(cdiHandler),
 		plugin.WithConfig(config),
 		plugin.WithDeviceListStrategies(deviceListStrategies),

internal/plugin/factory.go

@@ -17,6 +17,7 @@
 package plugin
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/NVIDIA/go-nvlib/pkg/nvlib/device"
@@ -46,7 +47,7 @@ type options struct {
 }
 
 // New a new set of plugins with the supplied options.
-func New(infolib info.Interface, nvmllib nvml.Interface, devicelib device.Interface, opts ...Option) ([]Interface, error) {
+func New(ctx context.Context, infolib info.Interface, nvmllib nvml.Interface, devicelib device.Interface, opts ...Option) ([]Interface, error) {
 	o := &options{
 		infolib:   infolib,
 		nvmllib:   nvmllib,
@@ -72,7 +73,7 @@ func New(infolib info.Interface, nvmllib nvml.Interface, devicelib device.Interf
 
 	var plugins []Interface
 	for _, resourceManager := range resourceManagers {
-		plugin, err := o.devicePluginForResource(resourceManager)
+		plugin, err := o.devicePluginForResource(ctx, resourceManager)
 		if err != nil {
 			return nil, fmt.Errorf("failed to create plugin: %w", err)
 		}

internal/plugin/server.go

@@ -50,6 +50,7 @@ const (
 
 // nvidiaDevicePlugin implements the Kubernetes device plugin API
 type nvidiaDevicePlugin struct {
+	ctx                  context.Context
 	rm                   rm.ResourceManager
 	config               *spec.Config
 	deviceListStrategies spec.DeviceListStrategies
@@ -68,13 +69,14 @@ type nvidiaDevicePlugin struct {
 }
 
 // devicePluginForResource creates a device plugin for the specified resource.
-func (o *options) devicePluginForResource(resourceManager rm.ResourceManager) (Interface, error) {
+func (o *options) devicePluginForResource(ctx context.Context, resourceManager rm.ResourceManager) (Interface, error) {
 	mpsOptions, err := o.getMPSOptions(resourceManager)
 	if err != nil {
 		return nil, err
 	}
 
 	plugin := nvidiaDevicePlugin{
+		ctx:                  ctx,
 		rm:                   resourceManager,
 		config:               o.config,
 		deviceListStrategies: o.deviceListStrategies,
@@ -245,7 +247,7 @@ func (plugin *nvidiaDevicePlugin) Register(kubeletSocket string) error {
 		},
 	}
 
-	_, err = client.Register(context.Background(), reqt)
+	_, err = client.Register(plugin.ctx, reqt)
 	if err != nil {
 		return err
 	}
@@ -432,7 +434,7 @@ func (plugin *nvidiaDevicePlugin) PreStartContainer(context.Context, *pluginapi.
 
 // dial establishes the gRPC communication with the registered device plugin.
 func (plugin *nvidiaDevicePlugin) dial(unixSocketPath string, timeout time.Duration) (*grpc.ClientConn, error) {
-	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	ctx, cancel := context.WithTimeout(plugin.ctx, timeout)
 	defer cancel()
 	//nolint:staticcheck  // TODO: Switch to grpc.NewClient
 	c, err := grpc.DialContext(ctx, unixSocketPath,