优化缓存Hash查询速度

internal/caches/list_file_hash_map.go

@@ -9,18 +9,34 @@ import (
 	"sync"
 )
 
+const HashMapSharding = 11
+
+var bigIntPool = sync.Pool{New: func() any {
+	return big.NewInt(0)
+}}
+
 // FileListHashMap 文件Hash列表
 type FileListHashMap struct {
-	m map[uint64]zero.Zero
+	m []map[uint64]zero.Zero
+
+	lockers []*sync.RWMutex
 
-	locker      sync.RWMutex
 	isAvailable bool
 	isReady     bool
 }
 
 func NewFileListHashMap() *FileListHashMap {
+	var m = make([]map[uint64]zero.Zero, HashMapSharding)
+	var lockers = make([]*sync.RWMutex, HashMapSharding)
+
+	for i := 0; i < HashMapSharding; i++ {
+		m[i] = map[uint64]zero.Zero{}
+		lockers[i] = &sync.RWMutex{}
+	}
+
 	return &FileListHashMap{
-		m:           map[uint64]zero.Zero{},
+		m:           m,
+		lockers:     lockers,
 		isAvailable: false,
 		isReady:     false,
 	}
@@ -56,9 +72,11 @@ func (this *FileListHashMap) Add(hash string) {
 		return
 	}
 
-	this.locker.Lock()
-	this.m[this.bigInt(hash)] = zero.New()
-	this.locker.Unlock()
+	hashInt, index := this.bigInt(hash)
+
+	this.lockers[index].Lock()
+	this.m[index][hashInt] = zero.New()
+	this.lockers[index].Unlock()
 }
 
 func (this *FileListHashMap) AddHashes(hashes []string) {
@@ -66,11 +84,12 @@ func (this *FileListHashMap) AddHashes(hashes []string) {
 		return
 	}
 
-	this.locker.Lock()
 	for _, hash := range hashes {
-		this.m[this.bigInt(hash)] = zero.New()
+		hashInt, index := this.bigInt(hash)
+		this.lockers[index].Lock()
+		this.m[index][hashInt] = zero.New()
+		this.lockers[index].Unlock()
 	}
-	this.locker.Unlock()
 }
 
 func (this *FileListHashMap) Delete(hash string) {
@@ -78,9 +97,10 @@ func (this *FileListHashMap) Delete(hash string) {
 		return
 	}
 
-	this.locker.Lock()
-	delete(this.m, this.bigInt(hash))
-	this.locker.Unlock()
+	hashInt, index := this.bigInt(hash)
+	this.lockers[index].Lock()
+	delete(this.m[index], hashInt)
+	this.lockers[index].Unlock()
 }
 
 func (this *FileListHashMap) Exist(hash string) bool {
@@ -91,16 +111,25 @@ func (this *FileListHashMap) Exist(hash string) bool {
 		// 只有完全Ready时才能判断是否为false
 		return true
 	}
-	this.locker.RLock()
-	_, ok := this.m[this.bigInt(hash)]
-	this.locker.RUnlock()
+
+	hashInt, index := this.bigInt(hash)
+
+	this.lockers[index].RLock()
+	_, ok := this.m[index][hashInt]
+	this.lockers[index].RUnlock()
 	return ok
 }
 
 func (this *FileListHashMap) Clean() {
-	this.locker.Lock()
-	this.m = map[uint64]zero.Zero{}
-	this.locker.Unlock()
+	for i := 0; i < HashMapSharding; i++ {
+		this.lockers[i].Lock()
+	}
+
+	this.m = make([]map[uint64]zero.Zero, HashMapSharding)
+
+	for i := HashMapSharding - 1; i >= 0; i-- {
+		this.lockers[i].Unlock()
+	}
 }
 
 func (this *FileListHashMap) IsReady() bool {
@@ -108,17 +137,36 @@ func (this *FileListHashMap) IsReady() bool {
 }
 
 func (this *FileListHashMap) Len() int {
-	this.locker.Lock()
-	defer this.locker.Unlock()
-	return len(this.m)
+	for i := 0; i < HashMapSharding; i++ {
+		this.lockers[i].Lock()
+	}
+
+	var count = 0
+	for _, shard := range this.m {
+		count += len(shard)
+	}
+
+	for i := HashMapSharding - 1; i >= 0; i-- {
+		this.lockers[i].Unlock()
+	}
+
+	return count
 }
 
 func (this *FileListHashMap) SetIsAvailable(isAvailable bool) {
 	this.isAvailable = isAvailable
 }
 
-func (this *FileListHashMap) bigInt(hash string) uint64 {
-	var bigInt = big.NewInt(0)
-	bigInt.SetString(hash, 16)
-	return bigInt.Uint64()
+func (this *FileListHashMap) SetIsReady(isReady bool) {
+	this.isReady = isReady
+}
+
+func (this *FileListHashMap) bigInt(hash string) (hashInt uint64, index int) {
+	var bigInt = bigIntPool.Get().(*big.Int)
+	bigInt.SetString(hash, 16)
+	hashInt = bigInt.Uint64()
+	bigIntPool.Put(bigInt)
+
+	index = int(hashInt % HashMapSharding)
+	return
 }

internal/caches/list_file_hash_map_test.go

@@ -6,6 +6,7 @@ import (
 	"github.com/TeaOSLab/EdgeNode/internal/caches"
 	"github.com/TeaOSLab/EdgeNode/internal/zero"
 	"github.com/iwind/TeaGo/Tea"
+	"github.com/iwind/TeaGo/rands"
 	"github.com/iwind/TeaGo/types"
 	stringutil "github.com/iwind/TeaGo/utils/string"
 	"math/big"
@@ -52,12 +53,34 @@ func TestFileListHashMap_Memory2(t *testing.T) {
 }
 
 func TestFileListHashMap_BigInt(t *testing.T) {
+	var bigInt = big.NewInt(0)
+
 	for _, s := range []string{"1", "2", "3", "123", "123456"} {
 		var hash = stringutil.Md5(s)
 
-		var bigInt = big.NewInt(0)
+		var bigInt1 = big.NewInt(0)
+		bigInt1.SetString(hash, 16)
+
 		bigInt.SetString(hash, 16)
-		t.Log(s, "=>", bigInt.Uint64(), "hash:", hash, "format:", strconv.FormatUint(bigInt.Uint64(), 16))
+
+		t.Log(s, "=>", bigInt1.Uint64(), "hash:", hash, "format:", strconv.FormatUint(bigInt1.Uint64(), 16), strconv.FormatUint(bigInt.Uint64(), 16))
+
+		if strconv.FormatUint(bigInt1.Uint64(), 16) != strconv.FormatUint(bigInt.Uint64(), 16) {
+			t.Fatal("not equal")
+		}
+	}
+
+	for i := 0; i < 1_000_000; i++ {
+		var hash = stringutil.Md5(types.String(i))
+
+		var bigInt1 = big.NewInt(0)
+		bigInt1.SetString(hash, 16)
+
+		bigInt.SetString(hash, 16)
+
+		if bigInt1.Uint64() != bigInt.Uint64() {
+			t.Fatal(i, "not equal")
+		}
 	}
 }
 
@@ -99,3 +122,23 @@ func Benchmark_BigInt(b *testing.B) {
 		_ = bigInt.Uint64()
 	}
 }
+
+func BenchmarkFileListHashMap_Exist(b *testing.B) {
+	var m = caches.NewFileListHashMap()
+	m.SetIsAvailable(true)
+	m.SetIsReady(true)
+
+	for i := 0; i < 1_000_000; i++ {
+		m.Add(types.String(i))
+	}
+
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			m.Add(types.String(rands.Int64()))
+			_ = m.Exist(types.String(rands.Int64()))
+		}
+	})
+}