新增安全map功能

util/example_test.go

@@ -2,91 +2,34 @@ package util_test
 
 import (
 	"fmt"
-	"github.com/name5566/leaf/util"
+	"testing"
+
+	"github.com/duanhf2012/origin/util"
 )
 
-func ExampleMap() {
+func ExampleMap(t *testing.T) {
-	m := new(util.Map)
+	maps := util.NewMapEx()
-
+	for i := 0; i < 10000; i++ {
-	fmt.Println(m.Get("key"))
+		maps.Set(i, i)
-	m.Set("key", "value")
-	fmt.Println(m.Get("key"))
-	m.Del("key")
-	fmt.Println(m.Get("key"))
-
-	m.Set(1, "1")
-	m.Set(2, 2)
-	m.Set("3", 3)
-
-	fmt.Println(m.Len())
-
-	// Output:
-	// <nil>
-	// value
-	// <nil>
-	// 3
-}
-
-func ExampleRandGroup() {
-	i := util.RandGroup(0, 0, 50, 50)
-	switch i {
-	case 2, 3:
-		fmt.Println("ok")
 	}
 
-	// Output:
+	for i := 0; i < 10000; i++ {
-	// ok
+		ret := maps.Get(i)
-}
+		if ret.(int) != i {
-
+			fmt.Printf("cannot find i:%d\n", i)
-func ExampleRandInterval() {
+		}
-	v := util.RandInterval(-1, 1)
-	switch v {
-	case -1, 0, 1:
-		fmt.Println("ok")
 	}
 
-	// Output:
+	for i := 0; i < 10000; i++ {
-	// ok
+		maps.LockSet(i, func(key interface{}, val interface{}) interface{} {
-}
+			return val.(int) + 1
-
+		})
-func ExampleRandIntervalN() {
-	r := util.RandIntervalN(-1, 0, 2)
-	if r[0] == -1 && r[1] == 0 ||
-		r[0] == 0 && r[1] == -1 {
-		fmt.Println("ok")
 	}
 
-	// Output:
+	for i := 0; i < 10000; i++ {
-	// ok
+		ret := maps.Get(i)
-}
+		if ret.(int) != i {
-
+			fmt.Printf("cannot find i:%d\n", i)
-func ExampleDeepCopy() {
+		}
-	src := []int{1, 2, 3}
-
-	var dst []int
-	util.DeepCopy(&dst, &src)
-
-	for _, v := range dst {
-		fmt.Println(v)
 	}
-
-	// Output:
-	// 1
-	// 2
-	// 3
-}
-
-func ExampleDeepClone() {
-	src := []int{1, 2, 3}
-
-	dst := util.DeepClone(src).([]int)
-
-	for _, v := range dst {
-		fmt.Println(v)
-	}
-
-	// Output:
-	// 1
-	// 2
-	// 3
 }

util/mapex.go

@@ -2,92 +2,202 @@ package util
 
 import (
 	"fmt"
+	"sync"
 
 	"github.com/duanhf2012/origin/util/hash"
 )
 
 const (
-	DEFAULT_MAX_HASH_NUM = 100
+	DEFAULT_SAFE_MAP_MAX_HASH_NUM = 10
 )
 
 type MapEx struct {
-	m          []Map
+	sync.RWMutex
-	hashMapNum uint
+	m          []map[interface{}]interface{}
+	l          []sync.RWMutex
+	hashMapNum int
+}
+
+func (m *MapEx) Init(hashMapNum int) {
+	m.hashMapNum = hashMapNum
+
+	m.m = []map[interface{}]interface{}{}
+	m.l = []sync.RWMutex{}
+
+	for i := 0; i < hashMapNum; i++ {
+		m.m = append(m.m, make(map[interface{}]interface{}))
+		m.l = append(m.l, sync.RWMutex{})
+	}
 }
 
 func NewMapEx() *MapEx {
 	mapEx := MapEx{}
-	mapEx.Init(DEFAULT_MAX_HASH_NUM)
+	mapEx.Init(DEFAULT_SAFE_MAP_MAX_HASH_NUM)
 	return &mapEx
 }
 
-func (m *MapEx) Init(hashMapNum uint) {
-	var i uint
-	for i = 0; i < hashMapNum; i++ {
-		m.m = append(m.m, Map{})
-	}
-
-	m.hashMapNum = hashMapNum
-}
-
 func (m *MapEx) ClearMap() {
-	var i uint
+	for i := 0; i < DEFAULT_SAFE_MAP_MAX_HASH_NUM; i++ {
-	for i = 0; i < m.hashMapNum; i++ {
+		m.l[i].Lock()
-		m.m[i].ClearMap()
+		m.m[i] = map[interface{}]interface{}{}
+		m.l[i].Unlock()
 	}
 }
 
-func (m *MapEx) GetHashCode(key interface{}) uint {
+func (m *MapEx) GetHashCode(key interface{}) int {
-	return hash.HashNumber(fmt.Sprint(key))
+	return int(hash.HashNumber(fmt.Sprint(key)))
 }
 
-func (m *MapEx) GetMapByKey(key interface{}) *Map {
+func (m *MapEx) GetArrayIdByKey(key interface{}) int {
 	idx := m.GetHashCode(key) % m.hashMapNum
+	if idx > m.hashMapNum {
+		return -1
+	}
+
+	return idx
+}
+
+func (m *MapEx) GetMapByKey(key interface{}) map[interface{}]interface{} {
+	idx := m.GetArrayIdByKey(key)
 	if idx < 0 || idx > m.hashMapNum {
 		return nil
 	}
 
-	return &m.m[idx]
+	return m.m[idx]
 }
 
-func (m *MapEx) Get(key interface{}) interface{} {
+func (m *MapEx) UnsafeGet(key interface{}) interface{} {
+
 	mapData := m.GetMapByKey(key)
 	if mapData == nil {
 		return nil
 	}
-	return mapData.Get(key)
+
+	val, ok := mapData[key]
+	if ok == false {
+		return nil
+	}
+
+	return val
+}
+
+func (m *MapEx) Get(key interface{}) interface{} {
+	idx := m.GetArrayIdByKey(key)
+	if idx < 0 || idx > m.hashMapNum {
+		return nil
+	}
+
+	m.l[idx].RLock()
+	defer m.l[idx].RUnlock()
+
+	val := m.m[idx]
+	ret, ok := val[key]
+	if ok == false {
+		return nil
+	}
+	return ret
 }
 
 func (m *MapEx) Set(key interface{}, value interface{}) {
+	idx := m.GetArrayIdByKey(key)
+	if idx < 0 || idx > m.hashMapNum {
+		return
+	}
+
+	m.l[idx].Lock()
+	defer m.l[idx].Unlock()
+
+	val := m.m[idx]
+	val[key] = value
+}
+
+func (m *MapEx) UnsafeDel(key interface{}) {
 	mapData := m.GetMapByKey(key)
 	if mapData == nil {
 		return
 	}
 
-	mapData.Set(key, value)
+	delete(mapData, key)
 }
 
 func (m *MapEx) Del(key interface{}) {
-
+	idx := m.GetArrayIdByKey(key)
-	mapData := m.GetMapByKey(key)
+	if idx < 0 || idx > m.hashMapNum {
-	if mapData == nil {
 		return
 	}
 
-	mapData.Del(key)
+	m.l[idx].Lock()
+	defer m.l[idx].Unlock()
+
+	val := m.m[idx]
+	delete(val, key)
+}
+
+func (m *MapEx) Len() int {
+	lens := 0
+	for i := 0; i < m.hashMapNum; i++ {
+		m.l[i].RLock()
+		lens += len(m.m[i])
+		m.l[i].RUnlock()
+	}
+
+	return lens
 }
 
 func (m *MapEx) RLockRange(f func(key interface{}, value interface{})) {
-	var i uint
+	for i := 0; i < m.hashMapNum; i++ {
-	for i = 0; i < m.hashMapNum; i++ {
+		m.l[i].RLock()
-		m.m[i].RLockRange(f)
+		for key, val := range m.m[i] {
+			f(key, val)
+		}
+		m.l[i].RUnlock()
 	}
-
 }
 
 func (m *MapEx) LockRange(f func(key interface{}, value interface{})) {
-	var i uint
+	for i := 0; i < m.hashMapNum; i++ {
-	for i = 0; i < m.hashMapNum; i++ {
+		m.l[i].Lock()
-		m.m[i].LockRange(f)
+		for key, val := range m.m[i] {
+			f(key, val)
+		}
+		m.l[i].Unlock()
 	}
 }
+
+func (m *MapEx) LockGet(key interface{}, f func(value interface{})) {
+	idx := m.GetArrayIdByKey(key)
+	if idx < 0 || idx > m.hashMapNum {
+		f(nil)
+		return
+	}
+
+	m.l[idx].Lock()
+	val := m.m[idx]
+	ret, ok := val[key]
+	if ok == false {
+		f(nil)
+	} else {
+		f(ret)
+	}
+	m.l[idx].Unlock()
+}
+
+func (m *MapEx) LockSet(key interface{}, f func(value interface{}) interface{}) {
+	idx := m.GetArrayIdByKey(key)
+	if idx < 0 || idx > m.hashMapNum {
+		f(nil)
+		return
+	}
+
+	m.l[idx].Lock()
+	val := m.m[idx]
+	ret, ok := val[key]
+
+	if ok == false {
+		val[key] = f(nil)
+	} else {
+		val[key] = f(ret)
+	}
+
+	m.l[idx].Unlock()
+}