From 751deb24f0f7fcdd0a1d6d6055122b8a72f5730f Mon Sep 17 00:00:00 2001
From: Luca Berneking <l.berneking@mittwald.de>
Date: Fri, 4 Oct 2024 16:00:57 +0200
Subject: [PATCH] Use single mutex for volume locks to prevent memory leak

Signed-off-by: Luca Berneking <l.berneking@mittwald.de>
---
 pkg/driver/controller.go  | 30 ++++++------------------
 pkg/driver/volume_lock.go | 49 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+), 23 deletions(-)
 create mode 100644 pkg/driver/volume_lock.go

diff --git a/pkg/driver/controller.go b/pkg/driver/controller.go
index f1d98e13..0308ffca 100644
--- a/pkg/driver/controller.go
+++ b/pkg/driver/controller.go
@@ -4,7 +4,6 @@ import (
 	"fmt"
 	"strconv"
 	"strings"
-	"sync"
 	"time"
 
 	"github.com/container-storage-interface/spec/lib/go/csi"
@@ -65,23 +64,7 @@ type controller struct {
 	k8sNodeInformer cache.SharedIndexInformer
 	zfsNodeInformer cache.SharedIndexInformer
 
-	volMutexes sync.Map
-}
-
-func (cs *controller) LockVolume(volume string) func() {
-	value, _ := cs.volMutexes.LoadOrStore(volume, &sync.Mutex{})
-	mtx := value.(*sync.Mutex)
-	mtx.Lock()
-	return func() { mtx.Unlock() }
-}
-
-func (cs *controller) LockVolumeWithSnapshot(volume string, snapshot string) func() {
-	unlockVol := cs.LockVolume(volume)
-	unlockSnap := cs.LockVolume(snapshot)
-	return func() {
-		unlockVol()
-		unlockSnap()
-	}
+	volumeLock *volumeLock
 }
 
 // NewController returns a new instance
@@ -90,6 +73,7 @@ func NewController(d *CSIDriver) csi.ControllerServer {
 	ctrl := &controller{
 		driver:       d,
 		capabilities: newControllerCapabilities(),
+		volumeLock:   newVolumeLock(),
 	}
 	if err := ctrl.init(); err != nil {
 		klog.Fatalf("init controller: %v", err)
@@ -467,7 +451,7 @@ func (cs *controller) CreateVolume(
 	contentSource := req.GetVolumeContentSource()
 	pvcName := helpers.GetInsensitiveParameter(&parameters, "csi.storage.k8s.io/pvc/name")
 
-	unlock := cs.LockVolume(volName)
+	unlock := cs.volumeLock.LockVolume(volName)
 	defer unlock()
 
 	if contentSource != nil && contentSource.GetSnapshot() != nil {
@@ -513,7 +497,7 @@ func (cs *controller) DeleteVolume(
 	}
 
 	volumeID := strings.ToLower(req.GetVolumeId())
-	unlock := cs.LockVolume(volumeID)
+	unlock := cs.volumeLock.LockVolume(volumeID)
 	defer unlock()
 
 	// verify if the volume has already been deleted
@@ -633,7 +617,7 @@ func (cs *controller) ControllerExpandVolume(
 			"ControllerExpandVolume: no volumeID provided",
 		)
 	}
-	unlock := cs.LockVolume(volumeID)
+	unlock := cs.volumeLock.LockVolume(volumeID)
 	defer unlock()
 
 	/* round off the new size */
@@ -731,7 +715,7 @@ func (cs *controller) CreateSnapshot(
 	if err != nil {
 		return nil, err
 	}
-	unlock := cs.LockVolumeWithSnapshot(volumeID, snapName)
+	unlock := cs.volumeLock.LockVolumeWithSnapshot(volumeID, snapName)
 	defer unlock()
 
 	snapTimeStamp := time.Now().Unix()
@@ -829,7 +813,7 @@ func (cs *controller) DeleteSnapshot(
 		// should succeed when an invalid snapshot id is used
 		return &csi.DeleteSnapshotResponse{}, nil
 	}
-	unlock := cs.LockVolumeWithSnapshot(snapshotID[0], snapshotID[1])
+	unlock := cs.volumeLock.LockVolumeWithSnapshot(snapshotID[0], snapshotID[1])
 	defer unlock()
 	if err := zfs.DeleteSnapshot(snapshotID[1]); err != nil {
 		return nil, status.Errorf(
diff --git a/pkg/driver/volume_lock.go b/pkg/driver/volume_lock.go
new file mode 100644
index 00000000..c3b140a8
--- /dev/null
+++ b/pkg/driver/volume_lock.go
@@ -0,0 +1,49 @@
+package driver
+
+import (
+	"sync"
+)
+
+type volumeLock struct {
+	cond   sync.Cond
+	locked map[string]struct{}
+}
+
+func newVolumeLock() *volumeLock {
+	return &volumeLock{
+		cond:   *sync.NewCond(&sync.Mutex{}),
+		locked: map[string]struct{}{},
+	}
+}
+
+func (l *volumeLock) LockVolume(volume string) func() {
+	l.cond.L.Lock()
+	defer l.cond.L.Unlock()
+
+	for {
+		if _, locked := l.locked[volume]; !locked {
+			break
+		}
+
+		l.cond.Wait()
+	}
+
+	l.locked[volume] = struct{}{}
+
+	return func() {
+		l.cond.L.Lock()
+		defer l.cond.L.Unlock()
+
+		delete(l.locked, volume)
+		l.cond.Broadcast()
+	}
+}
+
+func (l *volumeLock) LockVolumeWithSnapshot(volume string, snapshot string) func() {
+	unlockVol := l.LockVolume(volume)
+	unlockSnap := l.LockVolume(snapshot)
+	return func() {
+		unlockVol()
+		unlockSnap()
+	}
+}