Merge pull request #1370 from ydb-platform/xsync

internal/xsync improvements

Author: asmyasnikov

internal/xsync/map.go

@@ -3,13 +3,15 @@ package xsync
 import (
 	"fmt"
 	"sync"
+	"sync/atomic"
 )
 
 type Map[K comparable, V any] struct {
-	m sync.Map
+	m    sync.Map
+	size atomic.Int32
 }
 
-func (m *Map[K, V]) Load(key K) (value V, ok bool) {
+func (m *Map[K, V]) Get(key K) (value V, ok bool) {
 	v, ok := m.m.Load(key)
 	if !ok {
 		return value, false
@@ -38,22 +40,36 @@ func (m *Map[K, V]) Has(key K) bool {
 	return ok
 }
 
-func (m *Map[K, V]) Store(key K, value V) {
-	m.m.Store(key, value)
+func (m *Map[K, V]) Set(key K, value V) {
+	_, exists := m.m.LoadOrStore(key, value)
+
+	if !exists {
+		m.size.Add(1)
+	}
 }
 
-func (m *Map[K, V]) Delete(key K) (ok bool) {
-	_, ok = m.LoadAndDelete(key)
+func (m *Map[K, V]) Delete(key K) bool {
+	_, exists := m.Extract(key)
 
-	return ok
+	if !exists {
+		m.size.Add(-1)
+	}
+
+	return exists
 }
 
-func (m *Map[K, V]) LoadAndDelete(key K) (value V, ok bool) {
-	v, ok := m.m.LoadAndDelete(key)
-	if !ok {
+func (m *Map[K, V]) Size() int {
+	return int(m.size.Load())
+}
+
+func (m *Map[K, V]) Extract(key K) (value V, ok bool) {
+	v, exists := m.m.LoadAndDelete(key)
+	if !exists {
 		return value, false
 	}
 
+	m.size.Add(-1)
+
 	value, ok = v.(V)
 
 	return value, ok
@@ -65,10 +81,16 @@ func (m *Map[K, V]) Range(f func(key K, value V) bool) {
 	})
 }
 
-func (m *Map[K, V]) Clear() {
+func (m *Map[K, V]) Clear() (removed int) {
 	m.m.Range(func(k, v any) bool {
+		removed++
+
 		m.m.Delete(k)
 
 		return true
 	})
+
+	m.size.Add(int32(-removed))
+
+	return removed
 }

internal/xsync/map_test.go

@@ -8,9 +8,9 @@ import (
 
 func TestMap(t *testing.T) {
 	var m Map[int, string]
-	v, ok := m.Load(1)
+	v, ok := m.Get(1)
 	require.False(t, ok)
-	m.Store(1, "one")
+	m.Set(1, "one")
 	require.NotPanics(t, func() {
 		v = m.Must(1)
 		require.Equal(t, "one", v)
@@ -23,25 +23,26 @@ func TestMap(t *testing.T) {
 		v = m.Must(2)
 	})
 	m.m.Delete(2)
-	v, ok = m.LoadAndDelete(2)
+	v, ok = m.Extract(2)
 	require.False(t, ok)
 	require.Equal(t, "", v)
-	m.Store(2, "two")
-	v, ok = m.Load(2)
+	m.Set(2, "two")
+	v, ok = m.Get(2)
 	require.True(t, ok)
 	require.Equal(t, "two", v)
-	v, ok = m.LoadAndDelete(1)
+	v, ok = m.Extract(1)
 	require.True(t, ok)
 	require.Equal(t, "one", v)
 	require.False(t, m.Has(1))
-	m.Store(3, "three")
-	v, ok = m.Load(3)
+	m.Set(3, "three")
+	v, ok = m.Get(3)
 	require.True(t, ok)
 	require.Equal(t, "three", v)
 	exp := map[int]string{
 		2: "two",
 		3: "three",
 	}
+	require.Equal(t, 2, m.Size())
 	var unexp map[int]string
 	m.Range(func(key int, value string) bool {
 		if v, ok := exp[key]; ok && v == value {
@@ -55,6 +56,7 @@ func TestMap(t *testing.T) {
 	require.Empty(t, exp)
 	require.Empty(t, unexp)
 	m.Clear()
+	require.Zero(t, m.Size())
 	empty := true
 	m.Range(func(key int, value string) bool {
 		empty = false

internal/xsync/set.go

@@ -2,38 +2,53 @@ package xsync
 
 import (
 	"sync"
+	"sync/atomic"
 )
 
 type Set[K comparable] struct {
-	m sync.Map
+	m    sync.Map
+	size atomic.Int32
 }
 
-func (m *Set[K]) Has(key K) bool {
-	_, ok := m.m.Load(key)
+func (s *Set[K]) Has(key K) bool {
+	_, ok := s.m.Load(key)
 
 	return ok
 }
 
-func (m *Set[K]) Add(key K) bool {
-	_, loaded := m.m.LoadOrStore(key, struct{}{})
+func (s *Set[K]) Add(key K) bool {
+	_, exists := s.m.LoadOrStore(key, struct{}{})
 
-	return !loaded
+	if !exists {
+		s.size.Add(1)
+	}
+
+	return !exists
 }
 
-func (m *Set[K]) Remove(key K) (ok bool) {
-	_, ok = m.m.LoadAndDelete(key)
+func (s *Set[K]) Size() int {
+	return int(s.size.Load())
+}
 
-	return ok
+func (s *Set[K]) Remove(key K) bool {
+	_, exists := s.m.LoadAndDelete(key)
+
+	if exists {
+		s.size.Add(-1)
+	}
+
+	return exists
 }
 
-func (m *Set[K]) Clear() (removed int) {
-	m.m.Range(func(k, v any) bool {
+func (s *Set[K]) Clear() (removed int) {
+	s.m.Range(func(k, v any) bool {
 		removed++
 
-		m.m.Delete(k)
+		s.m.Delete(k)
 
 		return true
 	})
+	s.size.Add(int32(-removed))
 
 	return removed
 }

internal/xsync/set_test.go

@@ -15,5 +15,7 @@ func TestSet(t *testing.T) {
 	require.True(t, s.Add(2))
 	require.True(t, s.Remove(1))
 	require.False(t, s.Has(1))
+	require.Equal(t, 1, s.Size())
 	require.Equal(t, 1, s.Clear())
+	require.Zero(t, s.Size())
 }