Add map lock on Publish/Unpublish volume

pkg/csi_driver/node.go

@@ -52,11 +52,7 @@ type nodeServer struct {
 	volumeLocks           *util.VolumeLocks
 	k8sClients            clientset.Interface
 	limiter               rate.Limiter
-	volumeStateStore      map[string]*volumeState
-}
-
-type volumeState struct {
-	bucketAccessCheckPassed bool
+	volumeStateStore      *util.VolumeStateStore
 }
 
 func newNodeServer(driver *GCSDriver, mounter mount.Interface) csi.NodeServer {
@@ -67,7 +63,7 @@ func newNodeServer(driver *GCSDriver, mounter mount.Interface) csi.NodeServer {
 		volumeLocks:           util.NewVolumeLocks(),
 		k8sClients:            driver.config.K8sClients,
 		limiter:               *rate.NewLimiter(rate.Every(time.Second), 10),
-		volumeStateStore:      make(map[string]*volumeState),
+		volumeStateStore:      util.NewVolumeStateStore(),
 	}
 }
 
@@ -113,13 +109,13 @@ func (s *nodeServer) NodePublishVolume(ctx context.Context, req *csi.NodePublish
 	if bucketName != "_" && !skipBucketAccessCheck {
 		// Use target path as an volume identifier because it corresponds to Pods and volumes.
 		// Pods may belong to different namespaces and would need their own access check.
-		vs, ok := s.volumeStateStore[targetPath]
+		vs, ok := s.volumeStateStore.Load(targetPath)
 		if !ok {
-			s.volumeStateStore[targetPath] = &volumeState{}
-			vs = s.volumeStateStore[targetPath]
+			s.volumeStateStore.Store(targetPath, &util.VolumeState{})
+			vs, _ = s.volumeStateStore.Load(targetPath)
 		}
 
-		if !vs.bucketAccessCheckPassed {
+		if !vs.BucketAccessCheckPassed {
 			storageService, err := s.prepareStorageService(ctx, vc)
 			if err != nil {
 				return nil, status.Errorf(codes.Unauthenticated, "failed to prepare storage service: %v", err)
@@ -130,7 +126,7 @@ func (s *nodeServer) NodePublishVolume(ctx context.Context, req *csi.NodePublish
 				return nil, status.Errorf(storage.ParseErrCode(err), "failed to get GCS bucket %q: %v", bucketName, err)
 			}
 
-			vs.bucketAccessCheckPassed = true
+			vs.BucketAccessCheckPassed = true
 		}
 	}
 
@@ -239,7 +235,7 @@ func (s *nodeServer) NodeUnpublishVolume(_ context.Context, req *csi.NodeUnpubli
 		s.driver.config.MetricsManager.UnregisterMetricsCollector(targetPath)
 	}
 
-	delete(s.volumeStateStore, targetPath)
+	s.volumeStateStore.Delete(targetPath)
 
 	// Check if the target path is already mounted
 	if mounted, err := s.isDirMounted(targetPath); mounted || err != nil {

pkg/csi_driver/node_test.go

@@ -21,12 +21,14 @@ import (
 	"errors"
 	"os"
 	"path/filepath"
+	"sync"
 	"testing"
 
 	csi "github.com/container-storage-interface/spec/lib/go/csi"
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
 	"github.com/googlecloudplatform/gcs-fuse-csi-driver/pkg/cloud_provider/storage"
+	"github.com/googlecloudplatform/gcs-fuse-csi-driver/pkg/util"
 	"golang.org/x/net/context"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -279,3 +281,32 @@ func validateMountPoint(t *testing.T, name string, fm *mount.FakeMounter, e *mou
 		t.Errorf("unexpected options args (-got, +want)\n%s", diff)
 	}
 }
+
+func TestConcurrentMapWrites(t *testing.T) {
+	t.Parallel()
+	// Create a shared map for the test
+	sharedVSS := util.NewVolumeStateStore()
+	// Number of concurrent writes we want to simulate
+	numWrites := 2000
+
+	// Use a WaitGroup to wait for all goroutines to finish
+	var wg sync.WaitGroup
+
+	// Run concurrent tests manually using goroutines
+	for i := range numWrites {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			// Simulate concurrent writes to the shared map
+			sharedVSS.Store(string(rune(i)), &util.VolumeState{})
+		}()
+	}
+
+	// Wait for all goroutines to finish
+	wg.Wait()
+
+	// validate correct number of writes occurred
+	if int(sharedVSS.Size()) != numWrites {
+		t.Errorf("expected %d entries in the map, got %d", numWrites, sharedVSS.Size())
+	}
+}

pkg/util/volume_state_store_lock.go

@@ -0,0 +1,66 @@
+/*
+Copyright 2018 The Kubernetes Authors.
+Copyright 2022 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package util
+
+import (
+	"sync"
+	"sync/atomic"
+)
+
+// VolumeStateStore provides a thread-safe map for storing volume states.
+type VolumeStateStore struct {
+	store sync.Map // Concurrent-safe map
+	size  int32
+}
+
+type VolumeState struct {
+	BucketAccessCheckPassed bool
+}
+
+// NewVolumeStateStore initializes the volume state store.
+func NewVolumeStateStore() *VolumeStateStore {
+	return &VolumeStateStore{}
+}
+
+// NewVolumeStateStore initializes the volume state store.
+func (vss *VolumeStateStore) Size() int32 {
+	return vss.size
+}
+
+// Store adds or updates a volume state.
+func (vss *VolumeStateStore) Store(volumeID string, state *VolumeState) {
+	vss.store.Store(volumeID, state)
+	atomic.AddInt32(&vss.size, 1)
+}
+
+// Load retrieves the state of a volume.
+func (vss *VolumeStateStore) Load(volumeID string) (*VolumeState, bool) {
+	if value, ok := vss.store.Load(volumeID); ok {
+		if v, ok := value.(*VolumeState); ok {
+			return v, true
+		}
+	}
+
+	return nil, false
+}
+
+// Delete removes a volume from the store.
+func (vss *VolumeStateStore) Delete(volumeID string) {
+	vss.store.Delete(volumeID)
+	atomic.AddInt32(&vss.size, -1)
+}