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

@@ -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

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.
 	//

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, "")
 }

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)