Merge pull request kubernetes#129832 from pohly/dra-seamless-upgrade

DRA: seamless driver upgrades

Author: k8s-ci-robot

pkg/kubelet/cm/devicemanager/plugin/v1beta1/handler.go

@@ -43,8 +43,8 @@ func (s *server) RegisterPlugin(pluginName string, endpoint string, versions []s
 	return s.connectClient(pluginName, endpoint)
 }
 
-func (s *server) DeRegisterPlugin(pluginName string) {
-	klog.V(2).InfoS("Deregistering plugin", "plugin", pluginName)
+func (s *server) DeRegisterPlugin(pluginName, endpoint string) {
+	klog.V(2).InfoS("Deregistering plugin", "plugin", pluginName, "endpoint", endpoint)
 	client := s.getClient(pluginName)
 	if client != nil {
 		s.disconnectClient(pluginName, client)
@@ -88,7 +88,6 @@ func (s *server) disconnectClient(name string, c Client) error {
 	s.deregisterClient(name)
 	return c.Disconnect()
 }
-
 func (s *server) registerClient(name string, c Client) {
 	s.mutex.Lock()
 	defer s.mutex.Unlock()

pkg/kubelet/cm/dra/manager.go

@@ -98,7 +98,20 @@ func NewManagerImpl(kubeClient clientset.Interface, stateFileDirectory string, n
 }
 
 func (m *ManagerImpl) GetWatcherHandler() cache.PluginHandler {
-	return cache.PluginHandler(dra.NewRegistrationHandler(m.kubeClient, m.getNode))
+	// The time that DRA drivers have to come back after being unregistered
+	// before the kubelet removes their ResourceSlices.
+	//
+	// This must be long enough to actually allow stopping a pod and
+	// starting the replacement (otherwise ResourceSlices get deleted
+	// unnecessarily) and not too long (otherwise the time window were
+	// pods might still get scheduled to the node after removal of a
+	// driver is too long).
+	//
+	// 30 seconds might be long enough for a simple container restart.
+	// 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))
 }
 
 // Start starts the reconcile loop of the manager.

pkg/kubelet/cm/dra/manager_test.go

@@ -580,11 +580,11 @@ func TestPrepareResources(t *testing.T) {
 			}
 			defer draServerInfo.teardownFn()
 
-			plg := plugin.NewRegistrationHandler(nil, getFakeNode)
+			plg := plugin.NewRegistrationHandler(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) // for sake of next tests
+			defer plg.DeRegisterPlugin(test.driverName, draServerInfo.socketName) // for sake of next tests
 
 			if test.claimInfo != nil {
 				manager.cache.add(test.claimInfo)
@@ -717,11 +717,11 @@ func TestUnprepareResources(t *testing.T) {
 			}
 			defer draServerInfo.teardownFn()
 
-			plg := plugin.NewRegistrationHandler(nil, getFakeNode)
+			plg := plugin.NewRegistrationHandler(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) // for sake of next tests
+			defer plg.DeRegisterPlugin(test.driverName, draServerInfo.socketName) // for sake of next tests
 
 			manager := &ManagerImpl{
 				kubeClient: fakeKubeClient,
@@ -887,11 +887,11 @@ func TestParallelPrepareUnprepareResources(t *testing.T) {
 	}
 	defer draServerInfo.teardownFn()
 
-	plg := plugin.NewRegistrationHandler(nil, getFakeNode)
+	plg := plugin.NewRegistrationHandler(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)
+	defer plg.DeRegisterPlugin(driverName, draServerInfo.socketName)
 
 	// Create ClaimInfo cache
 	cache, err := newClaimInfoCache(t.TempDir(), draManagerStateFileName)

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

@@ -136,7 +136,7 @@ func TestGRPCConnIsReused(t *testing.T) {
 
 	// ensure the plugin we are using is registered
 	draPlugins.add(p)
-	defer draPlugins.delete(pluginName)
+	defer draPlugins.remove(pluginName, addr)
 
 	// we call `NodePrepareResource` 2 times and check whether a new connection is created or the same is reused
 	for i := 0; i < 2; i++ {
@@ -210,7 +210,7 @@ func TestNewDRAPluginClient(t *testing.T) {
 			setup: func(name string) tearDown {
 				draPlugins.add(&Plugin{name: name})
 				return func() {
-					draPlugins.delete(name)
+					draPlugins.remove(name, "")
 				}
 			},
 			pluginName: "dummy-plugin",
@@ -298,7 +298,7 @@ func TestGRPCMethods(t *testing.T) {
 			}
 
 			draPlugins.add(p)
-			defer draPlugins.delete(pluginName)
+			defer draPlugins.remove(pluginName, addr)
 
 			client, err := NewDRAPluginClient(pluginName)
 			if err != nil {

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

@@ -18,13 +18,16 @@ package plugin
 
 import (
 	"errors"
+	"fmt"
+	"slices"
 	"sync"
 )
 
 // PluginsStore holds a list of DRA Plugins.
 type pluginsStore struct {
 	sync.RWMutex
-	store map[string]*Plugin
+	// 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
@@ -37,43 +40,57 @@ func (s *pluginsStore) get(pluginName string) *Plugin {
 	s.RLock()
 	defer s.RUnlock()
 
-	return s.store[pluginName]
+	instances := s.store[pluginName]
+	if len(instances) == 0 {
+		return nil
+	}
+	// Heuristic: pick the most recent one. It's most likely
+	// the newest, except when kubelet got restarted and registered
+	// all running plugins in random order.
+	return instances[len(instances)-1]
 }
 
 // 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) (replacedPlugin *Plugin, replaced bool) {
+func (s *pluginsStore) add(p *Plugin) error {
 	s.Lock()
 	defer s.Unlock()
 
 	if s.store == nil {
-		s.store = make(map[string]*Plugin)
+		s.store = make(map[string][]*Plugin)
 	}
-
-	replacedPlugin, exists := s.store[p.name]
-	s.store[p.name] = p
-
-	if replacedPlugin != nil && replacedPlugin.cancel != nil {
-		replacedPlugin.cancel(errors.New("plugin got replaced"))
+	for _, oldP := range s.store[p.name] {
+		if oldP.endpoint == p.endpoint {
+			// One plugin instance cannot hijack the endpoint of another instance.
+			return fmt.Errorf("endpoint %s already registered for plugin %s", p.endpoint, p.name)
+		}
 	}
-
-	return replacedPlugin, exists
+	s.store[p.name] = append(s.store[p.name], p)
+	return nil
 }
 
-// Delete lets you delete a DRA Plugin by name.
-// This method is protected by a mutex.
-func (s *pluginsStore) delete(pluginName string) *Plugin {
+// 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) {
 	s.Lock()
 	defer s.Unlock()
 
-	p, exists := s.store[pluginName]
-	if !exists {
-		return nil
+	instances := s.store[pluginName]
+	i := slices.IndexFunc(instances, func(p *Plugin) bool { return p.endpoint == endpoint })
+	if i == -1 {
+		return nil, false
+	}
+	p := instances[i]
+	last := len(instances) == 1
+	if last {
+		delete(s.store, pluginName)
+	} else {
+		s.store[pluginName] = slices.Delete(instances, i, i+1)
 	}
+
 	if p.cancel != nil {
 		p.cancel(errors.New("plugin got removed"))
 	}
-	delete(s.store, pluginName)
-
-	return p
+	return p, last
 }

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

@@ -22,6 +22,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestAddSameName(t *testing.T) {
@@ -30,26 +31,35 @@ func TestAddSameName(t *testing.T) {
 
 	firstWasCancelled := false
 	p := &Plugin{
-		name:   pluginName,
-		cancel: func(err error) { firstWasCancelled = true },
+		name:     pluginName,
+		endpoint: "old",
+		cancel:   func(err error) { firstWasCancelled = true },
 	}
 
 	// ensure the plugin we are using is registered
-	draPlugins.add(p)
-	defer draPlugins.delete(p.name)
+	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")
 
+	// Same name, same endpoint -> error.
+	require.Error(t, draPlugins.add(p))
+
 	secondWasCancelled := false
 	p2 := &Plugin{
-		name:   pluginName,
-		cancel: func(err error) { secondWasCancelled = true },
+		name:     pluginName,
+		endpoint: "new",
+		cancel:   func(err error) { secondWasCancelled = true },
 	}
+	require.NoError(t, draPlugins.add(p2))
+	defer draPlugins.remove(p2.name, p2.endpoint)
 
-	draPlugins.add(p2)
-	defer draPlugins.delete(p2.name)
+	assert.False(t, firstWasCancelled, "should not cancel context after registering the second instance")
+	assert.False(t, secondWasCancelled, "should not cancel context of a new plugin")
 
-	assert.True(t, firstWasCancelled, "should cancel context after the second call")
+	// Remove old plugin.
+	draPlugins.remove(p.name, p.endpoint)
+	assert.True(t, firstWasCancelled, "should have canceled context after the explicit removal")
 	assert.False(t, secondWasCancelled, "should not cancel context of a new plugin")
 }
 
@@ -65,7 +75,7 @@ func TestDelete(t *testing.T) {
 	// ensure the plugin we are using is registered
 	draPlugins.add(p)
 
-	draPlugins.delete(p.name)
+	draPlugins.remove(p.name, "")
 
 	assert.True(t, wasCancelled, "should cancel context after the second call")
 }