fix: resolve channel exclusive mode lost and vchannel list bug (#47701)

issue: #47700

Move DescribeCollection before SpawnReplicasWithReplicaConfig in
LoadCollectionJob.Execute to use collInfo.GetVirtualChannelNames()
instead of broadcast message vchannels, and handle ErrCollectionNotFound
gracefully. Add tryBalanceNodeForChannel call in
TryEnableChannelExclusiveMode to ensure nodes are balanced across
channels when exclusive mode is enabled.

Fix querycoord test suite: initialize AssignPolicyFactory, add
VirtualChannelNames to DescribeCollection mock, and remove stale mock
overrides. Add unit tests for both fixes.

Signed-off-by: chyezh <chyezh@outlook.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

internal/querycoordv2/job/job_load.go

@@ -37,6 +37,7 @@ import (
 	"github.com/milvus-io/milvus/pkg/v2/metrics"
 	"github.com/milvus-io/milvus/pkg/v2/proto/querypb"
 	"github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
+	"github.com/milvus-io/milvus/pkg/v2/util/merr"
 	"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
 )
 
@@ -84,25 +85,26 @@ func NewLoadCollectionJob(
 
 func (job *LoadCollectionJob) Execute() error {
 	req := job.result.Message.Header()
-	vchannels := job.result.GetVChannelsWithoutControlChannel()
-
 	log := log.Ctx(job.ctx).With(zap.Int64("collectionID", req.GetCollectionId()))
 	meta.GlobalFailedLoadCache.Remove(req.GetCollectionId())
 
+	collInfo, err := job.broker.DescribeCollection(job.ctx, req.GetCollectionId())
+	if errors.Is(err, merr.ErrCollectionNotFound) {
+		return nil
+	}
+	if err != nil {
+		return err
+	}
+
 	// 1. create replica if not exist
 	if _, err := utils.SpawnReplicasWithReplicaConfig(job.ctx, job.meta, meta.SpawnWithReplicaConfigParams{
 		CollectionID: req.GetCollectionId(),
-		Channels:     vchannels,
+		Channels:     collInfo.GetVirtualChannelNames(),
 		Configs:      req.GetReplicas(),
 	}); err != nil {
 		return err
 	}
 
-	collInfo, err := job.broker.DescribeCollection(job.ctx, req.GetCollectionId())
-	if err != nil {
-		return err
-	}
-
 	// 2. put load info meta
 	fieldIndexIDs := make(map[int64]int64, len(req.GetLoadFields()))
 	fieldIDs := make([]int64, 0, len(req.GetLoadFields()))

internal/querycoordv2/job/job_load_test.go

@@ -0,0 +1,132 @@
+// Licensed to the LF AI & Data foundation under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package job
+
+import (
+	"context"
+	"testing"
+
+	"github.com/cockroachdb/errors"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/suite"
+
+	"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
+	"github.com/milvus-io/milvus/internal/querycoordv2/meta"
+	"github.com/milvus-io/milvus/pkg/v2/proto/messagespb"
+	"github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
+	"github.com/milvus-io/milvus/pkg/v2/util/merr"
+	"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
+)
+
+type LoadCollectionJobSuite struct {
+	suite.Suite
+}
+
+func (suite *LoadCollectionJobSuite) SetupSuite() {
+	paramtable.Init()
+}
+
+func (suite *LoadCollectionJobSuite) SetupTest() {
+	meta.GlobalFailedLoadCache = meta.NewFailedLoadCache()
+}
+
+func (suite *LoadCollectionJobSuite) buildBroadcastResult(collectionID int64, partitionIDs []int64) message.BroadcastResultAlterLoadConfigMessageV2 {
+	controlChannel := "_ctrl_channel"
+	replicas := []*messagespb.LoadReplicaConfig{
+		{ReplicaId: 1, ResourceGroupName: "__default_resource_group"},
+	}
+	broadcastMsg := message.NewAlterLoadConfigMessageBuilderV2().
+		WithHeader(&messagespb.AlterLoadConfigMessageHeader{
+			CollectionId: collectionID,
+			PartitionIds: partitionIDs,
+			Replicas:     replicas,
+		}).
+		WithBody(&messagespb.AlterLoadConfigMessageBody{}).
+		WithBroadcast([]string{controlChannel}).
+		MustBuildBroadcast()
+
+	specializedMsg := message.MustAsBroadcastAlterLoadConfigMessageV2(broadcastMsg)
+	return message.BroadcastResultAlterLoadConfigMessageV2{
+		Message: specializedMsg,
+		Results: map[string]*message.AppendResult{
+			controlChannel: {},
+		},
+	}
+}
+
+// TestDescribeCollectionNotFound tests that Execute returns nil when the collection is not found.
+func (suite *LoadCollectionJobSuite) TestDescribeCollectionNotFound() {
+	ctx := context.Background()
+	collectionID := int64(1000)
+
+	broker := meta.NewMockBroker(suite.T())
+	broker.EXPECT().DescribeCollection(mock.Anything, collectionID).
+		Return(nil, merr.WrapErrCollectionNotFound(collectionID))
+
+	result := suite.buildBroadcastResult(collectionID, []int64{100, 101})
+	job := NewLoadCollectionJob(ctx, result, nil, nil, broker, nil, nil, nil, nil, nil)
+
+	err := job.Execute()
+	suite.NoError(err)
+}
+
+// TestDescribeCollectionOtherError tests that Execute returns the error when DescribeCollection fails.
+func (suite *LoadCollectionJobSuite) TestDescribeCollectionOtherError() {
+	ctx := context.Background()
+	collectionID := int64(1001)
+
+	expectedErr := errors.New("broker unavailable")
+	broker := meta.NewMockBroker(suite.T())
+	broker.EXPECT().DescribeCollection(mock.Anything, collectionID).
+		Return(nil, expectedErr)
+
+	result := suite.buildBroadcastResult(collectionID, []int64{200, 201})
+	job := NewLoadCollectionJob(ctx, result, nil, nil, broker, nil, nil, nil, nil, nil)
+
+	err := job.Execute()
+	suite.Error(err)
+	suite.True(errors.Is(err, expectedErr))
+}
+
+// TestDescribeCollectionSuccess tests that Execute proceeds with VirtualChannelNames from DescribeCollection.
+func (suite *LoadCollectionJobSuite) TestDescribeCollectionSuccess() {
+	ctx := context.Background()
+	collectionID := int64(1002)
+	channels := []string{"ch1", "ch2"}
+
+	broker := meta.NewMockBroker(suite.T())
+	broker.EXPECT().DescribeCollection(mock.Anything, collectionID).
+		Return(&milvuspb.DescribeCollectionResponse{
+			CollectionID:        collectionID,
+			VirtualChannelNames: channels,
+		}, nil)
+
+	result := suite.buildBroadcastResult(collectionID, []int64{300, 301})
+	// We pass nil for meta to test that DescribeCollection is called before SpawnReplicasWithReplicaConfig.
+	// SpawnReplicasWithReplicaConfig will panic on nil meta, proving that DescribeCollection was called first.
+	job := NewLoadCollectionJob(ctx, result, nil, nil, broker, nil, nil, nil, nil, nil)
+
+	// This should panic at SpawnReplicasWithReplicaConfig because meta is nil,
+	// but this proves DescribeCollection was called and returned successfully first.
+	suite.Panics(func() {
+		job.Execute()
+	})
+}
+
+func TestLoadCollectionJob(t *testing.T) {
+	suite.Run(t, new(LoadCollectionJobSuite))
+}

internal/querycoordv2/meta/replica.go

@@ -393,6 +393,7 @@ func (replica *mutableReplica) TryEnableChannelExclusiveMode(channelNames ...str
 	if replica.exclusiveRWNodeToChannel == nil {
 		replica.exclusiveRWNodeToChannel = make(map[int64]string)
 	}
+	replica.tryBalanceNodeForChannel()
 }
 
 func (replica *mutableReplica) DisableChannelExclusiveMode() {

internal/querycoordv2/meta/replica_test.go

@@ -728,6 +728,117 @@ func (suite *ReplicaSuite) TestCalculateOptimalAssignments() {
 	suite.Equal(1, countsOfThree) // 1 channel gets 3 nodes
 }
 
+// TestTryEnableChannelExclusiveModeTriggersBalance tests that TryEnableChannelExclusiveMode
+// calls tryBalanceNodeForChannel to balance nodes across channels.
+func (suite *ReplicaSuite) TestTryEnableChannelExclusiveModeTriggersBalance() {
+	paramtable.Get().Save(paramtable.Get().QueryCoordCfg.Balancer.Key, ChannelLevelScoreBalancerName)
+	paramtable.Get().Save(paramtable.Get().QueryCoordCfg.ChannelExclusiveNodeFactor.Key, "1")
+	defer func() {
+		paramtable.Get().Reset(paramtable.Get().QueryCoordCfg.Balancer.Key)
+		paramtable.Get().Reset(paramtable.Get().QueryCoordCfg.ChannelExclusiveNodeFactor.Key)
+	}()
+
+	// Create a replica with nodes but no ChannelNodeInfos (nil) to trigger initialization path
+	r := newReplica(&querypb.Replica{
+		ID:            1,
+		CollectionID:  2,
+		ResourceGroup: DefaultResourceGroupName,
+		Nodes:         []int64{1, 2, 3, 4, 5, 6},
+	})
+
+	mutableReplica := r.CopyForWrite()
+	// Verify ChannelNodeInfos is nil before calling TryEnableChannelExclusiveMode
+	suite.Nil(mutableReplica.replicaPB.ChannelNodeInfos)
+
+	// Call TryEnableChannelExclusiveMode with channel names
+	mutableReplica.TryEnableChannelExclusiveMode("channel1", "channel2", "channel3")
+
+	newR := mutableReplica.IntoReplica()
+
+	// Verify that ChannelNodeInfos was created
+	suite.NotNil(newR.replicaPB.GetChannelNodeInfos())
+	suite.Equal(3, len(newR.replicaPB.GetChannelNodeInfos()))
+
+	// Verify that tryBalanceNodeForChannel was called and nodes were balanced
+	// 6 nodes / 3 channels = 2 nodes per channel
+	totalAssignedNodes := 0
+	for _, channelNodeInfo := range newR.replicaPB.GetChannelNodeInfos() {
+		suite.Equal(2, len(channelNodeInfo.GetRwNodes()))
+		totalAssignedNodes += len(channelNodeInfo.GetRwNodes())
+	}
+	suite.Equal(6, totalAssignedNodes)
+}
+
+// TestTryEnableChannelExclusiveModeExistingChannelNodeInfos tests that TryEnableChannelExclusiveMode
+// does not overwrite existing ChannelNodeInfos but still triggers balance.
+func (suite *ReplicaSuite) TestTryEnableChannelExclusiveModeExistingChannelNodeInfos() {
+	paramtable.Get().Save(paramtable.Get().QueryCoordCfg.Balancer.Key, ChannelLevelScoreBalancerName)
+	paramtable.Get().Save(paramtable.Get().QueryCoordCfg.ChannelExclusiveNodeFactor.Key, "1")
+	defer func() {
+		paramtable.Get().Reset(paramtable.Get().QueryCoordCfg.Balancer.Key)
+		paramtable.Get().Reset(paramtable.Get().QueryCoordCfg.ChannelExclusiveNodeFactor.Key)
+	}()
+
+	// Create a replica with existing ChannelNodeInfos but no balanced nodes
+	r := newReplica(&querypb.Replica{
+		ID:            1,
+		CollectionID:  2,
+		ResourceGroup: DefaultResourceGroupName,
+		Nodes:         []int64{1, 2, 3, 4},
+		ChannelNodeInfos: map[string]*querypb.ChannelNodeInfo{
+			"channel1": {},
+			"channel2": {},
+		},
+	})
+
+	mutableReplica := r.CopyForWrite()
+	// ChannelNodeInfos is not nil, so TryEnableChannelExclusiveMode should not overwrite
+	mutableReplica.TryEnableChannelExclusiveMode("channel1", "channel2")
+
+	newR := mutableReplica.IntoReplica()
+
+	// Verify existing channels are preserved (not overwritten with new ones)
+	suite.Equal(2, len(newR.replicaPB.GetChannelNodeInfos()))
+
+	// Verify that tryBalanceNodeForChannel was still called and nodes were balanced
+	// 4 nodes / 2 channels = 2 nodes per channel
+	totalAssignedNodes := 0
+	for _, channelNodeInfo := range newR.replicaPB.GetChannelNodeInfos() {
+		suite.Equal(2, len(channelNodeInfo.GetRwNodes()))
+		totalAssignedNodes += len(channelNodeInfo.GetRwNodes())
+	}
+	suite.Equal(4, totalAssignedNodes)
+}
+
+// TestTryEnableChannelExclusiveModeInsufficientNodes tests that TryEnableChannelExclusiveMode
+// properly handles the case where there are not enough nodes for exclusive mode.
+func (suite *ReplicaSuite) TestTryEnableChannelExclusiveModeInsufficientNodes() {
+	paramtable.Get().Save(paramtable.Get().QueryCoordCfg.Balancer.Key, ChannelLevelScoreBalancerName)
+	paramtable.Get().Save(paramtable.Get().QueryCoordCfg.ChannelExclusiveNodeFactor.Key, "3")
+	defer func() {
+		paramtable.Get().Reset(paramtable.Get().QueryCoordCfg.Balancer.Key)
+		paramtable.Get().Reset(paramtable.Get().QueryCoordCfg.ChannelExclusiveNodeFactor.Key)
+	}()
+
+	// 2 nodes for 2 channels with factor 3: need 6 nodes, only have 2
+	r := newReplica(&querypb.Replica{
+		ID:            1,
+		CollectionID:  2,
+		ResourceGroup: DefaultResourceGroupName,
+		Nodes:         []int64{1, 2},
+	})
+
+	mutableReplica := r.CopyForWrite()
+	mutableReplica.TryEnableChannelExclusiveMode("channel1", "channel2")
+
+	newR := mutableReplica.IntoReplica()
+
+	// With insufficient nodes, tryBalanceNodeForChannel should disable exclusive mode
+	for _, channelNodeInfo := range newR.replicaPB.GetChannelNodeInfos() {
+		suite.Equal(0, len(channelNodeInfo.GetRwNodes()))
+	}
+}
+
 func TestReplica(t *testing.T) {
 	suite.Run(t, new(ReplicaSuite))
 }

internal/querycoordv2/services_test.go

@@ -43,6 +43,7 @@ import (
 	"github.com/milvus-io/milvus/internal/mocks/distributed/mock_streaming"
 	"github.com/milvus-io/milvus/internal/mocks/streamingcoord/server/mock_balancer"
 	"github.com/milvus-io/milvus/internal/mocks/streamingcoord/server/mock_broadcaster"
+	"github.com/milvus-io/milvus/internal/querycoordv2/assign"
 	"github.com/milvus-io/milvus/internal/querycoordv2/balance"
 	"github.com/milvus-io/milvus/internal/querycoordv2/checkers"
 	"github.com/milvus-io/milvus/internal/querycoordv2/dist"
@@ -244,6 +245,8 @@ func (suite *ServiceSuite) SetupTest() {
 	suite.taskScheduler.EXPECT().GetSegmentTaskDelta(mock.Anything, mock.Anything).Return(0).Maybe()
 	suite.taskScheduler.EXPECT().GetChannelTaskDelta(mock.Anything, mock.Anything).Return(0).Maybe()
 	suite.jobScheduler.Start()
+	assign.ResetGlobalAssignPolicyFactoryForTest()
+	assign.InitGlobalAssignPolicyFactory(suite.taskScheduler, suite.nodeMgr, suite.dist, suite.meta, suite.targetMgr)
 	suite.balancer = balance.NewRowCountBasedBalancer(
 		suite.taskScheduler,
 		suite.nodeMgr,
@@ -304,10 +307,11 @@ func (suite *ServiceSuite) SetupTest() {
 		for _, collection := range suite.collections {
 			if collection == collectionID {
 				return &milvuspb.DescribeCollectionResponse{
-					DbName:         util.DefaultDBName,
-					DbId:           1,
-					CollectionID:   collectionID,
-					CollectionName: fmt.Sprintf("collection_%d", collectionID),
+					DbName:              util.DefaultDBName,
+					DbId:                1,
+					CollectionID:        collectionID,
+					CollectionName:      fmt.Sprintf("collection_%d", collectionID),
+					VirtualChannelNames: suite.channels[collectionID],
 					Schema: &schemapb.CollectionSchema{
 						Fields: []*schemapb.FieldSchema{
 							{FieldID: 100},
@@ -456,8 +460,6 @@ func (suite *ServiceSuite) TestLoadCollection() {
 
 	// Test load all collections
 	for _, collection := range suite.collections {
-		suite.broker.EXPECT().DescribeCollection(mock.Anything, mock.Anything).
-			Return(nil, nil)
 		suite.expectGetRecoverInfo(collection)
 
 		req := &querypb.LoadCollectionRequest{
@@ -992,8 +994,6 @@ func (suite *ServiceSuite) TestLoadPartition() {
 
 	// Test load all partitions
 	for _, collection := range suite.collections {
-		suite.broker.EXPECT().DescribeCollection(mock.Anything, mock.Anything).
-			Return(nil, nil)
 		suite.expectGetRecoverInfo(collection)
 
 		req := &querypb.LoadPartitionsRequest{
@@ -2079,8 +2079,6 @@ func (suite *ServiceSuite) expectGetRecoverInfo(collection int64) {
 }
 
 func (suite *ServiceSuite) expectLoadMetaRPCs() {
-	suite.broker.EXPECT().DescribeCollection(mock.Anything, mock.Anything).
-		Return(nil, nil).Maybe()
 	suite.broker.EXPECT().ListIndexes(mock.Anything, mock.Anything).
 		Return(nil, nil).Maybe()
 }
@@ -2265,6 +2263,7 @@ func (suite *ServiceSuite) TearDownTest() {
 	suite.targetObserver.Stop()
 	suite.collectionObserver.Stop()
 	suite.jobScheduler.Stop()
+	assign.ResetGlobalAssignPolicyFactoryForTest()
 }
 
 func TestService(t *testing.T) {