fix: pin requery to preferred replica (#49826)

issue: #49742

### What changed
- Propagate preferred shard leader hints for search requery workloads.
- Add shard-level metrics for preferred node hit, miss, unavailable, and
rejected results.
- Keep preferred-node selection outside the balancer path so it does not
affect balancer state.
- Preserve fallback to other serviceable replicas when the preferred
node is unavailable or excluded.

### Tests
- go test ./internal/proxy/shardclient

---------

Signed-off-by: chyezh <chyezh@outlook.com>

internal/proxy/search_pipeline.go

@@ -1003,6 +1003,7 @@ type requeryOperator struct {
 	partitionIDs       []int64
 	primaryFieldSchema *schemapb.FieldSchema
 	queryChannelsTs    map[string]Timestamp
+	queryChannelsNode  map[string]int64
 	consistencyLevel   commonpb.ConsistencyLevel
 	guaranteeTimestamp uint64
 	namespace          *string
@@ -1034,6 +1035,13 @@ func newRequeryOperator(t *searchTask, _ map[string]any) (operator, error) {
 			outputFieldNames.Insert(highlightDynFields...)
 		}
 	}
+	queryChannelsNode := make(map[string]int64)
+	if t.queryChannelsNode != nil {
+		t.queryChannelsNode.Range(func(channel string, nodeID int64) bool {
+			queryChannelsNode[channel] = nodeID
+			return true
+		})
+	}
 	return &requeryOperator{
 		traceCtx:           t.TraceCtx(),
 		outputFieldNames:   outputFieldNames.Collect(),
@@ -1042,6 +1050,7 @@ func newRequeryOperator(t *searchTask, _ map[string]any) (operator, error) {
 		collectionName:     t.request.GetCollectionName(),
 		primaryFieldSchema: pkField,
 		queryChannelsTs:    t.queryChannelsTs,
+		queryChannelsNode:  queryChannelsNode,
 		consistencyLevel:   t.GetConsistencyLevel(),
 		guaranteeTimestamp: t.GetGuaranteeTimestamp(),
 		notReturnAllMeta:   t.request.GetNotReturnAllMeta(),
@@ -1091,6 +1100,10 @@ func (op *requeryOperator) requery(ctx context.Context, span trace.Span, ids *sc
 	for k, v := range op.queryChannelsTs {
 		channelsMvcc[k] = v
 	}
+	preferredNodes := make(map[string]int64)
+	for k, v := range op.queryChannelsNode {
+		preferredNodes[k] = v
+	}
 	qt := &queryTask{
 		ctx:       op.traceCtx,
 		Condition: NewTaskCondition(op.traceCtx),
@@ -1110,6 +1123,7 @@ func (op *requeryOperator) requery(ctx context.Context, span trace.Span, ids *sc
 		lb:             op.node.(*Proxy).lbPolicy,
 		shardclientMgr: op.node.(*Proxy).shardMgr,
 		channelsMvcc:   channelsMvcc,
+		preferredNodes: preferredNodes,
 		fastSkip:       true,
 		reQuery:        true,
 	}

internal/proxy/shardclient/lb_policy.go

@@ -29,6 +29,7 @@ import (
 	"github.com/milvus-io/milvus/internal/querycoordv2/params"
 	"github.com/milvus-io/milvus/internal/types"
 	"github.com/milvus-io/milvus/pkg/v3/log"
+	"github.com/milvus-io/milvus/pkg/v3/metrics"
 	"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
 	"github.com/milvus-io/milvus/pkg/v3/util/merr"
 	"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
@@ -39,12 +40,13 @@ import (
 type ExecuteFunc func(context.Context, UniqueID, types.QueryNodeClient, string) error
 
 type ChannelWorkload struct {
-	Db             string
-	CollectionName string
-	CollectionID   int64
-	Channel        string
-	Nq             int64
-	Exec           ExecuteFunc
+	Db              string
+	CollectionName  string
+	CollectionID    int64
+	Channel         string
+	Nq              int64
+	Exec            ExecuteFunc
+	PreferredNodeID int64
 }
 
 type CollectionWorkLoad struct {
@@ -53,6 +55,7 @@ type CollectionWorkLoad struct {
 	CollectionID   int64
 	Nq             int64
 	Exec           ExecuteFunc
+	PreferredNodes map[string]int64
 }
 
 type LBPolicy interface {
@@ -132,19 +135,36 @@ func (lb *LBPolicyImpl) GetShardLeaderList(ctx context.Context, dbName string, c
 	return ret, err
 }
 
+func recordPreferredNodeSelection(status string) {
+	metrics.ProxyShardLeaderPreferredNodeCount.WithLabelValues(
+		status,
+	).Inc()
+}
+
+func preferredNodeID(workload CollectionWorkLoad, channel string) int64 {
+	if workload.PreferredNodes == nil {
+		return 0
+	}
+	nodeID := workload.PreferredNodes[channel]
+	if nodeID == 0 {
+		recordPreferredNodeSelection(metrics.PreferredNodeMissLabel)
+	}
+	return nodeID
+}
+
 // try to select the best node from the available nodes
-func (lb *LBPolicyImpl) selectNode(ctx context.Context, balancer LBBalancer, workload ChannelWorkload, excludeNodes *typeutil.UniqueSet) (NodeInfo, error) {
+func (lb *LBPolicyImpl) selectNode(ctx context.Context, balancer LBBalancer, workload ChannelWorkload, excludeNodes *typeutil.UniqueSet) (NodeInfo, bool, error) {
 	log := log.Ctx(ctx).With(
 		zap.Int64("collectionID", workload.CollectionID),
 		zap.String("channelName", workload.Channel),
 	)
 	// Select node using specified nodes
-	trySelectNode := func(withCache bool) (NodeInfo, error) {
+	trySelectNode := func(withCache bool) (NodeInfo, bool, error) {
 		shardLeaders, err := lb.GetShard(ctx, workload.Db, workload.CollectionName, workload.CollectionID, workload.Channel, withCache)
 		if err != nil {
 			log.Warn("failed to get shard delegator",
 				zap.Error(err))
-			return NodeInfo{}, err
+			return NodeInfo{}, false, err
 		}
 
 		// if all available delegator has been excluded even after refresh shard leader cache
@@ -192,10 +212,18 @@ func (lb *LBPolicyImpl) selectNode(ctx context.Context, balancer LBBalancer, wor
 		}
 		if len(candidateNodes) == 0 {
 			err = merr.WrapErrChannelNotAvailable(workload.Channel, "no available shard leaders")
-			return NodeInfo{}, err
+			return NodeInfo{}, false, err
+		}
+
+		if preferredNode, ok := serviceableNodes[workload.PreferredNodeID]; ok {
+			recordPreferredNodeSelection(metrics.PreferredNodeHitLabel)
+			return preferredNode, false, nil
+		} else if workload.PreferredNodeID != 0 {
+			recordPreferredNodeSelection(metrics.PreferredNodeUnavailableLabel)
 		}
 
 		balancer.RegisterNodeInfo(lo.Values(candidateNodes))
+
 		// prefer serviceable nodes
 		var targetNodeID int64
 		if len(serviceableNodes) > 0 {
@@ -204,30 +232,30 @@ func (lb *LBPolicyImpl) selectNode(ctx context.Context, balancer LBBalancer, wor
 			targetNodeID, err = balancer.SelectNode(ctx, lo.Keys(candidateNodes), workload.Nq)
 		}
 		if err != nil {
-			return NodeInfo{}, err
+			return NodeInfo{}, false, err
 		}
 
 		if _, ok := candidateNodes[targetNodeID]; !ok {
 			err = merr.WrapErrNodeNotAvailable(targetNodeID)
-			return NodeInfo{}, err
+			return NodeInfo{}, false, err
 		}
 
-		return candidateNodes[targetNodeID], nil
+		return candidateNodes[targetNodeID], true, nil
 	}
 
 	// First attempt with current shard leaders cache
 	withShardLeaderCache := true
-	targetNode, err := trySelectNode(withShardLeaderCache)
+	targetNode, selectedByBalancer, err := trySelectNode(withShardLeaderCache)
 	if err != nil {
 		// Second attempt with fresh shard leaders
 		withShardLeaderCache = false
-		targetNode, err = trySelectNode(withShardLeaderCache)
+		targetNode, selectedByBalancer, err = trySelectNode(withShardLeaderCache)
 		if err != nil {
-			return NodeInfo{}, err
+			return NodeInfo{}, false, err
 		}
 	}
 
-	return targetNode, nil
+	return targetNode, selectedByBalancer, nil
 }
 
 // ExecuteWithRetry will choose a qn to execute the workload, and retry if failed, until reach the max retryTimes.
@@ -268,7 +296,7 @@ func (lb *LBPolicyImpl) ExecuteWithRetry(ctx context.Context, workload ChannelWo
 		excludeNodes := typeutil.NewUniqueSet(blacklist...)
 		excludeNodes.Insert(requestExcludedNodes.Collect()...)
 		balancer := lb.getBalancer()
-		targetNode, err := lb.selectNode(ctx, balancer, workload, &excludeNodes)
+		targetNode, selectedByBalancer, err := lb.selectNode(ctx, balancer, workload, &excludeNodes)
 		if err != nil {
 			log.Warn("failed to select node for shard",
 				zap.Int64("nodeID", targetNode.NodeID),
@@ -281,7 +309,9 @@ func (lb *LBPolicyImpl) ExecuteWithRetry(ctx context.Context, workload ChannelWo
 			return true, err
 		}
 		// cancel work load which assign to the target node
-		defer balancer.CancelWorkload(targetNode.NodeID, workload.Nq)
+		if selectedByBalancer {
+			defer balancer.CancelWorkload(targetNode.NodeID, workload.Nq)
+		}
 
 		client, err := lb.clientMgr.GetClient(ctx, targetNode)
 		if err != nil {
@@ -347,12 +377,13 @@ func (lb *LBPolicyImpl) Execute(ctx context.Context, workload CollectionWorkLoad
 	// Single channel fast path: skip errgroup/goroutine overhead
 	if len(channelList) == 1 {
 		return lb.ExecuteWithRetry(ctx, ChannelWorkload{
-			Db:             workload.Db,
-			CollectionName: workload.CollectionName,
-			CollectionID:   workload.CollectionID,
-			Channel:        channelList[0],
-			Nq:             workload.Nq,
-			Exec:           workload.Exec,
+			Db:              workload.Db,
+			CollectionName:  workload.CollectionName,
+			CollectionID:    workload.CollectionID,
+			Channel:         channelList[0],
+			Nq:              workload.Nq,
+			Exec:            workload.Exec,
+			PreferredNodeID: preferredNodeID(workload, channelList[0]),
 		})
 	}
 
@@ -360,12 +391,13 @@ func (lb *LBPolicyImpl) Execute(ctx context.Context, workload CollectionWorkLoad
 	for _, channel := range channelList {
 		wg.Go(func() error {
 			return lb.ExecuteWithRetry(ctx, ChannelWorkload{
-				Db:             workload.Db,
-				CollectionName: workload.CollectionName,
-				CollectionID:   workload.CollectionID,
-				Channel:        channel,
-				Nq:             workload.Nq,
-				Exec:           workload.Exec,
+				Db:              workload.Db,
+				CollectionName:  workload.CollectionName,
+				CollectionID:    workload.CollectionID,
+				Channel:         channel,
+				Nq:              workload.Nq,
+				Exec:            workload.Exec,
+				PreferredNodeID: preferredNodeID(workload, channel),
 			})
 		})
 	}
@@ -383,12 +415,13 @@ func (lb *LBPolicyImpl) ExecuteOneChannel(ctx context.Context, workload Collecti
 	// let every request could retry at least twice, which could retry after update shard leader cache
 	for _, channel := range channelList {
 		return lb.ExecuteWithRetry(ctx, ChannelWorkload{
-			Db:             workload.Db,
-			CollectionName: workload.CollectionName,
-			CollectionID:   workload.CollectionID,
-			Channel:        channel,
-			Nq:             workload.Nq,
-			Exec:           workload.Exec,
+			Db:              workload.Db,
+			CollectionName:  workload.CollectionName,
+			CollectionID:    workload.CollectionID,
+			Channel:         channel,
+			Nq:              workload.Nq,
+			Exec:            workload.Exec,
+			PreferredNodeID: preferredNodeID(workload, channel),
 		})
 	}
 	return fmt.Errorf("no acitvate sheard leader exist for collection: %s", workload.CollectionName)

internal/proxy/shardclient/lb_policy_test.go

@@ -22,12 +22,15 @@ import (
 	"testing"
 
 	"github.com/cockroachdb/errors"
+	"github.com/prometheus/client_golang/prometheus/testutil"
+	"github.com/samber/lo"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/suite"
 	"go.uber.org/atomic"
 
 	"github.com/milvus-io/milvus/internal/mocks"
 	"github.com/milvus-io/milvus/internal/types"
+	"github.com/milvus-io/milvus/pkg/v3/metrics"
 	"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
 	"github.com/milvus-io/milvus/pkg/v3/util/merr"
 	"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
@@ -99,7 +102,7 @@ func (s *LBPolicySuite) TestSelectNode() {
 	s.lbBalancer.EXPECT().RegisterNodeInfo(mock.Anything)
 	s.lbBalancer.EXPECT().SelectNode(mock.Anything, mock.Anything, mock.Anything).Return(int64(5), nil)
 	excludeNodes := typeutil.NewUniqueSet()
-	targetNode, err := s.lbPolicy.selectNode(ctx, s.lbBalancer, ChannelWorkload{
+	targetNode, _, err := s.lbPolicy.selectNode(ctx, s.lbBalancer, ChannelWorkload{
 		Db:             s.dbName,
 		CollectionName: s.collectionName,
 		CollectionID:   s.collectionID,
@@ -121,7 +124,7 @@ func (s *LBPolicySuite) TestSelectNode() {
 	s.lbBalancer.EXPECT().RegisterNodeInfo(mock.Anything)
 	s.lbBalancer.EXPECT().SelectNode(mock.Anything, mock.Anything, mock.Anything).Return(int64(3), nil).Once()
 	excludeNodes = typeutil.NewUniqueSet()
-	targetNode, err = s.lbPolicy.selectNode(ctx, s.lbBalancer, ChannelWorkload{
+	targetNode, _, err = s.lbPolicy.selectNode(ctx, s.lbBalancer, ChannelWorkload{
 		Db:             s.dbName,
 		CollectionName: s.collectionName,
 		CollectionID:   s.collectionID,
@@ -139,7 +142,7 @@ func (s *LBPolicySuite) TestSelectNode() {
 	s.lbBalancer.EXPECT().RegisterNodeInfo(mock.Anything)
 	s.lbBalancer.EXPECT().SelectNode(mock.Anything, mock.Anything, mock.Anything).Return(int64(-1), merr.ErrNodeNotAvailable)
 	excludeNodes = typeutil.NewUniqueSet()
-	targetNode, err = s.lbPolicy.selectNode(ctx, s.lbBalancer, ChannelWorkload{
+	targetNode, _, err = s.lbPolicy.selectNode(ctx, s.lbBalancer, ChannelWorkload{
 		Db:             s.dbName,
 		CollectionName: s.collectionName,
 		CollectionID:   s.collectionID,
@@ -156,7 +159,7 @@ func (s *LBPolicySuite) TestSelectNode() {
 	s.mgr.EXPECT().GetShard(mock.Anything, false, s.dbName, s.collectionName, s.collectionID, s.channels[0]).Return(s.nodes, nil).Once()
 	s.lbBalancer.EXPECT().RegisterNodeInfo(mock.Anything)
 	s.lbBalancer.EXPECT().SelectNode(mock.Anything, mock.Anything, mock.Anything).Return(int64(-1), merr.ErrNodeNotAvailable)
-	targetNode, err = s.lbPolicy.selectNode(ctx, s.lbBalancer, ChannelWorkload{
+	targetNode, _, err = s.lbPolicy.selectNode(ctx, s.lbBalancer, ChannelWorkload{
 		Db:             s.dbName,
 		CollectionName: s.collectionName,
 		CollectionID:   s.collectionID,
@@ -171,7 +174,7 @@ func (s *LBPolicySuite) TestSelectNode() {
 	s.mgr.EXPECT().GetShard(mock.Anything, true, s.dbName, s.collectionName, s.collectionID, s.channels[0]).Return(nil, merr.ErrCollectionNotLoaded).Once()
 	s.mgr.EXPECT().GetShard(mock.Anything, false, s.dbName, s.collectionName, s.collectionID, s.channels[0]).Return(nil, merr.ErrCollectionNotLoaded).Once()
 	excludeNodes = typeutil.NewUniqueSet()
-	targetNode, err = s.lbPolicy.selectNode(ctx, s.lbBalancer, ChannelWorkload{
+	targetNode, _, err = s.lbPolicy.selectNode(ctx, s.lbBalancer, ChannelWorkload{
 		Db:             s.dbName,
 		CollectionName: s.collectionName,
 		CollectionID:   s.collectionID,
@@ -181,6 +184,170 @@ func (s *LBPolicySuite) TestSelectNode() {
 	s.ErrorIs(err, merr.ErrCollectionNotLoaded)
 }
 
+func (s *LBPolicySuite) TestPreferredNodeHint() {
+	ctx := context.Background()
+
+	s.mgr.EXPECT().GetShard(mock.Anything, true, s.dbName, s.collectionName, s.collectionID, s.channels[0]).Return(s.nodes, nil)
+	excludeNodes := typeutil.NewUniqueSet()
+	targetNode, selectedByBalancer, err := s.lbPolicy.selectNode(ctx, s.lbBalancer, ChannelWorkload{
+		Db:              s.dbName,
+		CollectionName:  s.collectionName,
+		CollectionID:    s.collectionID,
+		Channel:         s.channels[0],
+		Nq:              1,
+		PreferredNodeID: 3,
+	}, &excludeNodes)
+	s.NoError(err)
+	s.Equal(int64(3), targetNode.NodeID)
+	s.False(selectedByBalancer)
+}
+
+func (s *LBPolicySuite) TestPreferredNodeHintFallback() {
+	ctx := context.Background()
+	nodes := []NodeInfo{
+		{NodeID: 1, Address: "localhost", Serviceable: true},
+		{NodeID: 2, Address: "localhost", Serviceable: false},
+		{NodeID: 3, Address: "localhost", Serviceable: true},
+	}
+
+	s.mgr.EXPECT().GetShard(mock.Anything, true, s.dbName, s.collectionName, s.collectionID, s.channels[0]).Return(nodes, nil)
+	s.lbBalancer.EXPECT().RegisterNodeInfo(mock.Anything)
+	s.lbBalancer.EXPECT().SelectNode(mock.Anything, mock.MatchedBy(func(nodes []int64) bool {
+		return !lo.Contains(nodes, int64(2)) && lo.Contains(nodes, int64(1)) && lo.Contains(nodes, int64(3))
+	}), int64(1)).Return(int64(1), nil)
+	excludeNodes := typeutil.NewUniqueSet()
+	targetNode, selectedByBalancer, err := s.lbPolicy.selectNode(ctx, s.lbBalancer, ChannelWorkload{
+		Db:              s.dbName,
+		CollectionName:  s.collectionName,
+		CollectionID:    s.collectionID,
+		Channel:         s.channels[0],
+		Nq:              1,
+		PreferredNodeID: 2,
+	}, &excludeNodes)
+	s.NoError(err)
+	s.Equal(int64(1), targetNode.NodeID)
+	s.True(selectedByBalancer)
+}
+
+func (s *LBPolicySuite) TestExecuteUsesPreferredNodeHint() {
+	ctx := context.Background()
+
+	s.mgr.EXPECT().GetShardLeaderList(mock.Anything, s.dbName, s.collectionName, s.collectionID, true).Return([]string{s.channels[0]}, nil)
+	s.mgr.EXPECT().GetShard(mock.Anything, true, s.dbName, s.collectionName, s.collectionID, s.channels[0]).Return(s.nodes, nil)
+	s.mgr.EXPECT().GetClient(mock.Anything, mock.MatchedBy(func(node NodeInfo) bool {
+		return node.NodeID == 3
+	})).Return(s.qn, nil)
+
+	var executedNodeID int64
+	err := s.lbPolicy.Execute(ctx, CollectionWorkLoad{
+		Db:             s.dbName,
+		CollectionName: s.collectionName,
+		CollectionID:   s.collectionID,
+		Nq:             1,
+		PreferredNodes: map[string]int64{s.channels[0]: 3},
+		Exec: func(ctx context.Context, nodeID UniqueID, qn types.QueryNodeClient, channel string) error {
+			executedNodeID = nodeID
+			return nil
+		},
+	})
+	s.NoError(err)
+	s.Equal(int64(3), executedNodeID)
+}
+
+func (s *LBPolicySuite) TestPreferredNodeHintMetrics() {
+	ctx := context.Background()
+	collectionID := int64(99001)
+	channel := "preferred-metric-channel-hit"
+	before := testutil.ToFloat64(metrics.ProxyShardLeaderPreferredNodeCount.WithLabelValues(
+		metrics.PreferredNodeHitLabel,
+	))
+
+	s.mgr.EXPECT().GetShard(mock.Anything, true, s.dbName, s.collectionName, collectionID, channel).Return(s.nodes, nil)
+	excludeNodes := typeutil.NewUniqueSet()
+	targetNode, selectedByBalancer, err := s.lbPolicy.selectNode(ctx, s.lbBalancer, ChannelWorkload{
+		Db:              s.dbName,
+		CollectionName:  s.collectionName,
+		CollectionID:    collectionID,
+		Channel:         channel,
+		Nq:              1,
+		PreferredNodeID: 3,
+	}, &excludeNodes)
+	s.NoError(err)
+	s.Equal(int64(3), targetNode.NodeID)
+	s.False(selectedByBalancer)
+
+	after := testutil.ToFloat64(metrics.ProxyShardLeaderPreferredNodeCount.WithLabelValues(
+		metrics.PreferredNodeHitLabel,
+	))
+	s.Equal(float64(1), after-before)
+}
+
+func (s *LBPolicySuite) TestPreferredNodeHintMetricsDisabledForNormalWorkload() {
+	ctx := context.Background()
+	collectionID := int64(99002)
+	channel := "preferred-metric-channel-disabled"
+	before := testutil.ToFloat64(metrics.ProxyShardLeaderPreferredNodeCount.WithLabelValues(
+		metrics.PreferredNodeMissLabel,
+	))
+
+	s.mgr.EXPECT().GetShard(mock.Anything, true, s.dbName, s.collectionName, collectionID, channel).Return(s.nodes, nil)
+	s.lbBalancer.EXPECT().RegisterNodeInfo(mock.Anything)
+	s.lbBalancer.EXPECT().SelectNode(mock.Anything, mock.Anything, int64(1)).Return(int64(1), nil)
+	excludeNodes := typeutil.NewUniqueSet()
+	targetNode, _, err := s.lbPolicy.selectNode(ctx, s.lbBalancer, ChannelWorkload{
+		Db:             s.dbName,
+		CollectionName: s.collectionName,
+		CollectionID:   collectionID,
+		Channel:        channel,
+		Nq:             1,
+	}, &excludeNodes)
+	s.NoError(err)
+	s.Equal(int64(1), targetNode.NodeID)
+
+	after := testutil.ToFloat64(metrics.ProxyShardLeaderPreferredNodeCount.WithLabelValues(
+		metrics.PreferredNodeMissLabel,
+	))
+	s.Equal(float64(0), after-before)
+}
+
+func (s *LBPolicySuite) TestPreferredNodeFailureFallsBackToOtherReplica() {
+	ctx := context.Background()
+	channel := "preferred-node-fallback-channel"
+	nodes := []NodeInfo{
+		{NodeID: 1, Address: "localhost:9000", Serviceable: true},
+		{NodeID: 2, Address: "localhost:9001", Serviceable: true},
+	}
+	s.lbPolicy.retryOnReplica = 1
+
+	s.mgr.ExpectedCalls = nil
+	s.lbBalancer.ExpectedCalls = nil
+	s.mgr.EXPECT().GetShardLeaderList(mock.Anything, s.dbName, s.collectionName, s.collectionID, true).Return([]string{channel}, nil)
+	s.mgr.EXPECT().GetShard(mock.Anything, true, s.dbName, s.collectionName, s.collectionID, channel).Return(nodes, nil)
+	s.mgr.EXPECT().GetShard(mock.Anything, false, s.dbName, s.collectionName, s.collectionID, channel).Return(nodes, nil).Maybe()
+	s.mgr.EXPECT().GetClient(mock.Anything, mock.Anything).Return(s.qn, nil)
+	s.lbBalancer.EXPECT().RegisterNodeInfo(mock.Anything)
+	s.lbBalancer.EXPECT().SelectNode(mock.Anything, []int64{int64(2)}, int64(1)).Return(int64(2), nil).Once()
+	s.lbBalancer.EXPECT().CancelWorkload(mock.Anything, mock.Anything)
+
+	executedNodes := make([]int64, 0, 2)
+	err := s.lbPolicy.Execute(ctx, CollectionWorkLoad{
+		Db:             s.dbName,
+		CollectionName: s.collectionName,
+		CollectionID:   s.collectionID,
+		Nq:             1,
+		PreferredNodes: map[string]int64{channel: 1},
+		Exec: func(ctx context.Context, nodeID UniqueID, qn types.QueryNodeClient, channel string) error {
+			executedNodes = append(executedNodes, nodeID)
+			if nodeID == 1 {
+				return merr.ErrServiceUnavailable
+			}
+			return nil
+		},
+	})
+	s.NoError(err)
+	s.Equal([]int64{1, 2}, executedNodes)
+}
+
 func (s *LBPolicySuite) TestExecuteWithRetry() {
 	ctx := context.Background()
 
@@ -665,7 +832,7 @@ func (s *LBPolicySuite) TestSelectNodeEdgeCases() {
 	s.mgr.EXPECT().GetShard(mock.Anything, false, s.dbName, s.collectionName, s.collectionID, s.channels[0]).Return([]NodeInfo{}, nil).Once()
 
 	excludeNodes := typeutil.NewUniqueSet(s.nodeIDs...)
-	_, err := s.lbPolicy.selectNode(ctx, s.lbBalancer, ChannelWorkload{
+	_, _, err := s.lbPolicy.selectNode(ctx, s.lbBalancer, ChannelWorkload{
 		Db:             s.dbName,
 		CollectionName: s.collectionName,
 		CollectionID:   s.collectionID,
@@ -684,7 +851,7 @@ func (s *LBPolicySuite) TestSelectNodeEdgeCases() {
 	s.lbBalancer.EXPECT().SelectNode(mock.Anything, mock.Anything, mock.Anything).Return(int64(1), nil).Once()
 
 	excludeNodes = typeutil.NewUniqueSet(int64(1)) // Exclude the single node
-	targetNode, err := s.lbPolicy.selectNode(ctx, s.lbBalancer, ChannelWorkload{
+	targetNode, _, err := s.lbPolicy.selectNode(ctx, s.lbBalancer, ChannelWorkload{
 		Db:             s.dbName,
 		CollectionName: s.collectionName,
 		CollectionID:   s.collectionID,
@@ -711,7 +878,7 @@ func (s *LBPolicySuite) TestSelectNodeEdgeCases() {
 	}), mock.Anything).Return(int64(3), nil).Once()
 
 	excludeNodes = typeutil.NewUniqueSet(int64(1)) // Exclude node 1, node 3 should be available
-	targetNode, err = s.lbPolicy.selectNode(ctx, s.lbBalancer, ChannelWorkload{
+	targetNode, _, err = s.lbPolicy.selectNode(ctx, s.lbBalancer, ChannelWorkload{
 		Db:             s.dbName,
 		CollectionName: s.collectionName,
 		CollectionID:   s.collectionID,
@@ -772,7 +939,7 @@ func (s *LBPolicySuite) TestSelectNodeWithExcludeClearing() {
 	s.lbBalancer.EXPECT().SelectNode(mock.Anything, mock.Anything, mock.Anything).Return(int64(1), nil).Once()
 
 	excludeNodes := typeutil.NewUniqueSet(int64(1), int64(2)) // Exclude all available nodes
-	targetNode, err := s.lbPolicy.selectNode(ctx, s.lbBalancer, ChannelWorkload{
+	targetNode, _, err := s.lbPolicy.selectNode(ctx, s.lbBalancer, ChannelWorkload{
 		Db:             s.dbName,
 		CollectionName: s.collectionName,
 		CollectionID:   s.collectionID,
@@ -797,7 +964,7 @@ func (s *LBPolicySuite) TestSelectNodeWithExcludeClearing() {
 	s.lbBalancer.EXPECT().SelectNode(mock.Anything, mock.Anything, mock.Anything).Return(int64(2), nil).Once()
 
 	excludeNodes = typeutil.NewUniqueSet(int64(1)) // Only exclude node 1
-	targetNode, err = s.lbPolicy.selectNode(ctx, s.lbBalancer, ChannelWorkload{
+	targetNode, _, err = s.lbPolicy.selectNode(ctx, s.lbBalancer, ChannelWorkload{
 		Db:             s.dbName,
 		CollectionName: s.collectionName,
 		CollectionID:   s.collectionID,
@@ -816,7 +983,7 @@ func (s *LBPolicySuite) TestSelectNodeWithExcludeClearing() {
 	s.mgr.EXPECT().GetShard(mock.Anything, false, s.dbName, s.collectionName, s.collectionID, s.channels[0]).Return([]NodeInfo{}, nil).Once()
 
 	excludeNodes = typeutil.NewUniqueSet(int64(1))
-	_, err = s.lbPolicy.selectNode(ctx, s.lbBalancer, ChannelWorkload{
+	_, _, err = s.lbPolicy.selectNode(ctx, s.lbBalancer, ChannelWorkload{
 		Db:             s.dbName,
 		CollectionName: s.collectionName,
 		CollectionID:   s.collectionID,

internal/proxy/task_query.go

@@ -78,6 +78,7 @@ type queryTask struct {
 	shardclientMgr   shardclient.ShardClientMgr
 	lb               shardclient.LBPolicy
 	channelsMvcc     map[string]Timestamp
+	preferredNodes   map[string]int64
 	fastSkip         bool
 
 	reQuery              bool
@@ -915,6 +916,7 @@ func (t *queryTask) Execute(ctx context.Context) error {
 		CollectionName: t.collectionName,
 		Nq:             1,
 		Exec:           t.queryShard,
+		PreferredNodes: t.preferredNodes,
 	})
 	if err != nil {
 		log.Warn("fail to execute query", zap.Error(err))

internal/proxy/task_search.go

@@ -92,13 +92,14 @@ type searchTask struct {
 
 	partitionIDsSet *typeutil.ConcurrentSet[UniqueID]
 
-	mixCoord        types.MixCoordClient
-	node            types.ProxyComponent
-	lb              shardclient.LBPolicy
-	shardClientMgr  shardclient.ShardClientMgr
-	queryChannelsTs map[string]Timestamp
-	queryInfos      []*planpb.QueryInfo
-	relatedDataSize int64
+	mixCoord          types.MixCoordClient
+	node              types.ProxyComponent
+	lb                shardclient.LBPolicy
+	shardClientMgr    shardclient.ShardClientMgr
+	queryChannelsTs   map[string]Timestamp
+	queryChannelsNode *typeutil.ConcurrentMap[string, int64]
+	queryInfos        []*planpb.QueryInfo
+	relatedDataSize   int64
 
 	// Rerank configuration metadata (nil means no rerank)
 	rerankMeta rerankMeta
@@ -1076,6 +1077,7 @@ func (t *searchTask) Execute(ctx context.Context) error {
 	tr := timerecord.NewTimeRecorder(fmt.Sprintf("proxy execute search %d", t.ID()))
 	defer tr.CtxElapse(ctx, "done")
 
+	t.queryChannelsNode = typeutil.NewConcurrentMap[string, int64]()
 	err := t.lb.Execute(ctx, shardclient.CollectionWorkLoad{
 		Db:             t.request.GetDbName(),
 		CollectionID:   t.CollectionID,
@@ -1315,6 +1317,9 @@ func (t *searchTask) searchShard(ctx context.Context, nodeID int64, qn types.Que
 	if t.resultBuf != nil {
 		t.resultBuf.Insert(result)
 	}
+	if t.queryChannelsNode != nil {
+		t.queryChannelsNode.Insert(channel, nodeID)
+	}
 	t.lb.UpdateCostMetrics(nodeID, result.CostAggregation)
 
 	return nil

internal/proxy/task_search_test.go

@@ -4409,6 +4409,7 @@ func TestSearchTask_Requery(t *testing.T) {
 
 		lb := shardclient.NewMockLBPolicy(t)
 		lb.EXPECT().Execute(mock.Anything, mock.Anything).Run(func(ctx context.Context, workload shardclient.CollectionWorkLoad) {
+			assert.Equal(t, map[string]int64{"mock_qn": 1}, workload.PreferredNodes)
 			err = workload.Exec(ctx, 0, qn, "")
 			assert.NoError(t, err)
 		}).Return(nil)
@@ -4446,7 +4447,9 @@ func TestSearchTask_Requery(t *testing.T) {
 			node:                   node,
 			translatedOutputFields: outputFields,
 			shardClientMgr:         mgr,
+			queryChannelsNode:      typeutil.NewConcurrentMap[string, int64](),
 		}
+		qt.queryChannelsNode.Insert("mock_qn", 1)
 		op, err := newRequeryOperator(qt, nil)
 		assert.NoError(t, err)
 		queryResult, storageCost, err := op.(*requeryOperator).requery(ctx, nil, qt.result.Results.Ids, outputFields)
@@ -4588,6 +4591,47 @@ func TestSearchTask_Requery(t *testing.T) {
 	})
 }
 
+func TestSearchTask_SearchShardRecordsNodeHint(t *testing.T) {
+	ctx := context.Background()
+	const channel = "by-dev-rootcoord-dml_0_100v0"
+	const nodeID int64 = 101
+
+	qn := mocks.NewMockQueryNodeClient(t)
+	qn.EXPECT().Search(mock.Anything, mock.MatchedBy(func(req *querypb.SearchRequest) bool {
+		return len(req.GetDmlChannels()) == 1 && req.GetDmlChannels()[0] == channel
+	})).Return(&internalpb.SearchResults{
+		Status:          merr.Success(),
+		CostAggregation: &internalpb.CostAggregation{},
+	}, nil)
+
+	lb := shardclient.NewMockLBPolicy(t)
+	lb.EXPECT().UpdateCostMetrics(nodeID, mock.Anything).Return()
+
+	task := &searchTask{
+		ctx: ctx,
+		SearchRequest: &internalpb.SearchRequest{
+			Base: &commonpb.MsgBase{
+				MsgType:  commonpb.MsgType_Search,
+				SourceID: paramtable.GetNodeID(),
+			},
+		},
+		request: &milvuspb.SearchRequest{
+			DbName:         "default",
+			CollectionName: "test_search_shard_records_node_hint",
+		},
+		Condition:         NewTaskCondition(ctx),
+		lb:                lb,
+		resultBuf:         typeutil.NewConcurrentSet[*internalpb.SearchResults](),
+		queryChannelsNode: typeutil.NewConcurrentMap[string, int64](),
+	}
+
+	err := task.searchShard(ctx, nodeID, qn, channel)
+	require.NoError(t, err)
+	recordedNodeID, ok := task.queryChannelsNode.Get(channel)
+	require.True(t, ok)
+	assert.Equal(t, nodeID, recordedNodeID)
+}
+
 type GetPartitionIDsSuite struct {
 	suite.Suite
 

pkg/metrics/metrics.go

@@ -48,7 +48,16 @@ const (
 	CacheMissLabel   = "miss"
 	TimetickLabel    = "timetick"
 	AllLabel         = "all"
+)
 
+const (
+	PreferredNodeHitLabel         = "hit"
+	PreferredNodeMissLabel        = "miss"
+	PreferredNodeUnavailableLabel = "unavailable"
+	PreferredNodeRejectedLabel    = "rejected"
+)
+
+const (
 	UnissuedIndexTaskLabel   = "unissued"
 	InProgressIndexTaskLabel = "in-progress"
 	FinishedIndexTaskLabel   = "finished"

pkg/metrics/proxy_metrics.go

@@ -365,6 +365,17 @@ var (
 			nodeIDLabelName,
 		})
 
+	// ProxyShardLeaderPreferredNodeCount records preferred shard leader selection results.
+	ProxyShardLeaderPreferredNodeCount = prometheus.NewCounterVec(
+		prometheus.CounterOpts{
+			Namespace: milvusNamespace,
+			Subsystem: typeutil.ProxyRole,
+			Name:      "shard_leader_preferred_node_count",
+			Help:      "counter of preferred shard leader selection results",
+		}, []string{
+			statusLabelName,
+		})
+
 	// ProxyRateLimitReqCount integrates a counter monitoring metric for the rate-limit rpc requests.
 	ProxyRateLimitReqCount = prometheus.NewCounterVec(
 		prometheus.CounterOpts{
@@ -532,6 +543,7 @@ func RegisterProxy(registry *prometheus.Registry) {
 
 	registry.MustRegister(ProxyWorkLoadScore)
 	registry.MustRegister(ProxyExecutingTotalNq)
+	registry.MustRegister(ProxyShardLeaderPreferredNodeCount)
 	registry.MustRegister(ProxyRateLimitReqCount)
 
 	registry.MustRegister(ProxySlowQueryCount)