From 68df0c36202b52af9a260a83a5bd342de2de2320 Mon Sep 17 00:00:00 2001
From: Kevin Wan <wanjunfeng@gmail.com>
Date: Mon, 18 Sep 2023 11:01:46 +0800
Subject: [PATCH] chore: add more tests (#3575)

---
 core/collection/safemap.go      |  6 ++--
 core/collection/safemap_test.go | 62 +++++++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+), 2 deletions(-)

diff --git a/core/collection/safemap.go b/core/collection/safemap.go
index 6ab328a0..13c994a0 100644
--- a/core/collection/safemap.go
+++ b/core/collection/safemap.go
@@ -29,6 +29,8 @@ func NewSafeMap() *SafeMap {
 // Del deletes the value with the given key from m.
 func (m *SafeMap) Del(key any) {
 	m.lock.Lock()
+	defer m.lock.Unlock()
+
 	if _, ok := m.dirtyOld[key]; ok {
 		delete(m.dirtyOld, key)
 		m.deletionOld++
@@ -52,7 +54,6 @@ func (m *SafeMap) Del(key any) {
 		m.dirtyNew = make(map[any]any)
 		m.deletionNew = 0
 	}
-	m.lock.Unlock()
 }
 
 // Get gets the value with the given key from m.
@@ -89,6 +90,8 @@ func (m *SafeMap) Range(f func(key, val any) bool) {
 // Set sets the value into m with the given key.
 func (m *SafeMap) Set(key, value any) {
 	m.lock.Lock()
+	defer m.lock.Unlock()
+
 	if m.deletionOld <= maxDeletion {
 		if _, ok := m.dirtyNew[key]; ok {
 			delete(m.dirtyNew, key)
@@ -102,7 +105,6 @@ func (m *SafeMap) Set(key, value any) {
 		}
 		m.dirtyNew[key] = value
 	}
-	m.lock.Unlock()
 }
 
 // Size returns the size of m.
diff --git a/core/collection/safemap_test.go b/core/collection/safemap_test.go
index 26c734ee..dd54881d 100644
--- a/core/collection/safemap_test.go
+++ b/core/collection/safemap_test.go
@@ -147,3 +147,65 @@ func TestSafeMap_Range(t *testing.T) {
 	assert.Equal(t, m.dirtyNew, newMap.dirtyNew)
 	assert.Equal(t, m.dirtyOld, newMap.dirtyOld)
 }
+
+func TestSetManyTimes(t *testing.T) {
+	const iteration = maxDeletion * 2
+	m := NewSafeMap()
+	for i := 0; i < iteration; i++ {
+		m.Set(i, i)
+		if i%3 == 0 {
+			m.Del(i / 2)
+		}
+	}
+	var count int
+	m.Range(func(k, v any) bool {
+		count++
+		return count < maxDeletion/2
+	})
+	assert.Equal(t, maxDeletion/2, count)
+	for i := 0; i < iteration; i++ {
+		m.Set(i, i)
+		if i%3 == 0 {
+			m.Del(i / 2)
+		}
+	}
+	for i := 0; i < iteration; i++ {
+		m.Set(i, i)
+		if i%3 == 0 {
+			m.Del(i / 2)
+		}
+	}
+	for i := 0; i < iteration; i++ {
+		m.Set(i, i)
+		if i%3 == 0 {
+			m.Del(i / 2)
+		}
+	}
+
+	count = 0
+	m.Range(func(k, v any) bool {
+		count++
+		return count < maxDeletion
+	})
+	assert.Equal(t, maxDeletion, count)
+}
+
+func TestSetManyTimesNew(t *testing.T) {
+	m := NewSafeMap()
+	for i := 0; i < maxDeletion*3; i++ {
+		m.Set(i, i)
+	}
+	for i := 0; i < maxDeletion*2; i++ {
+		m.Del(i)
+	}
+	for i := 0; i < maxDeletion*3; i++ {
+		m.Set(i+maxDeletion*3, i+maxDeletion*3)
+	}
+	for i := 0; i < maxDeletion*2; i++ {
+		m.Del(i + maxDeletion*2)
+	}
+	for i := 0; i < maxDeletion-copyThreshold+1; i++ {
+		m.Del(i + maxDeletion*2)
+	}
+	assert.Equal(t, 0, len(m.dirtyNew))
+}