From f927cd0108858cc1688cf894d219b7058de95099 Mon Sep 17 00:00:00 2001 From: Patrick Ohly Date: Tue, 3 Jun 2025 12:23:57 +0200 Subject: [PATCH] DRA kubelet: simplify plugin creation and connection handling Instead of creating the gRPC connection on demand and forcing gRPC to connect, we establish it immediately and rely on gRPC to handle the underlying connection automatically like it usually does. It's not clear what benefit the one second connection timeout had. The way it is now, gRPC calls still fail when the underlying connection cannot be established. Having to have a separate context for establishing that connection just made the code more complex. The DRAPluginManager is the central component which manages plugins. Making it responsible for creating them reduces the number of places where a DRAPlugin struct needs to be initialized. Doing this in the DRAPluginManager instead of a stand-alone function simplifies the implementation of connection monitoring, because that will be something that is tied to the DRAPluginManager state. --- pkg/kubelet/cm/dra/plugin/dra_plugin.go | 76 ++----------- .../cm/dra/plugin/dra_plugin_manager.go | 103 +++++++++--------- .../cm/dra/plugin/dra_plugin_manager_test.go | 56 +++------- pkg/kubelet/cm/dra/plugin/dra_plugin_test.go | 50 ++------- 4 files changed, 88 insertions(+), 197 deletions(-) diff --git a/pkg/kubelet/cm/dra/plugin/dra_plugin.go b/pkg/kubelet/cm/dra/plugin/dra_plugin.go index ef1f52944f1..4b4b3092350 100644 --- a/pkg/kubelet/cm/dra/plugin/dra_plugin.go +++ b/pkg/kubelet/cm/dra/plugin/dra_plugin.go @@ -18,15 +18,10 @@ package plugin import ( "context" - "errors" "fmt" - "net" - "sync" "time" "google.golang.org/grpc" - "google.golang.org/grpc/connectivity" - "google.golang.org/grpc/credentials/insecure" "google.golang.org/grpc/status" "k8s.io/klog/v2" @@ -48,58 +43,13 @@ const defaultClientCallTimeout = 45 * time.Second // It implements the kubelet operations for preparing/unpreparing by calling // a gRPC interface that is implemented by the plugin. type DRAPlugin struct { - driverName string - backgroundCtx context.Context - cancel func(cause error) - - mutex sync.Mutex + driverName string conn *grpc.ClientConn endpoint string chosenService string // e.g. drapbv1beta1.DRAPluginService clientCallTimeout time.Duration } -func (p *DRAPlugin) getOrCreateGRPCConn() (*grpc.ClientConn, error) { - p.mutex.Lock() - defer p.mutex.Unlock() - - if p.conn != nil { - return p.conn, nil - } - - ctx := p.backgroundCtx - logger := klog.FromContext(ctx) - - network := "unix" - logger.V(4).Info("Creating new gRPC connection", "protocol", network, "endpoint", p.endpoint) - // grpc.Dial is deprecated. grpc.NewClient should be used instead. - // For now this gets ignored because this function is meant to establish - // the connection, with the one second timeout below. Perhaps that - // approach should be reconsidered? - //nolint:staticcheck - conn, err := grpc.Dial( - p.endpoint, - grpc.WithTransportCredentials(insecure.NewCredentials()), - grpc.WithContextDialer(func(ctx context.Context, target string) (net.Conn, error) { - return (&net.Dialer{}).DialContext(ctx, network, target) - }), - grpc.WithChainUnaryInterceptor(newMetricsInterceptor(p.driverName)), - ) - if err != nil { - return nil, err - } - - ctx, cancel := context.WithTimeout(context.Background(), time.Second) - defer cancel() - - if ok := conn.WaitForStateChange(ctx, connectivity.Connecting); !ok { - return nil, errors.New("timed out waiting for gRPC connection to be ready") - } - - p.conn = conn - return p.conn, nil -} - func (p *DRAPlugin) DriverName() string { return p.driverName } @@ -110,23 +60,21 @@ func (p *DRAPlugin) NodePrepareResources( opts ...grpc.CallOption, ) (*drapbv1beta1.NodePrepareResourcesResponse, error) { logger := klog.FromContext(ctx) + logger = klog.LoggerWithValues(logger, "driverName", p.driverName, "endpoint", p.endpoint) + ctx = klog.NewContext(ctx, logger) logger.V(4).Info("Calling NodePrepareResources rpc", "request", req) - conn, err := p.getOrCreateGRPCConn() - if err != nil { - return nil, err - } - ctx, cancel := context.WithTimeout(ctx, p.clientCallTimeout) defer cancel() + var err error var response *drapbv1beta1.NodePrepareResourcesResponse switch p.chosenService { case drapbv1beta1.DRAPluginService: - nodeClient := drapbv1beta1.NewDRAPluginClient(conn) + nodeClient := drapbv1beta1.NewDRAPluginClient(p.conn) response, err = nodeClient.NodePrepareResources(ctx, req) case drapbv1alpha4.NodeService: - nodeClient := drapbv1alpha4.NewNodeClient(conn) + nodeClient := drapbv1alpha4.NewNodeClient(p.conn) response, err = drapbv1alpha4.V1Alpha4ClientWrapper{NodeClient: nodeClient}.NodePrepareResources(ctx, req) default: // Shouldn't happen, validateSupportedServices should only @@ -144,22 +92,20 @@ func (p *DRAPlugin) NodeUnprepareResources( ) (*drapbv1beta1.NodeUnprepareResourcesResponse, error) { logger := klog.FromContext(ctx) logger.V(4).Info("Calling NodeUnprepareResource rpc", "request", req) - - conn, err := p.getOrCreateGRPCConn() - if err != nil { - return nil, err - } + logger = klog.LoggerWithValues(logger, "driverName", p.driverName, "endpoint", p.endpoint) + ctx = klog.NewContext(ctx, logger) ctx, cancel := context.WithTimeout(ctx, p.clientCallTimeout) defer cancel() + var err error var response *drapbv1beta1.NodeUnprepareResourcesResponse switch p.chosenService { case drapbv1beta1.DRAPluginService: - nodeClient := drapbv1beta1.NewDRAPluginClient(conn) + nodeClient := drapbv1beta1.NewDRAPluginClient(p.conn) response, err = nodeClient.NodeUnprepareResources(ctx, req) case drapbv1alpha4.NodeService: - nodeClient := drapbv1alpha4.NewNodeClient(conn) + nodeClient := drapbv1alpha4.NewNodeClient(p.conn) response, err = drapbv1alpha4.V1Alpha4ClientWrapper{NodeClient: nodeClient}.NodeUnprepareResources(ctx, req) default: // Shouldn't happen, validateSupportedServices should only diff --git a/pkg/kubelet/cm/dra/plugin/dra_plugin_manager.go b/pkg/kubelet/cm/dra/plugin/dra_plugin_manager.go index bbc35597a31..307edebb35f 100644 --- a/pkg/kubelet/cm/dra/plugin/dra_plugin_manager.go +++ b/pkg/kubelet/cm/dra/plugin/dra_plugin_manager.go @@ -24,6 +24,8 @@ import ( "sync" "time" + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/insecure" v1 "k8s.io/api/core/v1" resourceapi "k8s.io/api/resource/v1beta1" apierrors "k8s.io/apimachinery/pkg/api/errors" @@ -35,6 +37,7 @@ import ( drapbv1alpha4 "k8s.io/kubelet/pkg/apis/dra/v1alpha4" drapbv1beta1 "k8s.io/kubelet/pkg/apis/dra/v1beta1" "k8s.io/kubernetes/pkg/kubelet/pluginmanager/cache" + "k8s.io/utils/ptr" ) // DRAPluginManager keeps track of how to reach plugins registered for DRA drivers. @@ -44,7 +47,7 @@ import ( // [cache.PluginHandler] interface and needs to be added to the // plugin manager. // -// The null DRAPluginManager is not usable, use NewPluginManager. +// The null DRAPluginManager is not usable, use NewDRAPluginManager. type DRAPluginManager struct { // backgroundCtx is used for all future activities of the DRAPluginManager. // This is necessary because it implements APIs which don't @@ -58,7 +61,7 @@ type DRAPluginManager struct { wg sync.WaitGroup mutex sync.RWMutex - // driver name -> Plugin in the order in which they got added + // driver name -> DRAPlugin in the order in which they got added store map[string][]*DRAPlugin // pendingWipes maps a driver name to a cancel function for @@ -112,8 +115,20 @@ func NewDRAPluginManager(ctx context.Context, kubeClient kubernetes.Interface, g // Stop cancels any remaining background activities and blocks until all goroutines have stopped. func (pm *DRAPluginManager) Stop() { + defer pm.wg.Wait() // Must run after unlocking our mutex. + pm.mutex.Lock() + defer pm.mutex.Unlock() + pm.cancel(errors.New("Stop was called")) - pm.wg.Wait() + + // Close all connections, otherwise gRPC keeps doing things in the background. + for _, plugins := range pm.store { + for _, plugin := range plugins { + if err := plugin.conn.Close(); err != nil { + klog.FromContext(pm.backgroundCtx).Error(err, "Closing gRPC connection", "driverName", plugin.driverName, "endpoint", plugin.endpoint) + } + } + } } // wipeResourceSlices deletes ResourceSlices of the node, optionally just for a specific driver. @@ -196,7 +211,7 @@ func (pm *DRAPluginManager) GetPlugin(driverName string) (*DRAPlugin, error) { return plugin, nil } -// get lets you retrieve a DRA Plugin by name. +// get lets you retrieve a DRA DRAPlugin by name. func (pm *DRAPluginManager) get(driverName string) *DRAPlugin { pm.mutex.RLock() defer pm.mutex.RUnlock() @@ -223,46 +238,16 @@ func (pm *DRAPluginManager) get(driverName string) *DRAPlugin { // in advance which version to use resp. which optional services the plugin // supports. func (pm *DRAPluginManager) RegisterPlugin(driverName string, endpoint string, supportedServices []string, pluginClientTimeout *time.Duration) error { - // Prepare a context with its own logger for the plugin. - // - // The lifecycle of the plugin's background activities is tied to our - // root context, so canceling that will also cancel the plugin. - // - // The logger injects the driver name and endpoint as additional values - // into all log output related to the plugin. - ctx := pm.backgroundCtx - logger := klog.FromContext(ctx) - logger = klog.LoggerWithValues(logger, "driverName", driverName, "endpoint", endpoint) - ctx = klog.NewContext(ctx, logger) - chosenService, err := pm.validateSupportedServices(driverName, supportedServices) if err != nil { return fmt.Errorf("invalid supported gRPC versions of DRA driver plugin %s at endpoint %s: %w", driverName, endpoint, err) } - var timeout time.Duration - if pluginClientTimeout == nil { - timeout = defaultClientCallTimeout - } else { - timeout = *pluginClientTimeout - } + timeout := ptr.Deref(pluginClientTimeout, defaultClientCallTimeout) - ctx, cancel := context.WithCancelCause(ctx) - - plugin := &DRAPlugin{ - driverName: driverName, - backgroundCtx: ctx, - cancel: cancel, - conn: nil, - endpoint: endpoint, - chosenService: chosenService, - clientCallTimeout: timeout, - } - - // Storing endpoint of newly registered DRA Plugin into the map, where the DRA driver name will be the key + // Storing endpoint of newly registered DRA DRAPlugin into the map, where the DRA driver name will be the key // under which the manager will be able to get a plugin when it needs to call it. - if err := pm.add(plugin); err != nil { - cancel(err) + if err := pm.add(driverName, endpoint, chosenService, timeout); err != nil { // No wrapping, the error already contains details. return err } @@ -270,23 +255,44 @@ func (pm *DRAPluginManager) RegisterPlugin(driverName string, endpoint string, s return nil } -func (pm *DRAPluginManager) add(p *DRAPlugin) error { +func (pm *DRAPluginManager) add(driverName string, endpoint string, chosenService string, clientCallTimeout time.Duration) error { pm.mutex.Lock() defer pm.mutex.Unlock() if pm.store == nil { pm.store = make(map[string][]*DRAPlugin) } - for _, oldP := range pm.store[p.driverName] { - if oldP.endpoint == p.endpoint { + for _, oldP := range pm.store[driverName] { + if oldP.endpoint == endpoint { // One plugin instance cannot hijack the endpoint of another instance. - return fmt.Errorf("endpoint %s already registered for plugin %s", p.endpoint, p.driverName) + return fmt.Errorf("endpoint %s already registered for DRA driver plugin %s", endpoint, driverName) } } - logger := klog.FromContext(p.backgroundCtx) + logger := klog.FromContext(pm.backgroundCtx) + + // The gRPC connection gets created once. gRPC then connects to the gRPC server on demand. + target := "unix:" + endpoint + logger.V(4).Info("Creating new gRPC connection", "target", target) + conn, err := grpc.NewClient( + target, + grpc.WithTransportCredentials(insecure.NewCredentials()), + grpc.WithChainUnaryInterceptor(newMetricsInterceptor(driverName)), + ) + if err != nil { + return fmt.Errorf("create gRPC connection to DRA driver %s plugin at endpoint %s: %w", driverName, endpoint, err) + } + + p := &DRAPlugin{ + driverName: driverName, + endpoint: endpoint, + conn: conn, + chosenService: chosenService, + clientCallTimeout: clientCallTimeout, + } + pm.store[p.driverName] = append(pm.store[p.driverName], p) - logger.V(3).Info("Registered DRA plugin", "numInstances", len(pm.store[p.driverName])) + logger.V(3).Info("Registered DRA plugin", "driverName", p.driverName, "endpoint", p.endpoint, "chosenService", p.chosenService, "numPlugins", len(pm.store[p.driverName])) pm.sync(p.driverName) return nil } @@ -317,13 +323,8 @@ func (pm *DRAPluginManager) remove(driverName, endpoint string) { } else { pm.store[driverName] = slices.Delete(plugins, i, i+1) } - if p.cancel != nil { - // This cancels background attempts to establish a connection to the plugin. - // TODO: remove this in favor of non-blocking connection management. - p.cancel(errors.New("plugin got removed")) - } - logger := klog.FromContext(p.backgroundCtx) - logger.V(3).Info("Unregistered DRA plugin", "numInstances", len(pm.store[driverName])) + logger := klog.FromContext(pm.backgroundCtx) + logger.V(3).Info("Unregistered DRA plugin", "driverName", p.driverName, "endpoint", p.endpoint, "numPlugins", len(pm.store[driverName])) pm.sync(driverName) } @@ -350,7 +351,7 @@ func (pm *DRAPluginManager) sync(driverName string) { ctx, cancel := context.WithCancelCause(pm.backgroundCtx) ctx = klog.NewContext(ctx, logger) - // Clean up the ResourceSlices for the deleted Plugin since it + // Clean up the ResourceSlices for the deleted DRAPlugin since it // may have died without doing so itself and might never come // back. // diff --git a/pkg/kubelet/cm/dra/plugin/dra_plugin_manager_test.go b/pkg/kubelet/cm/dra/plugin/dra_plugin_manager_test.go index 2abc9869af9..c104c9a88fb 100644 --- a/pkg/kubelet/cm/dra/plugin/dra_plugin_manager_test.go +++ b/pkg/kubelet/cm/dra/plugin/dra_plugin_manager_test.go @@ -21,7 +21,6 @@ import ( "math/rand/v2" "testing" - "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "k8s.io/kubernetes/test/utils/ktesting" ) @@ -31,58 +30,39 @@ func TestAddSameName(t *testing.T) { // name will have a random value to avoid conflicts driverName := fmt.Sprintf("dummy-driver-%d", rand.IntN(10000)) - firstWasCancelled := false - p := &DRAPlugin{ - driverName: driverName, - backgroundCtx: tCtx, - endpoint: "old", - cancel: func(err error) { firstWasCancelled = true }, - } - // ensure the plugin we are using is registered draPlugins := NewDRAPluginManager(tCtx, nil, nil, 0) - require.NoError(t, draPlugins.add(p)) - - assert.False(t, firstWasCancelled, "should not cancel context after the first call") + tCtx.ExpectNoError(draPlugins.add(driverName, "old.sock", "", defaultClientCallTimeout), "add first plugin") + p, err := draPlugins.GetPlugin(driverName) + tCtx.ExpectNoError(err, "get first plugin") // Same name, same endpoint -> error. - require.Error(t, draPlugins.add(p)) + require.Error(tCtx, draPlugins.add(driverName, "old.sock", "", defaultClientCallTimeout)) - secondWasCancelled := false - p2 := &DRAPlugin{ - driverName: driverName, - backgroundCtx: tCtx, - endpoint: "new", - cancel: func(err error) { secondWasCancelled = true }, + tCtx.ExpectNoError(draPlugins.add(driverName, "new.sock", "", defaultClientCallTimeout), "add second plugin") + p2, err := draPlugins.GetPlugin(driverName) + tCtx.ExpectNoError(err, "get second plugin") + if p == p2 { + tCtx.Fatal("expected to get second plugin, got first one again") } - require.NoError(t, draPlugins.add(p2)) - defer draPlugins.remove(p2.driverName, p2.endpoint) - - 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") // Remove old plugin. draPlugins.remove(p.driverName, 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") + plugin, err := draPlugins.GetPlugin(driverName) + + // Now we should have p2 left. + tCtx.ExpectNoError(err, "get plugin") + if p2 != plugin { + tCtx.Fatal("expected to get second plugin again, got something else") + } } func TestDelete(t *testing.T) { tCtx := ktesting.Init(t) driverName := fmt.Sprintf("dummy-driver-%d", rand.IntN(10000)) - wasCancelled := false - p := &DRAPlugin{ - driverName: driverName, - backgroundCtx: tCtx, - cancel: func(err error) { wasCancelled = true }, - } - // ensure the plugin we are using is registered draPlugins := NewDRAPluginManager(tCtx, nil, nil, 0) - draPlugins.add(p) - - draPlugins.remove(p.driverName, "") - - assert.True(t, wasCancelled, "should cancel context after the second call") + tCtx.ExpectNoError(draPlugins.add(driverName, "dra.sock", "", defaultClientCallTimeout), "add plugin") + draPlugins.remove(driverName, "") } diff --git a/pkg/kubelet/cm/dra/plugin/dra_plugin_test.go b/pkg/kubelet/cm/dra/plugin/dra_plugin_test.go index 9627e45e8b4..2644661cd8a 100644 --- a/pkg/kubelet/cm/dra/plugin/dra_plugin_test.go +++ b/pkg/kubelet/cm/dra/plugin/dra_plugin_test.go @@ -117,28 +117,13 @@ func TestGRPCConnIsReused(t *testing.T) { m := sync.Mutex{} driverName := "dummy-driver" - p := &DRAPlugin{ - driverName: driverName, - backgroundCtx: tCtx, - endpoint: addr, - chosenService: service, - clientCallTimeout: defaultClientCallTimeout, - } - - conn, err := p.getOrCreateGRPCConn() - defer func() { - err := conn.Close() - if err != nil { - t.Error(err) - } - }() - if err != nil { - t.Fatal(err) - } // ensure the plugin we are using is registered draPlugins := NewDRAPluginManager(tCtx, nil, nil, 0) - tCtx.ExpectNoError(draPlugins.add(p), "add plugin") + tCtx.ExpectNoError(draPlugins.add(driverName, addr, service, defaultClientCallTimeout), "add plugin") + plugin, err := draPlugins.GetPlugin(driverName) + tCtx.ExpectNoError(err, "get plugin") + conn := plugin.conn // we call `NodePrepareResource` 2 times and check whether a new connection is created or the same is reused for i := 0; i < 2; i++ { @@ -164,9 +149,7 @@ func TestGRPCConnIsReused(t *testing.T) { _, err = plugin.NodePrepareResources(tCtx, req) assert.NoError(t, err) - plugin.mutex.Lock() conn := plugin.conn - plugin.mutex.Unlock() m.Lock() defer m.Unlock() @@ -196,14 +179,14 @@ func TestGetDRAPlugin(t *testing.T) { shouldError: true, }, { - description: "driver-name not found in the list", + description: "driver name not found in the list", driverName: "driver-name-not-found-in-the-list", shouldError: true, }, { description: "plugin exists", setup: func(draPlugins *DRAPluginManager) error { - return draPlugins.add(&DRAPlugin{backgroundCtx: draPlugins.backgroundCtx, driverName: "dummy-driver"}) + return draPlugins.add("dummy-driver", "/tmp/dra.sock", "", defaultClientCallTimeout) }, driverName: "dummy-driver", }, @@ -272,27 +255,8 @@ func TestGRPCMethods(t *testing.T) { defer teardown() driverName := "dummy-driver" - p := &DRAPlugin{ - driverName: driverName, - backgroundCtx: tCtx, - endpoint: addr, - chosenService: test.chosenService, - clientCallTimeout: defaultClientCallTimeout, - } - - conn, err := p.getOrCreateGRPCConn() - defer func() { - err := conn.Close() - if err != nil { - t.Error(err) - } - }() - if err != nil { - t.Fatal(err) - } - draPlugins := NewDRAPluginManager(tCtx, nil, nil, 0) - draPlugins.add(p) + tCtx.ExpectNoError(draPlugins.add(driverName, addr, test.chosenService, defaultClientCallTimeout)) plugin, err := draPlugins.GetPlugin(driverName) if err != nil { t.Fatal(err)