From f05dd9465a37c47e770403c3c916306c8be005a5 Mon Sep 17 00:00:00 2001
From: Hrishikesh Hiraskar <hrishi@enterpret.com>
Date: Sun, 8 Oct 2023 16:05:23 +0530
Subject: [PATCH] =?UTF-8?q?=F0=9F=A9=B9=20Fix:=20race=20condition=20in=20m?=
 =?UTF-8?q?emory=20storage?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This commit updates `Conn` method of memory storage to fix race conditions.

Fixes: #2400
---
 internal/storage/memory/memory.go      | 10 ++++----
 internal/storage/memory/memory_test.go | 34 ++++++++++++++++++++++++--
 2 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/internal/storage/memory/memory.go b/internal/storage/memory/memory.go
index b3f70e6c09..5ef9d2e8a1 100644
--- a/internal/storage/memory/memory.go
+++ b/internal/storage/memory/memory.go
@@ -137,9 +137,9 @@ func (s *Storage) gc() {
 	}
 }
 
-// Return database client
-func (s *Storage) Conn() map[string]entry {
-	s.mux.RLock()
-	defer s.mux.RUnlock()
-	return s.db
+// ConnLocked can be used to access the underlying db in a thread-safe manner.
+func (s *Storage) ConnLocked(f func(map[string]entry)) {
+	s.mux.Lock()
+	defer s.mux.Unlock()
+	f(s.db)
 }
diff --git a/internal/storage/memory/memory_test.go b/internal/storage/memory/memory_test.go
index 72692425cd..aa5841b781 100644
--- a/internal/storage/memory/memory_test.go
+++ b/internal/storage/memory/memory_test.go
@@ -124,9 +124,39 @@ func Test_Storage_Memory_Close(t *testing.T) {
 	utils.AssertEqual(t, nil, testStore.Close())
 }
 
-func Test_Storage_Memory_Conn(t *testing.T) {
+func Test_Storage_Memory_ConnLocked(t *testing.T) {
 	t.Parallel()
-	utils.AssertEqual(t, true, testStore.Conn() != nil)
+	assertConn := func(db map[string]entry) {
+		utils.AssertEqual(t, true, db != nil)
+	}
+	testStore.ConnLocked(assertConn)
+}
+
+func Test_Storage_Memory_ConnLocked_Race(t *testing.T) {
+	key := "test_race"
+	done := make(chan bool)
+	read := func(db map[string]entry) {
+		_ = db[key]
+	}
+	write := func(db map[string]entry) {
+		db[key] = entry{}
+	}
+	doUntilDone := func(fn func(map[string]entry)) {
+		for {
+			select {
+			case <-done:
+				return
+			default:
+				testStore.ConnLocked(fn)
+			}
+		}
+	}
+
+	go doUntilDone(read)
+	go doUntilDone(write)
+
+	time.Sleep(100 * time.Millisecond)
+	done <- true
 }
 
 // go test -v -run=^$ -bench=Benchmark_Storage_Memory -benchmem -count=4