enhance: Decouple shard client manager from shard cache

internal/proxy/lb_balancer.go

@@ -23,6 +23,7 @@ import (
 )
 
 type LBBalancer interface {
+	RegisterNodeInfo(nodeInfos []nodeInfo)
 	SelectNode(ctx context.Context, availableNodes []int64, nq int64) (int64, error)
 	CancelWorkload(node int64, nq int64)
 	UpdateCostMetrics(node int64, cost *internalpb.CostAggregation)

internal/proxy/lb_policy.go

@@ -39,7 +39,7 @@ type ChannelWorkload struct {
 	collectionName string
 	collectionID   int64
 	channel        string
-	shardLeaders   []int64
+	shardLeaders   []nodeInfo
 	nq             int64
 	exec           executeFunc
 	retryTimes     uint
@@ -116,9 +116,20 @@ func (lb *LBPolicyImpl) GetShardLeaders(ctx context.Context, dbName string, coll
 }
 
 // try to select the best node from the available nodes
-func (lb *LBPolicyImpl) selectNode(ctx context.Context, balancer LBBalancer, workload ChannelWorkload, excludeNodes typeutil.UniqueSet) (int64, error) {
-	availableNodes := lo.FilterMap(workload.shardLeaders, func(node int64, _ int) (int64, bool) { return node, !excludeNodes.Contain(node) })
-	targetNode, err := balancer.SelectNode(ctx, availableNodes, workload.nq)
+func (lb *LBPolicyImpl) selectNode(ctx context.Context, balancer LBBalancer, workload ChannelWorkload, excludeNodes typeutil.UniqueSet) (nodeInfo, error) {
+	filterDelegator := func(nodes []nodeInfo) map[int64]nodeInfo {
+		ret := make(map[int64]nodeInfo)
+		for _, node := range nodes {
+			if !excludeNodes.Contain(node.nodeID) {
+				ret[node.nodeID] = node
+			}
+		}
+		return ret
+	}
+
+	availableNodes := filterDelegator(workload.shardLeaders)
+	balancer.RegisterNodeInfo(lo.Values(availableNodes))
+	targetNode, err := balancer.SelectNode(ctx, lo.Keys(availableNodes), workload.nq)
 	if err != nil {
 		log := log.Ctx(ctx)
 		globalMetaCache.DeprecateShardCache(workload.db, workload.collectionName)
@@ -128,32 +139,33 @@ func (lb *LBPolicyImpl) selectNode(ctx context.Context, balancer LBBalancer, wor
 				zap.Int64("collectionID", workload.collectionID),
 				zap.String("channelName", workload.channel),
 				zap.Error(err))
-			return -1, err
+			return nodeInfo{}, err
 		}
 
-		availableNodes := lo.FilterMap(shardLeaders[workload.channel], func(node nodeInfo, _ int) (int64, bool) { return node.nodeID, !excludeNodes.Contain(node.nodeID) })
+		availableNodes = filterDelegator(shardLeaders[workload.channel])
 		if len(availableNodes) == 0 {
-			nodes := lo.Map(shardLeaders[workload.channel], func(node nodeInfo, _ int) int64 { return node.nodeID })
 			log.Warn("no available shard delegator found",
 				zap.Int64("collectionID", workload.collectionID),
 				zap.String("channelName", workload.channel),
-				zap.Int64s("nodes", nodes),
+				zap.Int64s("availableNodes", lo.Keys(availableNodes)),
 				zap.Int64s("excluded", excludeNodes.Collect()))
-			return -1, merr.WrapErrChannelNotAvailable("no available shard delegator found")
+			return nodeInfo{}, merr.WrapErrChannelNotAvailable("no available shard delegator found")
 		}
 
-		targetNode, err = balancer.SelectNode(ctx, availableNodes, workload.nq)
+		balancer.RegisterNodeInfo(lo.Values(availableNodes))
+		targetNode, err = balancer.SelectNode(ctx, lo.Keys(availableNodes), workload.nq)
 		if err != nil {
 			log.Warn("failed to select shard",
 				zap.Int64("collectionID", workload.collectionID),
 				zap.String("channelName", workload.channel),
-				zap.Int64s("availableNodes", availableNodes),
+				zap.Int64s("availableNodes", lo.Keys(availableNodes)),
+				zap.Int64s("excluded", excludeNodes.Collect()),
 				zap.Error(err))
-			return -1, err
+			return nodeInfo{}, err
 		}
 	}
 
-	return targetNode, nil
+	return availableNodes[targetNode], nil
 }
 
 // ExecuteWithRetry will choose a qn to execute the workload, and retry if failed, until reach the max retryTimes.
@@ -168,7 +180,7 @@ func (lb *LBPolicyImpl) ExecuteWithRetry(ctx context.Context, workload ChannelWo
 			log.Warn("failed to select node for shard",
 				zap.Int64("collectionID", workload.collectionID),
 				zap.String("channelName", workload.channel),
-				zap.Int64("nodeID", targetNode),
+				zap.Int64("nodeID", targetNode.nodeID),
 				zap.Error(err),
 			)
 			if lastErr != nil {
@@ -177,29 +189,30 @@ func (lb *LBPolicyImpl) ExecuteWithRetry(ctx context.Context, workload ChannelWo
 			return err
 		}
 		// cancel work load which assign to the target node
-		defer balancer.CancelWorkload(targetNode, workload.nq)
+		defer balancer.CancelWorkload(targetNode.nodeID, workload.nq)
 
 		client, err := lb.clientMgr.GetClient(ctx, targetNode)
 		if err != nil {
 			log.Warn("search/query channel failed, node not available",
 				zap.Int64("collectionID", workload.collectionID),
 				zap.String("channelName", workload.channel),
-				zap.Int64("nodeID", targetNode),
+				zap.Int64("nodeID", targetNode.nodeID),
 				zap.Error(err))
-			excludeNodes.Insert(targetNode)
+			excludeNodes.Insert(targetNode.nodeID)
 
 			lastErr = errors.Wrapf(err, "failed to get delegator %d for channel %s", targetNode, workload.channel)
 			return lastErr
 		}
+		defer lb.clientMgr.ReleaseClientRef(targetNode.nodeID)
 
-		err = workload.exec(ctx, targetNode, client, workload.channel)
+		err = workload.exec(ctx, targetNode.nodeID, client, workload.channel)
 		if err != nil {
 			log.Warn("search/query channel failed",
 				zap.Int64("collectionID", workload.collectionID),
 				zap.String("channelName", workload.channel),
-				zap.Int64("nodeID", targetNode),
+				zap.Int64("nodeID", targetNode.nodeID),
 				zap.Error(err))
-			excludeNodes.Insert(targetNode)
+			excludeNodes.Insert(targetNode.nodeID)
 			lastErr = errors.Wrapf(err, "failed to search/query delegator %d for channel %s", targetNode, workload.channel)
 			return lastErr
 		}
@@ -221,9 +234,9 @@ func (lb *LBPolicyImpl) Execute(ctx context.Context, workload CollectionWorkLoad
 	// let every request could retry at least twice, which could retry after update shard leader cache
 	retryTimes := Params.ProxyCfg.RetryTimesOnReplica.GetAsInt()
 	wg, ctx := errgroup.WithContext(ctx)
-	for channel, nodes := range dml2leaders {
-		channel := channel
-		nodes := lo.Map(nodes, func(node nodeInfo, _ int) int64 { return node.nodeID })
+	for k, v := range dml2leaders {
+		channel := k
+		nodes := v
 		channelRetryTimes := retryTimes
 		if len(nodes) > 0 {
 			channelRetryTimes *= len(nodes)

internal/proxy/lb_policy_test.go

@@ -49,7 +49,8 @@ type LBPolicySuite struct {
 	lbBalancer *MockLBBalancer
 	lbPolicy   *LBPolicyImpl
 
-	nodes    []int64
+	nodeIDs  []int64
+	nodes    []nodeInfo
 	channels []string
 	qnList   []*mocks.MockQueryNode
 
@@ -62,7 +63,14 @@ func (s *LBPolicySuite) SetupSuite() {
 }
 
 func (s *LBPolicySuite) SetupTest() {
-	s.nodes = []int64{1, 2, 3, 4, 5}
+	s.nodeIDs = make([]int64, 0)
+	for i := 1; i <= 5; i++ {
+		s.nodeIDs = append(s.nodeIDs, int64(i))
+		s.nodes = append(s.nodes, nodeInfo{
+			nodeID:  int64(i),
+			address: "localhost",
+		})
+	}
 	s.channels = []string{"channel1", "channel2"}
 	successStatus := commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}
 	qc := mocks.NewMockQueryCoordClient(s.T())
@@ -74,12 +82,12 @@ func (s *LBPolicySuite) SetupTest() {
 		Shards: []*querypb.ShardLeadersList{
 			{
 				ChannelName: s.channels[0],
-				NodeIds:     s.nodes,
+				NodeIds:     s.nodeIDs,
 				NodeAddrs:   []string{"localhost:9000", "localhost:9001", "localhost:9002", "localhost:9003", "localhost:9004"},
 			},
 			{
 				ChannelName: s.channels[1],
-				NodeIds:     s.nodes,
+				NodeIds:     s.nodeIDs,
 				NodeAddrs:   []string{"localhost:9000", "localhost:9001", "localhost:9002", "localhost:9003", "localhost:9004"},
 			},
 		},
@@ -96,7 +104,6 @@ func (s *LBPolicySuite) SetupTest() {
 	s.qn.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(nil, nil).Maybe()
 
 	s.mgr = NewMockShardClientManager(s.T())
-	s.mgr.EXPECT().UpdateShardLeaders(mock.Anything, mock.Anything).Return(nil).Maybe()
 	s.lbBalancer = NewMockLBBalancer(s.T())
 	s.lbBalancer.EXPECT().Start(context.Background()).Maybe()
 	s.lbPolicy = NewLBPolicyImpl(s.mgr)
@@ -164,6 +171,7 @@ func (s *LBPolicySuite) loadCollection() {
 
 func (s *LBPolicySuite) TestSelectNode() {
 	ctx := context.Background()
+	s.lbBalancer.EXPECT().RegisterNodeInfo(mock.Anything)
 	s.lbBalancer.EXPECT().SelectNode(mock.Anything, mock.Anything, mock.Anything).Return(5, nil)
 	targetNode, err := s.lbPolicy.selectNode(ctx, s.lbBalancer, ChannelWorkload{
 		db:             dbName,
@@ -174,39 +182,41 @@ func (s *LBPolicySuite) TestSelectNode() {
 		nq:             1,
 	}, typeutil.NewUniqueSet())
 	s.NoError(err)
-	s.Equal(int64(5), targetNode)
+	s.Equal(int64(5), targetNode.nodeID)
 
 	// test select node failed, then update shard leader cache and retry, expect success
 	s.lbBalancer.ExpectedCalls = nil
+	s.lbBalancer.EXPECT().RegisterNodeInfo(mock.Anything)
 	s.lbBalancer.EXPECT().SelectNode(mock.Anything, mock.Anything, mock.Anything).Return(-1, errors.New("fake err")).Times(1)
 	s.lbBalancer.EXPECT().SelectNode(mock.Anything, mock.Anything, mock.Anything).Return(3, nil)
 	targetNode, err = s.lbPolicy.selectNode(ctx, s.lbBalancer, ChannelWorkload{
 		db:             dbName,
 		collectionName: s.collectionName,
 		collectionID:   s.collectionID,
 		channel:        s.channels[0],
-		shardLeaders:   []int64{},
+		shardLeaders:   []nodeInfo{},
 		nq:             1,
 	}, typeutil.NewUniqueSet())
 	s.NoError(err)
-	s.Equal(int64(3), targetNode)
+	s.Equal(int64(3), targetNode.nodeID)
 
 	// test select node always fails, expected failure
 	s.lbBalancer.ExpectedCalls = nil
+	s.lbBalancer.EXPECT().RegisterNodeInfo(mock.Anything)
 	s.lbBalancer.EXPECT().SelectNode(mock.Anything, mock.Anything, mock.Anything).Return(-1, merr.ErrNodeNotAvailable)
 	targetNode, err = s.lbPolicy.selectNode(ctx, s.lbBalancer, ChannelWorkload{
 		db:             dbName,
 		collectionName: s.collectionName,
 		collectionID:   s.collectionID,
 		channel:        s.channels[0],
-		shardLeaders:   []int64{},
+		shardLeaders:   []nodeInfo{},
 		nq:             1,
 	}, typeutil.NewUniqueSet())
 	s.ErrorIs(err, merr.ErrNodeNotAvailable)
-	s.Equal(int64(-1), targetNode)
 
 	// test all nodes has been excluded, expected failure
 	s.lbBalancer.ExpectedCalls = nil
+	s.lbBalancer.EXPECT().RegisterNodeInfo(mock.Anything)
 	s.lbBalancer.EXPECT().SelectNode(mock.Anything, mock.Anything, mock.Anything).Return(-1, merr.ErrNodeNotAvailable)
 	targetNode, err = s.lbPolicy.selectNode(ctx, s.lbBalancer, ChannelWorkload{
 		db:             dbName,
@@ -215,12 +225,12 @@ func (s *LBPolicySuite) TestSelectNode() {
 		channel:        s.channels[0],
 		shardLeaders:   s.nodes,
 		nq:             1,
-	}, typeutil.NewUniqueSet(s.nodes...))
+	}, typeutil.NewUniqueSet(s.nodeIDs...))
 	s.ErrorIs(err, merr.ErrChannelNotAvailable)
-	s.Equal(int64(-1), targetNode)
 
 	// test get shard leaders failed, retry to select node failed
 	s.lbBalancer.ExpectedCalls = nil
+	s.lbBalancer.EXPECT().RegisterNodeInfo(mock.Anything)
 	s.lbBalancer.EXPECT().SelectNode(mock.Anything, mock.Anything, mock.Anything).Return(-1, merr.ErrNodeNotAvailable)
 	s.qc.ExpectedCalls = nil
 	s.qc.EXPECT().GetShardLeaders(mock.Anything, mock.Anything).Return(nil, merr.ErrServiceUnavailable)
@@ -233,15 +243,16 @@ func (s *LBPolicySuite) TestSelectNode() {
 		nq:             1,
 	}, typeutil.NewUniqueSet())
 	s.ErrorIs(err, merr.ErrServiceUnavailable)
-	s.Equal(int64(-1), targetNode)
 }
 
 func (s *LBPolicySuite) TestExecuteWithRetry() {
 	ctx := context.Background()
 
 	// test execute success
 	s.lbBalancer.ExpectedCalls = nil
+	s.mgr.EXPECT().ReleaseClientRef(mock.Anything)
 	s.mgr.EXPECT().GetClient(mock.Anything, mock.Anything).Return(s.qn, nil)
+	s.lbBalancer.EXPECT().RegisterNodeInfo(mock.Anything)
 	s.lbBalancer.EXPECT().SelectNode(mock.Anything, mock.Anything, mock.Anything).Return(1, nil)
 	s.lbBalancer.EXPECT().CancelWorkload(mock.Anything, mock.Anything)
 	err := s.lbPolicy.ExecuteWithRetry(ctx, ChannelWorkload{
@@ -260,6 +271,7 @@ func (s *LBPolicySuite) TestExecuteWithRetry() {
 
 	// test select node failed, expected error
 	s.lbBalancer.ExpectedCalls = nil
+	s.lbBalancer.EXPECT().RegisterNodeInfo(mock.Anything)
 	s.lbBalancer.EXPECT().SelectNode(mock.Anything, mock.Anything, mock.Anything).Return(-1, merr.ErrNodeNotAvailable)
 	err = s.lbPolicy.ExecuteWithRetry(ctx, ChannelWorkload{
 		db:             dbName,
@@ -277,8 +289,10 @@ func (s *LBPolicySuite) TestExecuteWithRetry() {
 
 	// test get client failed, and retry failed, expected success
 	s.mgr.ExpectedCalls = nil
+	s.mgr.EXPECT().ReleaseClientRef(mock.Anything)
 	s.mgr.EXPECT().GetClient(mock.Anything, mock.Anything).Return(nil, errors.New("fake error")).Times(1)
 	s.lbBalancer.ExpectedCalls = nil
+	s.lbBalancer.EXPECT().RegisterNodeInfo(mock.Anything)
 	s.lbBalancer.EXPECT().SelectNode(mock.Anything, mock.Anything, mock.Anything).Return(1, nil)
 	s.lbBalancer.EXPECT().CancelWorkload(mock.Anything, mock.Anything)
 	err = s.lbPolicy.ExecuteWithRetry(ctx, ChannelWorkload{
@@ -296,8 +310,10 @@ func (s *LBPolicySuite) TestExecuteWithRetry() {
 	s.Error(err)
 
 	s.mgr.ExpectedCalls = nil
+	s.mgr.EXPECT().ReleaseClientRef(mock.Anything)
 	s.mgr.EXPECT().GetClient(mock.Anything, mock.Anything).Return(nil, errors.New("fake error")).Times(1)
 	s.mgr.EXPECT().GetClient(mock.Anything, mock.Anything).Return(s.qn, nil)
+	s.lbBalancer.EXPECT().RegisterNodeInfo(mock.Anything)
 	s.lbBalancer.EXPECT().CancelWorkload(mock.Anything, mock.Anything)
 	err = s.lbPolicy.ExecuteWithRetry(ctx, ChannelWorkload{
 		db:             dbName,
@@ -315,8 +331,10 @@ func (s *LBPolicySuite) TestExecuteWithRetry() {
 
 	// test exec failed, then retry success
 	s.mgr.ExpectedCalls = nil
+	s.mgr.EXPECT().ReleaseClientRef(mock.Anything)
 	s.mgr.EXPECT().GetClient(mock.Anything, mock.Anything).Return(s.qn, nil)
 	s.lbBalancer.ExpectedCalls = nil
+	s.lbBalancer.EXPECT().RegisterNodeInfo(mock.Anything)
 	s.lbBalancer.EXPECT().SelectNode(mock.Anything, mock.Anything, mock.Anything).Return(1, nil)
 	s.lbBalancer.EXPECT().CancelWorkload(mock.Anything, mock.Anything)
 	counter := 0
@@ -341,6 +359,7 @@ func (s *LBPolicySuite) TestExecuteWithRetry() {
 	// test exec timeout
 	s.mgr.ExpectedCalls = nil
 	s.mgr.EXPECT().GetClient(mock.Anything, mock.Anything).Return(s.qn, nil)
+	s.mgr.EXPECT().ReleaseClientRef(mock.Anything)
 	s.lbBalancer.EXPECT().CancelWorkload(mock.Anything, mock.Anything)
 	s.qn.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(nil, nil).Maybe()
 	s.qn.EXPECT().Search(mock.Anything, mock.Anything).Return(nil, context.Canceled).Times(1)
@@ -365,7 +384,9 @@ func (s *LBPolicySuite) TestExecute() {
 	ctx := context.Background()
 	mockErr := errors.New("mock error")
 	// test  all channel success
+	s.mgr.EXPECT().ReleaseClientRef(mock.Anything)
 	s.mgr.EXPECT().GetClient(mock.Anything, mock.Anything).Return(s.qn, nil)
+	s.lbBalancer.EXPECT().RegisterNodeInfo(mock.Anything)
 	s.lbBalancer.EXPECT().SelectNode(mock.Anything, mock.Anything, mock.Anything).Return(1, nil)
 	s.lbBalancer.EXPECT().CancelWorkload(mock.Anything, mock.Anything)
 	err := s.lbPolicy.Execute(ctx, CollectionWorkLoad{