diff --git a/core/collection/safemap.go b/core/collection/safemap.go index 7be1e4af..4856f147 100644 --- a/core/collection/safemap.go +++ b/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() -} diff --git a/core/collection/safemap_test.go b/core/collection/safemap_test.go index 4e9ca1d7..24f42152 100644 --- a/core/collection/safemap_test.go +++ b/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) }