manager_test.go

Documentation: k8s.io/kubernetes/pkg/kubelet/cm/dra

     1  /*
     2  Copyright 2023 The Kubernetes Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package dra
    18  
    19  import (
    20  	"context"
    21  	"fmt"
    22  	"net"
    23  	"os"
    24  	"path/filepath"
    25  	"sync/atomic"
    26  	"testing"
    27  	"time"
    28  
    29  	"github.com/stretchr/testify/assert"
    30  	"google.golang.org/grpc"
    31  	v1 "k8s.io/api/core/v1"
    32  	resourcev1alpha2 "k8s.io/api/resource/v1alpha2"
    33  	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
    34  	"k8s.io/apimachinery/pkg/util/sets"
    35  	"k8s.io/client-go/kubernetes/fake"
    36  	"k8s.io/dynamic-resource-allocation/resourceclaim"
    37  	drapbv1 "k8s.io/kubelet/pkg/apis/dra/v1alpha3"
    38  	"k8s.io/kubernetes/pkg/kubelet/cm/dra/plugin"
    39  	"k8s.io/kubernetes/pkg/kubelet/cm/dra/state"
    40  	kubecontainer "k8s.io/kubernetes/pkg/kubelet/container"
    41  )
    42  
    43  const (
    44  	driverName      = "test-cdi-device"
    45  	driverClassName = "test"
    46  )
    47  
    48  type fakeDRADriverGRPCServer struct {
    49  	drapbv1.UnimplementedNodeServer
    50  	driverName             string
    51  	timeout                *time.Duration
    52  	prepareResourceCalls   atomic.Uint32
    53  	unprepareResourceCalls atomic.Uint32
    54  }
    55  
    56  func (s *fakeDRADriverGRPCServer) NodePrepareResources(ctx context.Context, req *drapbv1.NodePrepareResourcesRequest) (*drapbv1.NodePrepareResourcesResponse, error) {
    57  	s.prepareResourceCalls.Add(1)
    58  
    59  	if s.timeout != nil {
    60  		time.Sleep(*s.timeout)
    61  	}
    62  	deviceName := "claim-" + req.Claims[0].Uid
    63  	result := s.driverName + "/" + driverClassName + "=" + deviceName
    64  	return &drapbv1.NodePrepareResourcesResponse{Claims: map[string]*drapbv1.NodePrepareResourceResponse{req.Claims[0].Uid: {CDIDevices: []string{result}}}}, nil
    65  }
    66  
    67  func (s *fakeDRADriverGRPCServer) NodeUnprepareResources(ctx context.Context, req *drapbv1.NodeUnprepareResourcesRequest) (*drapbv1.NodeUnprepareResourcesResponse, error) {
    68  	s.unprepareResourceCalls.Add(1)
    69  
    70  	if s.timeout != nil {
    71  		time.Sleep(*s.timeout)
    72  	}
    73  	return &drapbv1.NodeUnprepareResourcesResponse{Claims: map[string]*drapbv1.NodeUnprepareResourceResponse{req.Claims[0].Uid: {}}}, nil
    74  }
    75  
    76  type tearDown func()
    77  
    78  type fakeDRAServerInfo struct {
    79  	// fake DRA server
    80  	server *fakeDRADriverGRPCServer
    81  	// fake DRA plugin socket name
    82  	socketName string
    83  	// teardownFn stops fake gRPC server
    84  	teardownFn tearDown
    85  }
    86  
    87  func setupFakeDRADriverGRPCServer(shouldTimeout bool, pluginClientTimeout *time.Duration) (fakeDRAServerInfo, error) {
    88  	socketDir, err := os.MkdirTemp("", "dra")
    89  	if err != nil {
    90  		return fakeDRAServerInfo{
    91  			server:     nil,
    92  			socketName: "",
    93  			teardownFn: nil,
    94  		}, err
    95  	}
    96  
    97  	socketName := filepath.Join(socketDir, "server.sock")
    98  	stopCh := make(chan struct{})
    99  
   100  	teardown := func() {
   101  		close(stopCh)
   102  		os.RemoveAll(socketName)
   103  	}
   104  
   105  	l, err := net.Listen("unix", socketName)
   106  	if err != nil {
   107  		teardown()
   108  		return fakeDRAServerInfo{
   109  			server:     nil,
   110  			socketName: "",
   111  			teardownFn: nil,
   112  		}, err
   113  	}
   114  
   115  	s := grpc.NewServer()
   116  	fakeDRADriverGRPCServer := &fakeDRADriverGRPCServer{
   117  		driverName: driverName,
   118  	}
   119  	if shouldTimeout {
   120  		timeout := *pluginClientTimeout * 2
   121  		fakeDRADriverGRPCServer.timeout = &timeout
   122  	}
   123  
   124  	drapbv1.RegisterNodeServer(s, fakeDRADriverGRPCServer)
   125  
   126  	go func() {
   127  		go s.Serve(l)
   128  		<-stopCh
   129  		s.GracefulStop()
   130  	}()
   131  
   132  	return fakeDRAServerInfo{
   133  		server:     fakeDRADriverGRPCServer,
   134  		socketName: socketName,
   135  		teardownFn: teardown,
   136  	}, nil
   137  }
   138  
   139  func TestNewManagerImpl(t *testing.T) {
   140  	kubeClient := fake.NewSimpleClientset()
   141  	for _, test := range []struct {
   142  		description        string
   143  		stateFileDirectory string
   144  		wantErr            bool
   145  	}{
   146  		{
   147  			description:        "invalid directory path",
   148  			stateFileDirectory: "",
   149  			wantErr:            true,
   150  		},
   151  		{
   152  			description:        "valid directory path",
   153  			stateFileDirectory: t.TempDir(),
   154  		},
   155  	} {
   156  		t.Run(test.description, func(t *testing.T) {
   157  			manager, err := NewManagerImpl(kubeClient, test.stateFileDirectory, "worker")
   158  			if test.wantErr {
   159  				assert.Error(t, err)
   160  				return
   161  			}
   162  
   163  			assert.NoError(t, err)
   164  			assert.NotNil(t, manager.cache)
   165  			assert.NotNil(t, manager.kubeClient)
   166  		})
   167  	}
   168  }
   169  
   170  func TestGetResources(t *testing.T) {
   171  	kubeClient := fake.NewSimpleClientset()
   172  	resourceClaimName := "test-pod-claim-1"
   173  
   174  	for _, test := range []struct {
   175  		description string
   176  		container   *v1.Container
   177  		pod         *v1.Pod
   178  		claimInfo   *ClaimInfo
   179  		wantErr     bool
   180  	}{
   181  		{
   182  			description: "claim info with annotations",
   183  			container: &v1.Container{
   184  				Name: "test-container",
   185  				Resources: v1.ResourceRequirements{
   186  					Claims: []v1.ResourceClaim{
   187  						{
   188  							Name: "test-pod-claim-1",
   189  						},
   190  					},
   191  				},
   192  			},
   193  			pod: &v1.Pod{
   194  				ObjectMeta: metav1.ObjectMeta{
   195  					Name:      "test-pod",
   196  					Namespace: "test-namespace",
   197  				},
   198  				Spec: v1.PodSpec{
   199  					ResourceClaims: []v1.PodResourceClaim{
   200  						{
   201  							Name:   "test-pod-claim-1",
   202  							Source: v1.ClaimSource{ResourceClaimName: &resourceClaimName},
   203  						},
   204  					},
   205  				},
   206  			},
   207  			claimInfo: &ClaimInfo{
   208  				annotations: map[string][]kubecontainer.Annotation{
   209  					"test-plugin": {
   210  						{
   211  							Name:  "test-annotation",
   212  							Value: "123",
   213  						},
   214  					},
   215  				},
   216  				ClaimInfoState: state.ClaimInfoState{
   217  					ClaimName: "test-pod-claim-1",
   218  					CDIDevices: map[string][]string{
   219  						driverName: {"123"},
   220  					},
   221  					Namespace: "test-namespace",
   222  				},
   223  			},
   224  		},
   225  		{
   226  			description: "claim info without annotations",
   227  			container: &v1.Container{
   228  				Name: "test-container",
   229  				Resources: v1.ResourceRequirements{
   230  					Claims: []v1.ResourceClaim{
   231  						{
   232  							Name: "test-pod-claim-1",
   233  						},
   234  					},
   235  				},
   236  			},
   237  			pod: &v1.Pod{
   238  				ObjectMeta: metav1.ObjectMeta{
   239  					Name:      "test-pod",
   240  					Namespace: "test-namespace",
   241  				},
   242  				Spec: v1.PodSpec{
   243  					ResourceClaims: []v1.PodResourceClaim{
   244  						{
   245  							Name:   "test-pod-claim-1",
   246  							Source: v1.ClaimSource{ResourceClaimName: &resourceClaimName},
   247  						},
   248  					},
   249  				},
   250  			},
   251  			claimInfo: &ClaimInfo{
   252  				ClaimInfoState: state.ClaimInfoState{
   253  					ClaimName: "test-pod-claim-1",
   254  					CDIDevices: map[string][]string{
   255  						driverName: {"123"},
   256  					},
   257  					Namespace: "test-namespace",
   258  				},
   259  			},
   260  		},
   261  		{
   262  			description: "no claim info",
   263  			container: &v1.Container{
   264  				Name: "test-container",
   265  				Resources: v1.ResourceRequirements{
   266  					Claims: []v1.ResourceClaim{
   267  						{
   268  							Name: "test-pod-claim-1",
   269  						},
   270  					},
   271  				},
   272  			},
   273  			pod: &v1.Pod{
   274  				ObjectMeta: metav1.ObjectMeta{
   275  					Name:      "test-pod",
   276  					Namespace: "test-namespace",
   277  				},
   278  				Spec: v1.PodSpec{
   279  					ResourceClaims: []v1.PodResourceClaim{
   280  						{
   281  							Name: "test-pod-claim-1",
   282  						},
   283  					},
   284  				},
   285  			},
   286  			wantErr: true,
   287  		},
   288  	} {
   289  		t.Run(test.description, func(t *testing.T) {
   290  			manager, err := NewManagerImpl(kubeClient, t.TempDir(), "worker")
   291  			assert.NoError(t, err)
   292  
   293  			if test.claimInfo != nil {
   294  				manager.cache.add(test.claimInfo)
   295  			}
   296  
   297  			containerInfo, err := manager.GetResources(test.pod, test.container)
   298  			if test.wantErr {
   299  				assert.Error(t, err)
   300  				return
   301  			}
   302  
   303  			assert.NoError(t, err)
   304  			assert.Equal(t, test.claimInfo.CDIDevices[driverName][0], containerInfo.CDIDevices[0].Name)
   305  		})
   306  	}
   307  }
   308  
   309  func getFakeNode() (*v1.Node, error) {
   310  	return &v1.Node{ObjectMeta: metav1.ObjectMeta{Name: "worker"}}, nil
   311  }
   312  
   313  func TestPrepareResources(t *testing.T) {
   314  	fakeKubeClient := fake.NewSimpleClientset()
   315  
   316  	for _, test := range []struct {
   317  		description          string
   318  		driverName           string
   319  		pod                  *v1.Pod
   320  		claimInfo            *ClaimInfo
   321  		resourceClaim        *resourcev1alpha2.ResourceClaim
   322  		wantErr              bool
   323  		wantTimeout          bool
   324  		wantResourceSkipped  bool
   325  		ExpectedPrepareCalls uint32
   326  	}{
   327  		{
   328  			description: "failed to fetch ResourceClaim",
   329  			driverName:  driverName,
   330  			pod: &v1.Pod{
   331  				ObjectMeta: metav1.ObjectMeta{
   332  					Name:      "test-pod",
   333  					Namespace: "test-namespace",
   334  					UID:       "test-reserved",
   335  				},
   336  				Spec: v1.PodSpec{
   337  					ResourceClaims: []v1.PodResourceClaim{
   338  						{
   339  							Name: "test-pod-claim-0",
   340  							Source: v1.ClaimSource{
   341  								ResourceClaimName: func() *string {
   342  									s := "test-pod-claim-0"
   343  									return &s
   344  								}(),
   345  							},
   346  						},
   347  					},
   348  				},
   349  			},
   350  			wantErr: true,
   351  		},
   352  		{
   353  			description: "plugin does not exist",
   354  			driverName:  "this-plugin-does-not-exist",
   355  			pod: &v1.Pod{
   356  				ObjectMeta: metav1.ObjectMeta{
   357  					Name:      "test-pod",
   358  					Namespace: "test-namespace",
   359  					UID:       "test-reserved",
   360  				},
   361  				Spec: v1.PodSpec{
   362  					ResourceClaims: []v1.PodResourceClaim{
   363  						{
   364  							Name: "test-pod-claim-1",
   365  							Source: v1.ClaimSource{
   366  								ResourceClaimName: func() *string {
   367  									s := "test-pod-claim-1"
   368  									return &s
   369  								}(),
   370  							},
   371  						},
   372  					},
   373  					Containers: []v1.Container{
   374  						{
   375  							Resources: v1.ResourceRequirements{
   376  								Claims: []v1.ResourceClaim{
   377  									{
   378  										Name: "test-pod-claim-1",
   379  									},
   380  								},
   381  							},
   382  						},
   383  					},
   384  				},
   385  			},
   386  			resourceClaim: &resourcev1alpha2.ResourceClaim{
   387  				ObjectMeta: metav1.ObjectMeta{
   388  					Name:      "test-pod-claim-1",
   389  					Namespace: "test-namespace",
   390  					UID:       "test-reserved",
   391  				},
   392  				Spec: resourcev1alpha2.ResourceClaimSpec{
   393  					ResourceClassName: "test-class",
   394  				},
   395  				Status: resourcev1alpha2.ResourceClaimStatus{
   396  					DriverName: driverName,
   397  					Allocation: &resourcev1alpha2.AllocationResult{
   398  						ResourceHandles: []resourcev1alpha2.ResourceHandle{
   399  							{Data: "test-data", DriverName: driverName},
   400  						},
   401  					},
   402  					ReservedFor: []resourcev1alpha2.ResourceClaimConsumerReference{
   403  						{UID: "test-reserved"},
   404  					},
   405  				},
   406  			},
   407  			wantErr: true,
   408  		},
   409  		{
   410  			description: "pod is not allowed to use resource claim",
   411  			driverName:  driverName,
   412  			pod: &v1.Pod{
   413  				ObjectMeta: metav1.ObjectMeta{
   414  					Name:      "test-pod",
   415  					Namespace: "test-namespace",
   416  					UID:       "test-reserved",
   417  				},
   418  				Spec: v1.PodSpec{
   419  					ResourceClaims: []v1.PodResourceClaim{
   420  						{
   421  							Name: "test-pod-claim-2",
   422  							Source: v1.ClaimSource{
   423  								ResourceClaimName: func() *string {
   424  									s := "test-pod-claim-2"
   425  									return &s
   426  								}(),
   427  							},
   428  						},
   429  					},
   430  				},
   431  			},
   432  			resourceClaim: &resourcev1alpha2.ResourceClaim{
   433  				ObjectMeta: metav1.ObjectMeta{
   434  					Name:      "test-pod-claim-2",
   435  					Namespace: "test-namespace",
   436  					UID:       "test-reserved",
   437  				},
   438  				Spec: resourcev1alpha2.ResourceClaimSpec{
   439  					ResourceClassName: "test-class",
   440  				},
   441  				Status: resourcev1alpha2.ResourceClaimStatus{
   442  					DriverName: driverName,
   443  					Allocation: &resourcev1alpha2.AllocationResult{
   444  						ResourceHandles: []resourcev1alpha2.ResourceHandle{
   445  							{Data: "test-data", DriverName: driverName},
   446  						},
   447  					},
   448  				},
   449  			},
   450  			wantErr: true,
   451  		},
   452  		{
   453  			description: "no container actually uses the claim",
   454  			driverName:  driverName,
   455  			pod: &v1.Pod{
   456  				ObjectMeta: metav1.ObjectMeta{
   457  					Name:      "test-pod",
   458  					Namespace: "test-namespace",
   459  					UID:       "test-reserved",
   460  				},
   461  				Spec: v1.PodSpec{
   462  					ResourceClaims: []v1.PodResourceClaim{
   463  						{
   464  							Name: "test-pod-claim-3",
   465  							Source: v1.ClaimSource{ResourceClaimName: func() *string {
   466  								s := "test-pod-claim-3"
   467  								return &s
   468  							}()},
   469  						},
   470  					},
   471  				},
   472  			},
   473  			resourceClaim: &resourcev1alpha2.ResourceClaim{
   474  				ObjectMeta: metav1.ObjectMeta{
   475  					Name:      "test-pod-claim-3",
   476  					Namespace: "test-namespace",
   477  					UID:       "test-reserved",
   478  				},
   479  				Spec: resourcev1alpha2.ResourceClaimSpec{
   480  					ResourceClassName: "test-class",
   481  				},
   482  				Status: resourcev1alpha2.ResourceClaimStatus{
   483  					DriverName: driverName,
   484  					Allocation: &resourcev1alpha2.AllocationResult{
   485  						ResourceHandles: []resourcev1alpha2.ResourceHandle{
   486  							{Data: "test-data", DriverName: driverName},
   487  						},
   488  					},
   489  					ReservedFor: []resourcev1alpha2.ResourceClaimConsumerReference{
   490  						{UID: "test-reserved"},
   491  					},
   492  				},
   493  			},
   494  			wantResourceSkipped: true,
   495  		},
   496  		{
   497  			description: "resource already prepared",
   498  			driverName:  driverName,
   499  			pod: &v1.Pod{
   500  				ObjectMeta: metav1.ObjectMeta{
   501  					Name:      "test-pod",
   502  					Namespace: "test-namespace",
   503  					UID:       "test-reserved",
   504  				},
   505  				Spec: v1.PodSpec{
   506  					ResourceClaims: []v1.PodResourceClaim{
   507  						{
   508  							Name: "test-pod-claim-4",
   509  							Source: v1.ClaimSource{ResourceClaimName: func() *string {
   510  								s := "test-pod-claim-4"
   511  								return &s
   512  							}()},
   513  						},
   514  					},
   515  					Containers: []v1.Container{
   516  						{
   517  							Resources: v1.ResourceRequirements{
   518  								Claims: []v1.ResourceClaim{
   519  									{
   520  										Name: "test-pod-claim-4",
   521  									},
   522  								},
   523  							},
   524  						},
   525  					},
   526  				},
   527  			},
   528  			claimInfo: &ClaimInfo{
   529  				ClaimInfoState: state.ClaimInfoState{
   530  					DriverName: driverName,
   531  					ClaimName:  "test-pod-claim-4",
   532  					Namespace:  "test-namespace",
   533  					PodUIDs:    sets.Set[string]{"test-another-pod-reserved": sets.Empty{}},
   534  				},
   535  				prepared: true,
   536  			},
   537  			resourceClaim: &resourcev1alpha2.ResourceClaim{
   538  				ObjectMeta: metav1.ObjectMeta{
   539  					Name:      "test-pod-claim-4",
   540  					Namespace: "test-namespace",
   541  					UID:       "test-reserved",
   542  				},
   543  				Spec: resourcev1alpha2.ResourceClaimSpec{
   544  					ResourceClassName: "test-class",
   545  				},
   546  				Status: resourcev1alpha2.ResourceClaimStatus{
   547  					DriverName: driverName,
   548  					Allocation: &resourcev1alpha2.AllocationResult{
   549  						ResourceHandles: []resourcev1alpha2.ResourceHandle{
   550  							{Data: "test-data", DriverName: driverName},
   551  						},
   552  					},
   553  					ReservedFor: []resourcev1alpha2.ResourceClaimConsumerReference{
   554  						{UID: "test-reserved"},
   555  					},
   556  				},
   557  			},
   558  			wantResourceSkipped: true,
   559  		},
   560  		{
   561  			description: "should timeout",
   562  			driverName:  driverName,
   563  			pod: &v1.Pod{
   564  				ObjectMeta: metav1.ObjectMeta{
   565  					Name:      "test-pod",
   566  					Namespace: "test-namespace",
   567  					UID:       "test-reserved",
   568  				},
   569  				Spec: v1.PodSpec{
   570  					ResourceClaims: []v1.PodResourceClaim{
   571  						{
   572  							Name: "test-pod-claim-5",
   573  							Source: v1.ClaimSource{ResourceClaimName: func() *string {
   574  								s := "test-pod-claim-5"
   575  								return &s
   576  							}()},
   577  						},
   578  					},
   579  					Containers: []v1.Container{
   580  						{
   581  							Resources: v1.ResourceRequirements{
   582  								Claims: []v1.ResourceClaim{
   583  									{
   584  										Name: "test-pod-claim-5",
   585  									},
   586  								},
   587  							},
   588  						},
   589  					},
   590  				},
   591  			},
   592  			resourceClaim: &resourcev1alpha2.ResourceClaim{
   593  				ObjectMeta: metav1.ObjectMeta{
   594  					Name:      "test-pod-claim-5",
   595  					Namespace: "test-namespace",
   596  					UID:       "test-reserved",
   597  				},
   598  				Spec: resourcev1alpha2.ResourceClaimSpec{
   599  					ResourceClassName: "test-class",
   600  				},
   601  				Status: resourcev1alpha2.ResourceClaimStatus{
   602  					DriverName: driverName,
   603  					Allocation: &resourcev1alpha2.AllocationResult{
   604  						ResourceHandles: []resourcev1alpha2.ResourceHandle{
   605  							{Data: "test-data", DriverName: driverName},
   606  						},
   607  					},
   608  					ReservedFor: []resourcev1alpha2.ResourceClaimConsumerReference{
   609  						{UID: "test-reserved"},
   610  					},
   611  				},
   612  			},
   613  			wantErr:              true,
   614  			wantTimeout:          true,
   615  			ExpectedPrepareCalls: 1,
   616  		},
   617  		{
   618  			description: "should prepare resource, claim not in cache",
   619  			driverName:  driverName,
   620  			pod: &v1.Pod{
   621  				ObjectMeta: metav1.ObjectMeta{
   622  					Name:      "test-pod",
   623  					Namespace: "test-namespace",
   624  					UID:       "test-reserved",
   625  				},
   626  				Spec: v1.PodSpec{
   627  					ResourceClaims: []v1.PodResourceClaim{
   628  						{
   629  							Name: "test-pod-claim-6",
   630  							Source: v1.ClaimSource{ResourceClaimName: func() *string {
   631  								s := "test-pod-claim-6"
   632  								return &s
   633  							}()},
   634  						},
   635  					},
   636  					Containers: []v1.Container{
   637  						{
   638  							Resources: v1.ResourceRequirements{
   639  								Claims: []v1.ResourceClaim{
   640  									{
   641  										Name: "test-pod-claim-6",
   642  									},
   643  								},
   644  							},
   645  						},
   646  					},
   647  				},
   648  			},
   649  			resourceClaim: &resourcev1alpha2.ResourceClaim{
   650  				ObjectMeta: metav1.ObjectMeta{
   651  					Name:      "test-pod-claim-6",
   652  					Namespace: "test-namespace",
   653  					UID:       "test-reserved",
   654  				},
   655  				Spec: resourcev1alpha2.ResourceClaimSpec{
   656  					ResourceClassName: "test-class",
   657  				},
   658  				Status: resourcev1alpha2.ResourceClaimStatus{
   659  					DriverName: driverName,
   660  					Allocation: &resourcev1alpha2.AllocationResult{
   661  						ResourceHandles: []resourcev1alpha2.ResourceHandle{
   662  							{Data: "test-data"},
   663  						},
   664  					},
   665  					ReservedFor: []resourcev1alpha2.ResourceClaimConsumerReference{
   666  						{UID: "test-reserved"},
   667  					},
   668  				},
   669  			},
   670  			ExpectedPrepareCalls: 1,
   671  		},
   672  		{
   673  			description: "should prepare resource. claim in cache, manager did not prepare resource",
   674  			driverName:  driverName,
   675  			pod: &v1.Pod{
   676  				ObjectMeta: metav1.ObjectMeta{
   677  					Name:      "test-pod",
   678  					Namespace: "test-namespace",
   679  					UID:       "test-reserved",
   680  				},
   681  				Spec: v1.PodSpec{
   682  					ResourceClaims: []v1.PodResourceClaim{
   683  						{
   684  							Name: "test-pod-claim",
   685  							Source: v1.ClaimSource{ResourceClaimName: func() *string {
   686  								s := "test-pod-claim"
   687  								return &s
   688  							}()},
   689  						},
   690  					},
   691  					Containers: []v1.Container{
   692  						{
   693  							Resources: v1.ResourceRequirements{
   694  								Claims: []v1.ResourceClaim{
   695  									{
   696  										Name: "test-pod-claim",
   697  									},
   698  								},
   699  							},
   700  						},
   701  					},
   702  				},
   703  			},
   704  			claimInfo: &ClaimInfo{
   705  				ClaimInfoState: state.ClaimInfoState{
   706  					DriverName: driverName,
   707  					ClassName:  "test-class",
   708  					ClaimName:  "test-pod-claim",
   709  					ClaimUID:   "test-reserved",
   710  					Namespace:  "test-namespace",
   711  					PodUIDs:    sets.Set[string]{"test-reserved": sets.Empty{}},
   712  					CDIDevices: map[string][]string{
   713  						driverName: {fmt.Sprintf("%s/%s=some-device", driverName, driverClassName)},
   714  					},
   715  					ResourceHandles: []resourcev1alpha2.ResourceHandle{{Data: "test-data"}},
   716  				},
   717  				annotations: make(map[string][]kubecontainer.Annotation),
   718  				prepared:    false,
   719  			},
   720  			resourceClaim: &resourcev1alpha2.ResourceClaim{
   721  				ObjectMeta: metav1.ObjectMeta{
   722  					Name:      "test-pod-claim",
   723  					Namespace: "test-namespace",
   724  					UID:       "test-reserved",
   725  				},
   726  				Spec: resourcev1alpha2.ResourceClaimSpec{
   727  					ResourceClassName: "test-class",
   728  				},
   729  				Status: resourcev1alpha2.ResourceClaimStatus{
   730  					DriverName: driverName,
   731  					Allocation: &resourcev1alpha2.AllocationResult{
   732  						ResourceHandles: []resourcev1alpha2.ResourceHandle{
   733  							{Data: "test-data"},
   734  						},
   735  					},
   736  					ReservedFor: []resourcev1alpha2.ResourceClaimConsumerReference{
   737  						{UID: "test-reserved"},
   738  					},
   739  				},
   740  			},
   741  			ExpectedPrepareCalls: 1,
   742  		},
   743  	} {
   744  		t.Run(test.description, func(t *testing.T) {
   745  			cache, err := newClaimInfoCache(t.TempDir(), draManagerStateFileName)
   746  			if err != nil {
   747  				t.Fatalf("failed to newClaimInfoCache, err:%v", err)
   748  			}
   749  
   750  			manager := &ManagerImpl{
   751  				kubeClient: fakeKubeClient,
   752  				cache:      cache,
   753  			}
   754  
   755  			if test.resourceClaim != nil {
   756  				if _, err := fakeKubeClient.ResourceV1alpha2().ResourceClaims(test.pod.Namespace).Create(context.Background(), test.resourceClaim, metav1.CreateOptions{}); err != nil {
   757  					t.Fatalf("failed to create ResourceClaim %s: %+v", test.resourceClaim.Name, err)
   758  				}
   759  			}
   760  
   761  			var pluginClientTimeout *time.Duration
   762  			if test.wantTimeout {
   763  				timeout := time.Millisecond * 20
   764  				pluginClientTimeout = &timeout
   765  			}
   766  
   767  			draServerInfo, err := setupFakeDRADriverGRPCServer(test.wantTimeout, pluginClientTimeout)
   768  			if err != nil {
   769  				t.Fatal(err)
   770  			}
   771  			defer draServerInfo.teardownFn()
   772  
   773  			plg := plugin.NewRegistrationHandler(nil, getFakeNode)
   774  			if err := plg.RegisterPlugin(test.driverName, draServerInfo.socketName, []string{"1.27"}, pluginClientTimeout); err != nil {
   775  				t.Fatalf("failed to register plugin %s, err: %v", test.driverName, err)
   776  			}
   777  			defer plg.DeRegisterPlugin(test.driverName) // for sake of next tests
   778  
   779  			if test.claimInfo != nil {
   780  				manager.cache.add(test.claimInfo)
   781  			}
   782  
   783  			err = manager.PrepareResources(test.pod)
   784  
   785  			assert.Equal(t, test.ExpectedPrepareCalls, draServerInfo.server.prepareResourceCalls.Load())
   786  
   787  			if test.wantErr {
   788  				assert.Error(t, err)
   789  				return // PrepareResources returned an error so stopping the subtest here
   790  			} else if test.wantResourceSkipped {
   791  				assert.NoError(t, err)
   792  				return // resource skipped so no need to continue
   793  			}
   794  
   795  			assert.NoError(t, err)
   796  			// check the cache contains the expected claim info
   797  			claimName, _, err := resourceclaim.Name(test.pod, &test.pod.Spec.ResourceClaims[0])
   798  			if err != nil {
   799  				t.Fatal(err)
   800  			}
   801  			claimInfo := manager.cache.get(*claimName, test.pod.Namespace)
   802  			if claimInfo == nil {
   803  				t.Fatalf("claimInfo not found in cache for claim %s", *claimName)
   804  			}
   805  			if claimInfo.DriverName != test.resourceClaim.Status.DriverName {
   806  				t.Fatalf("driverName mismatch: expected %s, got %s", test.resourceClaim.Status.DriverName, claimInfo.DriverName)
   807  			}
   808  			if claimInfo.ClassName != test.resourceClaim.Spec.ResourceClassName {
   809  				t.Fatalf("resourceClassName mismatch: expected %s, got %s", test.resourceClaim.Spec.ResourceClassName, claimInfo.ClassName)
   810  			}
   811  			if len(claimInfo.PodUIDs) != 1 || !claimInfo.PodUIDs.Has(string(test.pod.UID)) {
   812  				t.Fatalf("podUIDs mismatch: expected [%s], got %v", test.pod.UID, claimInfo.PodUIDs)
   813  			}
   814  			expectedResourceClaimDriverName := fmt.Sprintf("%s/%s=claim-%s", driverName, driverClassName, string(test.resourceClaim.Status.ReservedFor[0].UID))
   815  			if len(claimInfo.CDIDevices[test.resourceClaim.Status.DriverName]) != 1 || claimInfo.CDIDevices[test.resourceClaim.Status.DriverName][0] != expectedResourceClaimDriverName {
   816  				t.Fatalf("cdiDevices mismatch: expected [%s], got %v", []string{expectedResourceClaimDriverName}, claimInfo.CDIDevices[test.resourceClaim.Status.DriverName])
   817  			}
   818  		})
   819  	}
   820  }
   821  
   822  func TestUnprepareResources(t *testing.T) {
   823  	fakeKubeClient := fake.NewSimpleClientset()
   824  
   825  	for _, test := range []struct {
   826  		description            string
   827  		driverName             string
   828  		pod                    *v1.Pod
   829  		claimInfo              *ClaimInfo
   830  		wantErr                bool
   831  		wantTimeout            bool
   832  		wantResourceSkipped    bool
   833  		expectedUnprepareCalls uint32
   834  	}{
   835  		{
   836  			description: "plugin does not exist",
   837  			driverName:  "this-plugin-does-not-exist",
   838  			pod: &v1.Pod{
   839  				ObjectMeta: metav1.ObjectMeta{
   840  					Name:      "test-pod",
   841  					Namespace: "test-namespace",
   842  					UID:       "test-reserved",
   843  				},
   844  				Spec: v1.PodSpec{
   845  					ResourceClaims: []v1.PodResourceClaim{
   846  						{
   847  							Name: "another-claim-test",
   848  							Source: v1.ClaimSource{
   849  								ResourceClaimName: func() *string {
   850  									s := "another-claim-test"
   851  									return &s
   852  								}(),
   853  							},
   854  						},
   855  					},
   856  				},
   857  			},
   858  			claimInfo: &ClaimInfo{
   859  				ClaimInfoState: state.ClaimInfoState{
   860  					DriverName: driverName,
   861  					ClaimName:  "another-claim-test",
   862  					Namespace:  "test-namespace",
   863  					ResourceHandles: []resourcev1alpha2.ResourceHandle{
   864  						{
   865  							DriverName: driverName,
   866  							Data:       "test data",
   867  						},
   868  					},
   869  				},
   870  			},
   871  			wantErr: true,
   872  		},
   873  		{
   874  			description: "resource claim referenced by other pod(s)",
   875  			driverName:  driverName,
   876  			pod: &v1.Pod{
   877  				ObjectMeta: metav1.ObjectMeta{
   878  					Name:      "test-pod",
   879  					Namespace: "test-namespace",
   880  					UID:       "test-reserved",
   881  				},
   882  				Spec: v1.PodSpec{
   883  					ResourceClaims: []v1.PodResourceClaim{
   884  						{
   885  							Name: "test-pod-claim-1",
   886  							Source: v1.ClaimSource{ResourceClaimName: func() *string {
   887  								s := "test-pod-claim-1"
   888  								return &s
   889  							}()},
   890  						},
   891  					},
   892  					Containers: []v1.Container{
   893  						{
   894  							Resources: v1.ResourceRequirements{
   895  								Claims: []v1.ResourceClaim{
   896  									{
   897  										Name: "test-pod-claim-1",
   898  									},
   899  								},
   900  							},
   901  						},
   902  					},
   903  				},
   904  			},
   905  			claimInfo: &ClaimInfo{
   906  				ClaimInfoState: state.ClaimInfoState{
   907  					DriverName: driverName,
   908  					ClaimName:  "test-pod-claim-1",
   909  					Namespace:  "test-namespace",
   910  					PodUIDs:    sets.Set[string]{"test-reserved": sets.Empty{}, "test-reserved-2": sets.Empty{}},
   911  				},
   912  			},
   913  			wantResourceSkipped: true,
   914  		},
   915  		{
   916  			description: "should timeout",
   917  			driverName:  driverName,
   918  			pod: &v1.Pod{
   919  				ObjectMeta: metav1.ObjectMeta{
   920  					Name:      "test-pod",
   921  					Namespace: "test-namespace",
   922  					UID:       "test-reserved",
   923  				},
   924  				Spec: v1.PodSpec{
   925  					ResourceClaims: []v1.PodResourceClaim{
   926  						{
   927  							Name: "test-pod-claim-2",
   928  							Source: v1.ClaimSource{ResourceClaimName: func() *string {
   929  								s := "test-pod-claim-2"
   930  								return &s
   931  							}()},
   932  						},
   933  					},
   934  					Containers: []v1.Container{
   935  						{
   936  							Resources: v1.ResourceRequirements{
   937  								Claims: []v1.ResourceClaim{
   938  									{
   939  										Name: "test-pod-claim-2",
   940  									},
   941  								},
   942  							},
   943  						},
   944  					},
   945  				},
   946  			},
   947  			claimInfo: &ClaimInfo{
   948  				ClaimInfoState: state.ClaimInfoState{
   949  					DriverName: driverName,
   950  					ClaimName:  "test-pod-claim-2",
   951  					Namespace:  "test-namespace",
   952  					ResourceHandles: []resourcev1alpha2.ResourceHandle{
   953  						{
   954  							DriverName: driverName,
   955  							Data:       "test data",
   956  						},
   957  					},
   958  				},
   959  			},
   960  			wantErr:                true,
   961  			wantTimeout:            true,
   962  			expectedUnprepareCalls: 1,
   963  		},
   964  		{
   965  			description: "should unprepare resource, claim previously prepared by currently running manager",
   966  			driverName:  driverName,
   967  			pod: &v1.Pod{
   968  				ObjectMeta: metav1.ObjectMeta{
   969  					Name:      "test-pod",
   970  					Namespace: "test-namespace",
   971  					UID:       "test-reserved",
   972  				},
   973  				Spec: v1.PodSpec{
   974  					ResourceClaims: []v1.PodResourceClaim{
   975  						{
   976  							Name: "test-pod-claim-3",
   977  							Source: v1.ClaimSource{ResourceClaimName: func() *string {
   978  								s := "test-pod-claim-3"
   979  								return &s
   980  							}()},
   981  						},
   982  					},
   983  					Containers: []v1.Container{
   984  						{
   985  							Resources: v1.ResourceRequirements{
   986  								Claims: []v1.ResourceClaim{
   987  									{
   988  										Name: "test-pod-claim-3",
   989  									},
   990  								},
   991  							},
   992  						},
   993  					},
   994  				},
   995  			},
   996  			claimInfo: &ClaimInfo{
   997  				ClaimInfoState: state.ClaimInfoState{
   998  					DriverName: driverName,
   999  					ClaimName:  "test-pod-claim-3",
  1000  					Namespace:  "test-namespace",
  1001  					ResourceHandles: []resourcev1alpha2.ResourceHandle{
  1002  						{
  1003  							DriverName: driverName,
  1004  							Data:       "test data",
  1005  						},
  1006  					},
  1007  				},
  1008  				prepared: true,
  1009  			},
  1010  			expectedUnprepareCalls: 1,
  1011  		},
  1012  		{
  1013  			description: "should unprepare resource, claim previously was not prepared by currently running manager",
  1014  			driverName:  driverName,
  1015  			pod: &v1.Pod{
  1016  				ObjectMeta: metav1.ObjectMeta{
  1017  					Name:      "test-pod",
  1018  					Namespace: "test-namespace",
  1019  					UID:       "test-reserved",
  1020  				},
  1021  				Spec: v1.PodSpec{
  1022  					ResourceClaims: []v1.PodResourceClaim{
  1023  						{
  1024  							Name: "test-pod-claim",
  1025  							Source: v1.ClaimSource{ResourceClaimName: func() *string {
  1026  								s := "test-pod-claim"
  1027  								return &s
  1028  							}()},
  1029  						},
  1030  					},
  1031  					Containers: []v1.Container{
  1032  						{
  1033  							Resources: v1.ResourceRequirements{
  1034  								Claims: []v1.ResourceClaim{
  1035  									{
  1036  										Name: "test-pod-claim",
  1037  									},
  1038  								},
  1039  							},
  1040  						},
  1041  					},
  1042  				},
  1043  			},
  1044  			claimInfo: &ClaimInfo{
  1045  				ClaimInfoState: state.ClaimInfoState{
  1046  					DriverName: driverName,
  1047  					ClaimName:  "test-pod-claim",
  1048  					Namespace:  "test-namespace",
  1049  					ResourceHandles: []resourcev1alpha2.ResourceHandle{
  1050  						{
  1051  							DriverName: driverName,
  1052  							Data:       "test data",
  1053  						},
  1054  					},
  1055  				},
  1056  				prepared: false,
  1057  			},
  1058  			expectedUnprepareCalls: 1,
  1059  		},
  1060  	} {
  1061  		t.Run(test.description, func(t *testing.T) {
  1062  			cache, err := newClaimInfoCache(t.TempDir(), draManagerStateFileName)
  1063  			if err != nil {
  1064  				t.Fatalf("failed to create a new instance of the claimInfoCache, err: %v", err)
  1065  			}
  1066  
  1067  			var pluginClientTimeout *time.Duration
  1068  			if test.wantTimeout {
  1069  				timeout := time.Millisecond * 20
  1070  				pluginClientTimeout = &timeout
  1071  			}
  1072  
  1073  			draServerInfo, err := setupFakeDRADriverGRPCServer(test.wantTimeout, pluginClientTimeout)
  1074  			if err != nil {
  1075  				t.Fatal(err)
  1076  			}
  1077  			defer draServerInfo.teardownFn()
  1078  
  1079  			plg := plugin.NewRegistrationHandler(nil, getFakeNode)
  1080  			if err := plg.RegisterPlugin(test.driverName, draServerInfo.socketName, []string{"1.27"}, pluginClientTimeout); err != nil {
  1081  				t.Fatalf("failed to register plugin %s, err: %v", test.driverName, err)
  1082  			}
  1083  			defer plg.DeRegisterPlugin(test.driverName) // for sake of next tests
  1084  
  1085  			manager := &ManagerImpl{
  1086  				kubeClient: fakeKubeClient,
  1087  				cache:      cache,
  1088  			}
  1089  
  1090  			if test.claimInfo != nil {
  1091  				manager.cache.add(test.claimInfo)
  1092  			}
  1093  
  1094  			err = manager.UnprepareResources(test.pod)
  1095  
  1096  			assert.Equal(t, test.expectedUnprepareCalls, draServerInfo.server.unprepareResourceCalls.Load())
  1097  
  1098  			if test.wantErr {
  1099  				assert.Error(t, err)
  1100  				return // UnprepareResources returned an error so stopping the subtest here
  1101  			} else if test.wantResourceSkipped {
  1102  				assert.NoError(t, err)
  1103  				return // resource skipped so no need to continue
  1104  			}
  1105  
  1106  			assert.NoError(t, err)
  1107  			// Check that the cache has been updated correctly
  1108  			claimName, _, err := resourceclaim.Name(test.pod, &test.pod.Spec.ResourceClaims[0])
  1109  			if err != nil {
  1110  				t.Fatal(err)
  1111  			}
  1112  			claimInfo := manager.cache.get(*claimName, test.pod.Namespace)
  1113  			if claimInfo != nil {
  1114  				t.Fatalf("claimInfo still found in cache after calling UnprepareResources")
  1115  			}
  1116  		})
  1117  	}
  1118  }
  1119  
  1120  func TestPodMightNeedToUnprepareResources(t *testing.T) {
  1121  	fakeKubeClient := fake.NewSimpleClientset()
  1122  
  1123  	cache, err := newClaimInfoCache(t.TempDir(), draManagerStateFileName)
  1124  	if err != nil {
  1125  		t.Fatalf("failed to newClaimInfoCache, err:%v", err)
  1126  	}
  1127  
  1128  	manager := &ManagerImpl{
  1129  		kubeClient: fakeKubeClient,
  1130  		cache:      cache,
  1131  	}
  1132  
  1133  	podUID := sets.Set[string]{}
  1134  	podUID.Insert("test-pod-uid")
  1135  	manager.cache.add(&ClaimInfo{
  1136  		ClaimInfoState: state.ClaimInfoState{PodUIDs: podUID, ClaimName: "test-claim", Namespace: "test-namespace"},
  1137  	})
  1138  
  1139  	testClaimInfo := manager.cache.get("test-claim", "test-namespace")
  1140  	testClaimInfo.addPodReference("test-pod-uid")
  1141  
  1142  	manager.PodMightNeedToUnprepareResources("test-pod-uid")
  1143  }
  1144  
  1145  func TestGetContainerClaimInfos(t *testing.T) {
  1146  	cache, err := newClaimInfoCache(t.TempDir(), draManagerStateFileName)
  1147  	if err != nil {
  1148  		t.Fatalf("error occur:%v", err)
  1149  	}
  1150  	manager := &ManagerImpl{
  1151  		cache: cache,
  1152  	}
  1153  
  1154  	resourceClaimName := "test-resource-claim-1"
  1155  	resourceClaimName2 := "test-resource-claim-2"
  1156  
  1157  	for i, test := range []struct {
  1158  		expectedClaimName string
  1159  		pod               *v1.Pod
  1160  		container         *v1.Container
  1161  		claimInfo         *ClaimInfo
  1162  	}{
  1163  		{
  1164  			expectedClaimName: resourceClaimName,
  1165  			pod: &v1.Pod{
  1166  				Spec: v1.PodSpec{
  1167  					ResourceClaims: []v1.PodResourceClaim{
  1168  						{
  1169  							Name:   "claim1",
  1170  							Source: v1.ClaimSource{ResourceClaimName: &resourceClaimName},
  1171  						},
  1172  					},
  1173  				},
  1174  			},
  1175  			container: &v1.Container{
  1176  				Resources: v1.ResourceRequirements{
  1177  					Claims: []v1.ResourceClaim{
  1178  						{
  1179  							Name: "claim1",
  1180  						},
  1181  					},
  1182  				},
  1183  			},
  1184  			claimInfo: &ClaimInfo{ClaimInfoState: state.ClaimInfoState{ClaimName: resourceClaimName}},
  1185  		},
  1186  		{
  1187  			expectedClaimName: resourceClaimName2,
  1188  			pod: &v1.Pod{
  1189  				Spec: v1.PodSpec{
  1190  					ResourceClaims: []v1.PodResourceClaim{
  1191  						{
  1192  							Name:   "claim2",
  1193  							Source: v1.ClaimSource{ResourceClaimName: &resourceClaimName2},
  1194  						},
  1195  					},
  1196  				},
  1197  			},
  1198  			container: &v1.Container{
  1199  				Resources: v1.ResourceRequirements{
  1200  					Claims: []v1.ResourceClaim{
  1201  						{
  1202  							Name: "claim2",
  1203  						},
  1204  					},
  1205  				},
  1206  			},
  1207  			claimInfo: &ClaimInfo{ClaimInfoState: state.ClaimInfoState{ClaimName: resourceClaimName2}},
  1208  		},
  1209  	} {
  1210  		t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) {
  1211  			manager.cache.add(test.claimInfo)
  1212  
  1213  			fakeClaimInfos, err := manager.GetContainerClaimInfos(test.pod, test.container)
  1214  			assert.NoError(t, err)
  1215  			assert.Equal(t, 1, len(fakeClaimInfos))
  1216  			assert.Equal(t, test.expectedClaimName, fakeClaimInfos[0].ClaimInfoState.ClaimName)
  1217  
  1218  			manager.cache.delete(test.pod.Spec.ResourceClaims[0].Name, "default")
  1219  			_, err = manager.GetContainerClaimInfos(test.pod, test.container)
  1220  			assert.NoError(t, err)
  1221  		})
  1222  	}
  1223  }
  1224
View as plain text