refactor shared state codes

pkg/azuredisk/azuredisk_v2.go

@@ -297,7 +297,7 @@ func (d *DriverV2) StartControllersAndDieOnExit(ctx context.Context) {
 		os.Exit(1)
 	}
 
-	sharedState := controller.NewSharedState()
+	sharedState := controller.NewSharedState(mgr.GetClient())
 
 	// Setup a new controller to clean-up AzDriverNodes
 	// objects for the nodes which get deleted
@@ -353,8 +353,8 @@ func (d *DriverV2) StartControllersAndDieOnExit(ctx context.Context) {
 	// Leader controller manager should recover CRI if possible and clean them up before exiting.
 	go func() {
 		<-mgr.Elected()
-		// add cached client to sharedState
-		sharedState.RegisterCachedClient(mgr.GetClient())
+		// add sharedState to crd provisioner
+		d.crdProvisioner.RegisterSharedState(ctx, sharedState)
 
 		// recover lost states if necessary
 		klog.Infof("Elected as leader; initiating CRI recovery...")

pkg/controller/common.go

@@ -146,20 +146,17 @@ type SharedState struct {
 	cachedClient     client.Client
 }
 
-func NewSharedState() *SharedState {
+func NewSharedState(cachedClient client.Client) *SharedState {
 	return &SharedState{
 		podToClaimsMap:   &sync.Map{},
 		claimToPodsMap:   &sync.Map{},
 		volumeToClaimMap: &sync.Map{},
 		claimToVolumeMap: &sync.Map{},
 		podLocks:         &sync.Map{},
+		cachedClient:     cachedClient,
 	}
 }
 
-func (c *SharedState) RegisterCachedClient(client client.Client) {
-	c.cachedClient = client
-}
-
 func (c *SharedState) GetCachedClient() client.Client {
 	return c.cachedClient
 }

pkg/controller/common_test.go

@@ -291,8 +291,8 @@ func createPod(podNamespace, podName string, pvcs []string) v1.Pod {
 	return testPod
 }
 
-func initState(objs ...runtime.Object) (c *SharedState) {
-	c = NewSharedState()
+func initState(client client.Client, objs ...runtime.Object) (c *SharedState) {
+	c = NewSharedState(client)
 
 	for _, obj := range objs {
 		switch target := obj.(type) {

pkg/controller/pod_test.go

@@ -36,8 +36,7 @@ import (
 
 func NewTestPodController(controller *gomock.Controller, namespace string, objects ...runtime.Object) *ReconcilePod {
 	diskv1alpha1Objs, kubeObjs := splitObjects(objects...)
-	sharedState := initState(objects...)
-	sharedState.cachedClient = mockclient.NewMockClient(controller)
+	sharedState := initState(mockclient.NewMockClient(controller), objects...)
 
 	return &ReconcilePod{
 		azVolumeClient: diskfakes.NewSimpleClientset(diskv1alpha1Objs...),

pkg/controller/pv_test.go

@@ -40,8 +40,7 @@ type testReconcilePV struct {
 
 func newTestPVController(controller *gomock.Controller, namespace string, objects ...runtime.Object) *testReconcilePV {
 	diskv1alpha1Objs, kubeObjs := splitObjects(objects...)
-	sharedState := initState(objects...)
-	sharedState.cachedClient = mockclient.NewMockClient(controller)
+	sharedState := initState(mockclient.NewMockClient(controller), objects...)
 
 	return &testReconcilePV{
 		ReconcilePV: ReconcilePV{

pkg/controller/replica_test.go

@@ -39,8 +39,7 @@ import (
 
 func NewTestReplicaController(controller *gomock.Controller, namespace string, objects ...runtime.Object) *ReconcileReplica {
 	diskv1alpha1Objs, kubeObjs := splitObjects(objects...)
-	sharedState := initState(objects...)
-	sharedState.cachedClient = mockclient.NewMockClient(controller)
+	sharedState := initState(mockclient.NewMockClient(controller), objects...)
 
 	return &ReconcileReplica{
 		azVolumeClient: diskfakes.NewSimpleClientset(diskv1alpha1Objs...),

pkg/provisioner/crdprovisioner.go

@@ -149,8 +149,7 @@ func (c *CrdProvisioner) CreateVolume(
 	validVolumeName := azureutils.GetValidDiskName(volumeName)
 	azVolumeName := strings.ToLower(validVolumeName)
 
-	useCache := c.sharedState != nil && c.sharedState.GetCachedClient() != nil
-	azVolumeInstance, err := azureutils.GetAzVolume(ctx, c.getCachedClient(), c.azDiskClient, azVolumeName, c.namespace, useCache)
+	azVolumeInstance, err := azureutils.GetAzVolume(ctx, c.getCachedClient(), c.azDiskClient, azVolumeName, c.namespace, c.useCache())
 	azV := c.azDiskClient.DiskV1alpha1().AzVolumes(c.namespace)
 
 	if err == nil {
@@ -225,8 +224,8 @@ func (c *CrdProvisioner) CreateVolume(
 	}
 
 	conditionFunc := func() (bool, error) {
-		useCache := c.sharedState != nil && c.sharedState.GetCachedClient() != nil
-		azVolumeInstance, err := azureutils.GetAzVolume(ctx, c.getCachedClient(), c.azDiskClient, azVolumeName, c.namespace, useCache)
+
+		azVolumeInstance, err := azureutils.GetAzVolume(ctx, c.getCachedClient(), c.azDiskClient, azVolumeName, c.namespace, c.useCache())
 
 		if err != nil {
 			return true, err
@@ -276,8 +275,7 @@ func (c *CrdProvisioner) DeleteVolume(ctx context.Context, volumeID string, secr
 
 	azVolumeName := strings.ToLower(volumeName)
 
-	useCache := c.sharedState != nil && c.sharedState.GetCachedClient() != nil
-	azVolume, err := azureutils.GetAzVolume(ctx, c.getCachedClient(), c.azDiskClient, azVolumeName, c.namespace, useCache)
+	azVolume, err := azureutils.GetAzVolume(ctx, c.getCachedClient(), c.azDiskClient, azVolumeName, c.namespace, c.useCache())
 	if err != nil {
 		if errors.IsNotFound(err) {
 			klog.Infof("Could not find the volume name (%s). Deletion succeeded", volumeName)
@@ -315,8 +313,8 @@ func (c *CrdProvisioner) DeleteVolume(ctx context.Context, volumeID string, secr
 
 	conditionFunc := func() (bool, error) {
 		// Verify if the azVolume is deleted
-		useCache := c.sharedState != nil && c.sharedState.GetCachedClient() != nil
-		azVolume, err := azureutils.GetAzVolume(ctx, c.getCachedClient(), c.azDiskClient, azVolumeName, c.namespace, useCache)
+
+		azVolume, err := azureutils.GetAzVolume(ctx, c.getCachedClient(), c.azDiskClient, azVolumeName, c.namespace, c.useCache())
 		if err != nil {
 			if errors.IsNotFound(err) {
 				return true, nil
@@ -351,8 +349,7 @@ func (c *CrdProvisioner) PublishVolume(
 	if volumeContext == nil {
 		volumeContext = map[string]string{}
 	}
-	useCache := c.sharedState != nil && c.sharedState.GetCachedClient() != nil
-	azVolumeAttachmentInstance, err := azureutils.GetAzVolumeAttachment(ctx, c.getCachedClient(), c.azDiskClient, attachmentName, c.namespace, useCache)
+	azVolumeAttachmentInstance, err := azureutils.GetAzVolumeAttachment(ctx, c.getCachedClient(), c.azDiskClient, attachmentName, c.namespace, c.useCache())
 	if err == nil {
 		// If there already exists a primary attachment with populated responseObject return
 		if azVolumeAttachmentInstance.Status.Detail != nil && azVolumeAttachmentInstance.Status.Detail.PublishContext != nil && azVolumeAttachmentInstance.Status.Detail.Role == v1alpha1.PrimaryRole {
@@ -415,8 +412,8 @@ func (c *CrdProvisioner) PublishVolume(
 	}
 
 	conditionFunc := func() (bool, error) {
-		useCache := c.sharedState != nil && c.sharedState.GetCachedClient() != nil
-		azVolumeAttachmentInstance, err := azureutils.GetAzVolumeAttachment(ctx, c.getCachedClient(), c.azDiskClient, attachmentName, c.namespace, useCache)
+
+		azVolumeAttachmentInstance, err := azureutils.GetAzVolumeAttachment(ctx, c.getCachedClient(), c.azDiskClient, attachmentName, c.namespace, c.useCache())
 		if err != nil {
 			return false, err
 		}
@@ -462,8 +459,7 @@ func (c *CrdProvisioner) UnpublishVolume(
 
 	attachmentName := azureutils.GetAzVolumeAttachmentName(volumeName, nodeID)
 
-	useCache := c.sharedState != nil && c.sharedState.GetCachedClient() != nil
-	azVolumeAttachmentInstance, err := azureutils.GetAzVolumeAttachment(ctx, c.getCachedClient(), c.azDiskClient, attachmentName, c.namespace, useCache)
+	azVolumeAttachmentInstance, err := azureutils.GetAzVolumeAttachment(ctx, c.getCachedClient(), c.azDiskClient, attachmentName, c.namespace, c.useCache())
 	if err != nil {
 		if errors.IsNotFound(err) {
 			klog.Infof("AzVolumeAttachment (%s) has already been deleted.", attachmentName)
@@ -508,8 +504,8 @@ func (c *CrdProvisioner) UnpublishVolume(
 
 	conditionFunc := func() (bool, error) {
 		// Verify if the azVolume is deleted
-		useCache := c.sharedState != nil && c.sharedState.GetCachedClient() != nil
-		azVolumeAttachmentInstance, err := azureutils.GetAzVolumeAttachment(ctx, c.getCachedClient(), c.azDiskClient, attachmentName, c.namespace, useCache)
+
+		azVolumeAttachmentInstance, err := azureutils.GetAzVolumeAttachment(ctx, c.getCachedClient(), c.azDiskClient, attachmentName, c.namespace, c.useCache())
 		if err != nil {
 			if errors.IsNotFound(err) {
 				return true, nil
@@ -537,8 +533,7 @@ func (c *CrdProvisioner) ExpandVolume(
 	}
 	azVolumeName = strings.ToLower(azVolumeName)
 
-	useCache := c.sharedState != nil && c.sharedState.GetCachedClient() != nil
-	azVolume, err := azureutils.GetAzVolume(ctx, c.getCachedClient(), c.azDiskClient, azVolumeName, c.namespace, useCache)
+	azVolume, err := azureutils.GetAzVolume(ctx, c.getCachedClient(), c.azDiskClient, azVolumeName, c.namespace, c.useCache())
 	if err != nil || azVolume == nil {
 		klog.Errorf("Failed to retrieve existing volume id (%s)", volumeID)
 		return nil, status.Error(codes.Internal, fmt.Sprintf("Failed to retrieve volume id (%s), error: %v", volumeID, err))
@@ -556,8 +551,8 @@ func (c *CrdProvisioner) ExpandVolume(
 	}
 
 	conditionFunc := func() (bool, error) {
-		useCache := c.sharedState != nil && c.sharedState.GetCachedClient() != nil
-		azVolume, err := azureutils.GetAzVolume(ctx, c.getCachedClient(), c.azDiskClient, azVolumeName, c.namespace, useCache)
+
+		azVolume, err := azureutils.GetAzVolume(ctx, c.getCachedClient(), c.azDiskClient, azVolumeName, c.namespace, c.useCache())
 		if err != nil {
 			return true, err
 		}
@@ -663,3 +658,7 @@ func (c *CrdProvisioner) getCachedClient() client.Client {
 	}
 	return c.sharedState.GetCachedClient()
 }
+
+func (c *CrdProvisioner) useCache() bool {
+	return c.sharedState != nil && c.sharedState.GetCachedClient() != nil
+}