Fix create collection failed to assign channels (#19551)

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

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

internal/rootcoord/create_collection_task.go

@@ -128,11 +128,15 @@ func (t *createCollectionTask) assignPartitionID() error {
 }
 
 func (t *createCollectionTask) assignChannels() error {
-	vchanNames := make([]string, t.Req.ShardsNum)
+	vchanNames := make([]string, t.Req.GetShardsNum())
 	//physical channel names
-	chanNames := t.core.chanTimeTick.getDmlChannelNames(int(t.Req.ShardsNum))
+	chanNames := t.core.chanTimeTick.getDmlChannelNames(int(t.Req.GetShardsNum()))
 
-	for i := int32(0); i < t.Req.ShardsNum; i++ {
+	if int32(len(chanNames)) < t.Req.GetShardsNum() {
+		return fmt.Errorf("no enough channels, want: %d, got: %d", t.Req.GetShardsNum(), len(chanNames))
+	}
+
+	for i := int32(0); i < t.Req.GetShardsNum(); i++ {
 		vchanNames[i] = fmt.Sprintf("%s_%dv%d", chanNames[i], t.collID, i)
 	}
 	t.channels = collectionChannels{

internal/rootcoord/create_collection_task_test.go

@@ -218,10 +218,6 @@ func Test_createCollectionTask_Prepare(t *testing.T) {
 		assert.Error(t, err)
 	})
 
-	t.Run("failed to assign channels", func(t *testing.T) {
-		// TODO: error won't happen here.
-	})
-
 	t.Run("normal case", func(t *testing.T) {
 		defer cleanTestEnv()
 
@@ -251,6 +247,10 @@ func Test_createCollectionTask_Prepare(t *testing.T) {
 				Schema:         marshaledSchema,
 			},
 		}
+		task.Req.ShardsNum = int32(Params.RootCoordCfg.DmlChannelNum + 1) // no enough channels.
+		err = task.Prepare(context.Background())
+		assert.Error(t, err)
+		task.Req.ShardsNum = 1
 		err = task.Prepare(context.Background())
 		assert.NoError(t, err)
 	})

internal/rootcoord/dml_channels.go

@@ -180,11 +180,11 @@ func newDmlChannels(ctx context.Context, factory msgstream.Factory, chanNamePref
 }
 
 func (d *dmlChannels) getChannelNames(count int) []string {
+	d.mut.Lock()
+	defer d.mut.Unlock()
 	if count > len(d.channelsHeap) {
 		return nil
 	}
-	d.mut.Lock()
-	defer d.mut.Unlock()
 	// get next count items from heap
 	items := make([]*dmlMsgStream, 0, count)
 	result := make([]string, 0, count)

internal/rootcoord/dml_channels_test.go

@@ -149,6 +149,9 @@ func TestDmlChannels(t *testing.T) {
 	dml.addChannels(chans1...)
 	assert.Equal(t, 2, dml.getChannelNum())
 
+	chans2 := dml.getChannelNames(totalDmlChannelNum + 1)
+	assert.Nil(t, chans2)
+
 	dml.removeChannels(chans1...)
 	assert.Equal(t, 2, dml.getChannelNum())