From d958056218e28dc263af0b61a53803922c11d0a7 Mon Sep 17 00:00:00 2001
From: Gordon <46924906+FGadvancer@users.noreply.github.com>
Date: Wed, 13 Dec 2023 21:15:07 +0800
Subject: [PATCH] fix: reduce the granularity of the lock.

---
 internal/msggateway/n_ws_server.go |  2 +-
 internal/msggateway/user_map.go    | 47 ++++++++++++++++++++----------
 2 files changed, 33 insertions(+), 16 deletions(-)

diff --git a/internal/msggateway/n_ws_server.go b/internal/msggateway/n_ws_server.go
index 58d9e4ebcd..fd43be37da 100644
--- a/internal/msggateway/n_ws_server.go
+++ b/internal/msggateway/n_ws_server.go
@@ -159,7 +159,7 @@ func NewWsServer(opts ...Option) (*WsServer, error) {
 		unregisterChan:  make(chan *Client, 1000),
 		kickHandlerChan: make(chan *kickHandler, 1000),
 		validate:        v,
-		clients:         newUserMap(),
+		clients:         newUserMap(10),
 		Compressor:      NewGzipCompressor(),
 		Encoder:         NewGobEncoder(),
 	}, nil
diff --git a/internal/msggateway/user_map.go b/internal/msggateway/user_map.go
index 052d7de2d9..069d6f2c69 100644
--- a/internal/msggateway/user_map.go
+++ b/internal/msggateway/user_map.go
@@ -16,6 +16,7 @@ package msggateway
 
 import (
 	"context"
+	"hash/fnv"
 	"sync"
 
 	"github.com/OpenIMSDK/tools/log"
@@ -23,15 +24,26 @@ import (
 )
 
 type UserMap struct {
-	m sync.Map
+	maps []sync.Map
 }
 
-func newUserMap() *UserMap {
-	return &UserMap{}
+func newUserMap(numMaps int) *UserMap {
+	um := &UserMap{
+		maps: make([]sync.Map, numMaps),
+	}
+	return um
+}
+
+func (u *UserMap) getMap(key string) *sync.Map {
+	hash := fnv.New32a()
+	hash.Write([]byte(key))
+	index := hash.Sum32() % uint32(len(u.maps))
+	return &u.maps[index]
 }
 
 func (u *UserMap) GetAll(key string) ([]*Client, bool) {
-	allClients, ok := u.m.Load(key)
+	sm := u.getMap(key)
+	allClients, ok := sm.Load(key)
 	if ok {
 		return allClients.([]*Client), ok
 	}
@@ -39,7 +51,8 @@ func (u *UserMap) GetAll(key string) ([]*Client, bool) {
 }
 
 func (u *UserMap) Get(key string, platformID int) ([]*Client, bool, bool) {
-	allClients, userExisted := u.m.Load(key)
+	sm := u.getMap(key)
+	allClients, userExisted := sm.Load(key)
 	if userExisted {
 		var clients []*Client
 		for _, client := range allClients.([]*Client) {
@@ -56,22 +69,24 @@ func (u *UserMap) Get(key string, platformID int) ([]*Client, bool, bool) {
 }
 
 func (u *UserMap) Set(key string, v *Client) {
-	allClients, existed := u.m.Load(key)
+	sm := u.getMap(key)
+	allClients, existed := sm.Load(key)
 	if existed {
 		log.ZDebug(context.Background(), "Set existed", "user_id", key, "client", *v)
 		oldClients := allClients.([]*Client)
 		oldClients = append(oldClients, v)
-		u.m.Store(key, oldClients)
+		sm.Store(key, oldClients)
 	} else {
 		log.ZDebug(context.Background(), "Set not existed", "user_id", key, "client", *v)
 		var clients []*Client
 		clients = append(clients, v)
-		u.m.Store(key, clients)
+		sm.Store(key, clients)
 	}
 }
 
 func (u *UserMap) delete(key string, connRemoteAddr string) (isDeleteUser bool) {
-	allClients, existed := u.m.Load(key)
+	sm := u.getMap(key)
+	allClients, existed := sm.Load(key)
 	if existed {
 		oldClients := allClients.([]*Client)
 		var a []*Client
@@ -81,10 +96,10 @@ func (u *UserMap) delete(key string, connRemoteAddr string) (isDeleteUser bool)
 			}
 		}
 		if len(a) == 0 {
-			u.m.Delete(key)
+			sm.Delete(key)
 			return true
 		} else {
-			u.m.Store(key, a)
+			sm.Store(key, a)
 			return false
 		}
 	}
@@ -92,10 +107,11 @@ func (u *UserMap) delete(key string, connRemoteAddr string) (isDeleteUser bool)
 }
 
 func (u *UserMap) deleteClients(key string, clients []*Client) (isDeleteUser bool) {
+	sm := u.getMap(key)
 	m := utils.SliceToMapAny(clients, func(c *Client) (string, struct{}) {
 		return c.ctx.GetRemoteAddr(), struct{}{}
 	})
-	allClients, existed := u.m.Load(key)
+	allClients, existed := sm.Load(key)
 	if existed {
 		oldClients := allClients.([]*Client)
 		var a []*Client
@@ -105,10 +121,10 @@ func (u *UserMap) deleteClients(key string, clients []*Client) (isDeleteUser boo
 			}
 		}
 		if len(a) == 0 {
-			u.m.Delete(key)
+			sm.Delete(key)
 			return true
 		} else {
-			u.m.Store(key, a)
+			sm.Store(key, a)
 			return false
 		}
 	}
@@ -116,5 +132,6 @@ func (u *UserMap) deleteClients(key string, clients []*Client) (isDeleteUser boo
 }
 
 func (u *UserMap) DeleteAll(key string) {
-	u.m.Delete(key)
+	sm := u.getMap(key)
+	sm.Delete(key)
 }