use thread safe map

2 files changed, 149 insertions, 0 deletions

diff --git a/pkg/utils/thread_safe_map.go b/pkg/utils/thread_safe_map.go
new file mode 100644
index 000000000..a70cefc7d
--- /dev/null
+++ b/pkg/utils/thread_safe_map.go
@@ -0,0 +1,90 @@
+package utils
+
+import "sync"
+
+type ThreadSafeMap[K comparable, V any] struct {
+	mutex sync.RWMutex
+
+	innerMap map[K]V
+}
+
+func NewThreadSafeMap[K comparable, V any]() *ThreadSafeMap[K, V] {
+	return &ThreadSafeMap[K, V]{
+		innerMap: make(map[K]V),
+	}
+}
+
+func (m *ThreadSafeMap[K, V]) Get(key K) (V, bool) {
+	m.mutex.RLock()
+	defer m.mutex.RUnlock()
+
+	value, ok := m.innerMap[key]
+	return value, ok
+}
+
+func (m *ThreadSafeMap[K, V]) Set(key K, value V) {
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+
+	m.innerMap[key] = value
+}
+
+func (m *ThreadSafeMap[K, V]) Delete(key K) {
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+
+	delete(m.innerMap, key)
+}
+
+func (m *ThreadSafeMap[K, V]) Keys() []K {
+	m.mutex.RLock()
+	defer m.mutex.RUnlock()
+
+	keys := make([]K, 0, len(m.innerMap))
+	for key := range m.innerMap {
+		keys = append(keys, key)
+	}
+
+	return keys
+}
+
+func (m *ThreadSafeMap[K, V]) Values() []V {
+	m.mutex.RLock()
+	defer m.mutex.RUnlock()
+
+	values := make([]V, 0, len(m.innerMap))
+	for _, value := range m.innerMap {
+		values = append(values, value)
+	}
+
+	return values
+}
+
+func (m *ThreadSafeMap[K, V]) Len() int {
+	m.mutex.RLock()
+	defer m.mutex.RUnlock()
+
+	return len(m.innerMap)
+}
+
+func (m *ThreadSafeMap[K, V]) Clear() {
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+
+	m.innerMap = make(map[K]V)
+}
+
+func (m *ThreadSafeMap[K, V]) IsEmpty() bool {
+	m.mutex.RLock()
+	defer m.mutex.RUnlock()
+
+	return len(m.innerMap) == 0
+}
+
+func (m *ThreadSafeMap[K, V]) Has(key K) bool {
+	m.mutex.RLock()
+	defer m.mutex.RUnlock()
+
+	_, ok := m.innerMap[key]
+	return ok
+}
diff --git a/pkg/utils/thread_safe_map_test.go b/pkg/utils/thread_safe_map_test.go
new file mode 100644
index 000000000..9676cfe5f
--- /dev/null
+++ b/pkg/utils/thread_safe_map_test.go
@@ -0,0 +1,59 @@
+package utils
+
+import (
+	"testing"
+)
+
+func TestThreadSafeMap(t *testing.T) {
+	m := NewThreadSafeMap[int, int]()
+
+	m.Set(1, 1)
+	m.Set(2, 2)
+	m.Set(3, 3)
+
+	if m.Len() != 3 {
+		t.Errorf("Expected length to be 3, got %d", m.Len())
+	}
+
+	if !m.Has(1) {
+		t.Errorf("Expected to have key 1")
+	}
+
+	if m.Has(4) {
+		t.Errorf("Expected to not have key 4")
+	}
+
+	if _, ok := m.Get(1); !ok {
+		t.Errorf("Expected to have key 1")
+	}
+
+	if _, ok := m.Get(4); ok {
+		t.Errorf("Expected to not have key 4")
+	}
+
+	m.Delete(1)
+
+	if m.Has(1) {
+		t.Errorf("Expected to not have key 1")
+	}
+
+	m.Clear()
+
+	if m.Len() != 0 {
+		t.Errorf("Expected length to be 0, got %d", m.Len())
+	}
+}
+
+func TestThreadSafeMapConcurrentReadWrite(t *testing.T) {
+	m := NewThreadSafeMap[int, int]()
+
+	go func() {
+		for i := 0; i < 10000; i++ {
+			m.Set(0, 0)
+		}
+	}()
+
+	for i := 0; i < 10000; i++ {
+		m.Get(0)
+	}
+}