Merge pull request #242 from andyzhangx/lock

andyzhangx · web-flow · commit ece81ded8cf3 · 2020-09-21T12:51:44.000+08:00
fix: add lock for prevent account list throttling
diff --git a/charts/latest/blob-csi-driver/templates/csi-blob-controller.yaml b/charts/latest/blob-csi-driver/templates/csi-blob-controller.yaml
@@ -33,7 +33,7 @@ spec:
             - "--csi-address=$(ADDRESS)"
             - "--enable-leader-election"
             - "--leader-election-type=leases"
-            - "--timeout=30s"
+            - "--timeout=60s"
           env:
             - name: ADDRESS
               value: /csi/csi.sock
diff --git a/deploy/csi-blob-controller.yaml b/deploy/csi-blob-controller.yaml
@@ -32,7 +32,7 @@ spec:
             - "--csi-address=$(ADDRESS)"
             - "--enable-leader-election"
             - "--leader-election-type=leases"
-            - "--timeout=30s"
+            - "--timeout=60s"
           env:
             - name: ADDRESS
               value: /csi/csi.sock
diff --git a/pkg/blob/blob.go b/pkg/blob/blob.go
@@ -20,6 +20,8 @@ import (
 	"fmt"
 	"strings"
 
+	"sigs.k8s.io/blob-csi-driver/pkg/util"
+
 	csicommon "sigs.k8s.io/blob-csi-driver/pkg/csi-common"
 
 	"github.com/container-storage-interface/spec/lib/go/csi"
@@ -68,8 +70,9 @@ var (
 // Driver implements all interfaces of CSI drivers
 type Driver struct {
 	csicommon.CSIDriver
-	cloud   *azure.Cloud
-	mounter *mount.SafeFormatAndMount
+	cloud      *azure.Cloud
+	mounter    *mount.SafeFormatAndMount
+	volLockMap *util.LockMap
 }
 
 // NewDriver Creates a NewCSIDriver object. Assumes vendor version is equal to driver version &
@@ -79,6 +82,7 @@ func NewDriver(nodeID string) *Driver {
 	driver.Name = DriverName
 	driver.Version = driverVersion
 	driver.NodeID = nodeID
+	driver.volLockMap = util.NewLockMap()
 	return &driver
 }
 
diff --git a/pkg/blob/controllerserver.go b/pkg/blob/controllerserver.go
@@ -89,19 +89,19 @@ func (d *Driver) CreateVolume(ctx context.Context, req *csi.CreateVolumeRequest)
 		return nil, status.Errorf(codes.InvalidArgument, "protocol(%s) is not supported, supported protocol list: %v", protocol, supportedProtocolList)
 	}
 
+	enableHTTPSTrafficOnly := true
 	if protocol == nfs {
 		if account == "" {
 			return nil, status.Errorf(codes.InvalidArgument, "storage account must be specified when provisioning nfs file share")
 		}
+		enableHTTPSTrafficOnly = false
 	}
 
 	accountKind := string(storage.StorageV2)
 	if strings.HasPrefix(strings.ToLower(storageAccountType), "premium") {
 		accountKind = string(storage.BlockBlobStorage)
 	}
 
-	enableHTTPSTrafficOnly := true
-
 	tags, err := azure.ConvertTagsToMap(customTags)
 	if err != nil {
 		return nil, err
@@ -119,6 +119,10 @@ func (d *Driver) CreateVolume(ctx context.Context, req *csi.CreateVolumeRequest)
 
 	var accountName, accountKey string
 	if len(req.GetSecrets()) == 0 { // check whether account is provided by secret
+		lockKey := account + storageAccountType + accountKind + resourceGroup + location
+		d.volLockMap.LockEntry(lockKey)
+		defer d.volLockMap.UnlockEntry(lockKey)
+
 		err = wait.ExponentialBackoff(d.cloud.RequestBackoff(), func() (bool, error) {
 			var retErr error
 			accountName, accountKey, retErr = d.cloud.EnsureStorageAccount(accountOptions, protocol)
diff --git a/pkg/blob/nodeserver.go b/pkg/blob/nodeserver.go
@@ -148,16 +148,9 @@ func (d *Driver) NodeStageVolume(ctx context.Context, req *csi.NodeStageVolumeRe
 
 		source := fmt.Sprintf("%s:/%s/%s", blobStorageEndPoint, accountName, containerName)
 		mountOptions := util.JoinMountOptions(mountFlags, []string{"sec=sys,vers=3,nolock"})
-		mountComplete := false
-		err := wait.PollImmediate(1*time.Second, 2*time.Minute, func() (bool, error) {
-			err := d.mounter.MountSensitive(source, targetPath, nfs, mountOptions, []string{})
-			mountComplete = true
-			return true, err
-		})
-		if !mountComplete {
-			return nil, status.Error(codes.Internal, fmt.Sprintf("volume(%s) mount %q on %q failed with timeout(2m)", volumeID, source, targetPath))
-		}
-		if err != nil {
+		if err := wait.PollImmediate(1*time.Second, 2*time.Minute, func() (bool, error) {
+			return true, d.mounter.MountSensitive(source, targetPath, nfs, mountOptions, []string{})
+		}); err != nil {
 			return nil, status.Error(codes.Internal, fmt.Sprintf("volume(%s) mount %q on %q failed with %v", volumeID, source, targetPath, err))
 		}
 		klog.V(2).Infof("volume(%s) mount %q on %q succeeded", volumeID, source, targetPath)
diff --git a/pkg/util/util.go b/pkg/util/util.go
@@ -18,6 +18,7 @@ package util
 
 import (
 	"os"
+	"sync"
 )
 
 const (
@@ -80,3 +81,51 @@ func MakeDir(pathname string) error {
 	}
 	return nil
 }
+
+// LockMap used to lock on entries
+type LockMap struct {
+	sync.Mutex
+	mutexMap map[string]*sync.Mutex
+}
+
+// NewLockMap returns a new lock map
+func NewLockMap() *LockMap {
+	return &LockMap{
+		mutexMap: make(map[string]*sync.Mutex),
+	}
+}
+
+// LockEntry acquires a lock associated with the specific entry
+func (lm *LockMap) LockEntry(entry string) {
+	lm.Lock()
+	// check if entry does not exists, then add entry
+	if _, exists := lm.mutexMap[entry]; !exists {
+		lm.addEntry(entry)
+	}
+
+	lm.Unlock()
+	lm.lockEntry(entry)
+}
+
+// UnlockEntry release the lock associated with the specific entry
+func (lm *LockMap) UnlockEntry(entry string) {
+	lm.Lock()
+	defer lm.Unlock()
+
+	if _, exists := lm.mutexMap[entry]; !exists {
+		return
+	}
+	lm.unlockEntry(entry)
+}
+
+func (lm *LockMap) addEntry(entry string) {
+	lm.mutexMap[entry] = &sync.Mutex{}
+}
+
+func (lm *LockMap) lockEntry(entry string) {
+	lm.mutexMap[entry].Lock()
+}
+
+func (lm *LockMap) unlockEntry(entry string) {
+	lm.mutexMap[entry].Unlock()
+}
diff --git a/pkg/util/util_test.go b/pkg/util/util_test.go
@@ -19,6 +19,7 @@ package util
 import (
 	"os"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 )
@@ -39,6 +40,77 @@ func TestRoundUpGiB(t *testing.T) {
 	}
 }
 
+func TestSimpleLockEntry(t *testing.T) {
+	testLockMap := NewLockMap()
+
+	callbackChan1 := make(chan interface{})
+	go testLockMap.lockAndCallback(t, "entry1", callbackChan1)
+	ensureCallbackHappens(t, callbackChan1)
+}
+
+func TestSimpleLockUnlockEntry(t *testing.T) {
+	testLockMap := NewLockMap()
+
+	callbackChan1 := make(chan interface{})
+	go testLockMap.lockAndCallback(t, "entry1", callbackChan1)
+	ensureCallbackHappens(t, callbackChan1)
+	testLockMap.UnlockEntry("entry1")
+}
+
+func TestConcurrentLockEntry(t *testing.T) {
+	testLockMap := NewLockMap()
+
+	callbackChan1 := make(chan interface{})
+	callbackChan2 := make(chan interface{})
+
+	go testLockMap.lockAndCallback(t, "entry1", callbackChan1)
+	ensureCallbackHappens(t, callbackChan1)
+
+	go testLockMap.lockAndCallback(t, "entry1", callbackChan2)
+	ensureNoCallback(t, callbackChan2)
+
+	testLockMap.UnlockEntry("entry1")
+	ensureCallbackHappens(t, callbackChan2)
+	testLockMap.UnlockEntry("entry1")
+}
+
+func (lm *LockMap) lockAndCallback(t *testing.T, entry string, callbackChan chan<- interface{}) {
+	lm.LockEntry(entry)
+	callbackChan <- true
+}
+
+var callbackTimeout = 2 * time.Second
+
+func ensureCallbackHappens(t *testing.T, callbackChan <-chan interface{}) bool {
+	select {
+	case <-callbackChan:
+		return true
+	case <-time.After(callbackTimeout):
+		t.Fatalf("timed out waiting for callback")
+		return false
+	}
+}
+
+func ensureNoCallback(t *testing.T, callbackChan <-chan interface{}) bool {
+	select {
+	case <-callbackChan:
+		t.Fatalf("unexpected callback")
+		return false
+	case <-time.After(callbackTimeout):
+		return true
+	}
+}
+
+func TestUnlockEntryNotExists(t *testing.T) {
+	testLockMap := NewLockMap()
+
+	callbackChan1 := make(chan interface{})
+	go testLockMap.lockAndCallback(t, "entry1", callbackChan1)
+	ensureCallbackHappens(t, callbackChan1)
+	// entry2 does not exist
+	testLockMap.UnlockEntry("entry2")
+	testLockMap.UnlockEntry("entry1")
+}
 func TestBytesToGiB(t *testing.T) {
 	var sizeInBytes int64 = 5 * GiB
 
