pass a single context throughout the device-plugin method call stack

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 <tibrahim@nvidia.com>

Author: tariq1890

cmd/nvidia-device-plugin/main.go

@@ -354,7 +354,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)
 
@@ -62,6 +63,7 @@ func GetPlugins(infolib info.Interface, nvmllib nvml.Interface, devicelib device
 	}
 
 	plugins, err := plugin.New(infolib, nvmllib, devicelib,
+		plugin.WithContext(ctx),
 		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"
@@ -31,6 +32,7 @@ import (
 )
 
 type options struct {
+	ctx       context.Context
 	infolib   info.Interface
 	nvmllib   nvml.Interface
 	devicelib device.Interface

internal/plugin/options.go

@@ -17,6 +17,8 @@
 package plugin
 
 import (
+	"context"
+
 	"github.com/NVIDIA/go-nvlib/pkg/nvlib/info"
 	"github.com/NVIDIA/go-nvml/pkg/nvml"
 
@@ -28,6 +30,12 @@ import (
 // Option is a function that configures a options
 type Option func(*options)
 
+func WithContext(ctx context.Context) Option {
+	return func(m *options) {
+		m.ctx = ctx
+	}
+}
+
 // WithCDIHandler sets the CDI handler for the options
 func WithCDIHandler(handler cdi.Interface) Option {
 	return func(m *options) {

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
@@ -75,6 +76,7 @@ func (o *options) devicePluginForResource(resourceManager rm.ResourceManager) (I
 	}
 
 	plugin := nvidiaDevicePlugin{
+		ctx:                  o.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,