refactor: sequential range over safemap (#2316)

core/collection/safemap.go

@@ -68,6 +68,24 @@ func (m *SafeMap) Get(key interface{}) (interface{}, bool) {
 	return val, ok
 }
 
+// Range calls f sequentially for each key and value present in the map.
+// If f returns false, range stops the iteration.
+func (m *SafeMap) Range(f func(key, val interface{}) bool) {
+	m.lock.RLock()
+	defer m.lock.RUnlock()
+
+	for k, v := range m.dirtyOld {
+		if !f(k, v) {
+			return
+		}
+	}
+	for k, v := range m.dirtyNew {
+		if !f(k, v) {
+			return
+		}
+	}
+}
+
 // Set sets the value into m with the given key.
 func (m *SafeMap) Set(key, value interface{}) {
 	m.lock.Lock()
@@ -94,31 +112,3 @@ func (m *SafeMap) Size() int {
 	m.lock.RUnlock()
 	return size
 }
-
-// Range calls f sequentially for each key and value present in the map.
-// If f returns false, range stops the iteration.
-// Range m .
-func (m *SafeMap) Range(f func(key, value interface{}) bool) {
-	m.lock.RLock()
-	defer m.lock.RUnlock()
-	var wg sync.WaitGroup
-
-	wg.Add(2)
-	go func() {
-		for k, v := range m.dirtyOld {
-			if !f(k, v) {
-				break
-			}
-		}
-		wg.Done()
-	}()
-	go func() {
-		for k, v := range m.dirtyNew {
-			if !f(k, v) {
-				break
-			}
-		}
-		wg.Done()
-	}()
-	wg.Wait()
-}

core/collection/safemap_test.go

@@ -1,7 +1,7 @@
 package collection
 
 import (
-	"go.uber.org/atomic"
+	"sync/atomic"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -115,8 +115,9 @@ func TestSafeMap_Range(t *testing.T) {
 		exception1 = 5
 		exception2 = 500
 	)
+
 	m := NewSafeMap()
-	m_new := NewSafeMap()
+	newMap := NewSafeMap()
 
 	for i := 0; i < size; i++ {
 		m.Set(i, i)
@@ -136,14 +137,13 @@ func TestSafeMap_Range(t *testing.T) {
 		}
 	}
 
-	count := atomic.Int32{}
+	var count int32
 	m.Range(func(k, v interface{}) bool {
-		count.Add(1)
-		m_new.Set(k, v)
+		atomic.AddInt32(&count, 1)
+		newMap.Set(k, v)
 		return true
 	})
-	assert.Equal(t, int(count.Load()), m.Size())
-	assert.Equal(t, m.dirtyNew, m_new.dirtyNew)
-	assert.Equal(t, m.dirtyOld, m_new.dirtyOld)
-
+	assert.Equal(t, int(atomic.LoadInt32(&count)), m.Size())
+	assert.Equal(t, m.dirtyNew, newMap.dirtyNew)
+	assert.Equal(t, m.dirtyOld, newMap.dirtyOld)
 }