From 986a0997bba5410920753050b094f7718c6b7c38 Mon Sep 17 00:00:00 2001
From: joshvanl <me@joshvanl.dev>
Date: Tue, 23 Jul 2024 16:23:18 +0100
Subject: [PATCH] MutexMap: Adds DeleteRUnlock and fixes RLock/RUnlock

Signed-off-by: joshvanl <me@joshvanl.dev>
---
 concurrency/mutexmap.go      | 16 ++++++++++++++--
 concurrency/mutexmap_test.go | 28 ++++++++++++++++++++--------
 2 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/concurrency/mutexmap.go b/concurrency/mutexmap.go
index 63f93c9..7a87ed9 100644
--- a/concurrency/mutexmap.go
+++ b/concurrency/mutexmap.go
@@ -30,6 +30,7 @@ import (
 // - Clear(): Removes all mutexes from the map.
 // - ItemCount() int: Returns the number of items (mutexes) in the map.
 // - DeleteUnlock(key T): Removes the mutex associated with the given key from the map and releases the lock.
+// - DeleteRUnlock(key T): Removes the mutex associated with the given key from the map and releases the read lock.
 type MutexMap[T comparable] interface {
 	Lock(key T)
 	Unlock(key T)
@@ -39,6 +40,7 @@ type MutexMap[T comparable] interface {
 	Clear()
 	ItemCount() int
 	DeleteUnlock(key T)
+	DeleteRUnlock(key T)
 }
 
 type mutexMap[T comparable] struct {
@@ -90,7 +92,7 @@ func (a *mutexMap[T]) RLock(key T) {
 		}
 		a.lock.Unlock()
 	}
-	mutex.Lock()
+	mutex.RLock()
 }
 
 func (a *mutexMap[T]) RUnlock(key T) {
@@ -98,7 +100,7 @@ func (a *mutexMap[T]) RUnlock(key T) {
 	mutex, ok := a.items[key]
 	a.lock.RUnlock()
 	if ok {
-		mutex.Unlock()
+		mutex.RUnlock()
 	}
 }
 
@@ -118,6 +120,16 @@ func (a *mutexMap[T]) DeleteUnlock(key T) {
 	}
 }
 
+func (a *mutexMap[T]) DeleteRUnlock(key T) {
+	a.lock.Lock()
+	mutex, ok := a.items[key]
+	delete(a.items, key)
+	a.lock.Unlock()
+	if ok {
+		mutex.RUnlock()
+	}
+}
+
 func (a *mutexMap[T]) Clear() {
 	a.lock.Lock()
 	clear(a.items)
diff --git a/concurrency/mutexmap_test.go b/concurrency/mutexmap_test.go
index 2db2e13..bdebcfc 100644
--- a/concurrency/mutexmap_test.go
+++ b/concurrency/mutexmap_test.go
@@ -15,8 +15,11 @@ package concurrency
 
 import (
 	"sync"
+	"sync/atomic"
 	"testing"
+	"time"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -37,7 +40,7 @@ func TestNewMutexMap_Add_Delete(t *testing.T) {
 	})
 
 	t.Run("Concurrently lock and unlock mutexes", func(t *testing.T) {
-		var counter int
+		var counter atomic.Int64
 		var wg sync.WaitGroup
 
 		numGoroutines := 10
@@ -48,13 +51,13 @@ func TestNewMutexMap_Add_Delete(t *testing.T) {
 			go func() {
 				defer wg.Done()
 				mm.Lock("key1")
-				counter++
+				counter.Add(1)
 				mm.Unlock("key1")
 			}()
 		}
 		wg.Wait()
 
-		require.Equal(t, 10, counter)
+		require.Equal(t, int64(10), counter.Load())
 	})
 
 	t.Run("RLock and RUnlock mutex", func(t *testing.T) {
@@ -65,24 +68,33 @@ func TestNewMutexMap_Add_Delete(t *testing.T) {
 	})
 
 	t.Run("Concurrently RLock and RUnlock mutexes", func(t *testing.T) {
-		var counter int
+		var counter atomic.Int64
 		var wg sync.WaitGroup
 
 		numGoroutines := 10
-		wg.Add(numGoroutines)
+		wg.Add(numGoroutines * 2)
 
 		// Concurrently RLock and RUnlock for each key
 		for i := 0; i < numGoroutines; i++ {
 			go func() {
 				defer wg.Done()
 				mm.RLock("key1")
-				counter++
+				counter.Add(1)
+			}()
+		}
+
+		assert.EventuallyWithT(t, func(ct *assert.CollectT) {
+			assert.Equal(ct, int64(10), counter.Load())
+		}, 5*time.Second, 10*time.Millisecond)
+
+		for i := 0; i < numGoroutines; i++ {
+			go func() {
+				defer wg.Done()
 				mm.RUnlock("key1")
 			}()
 		}
-		wg.Wait()
 
-		require.Equal(t, 10, counter)
+		wg.Wait()
 	})
 
 	t.Run("Delete mutex", func(t *testing.T) {