Fix race condition with map

map.go

@@ -2,19 +2,26 @@ package storage
 
 import (
 	"fmt"
+	"sync"
 )
 
-type Map map[string]map[string][]byte
+type Map struct {
+	m    map[string]map[string][]byte
+	lock *sync.RWMutex
+}
 
 func NewMap() *Map {
-	m := make(map[string]map[string][]byte)
-	n := Map(m)
-	return &n
+	return &Map{
+		m:    make(map[string]map[string][]byte),
+		lock: &sync.RWMutex{},
+	}
 }
 
 func (m *Map) String() string {
+	m.lock.RLock()
+	defer m.lock.RUnlock()
 	s := ""
-	for k, v := range *m {
+	for k, v := range m.m {
 		if s != "" {
 			s += ",\n"
 		}
@@ -24,20 +31,23 @@ func (m *Map) String() string {
 }
 
 func (m *Map) Close() error {
-	m = nil
+	m.lock.Lock()
+	defer m.lock.Unlock()
+	m.m = nil
 	return nil
 }
 
 func (m *Map) List(ns []string, limits ...string) ([]string, error) {
+	m.lock.RLock()
+	defer m.lock.RUnlock()
 	namespace := resolveNamespace(ns)
 	limits = resolveLimits(limits)
 
 	keys := []string{}
-	n, _ := (*m)[DefaultNamespace]
-	if v, ok := (*m)[namespace]; ok {
-		n = v
+	if _, ok := m.m[namespace]; !ok {
+		namespace = DefaultNamespace
 	}
-	for k := range n {
+	for k := range m.m[namespace] {
 		if k >= limits[0] && k <= limits[1] {
 			keys = append(keys, k)
 		}
@@ -47,29 +57,33 @@ func (m *Map) List(ns []string, limits ...string) ([]string, error) {
 }
 
 func (m *Map) Get(key string, ns ...string) ([]byte, error) {
+	m.lock.RLock()
+	defer m.lock.RUnlock()
 	namespace := resolveNamespace(ns)
-	if _, ok := (*m)[namespace]; !ok {
+	if _, ok := m.m[namespace]; !ok {
 		return nil, ErrNotFound
 	}
-	if _, ok := (*m)[namespace][key]; !ok {
+	if _, ok := m.m[namespace][key]; !ok {
 		return nil, ErrNotFound
 	}
-	return (*m)[namespace][key], nil
+	return m.m[namespace][key], nil
 }
 
 func (m *Map) Set(key string, value []byte, ns ...string) error {
+	m.lock.Lock()
+	defer m.lock.Unlock()
 	namespace := resolveNamespace(ns)
 	if value == nil {
-		if _, ok := (*m)[namespace]; !ok {
+		if _, ok := m.m[namespace]; !ok {
 			return nil
-		} else if _, ok := (*m)[namespace][key]; ok {
-			delete((*m)[namespace], key)
+		} else if _, ok := m.m[namespace][key]; ok {
+			delete(m.m[namespace], key)
 		}
 	} else {
-		if _, ok := (*m)[namespace]; !ok {
-			(*m)[namespace] = make(map[string][]byte)
+		if _, ok := m.m[namespace]; !ok {
+			m.m[namespace] = make(map[string][]byte)
 		}
-		(*m)[namespace][key] = value
+		m.m[namespace][key] = value
 	}
 	return nil
 }