diff --git a/core/collection/safemap.go b/core/collection/safemap.go index 58e36d4e..7be1e4af 100644 --- a/core/collection/safemap.go +++ b/core/collection/safemap.go @@ -94,3 +94,31 @@ 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 46e64749..4e9ca1d7 100644 --- a/core/collection/safemap_test.go +++ b/core/collection/safemap_test.go @@ -1,6 +1,7 @@ package collection import ( + "go.uber.org/atomic" "testing" "github.com/stretchr/testify/assert" @@ -107,3 +108,42 @@ func testSafeMapWithParameters(t *testing.T, size, exception int) { } } } + +func TestSafeMap_Range(t *testing.T) { + const ( + size = 100000 + exception1 = 5 + exception2 = 500 + ) + m := NewSafeMap() + m_new := NewSafeMap() + + for i := 0; i < size; i++ { + m.Set(i, i) + } + for i := 0; i < size; i++ { + if i%exception1 == 0 { + m.Del(i) + } + } + + for i := size; i < size<<1; i++ { + m.Set(i, i) + } + for i := size; i < size<<1; i++ { + if i%exception2 != 0 { + m.Del(i) + } + } + + count := atomic.Int32{} + m.Range(func(k, v interface{}) bool { + count.Add(1) + m_new.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) + +}