diff --git a/common/utils/collection.go b/common/utils/collection.go index 5608dcc8f..2ad5db38c 100644 --- a/common/utils/collection.go +++ b/common/utils/collection.go @@ -95,10 +95,14 @@ func (set *SyncSet[K]) ToSlice() []K { } func (set *SyncSet[K]) Range(fn func(val K)) { + snapshot := map[K]struct{}{} set.lock.RLock() - defer set.lock.RUnlock() - for k := range set.container { + snapshot[k] = struct{}{} + } + set.lock.RUnlock() + + for k := range snapshot { fn(k) } } @@ -188,9 +192,14 @@ func (s *SegmentMap[K, V]) Range(f func(k K, v V)) { lock := s.locks[i] solt := s.solts[i] func() { + snapshot := map[K]V{} lock.RLock() - defer lock.RUnlock() for k, v := range solt { + snapshot[k] = v + } + lock.RUnlock() + + for k, v := range snapshot { f(k, v) } }() @@ -268,10 +277,14 @@ func (s *SyncMap[K, V]) Store(key K, val V) { // Range func (s *SyncMap[K, V]) Range(f func(key K, val V) bool) { + snapshot := map[K]V{} s.lock.RLock() - defer s.lock.RUnlock() - for k, v := range s.m { + snapshot[k] = v + } + s.lock.RUnlock() + + for k, v := range snapshot { _ = f(k, v) } } diff --git a/common/utils/collection_test.go b/common/utils/collection_test.go index ef817500b..07058767e 100644 --- a/common/utils/collection_test.go +++ b/common/utils/collection_test.go @@ -92,4 +92,40 @@ func Test_SyncSegmentMap(t *testing.T) { break } } + + for { + count := 0 + segmentMap.Range(func(k string, v *SegmentMap[string, string]) { + v.Range(func(k string, _ string) { + v.Del(k) + }) + count++ + }) + if count == total { + break + } + } +} + +func Test_SyncMap(t *testing.T) { + syncMap := NewSyncMap[int, int]() + + for i := 0; i < 10; i++ { + syncMap.Store(i, i) + } + + assert.Equal(t, 10, syncMap.Len()) + + syncMap.Range(func(key, val int) bool { + syncMap.Delete(key) + syncMap.Store(key+100, key+100) + return true + }) + + assert.Equal(t, 10, syncMap.Len()) + + for i := 0; i < 10; i++ { + _, ok := syncMap.Load(i) + assert.False(t, ok) + } }