Replace insertChannelsMap with channelsMgr (#5453)

Signed-off-by: dragondriver <jiquan.long@zilliz.com>

internal/proxynode/channels_mgr.go

@@ -3,6 +3,7 @@ package proxynode
 import (
 	"context"
 	"fmt"
+	"math/rand"
 	"runtime"
 	"sort"
 	"sync"
@@ -64,6 +65,41 @@ type masterService interface {
 	GetChannels(collectionID UniqueID) (map[vChan]pChan, error)
 }
 
+type mockMaster struct {
+	collectionID2Channels map[UniqueID]map[vChan]pChan
+}
+
+func newMockMaster() *mockMaster {
+	return &mockMaster{
+		collectionID2Channels: make(map[UniqueID]map[vChan]pChan),
+	}
+}
+
+func genUniqueStr() string {
+	l := rand.Uint64()%100 + 1
+	b := make([]byte, l)
+	if _, err := rand.Read(b); err != nil {
+		return ""
+	}
+	return fmt.Sprintf("%X", b)
+}
+
+func (m *mockMaster) GetChannels(collectionID UniqueID) (map[vChan]pChan, error) {
+	channels, ok := m.collectionID2Channels[collectionID]
+	if ok {
+		return channels, nil
+	}
+
+	channels = make(map[vChan]pChan)
+	l := rand.Uint64()%10 + 1
+	for i := 0; uint64(i) < l; i++ {
+		channels[genUniqueStr()] = genUniqueStr()
+	}
+
+	m.collectionID2Channels[collectionID] = channels
+	return channels, nil
+}
+
 type channelsMgrImpl struct {
 	collectionID2VIDs map[UniqueID][]int // id are sorted
 	collMtx           sync.RWMutex

internal/proxynode/channels_mgr_test.go

 package proxynode
 
 import (
-	"fmt"
-	"math/rand"
 	"testing"
 
 	"github.com/milvus-io/milvus/internal/msgstream"
@@ -24,41 +22,6 @@ func TestNaiveUniqueIntGenerator_get(t *testing.T) {
 	}
 }
 
-type mockMaster struct {
-	collectionID2Channels map[UniqueID]map[vChan]pChan
-}
-
-func newMockMaster() *mockMaster {
-	return &mockMaster{
-		collectionID2Channels: make(map[UniqueID]map[vChan]pChan),
-	}
-}
-
-func genUniqueStr() string {
-	l := rand.Uint64()%100 + 1
-	b := make([]byte, l)
-	if _, err := rand.Read(b); err != nil {
-		return ""
-	}
-	return fmt.Sprintf("%X", b)
-}
-
-func (m *mockMaster) GetChannels(collectionID UniqueID) (map[vChan]pChan, error) {
-	channels, ok := m.collectionID2Channels[collectionID]
-	if ok {
-		return channels, nil
-	}
-
-	channels = make(map[vChan]pChan)
-	l := rand.Uint64()%10 + 1
-	for i := 0; uint64(i) < l; i++ {
-		channels[genUniqueStr()] = genUniqueStr()
-	}
-
-	m.collectionID2Channels[collectionID] = channels
-	return channels, nil
-}
-
 func TestChannelsMgrImpl_getChannels(t *testing.T) {
 	master := newMockMaster()
 	factory := msgstream.NewSimpleMsgStreamFactory()

internal/proxynode/impl.go

@@ -143,6 +143,7 @@ func (node *ProxyNode) DropCollection(ctx context.Context, request *milvuspb.Dro
 		Condition:             NewTaskCondition(ctx),
 		DropCollectionRequest: request,
 		masterService:         node.masterService,
+		chMgr:                 node.chMgr,
 	}
 
 	err := node.sched.DdQueue.Enqueue(dct)
@@ -1062,6 +1063,7 @@ func (node *ProxyNode) Insert(ctx context.Context, request *milvuspb.InsertReque
 		},
 		rowIDAllocator: node.idAllocator,
 		segIDAssigner:  node.segAssigner,
+		chMgr:          node.chMgr,
 	}
 	if len(it.PartitionName) <= 0 {
 		it.PartitionName = Params.DefaultPartitionName

internal/proxynode/proxy_node.go

@@ -54,6 +54,8 @@ type ProxyNode struct {
 	proxyService  types.ProxyService
 	queryService  types.QueryService
 
+	chMgr channelsMgr
+
 	sched *TaskScheduler
 	tick  *timeTick
 
@@ -224,6 +226,11 @@ func (node *ProxyNode) Init() error {
 	node.segAssigner = segAssigner
 	node.segAssigner.PeerID = Params.ProxyID
 
+	// TODO(dragondriver): use real master service instance
+	mockMasterIns := newMockMaster()
+	chMgr := newChannelsMgr(mockMasterIns, node.msFactory)
+	node.chMgr = chMgr
+
 	node.sched, err = NewTaskScheduler(node.ctx, node.idAllocator, node.tsoAllocator, node.msFactory)
 	if err != nil {
 		return err

internal/proxynode/task.go

@@ -106,6 +106,7 @@ type InsertTask struct {
 	result         *milvuspb.InsertResponse
 	rowIDAllocator *allocator.IDAllocator
 	segIDAssigner  *SegIDAssigner
+	chMgr          channelsMgr
 }
 
 func (it *InsertTask) TraceCtx() context.Context {
@@ -684,30 +685,14 @@ func (it *InsertTask) Execute(ctx context.Context) error {
 
 	msgPack.Msgs[0] = tsMsg
 
-	stream, err := globalInsertChannelsMap.GetInsertMsgStream(collID)
+	stream, err := it.chMgr.getDMLStream(collID)
 	if err != nil {
-		resp, _ := it.dataService.GetInsertChannels(ctx, &datapb.GetInsertChannelsRequest{
-			Base: &commonpb.MsgBase{
-				MsgType:   commonpb.MsgType_Insert, // todo
-				MsgID:     it.Base.MsgID,           // todo
-				Timestamp: 0,                       // todo
-				SourceID:  Params.ProxyID,
-			},
-			DbID:         0, // todo
-			CollectionID: collID,
-		})
-		if resp == nil {
-			return errors.New("get insert channels resp is nil")
-		}
-		if resp.Status.ErrorCode != commonpb.ErrorCode_Success {
-			return errors.New(resp.Status.Reason)
-		}
-		err = globalInsertChannelsMap.CreateInsertMsgStream(collID, resp.Values)
+		err = it.chMgr.createDMLMsgStream(collID)
 		if err != nil {
 			return err
 		}
 	}
-	stream, err = globalInsertChannelsMap.GetInsertMsgStream(collID)
+	stream, err = it.chMgr.getDMLStream(collID)
 	if err != nil {
 		it.result.Status.ErrorCode = commonpb.ErrorCode_UnexpectedError
 		it.result.Status.Reason = err.Error()
@@ -849,36 +834,7 @@ func (cct *CreateCollectionTask) PreExecute(ctx context.Context) error {
 func (cct *CreateCollectionTask) Execute(ctx context.Context) error {
 	var err error
 	cct.result, err = cct.masterService.CreateCollection(ctx, cct.CreateCollectionRequest)
-	if err != nil {
-		return err
-	}
-	if cct.result.ErrorCode == commonpb.ErrorCode_Success {
-		collID, err := globalMetaCache.GetCollectionID(ctx, cct.CollectionName)
-		if err != nil {
-			return err
-		}
-		resp, _ := cct.dataServiceClient.GetInsertChannels(ctx, &datapb.GetInsertChannelsRequest{
-			Base: &commonpb.MsgBase{
-				MsgType:   commonpb.MsgType_Insert, // todo
-				MsgID:     cct.Base.MsgID,          // todo
-				Timestamp: 0,                       // todo
-				SourceID:  Params.ProxyID,
-			},
-			DbID:         0, // todo
-			CollectionID: collID,
-		})
-		if resp == nil {
-			return errors.New("get insert channels resp is nil")
-		}
-		if resp.Status.ErrorCode != commonpb.ErrorCode_Success {
-			return errors.New(resp.Status.Reason)
-		}
-		err = globalInsertChannelsMap.CreateInsertMsgStream(collID, resp.Values)
-		if err != nil {
-			return err
-		}
-	}
-	return nil
+	return err
 }
 
 func (cct *CreateCollectionTask) PostExecute(ctx context.Context) error {
@@ -891,6 +847,7 @@ type DropCollectionTask struct {
 	ctx           context.Context
 	masterService types.MasterService
 	result        *commonpb.Status
+	chMgr         channelsMgr
 }
 
 func (dct *DropCollectionTask) TraceCtx() context.Context {
@@ -951,7 +908,7 @@ func (dct *DropCollectionTask) Execute(ctx context.Context) error {
 		return err
 	}
 
-	err = globalInsertChannelsMap.CloseInsertMsgStream(collID)
+	err = dct.chMgr.removeDMLStream(collID)
 	if err != nil {
 		return err
 	}