Use RWMutex in NSMap and reduce lock area

djdongjin · djdongjin · commit 9e09bfb59093 · 2023-06-14T17:50:54.000Z
Signed-off-by: Jin Dong &lt;djdongjin95@gmail.com&gt;
diff --git a/runtime/nsmap.go b/runtime/nsmap.go
@@ -31,7 +31,7 @@ type object interface {
 
 // NSMap extends Map type with a notion of namespaces passed via Context.
 type NSMap[T object] struct {
-	mu      sync.Mutex
+	mu      sync.RWMutex
 	objects map[string]map[string]T
 }
 
@@ -44,13 +44,14 @@ func NewNSMap[T object]() *NSMap[T] {
 
 // Get a task
 func (m *NSMap[T]) Get(ctx context.Context, id string) (T, error) {
-	m.mu.Lock()
-	defer m.mu.Unlock()
 	namespace, err := namespaces.NamespaceRequired(ctx)
 	var t T
 	if err != nil {
 		return t, err
 	}
+
+	m.mu.RLock()
+	defer m.mu.RUnlock()
 	tasks, ok := m.objects[namespace]
 	if !ok {
 		return t, errdefs.ErrNotFound
@@ -64,8 +65,8 @@ func (m *NSMap[T]) Get(ctx context.Context, id string) (T, error) {
 
 // GetAll objects under a namespace
 func (m *NSMap[T]) GetAll(ctx context.Context, noNS bool) ([]T, error) {
-	m.mu.Lock()
-	defer m.mu.Unlock()
+	m.mu.RLock()
+	defer m.mu.RUnlock()
 	var o []T
 	if noNS {
 		for ns := range m.objects {
@@ -100,10 +101,10 @@ func (m *NSMap[T]) Add(ctx context.Context, t T) error {
 
 // AddWithNamespace adds a task with the provided namespace
 func (m *NSMap[T]) AddWithNamespace(namespace string, t T) error {
+	id := t.ID()
+
 	m.mu.Lock()
 	defer m.mu.Unlock()
-
-	id := t.ID()
 	if _, ok := m.objects[namespace]; !ok {
 		m.objects[namespace] = make(map[string]T)
 	}
@@ -116,21 +117,22 @@ func (m *NSMap[T]) AddWithNamespace(namespace string, t T) error {
 
 // Delete a task
 func (m *NSMap[T]) Delete(ctx context.Context, id string) {
-	m.mu.Lock()
-	defer m.mu.Unlock()
 	namespace, err := namespaces.NamespaceRequired(ctx)
 	if err != nil {
 		return
 	}
+
+	m.mu.Lock()
+	defer m.mu.Unlock()
 	tasks, ok := m.objects[namespace]
 	if ok {
 		delete(tasks, id)
 	}
 }
 
 func (m *NSMap[T]) IsEmpty() bool {
-	m.mu.Lock()
-	defer m.mu.Unlock()
+	m.mu.RLock()
+	defer m.mu.RUnlock()
 
 	for ns := range m.objects {
 		if len(m.objects[ns]) > 0 {
