From 344df7835cdc0e5d1b9008dda4fa76ac223657bf Mon Sep 17 00:00:00 2001 From: Abu Kashem Date: Tue, 1 Jul 2025 13:32:45 -0400 Subject: [PATCH 1/2] refactor dra fake grpc server --- pkg/kubelet/cm/dra/plugin/dra_plugin_test.go | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/pkg/kubelet/cm/dra/plugin/dra_plugin_test.go b/pkg/kubelet/cm/dra/plugin/dra_plugin_test.go index 9a7c856bdb4..5779b25a56a 100644 --- a/pkg/kubelet/cm/dra/plugin/dra_plugin_test.go +++ b/pkg/kubelet/cm/dra/plugin/dra_plugin_test.go @@ -37,6 +37,13 @@ import ( "k8s.io/kubernetes/test/utils/ktesting" ) +// this interface satisfies what setupGRPCServerWithFake needs +type fakeGRPCServerInterface interface { + drapbv1.DRAPluginServer + drahealthv1alpha1.DRAResourceHealthServer + drapbv1.UnsafeDRAPluginServer +} + type fakeGRPCServer struct { drapbv1beta1.UnimplementedDRAPluginServer drahealthv1alpha1.UnimplementedDRAResourceHealthServer @@ -44,6 +51,7 @@ type fakeGRPCServer struct { } var _ drapbv1.DRAPluginServer = &fakeGRPCServer{} +var _ fakeGRPCServerInterface = &fakeGRPCServer{} func (f *fakeGRPCServer) NodePrepareResources(ctx context.Context, in *drapbv1.NodePrepareResourcesRequest) (*drapbv1.NodePrepareResourcesResponse, error) { return &drapbv1.NodePrepareResourcesResponse{Claims: map[string]*drapbv1.NodePrepareResourceResponse{"claim-uid": { @@ -82,7 +90,7 @@ func (f *fakeGRPCServer) NodeWatchResources(in *drahealthv1alpha1.NodeWatchResou // tearDown is an idempotent cleanup function. type tearDown func() -func setupFakeGRPCServer(service, addr string) (tearDown, error) { +func setupGRPCServerWithFake(service, addr string, fakeGRPCServer fakeGRPCServerInterface) (tearDown, error) { ctx, cancel := context.WithCancel(context.Background()) teardown := func() { cancel() @@ -95,7 +103,6 @@ func setupFakeGRPCServer(service, addr string) (tearDown, error) { } s := grpc.NewServer() - fakeGRPCServer := &fakeGRPCServer{} switch service { case drapbv1.DRAPluginService: @@ -127,6 +134,10 @@ func setupFakeGRPCServer(service, addr string) (tearDown, error) { return teardown, nil } +func setupFakeGRPCServer(service, addr string) (tearDown, error) { + return setupGRPCServerWithFake(service, addr, &fakeGRPCServer{}) +} + func TestGRPCConnIsReused(t *testing.T) { tCtx := ktesting.Init(t) service := drapbv1.DRAPluginService From 11e3147fdd2a55b59902a8513ecf98a2260790a2 Mon Sep 17 00:00:00 2001 From: Abu Kashem Date: Tue, 1 Jul 2025 13:33:43 -0400 Subject: [PATCH 2/2] kubelet: add a test to verify that dra plugin enforces timeout --- pkg/kubelet/cm/dra/plugin/dra_plugin_test.go | 113 +++++++++++++++++++ 1 file changed, 113 insertions(+) diff --git a/pkg/kubelet/cm/dra/plugin/dra_plugin_test.go b/pkg/kubelet/cm/dra/plugin/dra_plugin_test.go index 5779b25a56a..1465d0dbcfb 100644 --- a/pkg/kubelet/cm/dra/plugin/dra_plugin_test.go +++ b/pkg/kubelet/cm/dra/plugin/dra_plugin_test.go @@ -35,6 +35,9 @@ import ( drapbv1 "k8s.io/kubelet/pkg/apis/dra/v1" drapbv1beta1 "k8s.io/kubelet/pkg/apis/dra/v1beta1" "k8s.io/kubernetes/test/utils/ktesting" + + grpccodes "google.golang.org/grpc/codes" + grpcstatus "google.golang.org/grpc/status" ) // this interface satisfies what setupGRPCServerWithFake needs @@ -336,6 +339,116 @@ func TestGRPCMethods(t *testing.T) { } } +func TestGRPCWithTimeoutEnforced(t *testing.T) { + tCtx := ktesting.Init(t) + + service := drapbv1.DRAPluginService + addr := path.Join(t.TempDir(), "dra.sock") + + // we force NodePrepareResources to block, this will cause the timeout + // to elapse, as a consequence the client should abort. + // once the client aborts, we unblock NodePrepareResources + blocked := make(chan struct{}) + server := &timeoutFakeGRPCServer{ + t: t, + fakeGRPCServer: &fakeGRPCServer{}, + blocked: blocked, + done: make(chan struct{}), + } + teardown, err := setupGRPCServerWithFake(service, addr, server) + require.NoError(t, err, "failed to setup grpc server") + defer teardown() + + driverName := "dummy-driver" + timeout := time.Second + manager := NewDRAPluginManager(tCtx, nil, nil, nil, 0) + err = manager.add(driverName, addr, service, timeout) + require.NoError(t, err, "unexpected error while adding the plugin") + + plugin, err := manager.GetPlugin(driverName) + require.NoError(t, err, "unexpected error while retrieving the plugin") + + // we will invoke the method on a new gorouinte, in case there + // is no timeout enforced it might block forever + errCh := make(chan error, 1) + go func() { + defer close(errCh) + req := &drapbv1.NodePrepareResourcesRequest{ + Claims: []*drapbv1.Claim{ + { + Namespace: "dummy-namespace", + Uid: "dummy-uid", + Name: "dummy-claim", + }, + }, + } + _, err := plugin.NodePrepareResources(tCtx, req) + errCh <- err + }() + + // wait for the grpc caller to timeout + select { + // not using wait.ForeverTestTimeout, we will wait at most + // 3*timeout to account for flakes in CI + case <-time.After(3 * timeout): + t.Errorf("expected the grpc caller to return after the timeout had elapsed") + case err = <-errCh: + // the grpc call returned + } + + // unblock the request handler on the server, and wait for it to + // return, this ensures that the server method was invoked. + // if the timeout is not enforced, we don't leak any goroutine + close(blocked) + select { + case <-server.done: + // not using wait.ForeverTestTimeout, we will wait at most + // 3*timeout to account for flakes in CI + case <-time.After(3 * timeout): + t.Errorf("expected the grpc method to have been invoked") + } + + require.Error(t, err, "expected the grpc method to return an error") + status, ok := grpcstatus.FromError(err) + // if it is not a gRPC error then the operation may have failed before + // the gRPC method was called, otherwise we would get gRPC error. + require.True(t, ok, "expected an error of type: %T, but got: %T", &grpcstatus.Status{}, err) + + // DeadlineExceeded means operation expired before completion. The gRPC + // framework will generate this error code when the deadline is exceeded. + // More here: https://github.com/grpc/grpc/blob/master/doc/statuscodes.md + assert.Equal(t, grpccodes.DeadlineExceeded, status.Code(), "expected status code to match") +} + +// it embeds a fakeGRPCServer instance and overrides the NodePrepareResources +// method to simulate a timeout scenario. +// NodePrepareResources will block until the test explicitly closes the blocked +// channel to unblock it, and it will close the done channel when it completes +// so the test can wait for it to finish. +type timeoutFakeGRPCServer struct { + *fakeGRPCServer + + t *testing.T + blocked <-chan struct{} + done chan struct{} +} + +func (f *timeoutFakeGRPCServer) NodePrepareResources(ctx context.Context, in *drapbv1.NodePrepareResourcesRequest) (*drapbv1.NodePrepareResourcesResponse, error) { + defer close(f.done) + now := time.Now() + deadline, ok := ctx.Deadline() + f.t.Logf("request context has deadline: %t, after: %s", ok, deadline.Sub(now)) + if !ok { + f.t.Errorf("expected the request context to have a deadline") + } + f.t.Logf("NodePrepareResources: blocking so the client times out before the server sends a reply") + + <-f.blocked + + f.t.Logf("NodePrepareResources: blocked for: %s", time.Since(now)) + return &drapbv1.NodePrepareResourcesResponse{}, nil +} + func assertError(t *testing.T, expectError string, err error) { t.Helper() switch {