DRA kubelet: code cleanup

Replaced the Manager interface with a simple struct. There was no need for the
interface.

Replaced the global instance of the DRA plugins store with a fresh instanced
owned by the manager. This makes unit testing a bit easier (no need to restore
state, would enable parallelizing long-running tests).

Simplified the plugin.PluginsStore type to just "plugin.Store" because it
stuttered. The Plugin type is kept because having one struct named after its
package is one common exception from the "don't stutter" guideline.
For the sake of clarify, the package gets imported as "draplugin" (=
<parent dir>/<package>).

Removed unused NewManager "node" parameter and replaced direct construction of
cache and manager with calls to NewManager because that is how the manager
should get constructed (less code, too).

Fixed incorrect description of Manager: the plugin store is the entity which
manages drivers, not the manager. The manager is focused on the DRA logic
around ResourceClaims.

pkg/kubelet/cm/container_manager_linux.go

@@ -131,8 +131,8 @@ type containerManagerImpl struct {
 	memoryManager memorymanager.Manager
 	// Interface for Topology resource co-ordination
 	topologyManager topologymanager.Manager
-	// Interface for Dynamic Resource Allocation management.
-	draManager dra.Manager
+	// Implementation of Dynamic Resource Allocation (DRA).
+	draManager *dra.Manager
 	// kubeClient is the interface to the Kubernetes API server. May be nil if the kubelet is running in standalone mode.
 	kubeClient clientset.Interface
 }
