From e9fef500b70355d9e3ed53079a230a8849891e29 Mon Sep 17 00:00:00 2001
From: sunnyyssh <blit228@mail.ru>
Date: Wed, 25 Dec 2024 13:16:20 +0300
Subject: [PATCH] =?UTF-8?q?=F0=9F=A9=B9=20Fix:=20Add=20key=20removal=20in?=
 =?UTF-8?q?=20MemoryLock?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 middleware/idempotency/locker.go | 32 +++++++++++++++++++++++---------
 1 file changed, 23 insertions(+), 9 deletions(-)

diff --git a/middleware/idempotency/locker.go b/middleware/idempotency/locker.go
index 2c3348b8f3..c3417dd6ca 100644
--- a/middleware/idempotency/locker.go
+++ b/middleware/idempotency/locker.go
@@ -10,42 +10,56 @@ type Locker interface {
 	Unlock(key string) error
 }
 
+type countedLock struct {
+	mu     sync.Mutex
+	locked int
+}
+
 type MemoryLock struct {
-	keys map[string]*sync.Mutex
+	keys map[string]*countedLock
 	mu   sync.Mutex
 }
 
 func (l *MemoryLock) Lock(key string) error {
 	l.mu.Lock()
-	mu, ok := l.keys[key]
+	lock, ok := l.keys[key]
 	if !ok {
-		mu = new(sync.Mutex)
-		l.keys[key] = mu
+		lock = new(countedLock)
+		l.keys[key] = lock
 	}
+	lock.locked++
 	l.mu.Unlock()
 
-	mu.Lock()
+	lock.mu.Lock()
 
 	return nil
 }
 
 func (l *MemoryLock) Unlock(key string) error {
 	l.mu.Lock()
-	mu, ok := l.keys[key]
-	l.mu.Unlock()
+	lock, ok := l.keys[key]
 	if !ok {
 		// This happens if we try to unlock an unknown key
+		l.mu.Unlock()
 		return nil
 	}
 
-	mu.Unlock()
+	lock.locked--
+	if lock.locked <= 0 {
+		// This happens if countedLock is used to Lock and Unlock the same number of times
+		// So, we can delete the key to prevent memory leak
+		delete(l.keys, key)
+	}
+	l.mu.Unlock()
+
+	lock.mu.Unlock()
 
 	return nil
 }
 
 func NewMemoryLock() *MemoryLock {
 	return &MemoryLock{
-		keys: make(map[string]*sync.Mutex),
+		keys: make(map[string]*countedLock),
 	}
 }