kubelet: add DRA plugin name to the Plugin struct

Author: bart0sh

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

@@ -51,6 +51,7 @@ func NewDRAPluginClient(pluginName string) (*Plugin, error) {
 }
 
 type Plugin struct {
+	name          string
 	backgroundCtx context.Context
 	cancel        func(cause error)
 

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

@@ -114,7 +114,9 @@ func TestGRPCConnIsReused(t *testing.T) {
 	wg := sync.WaitGroup{}
 	m := sync.Mutex{}
 
+	pluginName := "dummy-plugin"
 	p := &Plugin{
+		name:              pluginName,
 		backgroundCtx:     tCtx,
 		endpoint:          addr,
 		clientCallTimeout: defaultClientCallTimeout,
@@ -132,15 +134,15 @@ func TestGRPCConnIsReused(t *testing.T) {
 	}
 
 	// ensure the plugin we are using is registered
-	draPlugins.add("dummy-plugin", p)
-	defer draPlugins.delete("dummy-plugin")
+	draPlugins.add(p)
+	defer draPlugins.delete(pluginName)
 
 	// 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("dummy-plugin")
+			client, err := NewDRAPluginClient(pluginName)
 			if err != nil {
 				t.Error(err)
 				return
@@ -205,7 +207,7 @@ func TestNewDRAPluginClient(t *testing.T) {
 		{
 			description: "plugin exists",
 			setup: func(name string) tearDown {
-				draPlugins.add(name, &Plugin{})
+				draPlugins.add(&Plugin{name: name})
 				return func() {
 					draPlugins.delete(name)
 				}
@@ -251,7 +253,9 @@ func TestNodeUnprepareResources(t *testing.T) {
 			}
 			defer teardown()
 
+			pluginName := "dummy-plugin"
 			p := &Plugin{
+				name:              pluginName,
 				backgroundCtx:     tCtx,
 				endpoint:          addr,
 				clientCallTimeout: defaultClientCallTimeout,
@@ -268,10 +272,10 @@ func TestNodeUnprepareResources(t *testing.T) {
 				t.Fatal(err)
 			}
 
-			draPlugins.add("dummy-plugin", p)
-			defer draPlugins.delete("dummy-plugin")
+			draPlugins.add(p)
+			defer draPlugins.delete(pluginName)
 
-			client, err := NewDRAPluginClient("dummy-plugin")
+			client, err := NewDRAPluginClient(pluginName)
 			if err != nil {
 				t.Fatal(err)
 			}

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

@@ -42,16 +42,16 @@ 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(pluginName string, p *Plugin) (replaced bool) {
+func (s *pluginsStore) add(p *Plugin) (replaced bool) {
 	s.Lock()
 	defer s.Unlock()
 
 	if s.store == nil {
 		s.store = make(map[string]*Plugin)
 	}
 
-	_, exists := s.store[pluginName]
-	s.store[pluginName] = p
+	_, exists := s.store[p.name]
+	s.store[p.name] = p
 	return exists
 }
 

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

@@ -160,6 +160,7 @@ func (h *RegistrationHandler) RegisterPlugin(pluginName string, endpoint string,
 	ctx, cancel := context.WithCancelCause(ctx)
 
 	pluginInstance := &Plugin{
+		name:                    pluginName,
 		backgroundCtx:           ctx,
 		cancel:                  cancel,
 		conn:                    nil,
@@ -170,7 +171,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 draPlugins.add(pluginName, pluginInstance) {
+	if draPlugins.add(pluginInstance) {
 		logger.V(1).Info("Already registered, previous plugin was replaced")
 	}