From 0ae19c82a6c22e321fcd8b739fd4f54fa447ac10 Mon Sep 17 00:00:00 2001
From: huweiwen <huweiwen.hww@alibaba-inc.com>
Date: Tue, 30 Jan 2024 21:42:15 +0800
Subject: [PATCH] fix data race reported by `go test ./... -race`

---
 pkg/csi/mock_client.go                     | 29 +++++++++++-----------
 pkg/modifycontroller/controller.go         | 29 ++++++++++++++--------
 pkg/modifycontroller/modify_status_test.go |  5 +++-
 3 files changed, 36 insertions(+), 27 deletions(-)

diff --git a/pkg/csi/mock_client.go b/pkg/csi/mock_client.go
index f367d858c..57e85299d 100644
--- a/pkg/csi/mock_client.go
+++ b/pkg/csi/mock_client.go
@@ -3,6 +3,7 @@ package csi
 import (
 	"context"
 	"fmt"
+	"sync/atomic"
 
 	"github.com/container-storage-interface/spec/lib/go/csi"
 	"github.com/kubernetes-csi/csi-lib-utils/connection"
@@ -20,8 +21,6 @@ func NewMockClient(
 		supportsNodeResize:                      supportsNodeResize,
 		supportsControllerResize:                supportsControllerResize,
 		supportsControllerModify:                supportsControllerModify,
-		expandCalled:                            0,
-		modifyCalled:                            0,
 		supportsPluginControllerService:         supportsPluginControllerService,
 		supportsControllerSingleNodeMultiWriter: supportsControllerSingleNodeMultiWriter,
 	}
@@ -34,13 +33,13 @@ type MockClient struct {
 	supportsControllerModify                bool
 	supportsPluginControllerService         bool
 	supportsControllerSingleNodeMultiWriter bool
-	expandCalled                            int
-	modifyCalled                            int
+	expandCalled                            atomic.Int32
+	modifyCalled                            atomic.Int32
 	expansionFailed                         bool
 	modifyFailed                            bool
 	checkMigratedLabel                      bool
-	usedSecrets                             map[string]string
-	usedCapability                          *csi.VolumeCapability
+	usedSecrets                             atomic.Pointer[map[string]string]
+	usedCapability                          atomic.Pointer[csi.VolumeCapability]
 }
 
 func (c *MockClient) GetDriverName(context.Context) (string, error) {
@@ -87,7 +86,7 @@ func (c *MockClient) Expand(
 	capability *csi.VolumeCapability) (int64, bool, error) {
 	// TODO: Determine whether the operation succeeds or fails by parameters.
 	if c.expansionFailed {
-		c.expandCalled++
+		c.expandCalled.Add(1)
 		return requestBytes, c.supportsNodeResize, fmt.Errorf("expansion failed")
 	}
 	if c.checkMigratedLabel {
@@ -99,27 +98,27 @@ func (c *MockClient) Expand(
 			return requestBytes, c.supportsNodeResize, err
 		}
 	}
-	c.expandCalled++
-	c.usedSecrets = secrets
-	c.usedCapability = capability
+	c.expandCalled.Add(1)
+	c.usedSecrets.Store(&secrets)
+	c.usedCapability.Store(capability)
 	return requestBytes, c.supportsNodeResize, nil
 }
 
 func (c *MockClient) GetExpandCount() int {
-	return c.expandCalled
+	return int(c.expandCalled.Load())
 }
 
 func (c *MockClient) GetModifyCount() int {
-	return c.modifyCalled
+	return int(c.modifyCalled.Load())
 }
 
 func (c *MockClient) GetCapability() *csi.VolumeCapability {
-	return c.usedCapability
+	return c.usedCapability.Load()
 }
 
 // GetSecrets returns secrets used for volume expansion
 func (c *MockClient) GetSecrets() map[string]string {
-	return c.usedSecrets
+	return *c.usedSecrets.Load()
 }
 
 func (c *MockClient) CloseConnection() {
@@ -131,7 +130,7 @@ func (c *MockClient) Modify(
 	volumeID string,
 	secrets map[string]string,
 	mutableParameters map[string]string) error {
-	c.modifyCalled++
+	c.modifyCalled.Add(1)
 	if c.modifyFailed {
 		return fmt.Errorf("modify failed")
 	}
diff --git a/pkg/modifycontroller/controller.go b/pkg/modifycontroller/controller.go
index 02b58bb25..0b2fde85d 100644
--- a/pkg/modifycontroller/controller.go
+++ b/pkg/modifycontroller/controller.go
@@ -196,21 +196,13 @@ func isFirstTimeModifyVolumeWithPVC(pvc *v1.PersistentVolumeClaim, pv *v1.Persis
 	return false
 }
 
-// Run starts the controller.
-func (ctrl *modifyController) Run(
-	workers int, ctx context.Context) {
-	defer ctrl.claimQueue.ShutDown()
-
-	klog.InfoS("Starting external resizer for modify volume", "controller", ctrl.name)
-	defer klog.InfoS("Shutting down external resizer", "controller", ctrl.name)
-
-	stopCh := ctx.Done()
+func (ctrl *modifyController) init(ctx context.Context) bool {
 	informersSyncd := []cache.InformerSynced{ctrl.pvListerSynced, ctrl.pvcListerSynced}
 	informersSyncd = append(informersSyncd, ctrl.vacListerSynced)
 
-	if !cache.WaitForCacheSync(stopCh, informersSyncd...) {
+	if !cache.WaitForCacheSync(ctx.Done(), informersSyncd...) {
 		klog.ErrorS(nil, "Cannot sync pod, pv, pvc or vac caches")
-		return
+		return false
 	}
 
 	// Cache all the InProgress/Infeasible PVCs as Uncertain for ModifyVolume
@@ -218,7 +210,22 @@ func (ctrl *modifyController) Run(
 	if err != nil {
 		klog.ErrorS(err, "Failed to initialize uncertain pvcs")
 	}
+	return true
+}
 
+// Run starts the controller.
+func (ctrl *modifyController) Run(
+	workers int, ctx context.Context) {
+	defer ctrl.claimQueue.ShutDown()
+
+	klog.InfoS("Starting external resizer for modify volume", "controller", ctrl.name)
+	defer klog.InfoS("Shutting down external resizer", "controller", ctrl.name)
+
+	if !ctrl.init(ctx) {
+		return
+	}
+
+	stopCh := ctx.Done()
 	for i := 0; i < workers; i++ {
 		go wait.Until(ctrl.sync, 0, stopCh)
 	}
diff --git a/pkg/modifycontroller/modify_status_test.go b/pkg/modifycontroller/modify_status_test.go
index 113dd84af..e0850a9bc 100644
--- a/pkg/modifycontroller/modify_status_test.go
+++ b/pkg/modifycontroller/modify_status_test.go
@@ -321,7 +321,10 @@ func TestRemovePVCFromModifyVolumeUncertainCache(t *testing.T) {
 
 			ctx := context.TODO()
 			defer ctx.Done()
-			go controller.Run(1, ctx)
+			success := ctrlInstance.init(ctx)
+			if !success {
+				t.Fatal("failed to init controller")
+			}
 
 			for _, obj := range initialObjects {
 				switch obj.(type) {