@@ -310,7 +310,7 @@ func NewContainerManager(mountUtil mount.Interface, cadvisorInterface cadvisor.I
 	// Initialize DRA manager
 	if utilfeature.DefaultFeatureGate.Enabled(kubefeatures.DynamicResourceAllocation) {
 		klog.InfoS("Creating Dynamic Resource Allocation (DRA) manager")
-		cm.draManager, err = dra.NewManagerImpl(kubeClient, nodeConfig.KubeletRootDir, nodeConfig.NodeName)
+		cm.draManager, err = dra.NewManager(kubeClient, nodeConfig.KubeletRootDir)
 		if err != nil {
 			return nil, err
 		}

pkg/kubelet/cm/dra/manager.go

@@ -32,7 +32,7 @@ import (
 	"k8s.io/dynamic-resource-allocation/resourceclaim"
 	"k8s.io/klog/v2"
 	drapb "k8s.io/kubelet/pkg/apis/dra/v1beta1"
-	dra "k8s.io/kubernetes/pkg/kubelet/cm/dra/plugin"
+	draplugin "k8s.io/kubernetes/pkg/kubelet/cm/dra/plugin"
 	"k8s.io/kubernetes/pkg/kubelet/cm/dra/state"
 	"k8s.io/kubernetes/pkg/kubelet/config"
 	kubecontainer "k8s.io/kubernetes/pkg/kubelet/container"
@@ -52,8 +52,14 @@ type ActivePodsFunc func() []*v1.Pod
 // GetNodeFunc is a function that returns the node object using the kubelet's node lister.
 type GetNodeFunc func() (*v1.Node, error)
 
-// ManagerImpl is the structure in charge of managing DRA drivers.
-type ManagerImpl struct {
+// Manager is responsible for managing ResourceClaims.
+// It ensures that they are prepared before starting pods
+// and that they are unprepared before the last consuming
+// pod is declared as terminated.
+type Manager struct {
+	// draPlugins manages the registered plugins.
+	draPlugins *draplugin.Store
+
 	// cache contains cached claim info
 	cache *claimInfoCache
 
@@ -75,7 +81,7 @@ type ManagerImpl struct {
 	getNode GetNodeFunc
 }
 
-// NewManagerImpl creates a new manager.
+// NewManager creates a new DRA manager.
 //
 // Most errors returned by the manager show up in the context of a pod.
 // They try to adhere to the following convention:
@@ -84,7 +90,7 @@ type ManagerImpl struct {
 // - Don't include the namespace, it can be inferred from the context.
 // - Avoid repeated "failed to ...: failed to ..." when wrapping errors.
 // - Avoid wrapping when it does not provide relevant additional information to keep the user-visible error short.
-func NewManagerImpl(kubeClient clientset.Interface, stateFileDirectory string, nodeName types.NodeName) (*ManagerImpl, error) {
+func NewManager(kubeClient clientset.Interface, stateFileDirectory string) (*Manager, error) {
 	claimInfoCache, err := newClaimInfoCache(stateFileDirectory, draManagerStateFileName)
 	if err != nil {
 		return nil, fmt.Errorf("create ResourceClaim cache: %w", err)
@@ -94,8 +100,9 @@ func NewManagerImpl(kubeClient clientset.Interface, stateFileDirectory string, n
 	// We should consider making it configurable in the future.
 	reconcilePeriod := defaultReconcilePeriod
 
-	manager := &ManagerImpl{
+	manager := &Manager{
 		cache:           claimInfoCache,
+		draPlugins:      new(draplugin.Store),
 		kubeClient:      kubeClient,
 		reconcilePeriod: reconcilePeriod,
 		activePods:      nil,
@@ -105,7 +112,7 @@ func NewManagerImpl(kubeClient clientset.Interface, stateFileDirectory string, n
 	return manager, nil
 }
 
-func (m *ManagerImpl) GetWatcherHandler() cache.PluginHandler {
+func (m *Manager) GetWatcherHandler() cache.PluginHandler {
 	// The time that DRA drivers have to come back after being unregistered
 	// before the kubelet removes their ResourceSlices.
 	//
@@ -119,11 +126,12 @@ func (m *ManagerImpl) GetWatcherHandler() cache.PluginHandler {
 	// If a DRA driver wants to be sure that slices don't get wiped,
 	// it should use rolling updates.
 	wipingDelay := 30 * time.Second
-	return cache.PluginHandler(dra.NewRegistrationHandler(m.kubeClient, m.getNode, wipingDelay))
+	/* TODO: figure out where to get a context, NewRegistrationHandler runs with context.Background() */
+	return cache.PluginHandler(draplugin.NewRegistrationHandler(m.draPlugins, m.kubeClient, m.getNode, wipingDelay))
 }
 
 // Start starts the reconcile loop of the manager.
-func (m *ManagerImpl) Start(ctx context.Context, activePods ActivePodsFunc, getNode GetNodeFunc, sourcesReady config.SourcesReady) error {
+func (m *Manager) Start(ctx context.Context, activePods ActivePodsFunc, getNode GetNodeFunc, sourcesReady config.SourcesReady) error {
 	m.activePods = activePods
 	m.getNode = getNode
 	m.sourcesReady = sourcesReady
@@ -132,7 +140,7 @@ func (m *ManagerImpl) Start(ctx context.Context, activePods ActivePodsFunc, getN
 }
 
 // reconcileLoop ensures that any stale state in the manager's claimInfoCache gets periodically reconciled.
-func (m *ManagerImpl) reconcileLoop(ctx context.Context) {
+func (m *Manager) reconcileLoop(ctx context.Context) {
 	logger := klog.FromContext(ctx)
 	// Only once all sources are ready do we attempt to reconcile.
 	// This ensures that the call to m.activePods() below will succeed with
@@ -184,7 +192,7 @@ func (m *ManagerImpl) reconcileLoop(ctx context.Context) {
 // for the input container, issue NodePrepareResources rpc requests
 // for each new resource requirement, process their responses and update the cached
 // containerResources on success.
-func (m *ManagerImpl) PrepareResources(ctx context.Context, pod *v1.Pod) error {
+func (m *Manager) PrepareResources(ctx context.Context, pod *v1.Pod) error {
 	startTime := time.Now()
 	err := m.prepareResources(ctx, pod)
 	metrics.DRAOperationsDuration.WithLabelValues("PrepareResources", strconv.FormatBool(err == nil)).Observe(time.Since(startTime).Seconds())
@@ -194,10 +202,10 @@ func (m *ManagerImpl) PrepareResources(ctx context.Context, pod *v1.Pod) error {
 	return nil
 }
 
-func (m *ManagerImpl) prepareResources(ctx context.Context, pod *v1.Pod) error {
+func (m *Manager) prepareResources(ctx context.Context, pod *v1.Pod) error {
 	var err error
 	logger := klog.FromContext(ctx)
-	batches := make(map[*dra.Plugin][]*drapb.Claim)
+	batches := make(map[*draplugin.Plugin][]*drapb.Claim)
 	resourceClaims := make(map[types.UID]*resourceapi.ResourceClaim)
 
 	// Do a validation pass *without* changing the claim info cache.
@@ -213,7 +221,7 @@ func (m *ManagerImpl) prepareResources(ctx context.Context, pod *v1.Pod) error {
 		resourceClaim *resourceapi.ResourceClaim
 		podClaim      *v1.PodResourceClaim
 		claimInfo     *ClaimInfo
-		drivers       map[string]*dra.Plugin
+		drivers       map[string]*draplugin.Plugin
 	}, len(pod.Spec.ResourceClaims))
 	for i := range pod.Spec.ResourceClaims {
 		podClaim := &pod.Spec.ResourceClaims[i]
@@ -260,12 +268,12 @@ func (m *ManagerImpl) prepareResources(ctx context.Context, pod *v1.Pod) error {
 			return fmt.Errorf("ResourceClaim %s: %w", resourceClaim.Name, err)
 		}
 		infos[i].claimInfo = claimInfo
-		infos[i].drivers = make(map[string]*dra.Plugin, len(claimInfo.DriverState))
+		infos[i].drivers = make(map[string]*draplugin.Plugin, len(claimInfo.DriverState))
 		for driverName := range claimInfo.DriverState {
 			if plugin := infos[i].drivers[driverName]; plugin != nil {
 				continue
 			}
-			client, err := dra.NewDRAPluginClient(driverName)
+			client, err := m.draPlugins.NewDRAPluginClient(driverName)
 			if err != nil {
 				// No wrapping, error includes driver name already.
 				return err
@@ -421,7 +429,7 @@ func lookupClaimRequest(claims []*drapb.Claim, claimUID string) *drapb.Claim {
 
 // GetResources gets a ContainerInfo object from the claimInfo cache.
 // This information is used by the caller to update a container config.
-func (m *ManagerImpl) GetResources(pod *v1.Pod, container *v1.Container) (*ContainerInfo, error) {
+func (m *Manager) GetResources(pod *v1.Pod, container *v1.Container) (*ContainerInfo, error) {
 	cdiDevices := []kubecontainer.CDIDevice{}
 
 	for i := range pod.Spec.ResourceClaims {
@@ -466,7 +474,7 @@ func (m *ManagerImpl) GetResources(pod *v1.Pod, container *v1.Container) (*Conta
 // This function is idempotent and may be called multiple times against the same pod.
 // As such, calls to the underlying NodeUnprepareResource API are skipped for claims that have
 // already been successfully unprepared.
-func (m *ManagerImpl) UnprepareResources(ctx context.Context, pod *v1.Pod) error {
+func (m *Manager) UnprepareResources(ctx context.Context, pod *v1.Pod) error {
 	startTime := time.Now()
 	err := m.unprepareResourcesForPod(ctx, pod)
 	metrics.DRAOperationsDuration.WithLabelValues("UnprepareResources", strconv.FormatBool(err == nil)).Observe(time.Since(startTime).Seconds())
@@ -476,7 +484,7 @@ func (m *ManagerImpl) UnprepareResources(ctx context.Context, pod *v1.Pod) error
 	return nil
 }
 
-func (m *ManagerImpl) unprepareResourcesForPod(ctx context.Context, pod *v1.Pod) error {
+func (m *Manager) unprepareResourcesForPod(ctx context.Context, pod *v1.Pod) error {
 	var claimNames []string
 	for i := range pod.Spec.ResourceClaims {
 		claimName, _, err := resourceclaim.Name(pod, &pod.Spec.ResourceClaims[i])
@@ -495,7 +503,7 @@ func (m *ManagerImpl) unprepareResourcesForPod(ctx context.Context, pod *v1.Pod)
 	return m.unprepareResources(ctx, pod.UID, pod.Namespace, claimNames)
 }
 
-func (m *ManagerImpl) unprepareResources(ctx context.Context, podUID types.UID, namespace string, claimNames []string) error {
+func (m *Manager) unprepareResources(ctx context.Context, podUID types.UID, namespace string, claimNames []string) error {
 	logger := klog.FromContext(ctx)
 	batches := make(map[string][]*drapb.Claim)
 	claimNamesMap := make(map[types.UID]string)
@@ -549,7 +557,7 @@ func (m *ManagerImpl) unprepareResources(ctx context.Context, podUID types.UID,
 	// We could try to continue, but that would make the code more complex.
 	for driverName, claims := range batches {
 		// Call NodeUnprepareResources RPC for all resource handles.
-		client, err := dra.NewDRAPluginClient(driverName)
+		client, err := m.draPlugins.NewDRAPluginClient(driverName)
 		if client == nil {
 			// No wrapping, error includes driver name already.
 			return err
@@ -601,14 +609,14 @@ func (m *ManagerImpl) unprepareResources(ctx context.Context, podUID types.UID,
 
 // PodMightNeedToUnprepareResources returns true if the pod might need to
 // unprepare resources
-func (m *ManagerImpl) PodMightNeedToUnprepareResources(uid types.UID) bool {
+func (m *Manager) PodMightNeedToUnprepareResources(uid types.UID) bool {
 	m.cache.Lock()
 	defer m.cache.Unlock()
 	return m.cache.hasPodReference(uid)
 }
 
 // GetContainerClaimInfos gets Container's ClaimInfo
-func (m *ManagerImpl) GetContainerClaimInfos(pod *v1.Pod, container *v1.Container) ([]*ClaimInfo, error) {
+func (m *Manager) GetContainerClaimInfos(pod *v1.Pod, container *v1.Container) ([]*ClaimInfo, error) {
 	claimInfos := make([]*ClaimInfo, 0, len(pod.Spec.ResourceClaims))
 
 	for i, podResourceClaim := range pod.Spec.ResourceClaims {

pkg/kubelet/cm/dra/manager_test.go

@@ -199,7 +199,7 @@ func TestNewManagerImpl(t *testing.T) {
 		},
 	} {
 		t.Run(test.description, func(t *testing.T) {
-			manager, err := NewManagerImpl(kubeClient, test.stateFileDirectory, "worker")
+			manager, err := NewManager(kubeClient, test.stateFileDirectory)
 			if test.wantErr {
 				assert.Error(t, err)
 				return
@@ -363,7 +363,7 @@ func TestGetResources(t *testing.T) {
 		},
 	} {
 		t.Run(test.description, func(t *testing.T) {
-			manager, err := NewManagerImpl(kubeClient, t.TempDir(), "worker")
+			manager, err := NewManager(kubeClient, t.TempDir())
 			require.NoError(t, err)
 
 			if test.claimInfo != nil {
@@ -558,15 +558,9 @@ func TestPrepareResources(t *testing.T) {
 	} {
 		t.Run(test.description, func(t *testing.T) {
 			tCtx := ktesting.Init(t)
-			cache, err := newClaimInfoCache(t.TempDir(), draManagerStateFileName)
-			if err != nil {
-				t.Fatalf("failed to newClaimInfoCache, err:%v", err)
-			}
 
-			manager := &ManagerImpl{
-				kubeClient: fakeKubeClient,
-				cache:      cache,
-			}
+			manager, err := NewManager(fakeKubeClient, t.TempDir())
+			require.NoError(t, err, "create DRA manager")
 
 			if test.claim != nil {
 				if _, err := fakeKubeClient.ResourceV1beta1().ResourceClaims(test.pod.Namespace).Create(tCtx, test.claim, metav1.CreateOptions{}); err != nil {
@@ -589,11 +583,10 @@ func TestPrepareResources(t *testing.T) {
 			}
 			defer draServerInfo.teardownFn()
 
-			plg := plugin.NewRegistrationHandler(nil, getFakeNode, time.Second /* very short wiping delay for testing */)
+			plg := plugin.NewRegistrationHandler(manager.draPlugins, nil, nil, time.Second /* very short wiping delay for testing */)
 			if err := plg.RegisterPlugin(test.driverName, draServerInfo.socketName, []string{drapb.DRAPluginService}, pluginClientTimeout); err != nil {
 				t.Fatalf("failed to register plugin %s, err: %v", test.driverName, err)
 			}
-			defer plg.DeRegisterPlugin(test.driverName, draServerInfo.socketName) // for sake of next tests
 
 			if test.claimInfo != nil {
 				manager.cache.add(test.claimInfo)
@@ -709,10 +702,6 @@ func TestUnprepareResources(t *testing.T) {
 	} {
 		t.Run(test.description, func(t *testing.T) {
 			tCtx := ktesting.Init(t)
-			cache, err := newClaimInfoCache(t.TempDir(), draManagerStateFileName)
-			if err != nil {
-				t.Fatalf("failed to create a new instance of the claimInfoCache, err: %v", err)
-			}
 
 			var pluginClientTimeout *time.Duration
 			if test.wantTimeout {
@@ -726,16 +715,13 @@ func TestUnprepareResources(t *testing.T) {
 			}
 			defer draServerInfo.teardownFn()
 
-			plg := plugin.NewRegistrationHandler(nil, getFakeNode, time.Second /* very short wiping delay for testing */)
+			manager, err := NewManager(fakeKubeClient, t.TempDir())
+			require.NoError(t, err, "create DRA manager")
+
+			plg := plugin.NewRegistrationHandler(manager.draPlugins, nil, getFakeNode, time.Second /* very short wiping delay for testing */)
 			if err := plg.RegisterPlugin(test.driverName, draServerInfo.socketName, []string{drapb.DRAPluginService}, pluginClientTimeout); err != nil {
 				t.Fatalf("failed to register plugin %s, err: %v", test.driverName, err)
 			}
-			defer plg.DeRegisterPlugin(test.driverName, draServerInfo.socketName) // for sake of next tests
-
-			manager := &ManagerImpl{
-				kubeClient: fakeKubeClient,
-				cache:      cache,
-			}
 
 			if test.claimInfo != nil {
 				manager.cache.add(test.claimInfo)
@@ -773,16 +759,8 @@ func TestUnprepareResources(t *testing.T) {
 
 func TestPodMightNeedToUnprepareResources(t *testing.T) {
 	fakeKubeClient := fake.NewSimpleClientset()
-
-	cache, err := newClaimInfoCache(t.TempDir(), draManagerStateFileName)
-	if err != nil {
-		t.Fatalf("failed to newClaimInfoCache, err:%v", err)
-	}
-
-	manager := &ManagerImpl{
-		kubeClient: fakeKubeClient,
-		cache:      cache,
-	}
+	manager, err := NewManager(fakeKubeClient, t.TempDir())
+	require.NoError(t, err, "create DRA manager")
 
 	claimInfo := &ClaimInfo{
 		ClaimInfoState: state.ClaimInfoState{PodUIDs: sets.New(podUID), ClaimName: claimName, Namespace: namespace},
@@ -855,13 +833,8 @@ func TestGetContainerClaimInfos(t *testing.T) {
 		},
 	} {
 		t.Run(test.description, func(t *testing.T) {
-			cache, err := newClaimInfoCache(t.TempDir(), draManagerStateFileName)
-			if err != nil {
-				t.Fatalf("error occur:%v", err)
-			}
-			manager := &ManagerImpl{
-				cache: cache,
-			}
+			manager, err := NewManager(nil, t.TempDir())
+			require.NoError(t, err, "create DRA manager")
 
 			if test.claimInfo != nil {
 				manager.cache.add(test.claimInfo)
@@ -896,22 +869,15 @@ func TestParallelPrepareUnprepareResources(t *testing.T) {
 	}
 	defer draServerInfo.teardownFn()
 
-	plg := plugin.NewRegistrationHandler(nil, getFakeNode, time.Second /* very short wiping delay for testing */)
+	// Create fake Kube client and DRA manager
+	fakeKubeClient := fake.NewSimpleClientset()
+	manager, err := NewManager(fakeKubeClient, t.TempDir())
+	require.NoError(t, err, "create DRA manager")
+
+	plg := plugin.NewRegistrationHandler(manager.draPlugins, nil, getFakeNode, time.Second /* very short wiping delay for testing */)
 	if err := plg.RegisterPlugin(driverName, draServerInfo.socketName, []string{drapb.DRAPluginService}, nil); err != nil {
 		t.Fatalf("failed to register plugin %s, err: %v", driverName, err)
 	}
-	defer plg.DeRegisterPlugin(driverName, draServerInfo.socketName)
-
-	// Create ClaimInfo cache
-	cache, err := newClaimInfoCache(t.TempDir(), draManagerStateFileName)
-	if err != nil {
-		t.Errorf("failed to newClaimInfoCache, err: %+v", err)
-		return
-	}
-
-	// Create fake Kube client and DRA manager
-	fakeKubeClient := fake.NewSimpleClientset()
-	manager := &ManagerImpl{kubeClient: fakeKubeClient, cache: cache}
 
 	// Call PrepareResources in parallel
 	var wgSync, wgStart sync.WaitGroup // groups to sync goroutines

pkg/kubelet/cm/dra/plugin/plugin.go

@@ -35,24 +35,6 @@ import (
 	"k8s.io/kubernetes/pkg/kubelet/metrics"
 )
 
-// NewDRAPluginClient returns a wrapper around those gRPC methods of a DRA
-// driver kubelet plugin which need to be called by kubelet. The wrapper
-// handles gRPC connection management and logging. Connections are reused
-// across different NewDRAPluginClient calls.
-//
-// It returns an informative error message including the driver name
-// with an explanation why the driver is not usable.
-func NewDRAPluginClient(driverName string) (*Plugin, error) {
-	if driverName == "" {
-		return nil, errors.New("DRA driver name is empty")
-	}
-	client := draPlugins.get(driverName)
-	if client == nil {
-		return nil, fmt.Errorf("DRA driver %s is not registered", driverName)
-	}
-	return client, nil
-}
-
 type Plugin struct {
 	name          string
 	backgroundCtx context.Context

pkg/kubelet/cm/dra/plugin/plugin_test.go

@@ -27,7 +27,9 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc"
+
 	drapbv1alpha4 "k8s.io/kubelet/pkg/apis/dra/v1alpha4"
 	drapbv1beta1 "k8s.io/kubelet/pkg/apis/dra/v1beta1"
 	"k8s.io/kubernetes/test/utils/ktesting"
@@ -135,15 +137,15 @@ func TestGRPCConnIsReused(t *testing.T) {
 	}
 
 	// ensure the plugin we are using is registered
-	draPlugins.add(p)
-	defer draPlugins.remove(pluginName, addr)
+	var draPlugins Store
+	tCtx.ExpectNoError(draPlugins.add(p), "add plugin")
 
 	// we call `NodePrepareResource` 2 times and check whether a new connection is created or the same is reused
 	for i := 0; i < 2; i++ {
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
-			client, err := NewDRAPluginClient(pluginName)
+			client, err := draPlugins.NewDRAPluginClient(pluginName)
 			if err != nil {
 				t.Error(err)
 				return
@@ -185,42 +187,33 @@ func TestGRPCConnIsReused(t *testing.T) {
 func TestNewDRAPluginClient(t *testing.T) {
 	for _, test := range []struct {
 		description string
-		setup       func(string) tearDown
+		setup       func(*Store) error
 		pluginName  string
 		shouldError bool
 	}{
 		{
 			description: "plugin name is empty",
-			setup: func(_ string) tearDown {
-				return func() {}
-			},
-			pluginName:  "",
 			shouldError: true,
 		},
 		{
 			description: "plugin name not found in the list",
-			setup: func(_ string) tearDown {
-				return func() {}
-			},
 			pluginName:  "plugin-name-not-found-in-the-list",
 			shouldError: true,
 		},
 		{
 			description: "plugin exists",
-			setup: func(name string) tearDown {
-				draPlugins.add(&Plugin{name: name})
-				return func() {
-					draPlugins.remove(name, "")
-				}
+			setup: func(draPlugins *Store) error {
+				return draPlugins.add(&Plugin{name: "dummy-plugin"})
 			},
 			pluginName: "dummy-plugin",
 		},
 	} {
 		t.Run(test.description, func(t *testing.T) {
-			teardown := test.setup(test.pluginName)
-			defer teardown()
-
-			client, err := NewDRAPluginClient(test.pluginName)
+			var draPlugins Store
+			if test.setup != nil {
+				require.NoError(t, test.setup(&draPlugins), "setup plugin")
+			}
+			client, err := draPlugins.NewDRAPluginClient(test.pluginName)
 			if test.shouldError {
 				assert.Nil(t, client)
 				assert.Error(t, err)
@@ -297,10 +290,9 @@ func TestGRPCMethods(t *testing.T) {
 				t.Fatal(err)
 			}
 
+			var draPlugins Store
 			draPlugins.add(p)
-			defer draPlugins.remove(pluginName, addr)
-
-			client, err := NewDRAPluginClient(pluginName)
+			client, err := draPlugins.NewDRAPluginClient(pluginName)
 			if err != nil {
 				t.Fatal(err)
 			}

pkg/kubelet/cm/dra/plugin/plugins_store.go

@@ -23,20 +23,37 @@ import (
 	"sync"
 )
 
-// PluginsStore holds a list of DRA Plugins.
-type pluginsStore struct {
+// Store keeps track of how to reach plugins registered for DRA drivers.
+// Each plugin has a gRPC endpoint. There may be more than one plugin per driver.
+//
+// The null Store is usable.
+type Store struct {
 	sync.RWMutex
 	// plugin name -> Plugin in the order in which they got added
 	store map[string][]*Plugin
 }
 
-// draPlugins map keeps track of all registered DRA plugins on the node
-// and their corresponding sockets.
-var draPlugins = &pluginsStore{}
+// NewDRAPluginClient returns a wrapper around those gRPC methods of a DRA
+// driver kubelet plugin which need to be called by kubelet. The wrapper
+// handles gRPC connection management and logging. Connections are reused
+// across different NewDRAPluginClient calls.
+//
+// It returns an informative error message including the driver name
+// with an explanation why the driver is not usable.
+func (s *Store) NewDRAPluginClient(driverName string) (*Plugin, error) {
+	if driverName == "" {
+		return nil, errors.New("DRA driver name is empty")
+	}
+	client := s.get(driverName)
+	if client == nil {
+		return nil, fmt.Errorf("DRA driver %s is not registered", driverName)
+	}
+	return client, nil
+}
 
 // Get lets you retrieve a DRA Plugin by name.
 // This method is protected by a mutex.
-func (s *pluginsStore) get(pluginName string) *Plugin {
+func (s *Store) get(pluginName string) *Plugin {
 	s.RLock()
 	defer s.RUnlock()
 
@@ -52,7 +69,7 @@ func (s *pluginsStore) get(pluginName string) *Plugin {
 
 // Set lets you save a DRA Plugin to the list and give it a specific name.
 // This method is protected by a mutex.
-func (s *pluginsStore) add(p *Plugin) error {
+func (s *Store) add(p *Plugin) error {
 	s.Lock()
 	defer s.Unlock()
 
@@ -72,7 +89,7 @@ func (s *pluginsStore) add(p *Plugin) error {
 // remove lets you remove one endpoint for a DRA Plugin.
 // This method is protected by a mutex. It returns the
 // plugin if found and true if that was the last instance
-func (s *pluginsStore) remove(pluginName, endpoint string) (*Plugin, bool) {
+func (s *Store) remove(pluginName, endpoint string) (*Plugin, bool) {
 	s.Lock()
 	defer s.Unlock()
 

pkg/kubelet/cm/dra/plugin/plugins_store_test.go

@@ -37,8 +37,8 @@ func TestAddSameName(t *testing.T) {
 	}
 
 	// ensure the plugin we are using is registered
+	var draPlugins Store
 	require.NoError(t, draPlugins.add(p))
-	defer draPlugins.remove(p.name, p.endpoint)
 
 	assert.False(t, firstWasCancelled, "should not cancel context after the first call")
 
@@ -73,6 +73,7 @@ func TestDelete(t *testing.T) {
 	}
 
 	// ensure the plugin we are using is registered
+	var draPlugins Store
 	draPlugins.add(p)
 
 	draPlugins.remove(p.name, "")

pkg/kubelet/cm/dra/plugin/registration.go

@@ -47,6 +47,7 @@ import (
 const defaultClientCallTimeout = 45 * time.Second
 
 // RegistrationHandler is the handler which is fed to the pluginwatcher API.
+// It updates a plugin store.
 type RegistrationHandler struct {
 	// backgroundCtx is used for all future activities of the handler.
 	// This is necessary because it implements APIs which don't
@@ -56,6 +57,7 @@ type RegistrationHandler struct {
 	kubeClient    kubernetes.Interface
 	getNode       func() (*v1.Node, error)
 	wipingDelay   time.Duration
+	draPlugins    *Store
 
 	wg    sync.WaitGroup
 	mutex sync.Mutex
@@ -72,17 +74,16 @@ type RegistrationHandler struct {
 
 var _ cache.PluginHandler = &RegistrationHandler{}
 
-// NewPluginHandler returns new registration handler.
+// NewRegistrationHandler creates a new registration handler.
 //
-// Must only be called once per process because it manages global state.
 // If a kubeClient is provided, then it synchronizes ResourceSlices
 // with the resource information provided by plugins.
-func NewRegistrationHandler(kubeClient kubernetes.Interface, getNode func() (*v1.Node, error), wipingDelay time.Duration) *RegistrationHandler {
+func NewRegistrationHandler(draPlugins *Store, kubeClient kubernetes.Interface, getNode func() (*v1.Node, error), wipingDelay time.Duration) *RegistrationHandler {
 	// The context and thus logger should come from the caller.
-	return newRegistrationHandler(context.TODO(), kubeClient, getNode, wipingDelay)
+	return newRegistrationHandler(context.TODO(), draPlugins, kubeClient, getNode, wipingDelay)
 }
 
-func newRegistrationHandler(ctx context.Context, kubeClient kubernetes.Interface, getNode func() (*v1.Node, error), wipingDelay time.Duration) *RegistrationHandler {
+func newRegistrationHandler(ctx context.Context, draPlugins *Store, kubeClient kubernetes.Interface, getNode func() (*v1.Node, error), wipingDelay time.Duration) *RegistrationHandler {
 	ctx, cancel := context.WithCancelCause(ctx)
 	handler := &RegistrationHandler{
 		backgroundCtx: klog.NewContext(ctx, klog.LoggerWithName(klog.FromContext(ctx), "DRA registration handler")),
@@ -91,6 +92,7 @@ func newRegistrationHandler(ctx context.Context, kubeClient kubernetes.Interface
 		getNode:       getNode,
 		wipingDelay:   wipingDelay,
 		pendingWipes:  make(map[string]*context.CancelCauseFunc),
+		draPlugins:    draPlugins,
 	}
 
 	// When kubelet starts up, no DRA driver has registered yet. None of
@@ -228,7 +230,7 @@ func (h *RegistrationHandler) RegisterPlugin(pluginName string, endpoint string,
 
 	// Storing endpoint of newly registered DRA Plugin into the map, where plugin name will be the key
 	// all other DRA components will be able to get the actual socket of DRA plugins by its name.
-	if err := draPlugins.add(pluginInstance); err != nil {
+	if err := h.draPlugins.add(pluginInstance); err != nil {
 		cancel(err)
 		// No wrapping, the error already contains details.
 		return err
@@ -280,7 +282,7 @@ func (h *RegistrationHandler) validateSupportedServices(pluginName string, suppo
 // DeRegisterPlugin is called when a plugin has removed its socket,
 // signaling it is no longer available.
 func (h *RegistrationHandler) DeRegisterPlugin(pluginName, endpoint string) {
-	if p, last := draPlugins.remove(pluginName, endpoint); p != nil {
+	if p, last := h.draPlugins.remove(pluginName, endpoint); p != nil {
 		// This logger includes endpoint and pluginName.
 		logger := klog.FromContext(p.backgroundCtx)
 		logger.V(3).Info("Deregister DRA plugin", "lastInstance", last)

pkg/kubelet/cm/dra/plugin/registration_test.go

@@ -167,7 +167,8 @@ func TestRegistrationHandler(t *testing.T) {
 			}
 
 			// The handler wipes all slices at startup.
-			handler := newRegistrationHandler(tCtx, client, getFakeNode, time.Second /* very short wiping delay for testing */)
+			var draPlugins Store
+			handler := newRegistrationHandler(tCtx, &draPlugins, client, getFakeNode, time.Second /* very short wiping delay for testing */)
 			tCtx.Cleanup(handler.Stop)
 			requireNoSlices := func() {
 				t.Helper()
@@ -213,7 +214,7 @@ func TestRegistrationHandler(t *testing.T) {
 				require.NoError(t, err)
 			}
 			plugin := draPlugins.get(test.pluginName)
-			assert.NotNil(t, plugin, "plugin should be registered")
+			require.NotNil(t, plugin, "plugin should be registered")
 			t.Cleanup(func() {
 				if client != nil {
 					// Create the slice as if the plugin had done that while it runs.

pkg/kubelet/cm/dra/types.go

@@ -17,43 +17,9 @@ limitations under the License.
 package dra
 
 import (
-	"context"
-
-	v1 "k8s.io/api/core/v1"
-	"k8s.io/apimachinery/pkg/types"
-	"k8s.io/kubernetes/pkg/kubelet/config"
 	kubecontainer "k8s.io/kubernetes/pkg/kubelet/container"
-	"k8s.io/kubernetes/pkg/kubelet/pluginmanager/cache"
 )
 
-// Manager manages all the DRA resource plugins running on a node.
-type Manager interface {
-	// GetWatcherHandler returns the plugin handler for the DRA.
-	GetWatcherHandler() cache.PluginHandler
-
-	// Start starts the reconcile loop of the manager.
-	// This will ensure that all claims are unprepared even if pods get deleted unexpectedly.
-	Start(ctx context.Context, activePods ActivePodsFunc, getNode GetNodeFunc, sourcesReady config.SourcesReady) error
-
-	// PrepareResources prepares resources for a pod.
-	// It communicates with the DRA resource plugin to prepare resources.
-	PrepareResources(ctx context.Context, pod *v1.Pod) error
-
-	// UnprepareResources calls NodeUnprepareResource GRPC from DRA plugin to unprepare pod resources
-	UnprepareResources(ctx context.Context, pod *v1.Pod) error
-
-	// GetResources gets a ContainerInfo object from the claimInfo cache.
-	// This information is used by the caller to update a container config.
-	GetResources(pod *v1.Pod, container *v1.Container) (*ContainerInfo, error)
-
-	// PodMightNeedToUnprepareResources returns true if the pod with the given UID
-	// might need to unprepare resources.
-	PodMightNeedToUnprepareResources(UID types.UID) bool
-
-	// GetContainerClaimInfos gets Container ClaimInfo objects
-	GetContainerClaimInfos(pod *v1.Pod, container *v1.Container) ([]*ClaimInfo, error)
-}
-
 // ContainerInfo contains information required by the runtime to consume prepared resources.
 type ContainerInfo struct {
 	// CDI Devices for the container