Merge pull request #128630 from SergeyKanzhelev/cancelOldDRAPluginContext

call cancel on plugin that is replaced by another plugin with the same name

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

@@ -42,7 +42,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) (replaced bool) {
+func (s *pluginsStore) add(p *Plugin) (replacedPlugin *Plugin, replaced bool) {
 	s.Lock()
 	defer s.Unlock()
 
@@ -50,9 +50,14 @@ func (s *pluginsStore) add(p *Plugin) (replaced bool) {
 		s.store = make(map[string]*Plugin)
 	}
 
-	_, exists := s.store[p.name]
+	replacedPlugin, exists := s.store[p.name]
 	s.store[p.name] = p
-	return exists
+
+	if replacedPlugin != nil && replacedPlugin.cancel != nil {
+		replacedPlugin.cancel(errors.New("plugin got replaced"))
+	}
+
+	return replacedPlugin, exists
 }
 
 // Delete lets you delete a DRA Plugin by name.

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

@@ -0,0 +1,71 @@
+/*
+Copyright 2024 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package plugin
+
+import (
+	"fmt"
+	"math/rand/v2"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestAddSameName(t *testing.T) {
+	// name will have a random value to avoid conflicts
+	pluginName := fmt.Sprintf("dummy-plugin-%d", rand.IntN(10000))
+
+	firstWasCancelled := false
+	p := &Plugin{
+		name:   pluginName,
+		cancel: func(err error) { firstWasCancelled = true },
+	}
+
+	// ensure the plugin we are using is registered
+	draPlugins.add(p)
+	defer draPlugins.delete(p.name)
+
+	assert.False(t, firstWasCancelled, "should not cancel context after the first call")
+
+	secondWasCancelled := false
+	p2 := &Plugin{
+		name:   pluginName,
+		cancel: func(err error) { secondWasCancelled = true },
+	}
+
+	draPlugins.add(p2)
+	defer draPlugins.delete(p2.name)
+
+	assert.True(t, firstWasCancelled, "should cancel context after the second call")
+	assert.False(t, secondWasCancelled, "should not cancel context of a new plugin")
+}
+
+func TestDelete(t *testing.T) {
+	pluginName := fmt.Sprintf("dummy-plugin-%d", rand.IntN(10000))
+
+	wasCancelled := false
+	p := &Plugin{
+		name:   pluginName,
+		cancel: func(err error) { wasCancelled = true },
+	}
+
+	// ensure the plugin we are using is registered
+	draPlugins.add(p)
+
+	draPlugins.delete(p.name)
+
+	assert.True(t, wasCancelled, "should cancel context after the second call")
+}

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

@@ -179,8 +179,9 @@ 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(pluginInstance) {
-		logger.V(1).Info("Already registered, previous plugin was replaced")
+
+	if oldPlugin, replaced := draPlugins.add(pluginInstance); replaced {
+		logger.V(1).Info("DRA plugin already registered, the old plugin was replaced and will be forgotten by the kubelet till the next kubelet restart", "oldEndpoint", oldPlugin.endpoint)
 	}
 
 	return nil