Use map&mutex for dispatcher client (#26380)

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>

pkg/mq/msgdispatcher/client.go

@@ -17,6 +17,8 @@
 package msgdispatcher
 
 import (
+	"sync"
+
 	"github.com/milvus-io/milvus-proto/go-api/v2/msgpb"
 	"go.uber.org/zap"
 
@@ -24,7 +26,6 @@ import (
 	"github.com/milvus-io/milvus/pkg/mq/msgstream"
 	"github.com/milvus-io/milvus/pkg/mq/msgstream/mqwrapper"
 	"github.com/milvus-io/milvus/pkg/util/funcutil"
-	"github.com/milvus-io/milvus/pkg/util/typeutil"
 )
 
 type (
@@ -42,18 +43,20 @@ type Client interface {
 var _ Client = (*client)(nil)
 
 type client struct {
-	role     string
+	role       string
-	nodeID   int64
+	nodeID     int64
-	managers *typeutil.ConcurrentMap[string, DispatcherManager]
+	managers   map[string]DispatcherManager
-	factory  msgstream.Factory
+	managerMut sync.Mutex
+	factory    msgstream.Factory
 }
 
 func NewClient(factory msgstream.Factory, role string, nodeID int64) Client {
 	return &client{
-		role:     role,
+		role:    role,
-		nodeID:   nodeID,
+		nodeID:  nodeID,
-		factory:  factory,
+		factory: factory,
-		managers: typeutil.NewConcurrentMap[string, DispatcherManager](),
+		// managers: typeutil.NewConcurrentMap[string, DispatcherManager](),
+		managers: make(map[string]DispatcherManager),
 	}
 }
 
@@ -61,18 +64,20 @@ func (c *client) Register(vchannel string, pos *Pos, subPos SubPos) (<-chan *Msg
 	log := log.With(zap.String("role", c.role),
 		zap.Int64("nodeID", c.nodeID), zap.String("vchannel", vchannel))
 	pchannel := funcutil.ToPhysicalChannel(vchannel)
+	c.managerMut.Lock()
+	defer c.managerMut.Unlock()
 	var manager DispatcherManager
-	manager, ok := c.managers.Get(pchannel)
+	manager, ok := c.managers[pchannel]
 	if !ok {
 		manager = NewDispatcherManager(pchannel, c.role, c.nodeID, c.factory)
-		c.managers.Insert(pchannel, manager)
+		c.managers[pchannel] = manager
 		go manager.Run()
 	}
 	ch, err := manager.Add(vchannel, pos, subPos)
 	if err != nil {
 		if manager.Num() == 0 {
 			manager.Close()
-			c.managers.GetAndRemove(pchannel)
+			delete(c.managers, pchannel)
 		}
 		log.Error("register failed", zap.Error(err))
 		return nil, err
@@ -83,11 +88,13 @@ func (c *client) Register(vchannel string, pos *Pos, subPos SubPos) (<-chan *Msg
 
 func (c *client) Deregister(vchannel string) {
 	pchannel := funcutil.ToPhysicalChannel(vchannel)
-	if manager, ok := c.managers.Get(pchannel); ok {
+	c.managerMut.Lock()
+	defer c.managerMut.Unlock()
+	if manager, ok := c.managers[pchannel]; ok {
 		manager.Remove(vchannel)
 		if manager.Num() == 0 {
 			manager.Close()
-			c.managers.GetAndRemove(pchannel)
+			delete(c.managers, pchannel)
 		}
 		log.Info("deregister done", zap.String("role", c.role),
 			zap.Int64("nodeID", c.nodeID), zap.String("vchannel", vchannel))
@@ -97,11 +104,12 @@ func (c *client) Deregister(vchannel string) {
 func (c *client) Close() {
 	log := log.With(zap.String("role", c.role),
 		zap.Int64("nodeID", c.nodeID))
-	c.managers.Range(func(pchannel string, manager DispatcherManager) bool {
+	c.managerMut.Lock()
+	defer c.managerMut.Unlock()
+	for pchannel, manager := range c.managers {
 		log.Info("close manager", zap.String("channel", pchannel))
-		c.managers.GetAndRemove(pchannel)
+		delete(c.managers, pchannel)
 		manager.Close()
-		return true
+	}
-	})
 	log.Info("dispatcher client closed")
 }

pkg/mq/msgdispatcher/client_test.go

@@ -61,6 +61,9 @@ func TestClient_Concurrency(t *testing.T) {
 	wg.Wait()
 	expected := int(total - deregisterCount.Load())
 
-	n := client1.(*client).managers.Len()
+	c := client1.(*client)
+	c.managerMut.Lock()
+	n := len(c.managers)
+	c.managerMut.Unlock()
 	assert.Equal(t, expected, n)
 }