Clean flowGraph if watchChannel failed

internal/querynode/collection_replica.go

@@ -53,7 +53,7 @@ type ReplicaInterface interface {
 	// getCollectionIDs returns all collection ids in the collectionReplica
 	getCollectionIDs() []UniqueID
 	// addCollection creates a new collection and add it to collectionReplica
-	addCollection(collectionID UniqueID, schema *schemapb.CollectionSchema) error
+	addCollection(collectionID UniqueID, schema *schemapb.CollectionSchema) *Collection
 	// removeCollection removes the collection from collectionReplica
 	removeCollection(collectionID UniqueID) error
 	// getCollectionByID gets the collection which id is collectionID
@@ -199,18 +199,18 @@ func (colReplica *collectionReplica) getCollectionIDs() []UniqueID {
 }
 
 // addCollection creates a new collection and add it to collectionReplica
-func (colReplica *collectionReplica) addCollection(collectionID UniqueID, schema *schemapb.CollectionSchema) error {
+func (colReplica *collectionReplica) addCollection(collectionID UniqueID, schema *schemapb.CollectionSchema) *Collection {
 	colReplica.mu.Lock()
 	defer colReplica.mu.Unlock()
 
-	if ok := colReplica.hasCollectionPrivate(collectionID); ok {
-		return fmt.Errorf("collection has been loaded, id %d", collectionID)
+	if col, ok := colReplica.collections[collectionID]; ok {
+		return col
 	}
 
 	var newCollection = newCollection(collectionID, schema)
 	colReplica.collections[collectionID] = newCollection
 
-	return nil
+	return newCollection
 }
 
 // removeCollection removes the collection from collectionReplica

internal/querynode/data_sync_service.go

@@ -42,16 +42,18 @@ type dataSyncService struct {
 }
 
 // addFlowGraphsForDMLChannels add flowGraphs to dmlChannel2FlowGraph
-func (dsService *dataSyncService) addFlowGraphsForDMLChannels(collectionID UniqueID, dmlChannels []string) {
+func (dsService *dataSyncService) addFlowGraphsForDMLChannels(collectionID UniqueID, dmlChannels []string) map[string]*queryNodeFlowGraph {
 	dsService.mu.Lock()
 	defer dsService.mu.Unlock()
 
+	results := make(map[string]*queryNodeFlowGraph)
 	for _, channel := range dmlChannels {
-		if _, ok := dsService.dmlChannel2FlowGraph[channel]; ok {
+		if fg, ok := dsService.dmlChannel2FlowGraph[channel]; ok {
 			log.Warn("dml flow graph has been existed",
 				zap.Any("collectionID", collectionID),
 				zap.Any("channel", channel),
 			)
+			results[channel] = fg
 			continue
 		}
 		newFlowGraph := newQueryNodeFlowGraph(dsService.ctx,
@@ -64,20 +66,25 @@ func (dsService *dataSyncService) addFlowGraphsForDMLChannels(collectionID Uniqu
 		log.Debug("add DML flow graph",
 			zap.Any("collectionID", collectionID),
 			zap.Any("channel", channel))
+		results[channel] = newFlowGraph
 	}
+
+	return results
 }
 
 // addFlowGraphsForDeltaChannels add flowGraphs to deltaChannel2FlowGraph
-func (dsService *dataSyncService) addFlowGraphsForDeltaChannels(collectionID UniqueID, deltaChannels []string) {
+func (dsService *dataSyncService) addFlowGraphsForDeltaChannels(collectionID UniqueID, deltaChannels []string) map[string]*queryNodeFlowGraph {
 	dsService.mu.Lock()
 	defer dsService.mu.Unlock()
 
+	results := make(map[string]*queryNodeFlowGraph)
 	for _, channel := range deltaChannels {
-		if _, ok := dsService.deltaChannel2FlowGraph[channel]; ok {
+		if fg, ok := dsService.deltaChannel2FlowGraph[channel]; ok {
 			log.Warn("delta flow graph has been existed",
 				zap.Any("collectionID", collectionID),
 				zap.Any("channel", channel),
 			)
+			results[channel] = fg
 			continue
 		}
 		newFlowGraph := newQueryNodeDeltaFlowGraph(dsService.ctx,
@@ -90,7 +97,10 @@ func (dsService *dataSyncService) addFlowGraphsForDeltaChannels(collectionID Uni
 		log.Debug("add delta flow graph",
 			zap.Any("collectionID", collectionID),
 			zap.Any("channel", channel))
+		results[channel] = newFlowGraph
 	}
+
+	return results
 }
 
 // getFlowGraphByDMLChannel returns the DML flowGraph by channel

internal/querynode/mock_test.go

@@ -869,10 +869,7 @@ func genSimpleReplica() (ReplicaInterface, error) {
 	}
 	r := newCollectionReplica(kv)
 	schema := genSimpleSegCoreSchema()
-	err = r.addCollection(defaultCollectionID, schema)
-	if err != nil {
-		return nil, err
-	}
+	r.addCollection(defaultCollectionID, schema)
 	err = r.addPartition(defaultCollectionID, defaultPartitionID)
 	return r, err
 }

internal/querynode/query_collection_test.go

@@ -144,8 +144,7 @@ func TestQueryCollection_withoutVChannel(t *testing.T) {
 	historical := newHistorical(context.Background(), historicalReplica, tsReplica)
 
 	//add a segment to historical data
-	err = historical.replica.addCollection(0, schema)
-	assert.Nil(t, err)
+	historical.replica.addCollection(0, schema)
 	err = historical.replica.addPartition(0, 1)
 	assert.Nil(t, err)
 	err = historical.replica.addSegment(2, 1, 0, "testChannel", segmentTypeSealed, true)
@@ -168,8 +167,7 @@ func TestQueryCollection_withoutVChannel(t *testing.T) {
 
 	//create a streaming
 	streaming := newStreaming(ctx, streamingReplica, factory, etcdKV, tsReplica)
-	err = streaming.replica.addCollection(0, schema)
-	assert.Nil(t, err)
+	streaming.replica.addCollection(0, schema)
 	err = streaming.replica.addPartition(0, 1)
 	assert.Nil(t, err)
 

internal/querynode/query_node_test.go

@@ -147,8 +147,7 @@ func initTestMeta(t *testing.T, node *QueryNode, collectionID UniqueID, segmentI
 	}
 	collectionMeta := genTestCollectionMeta(collectionID, isBinary)
 
-	var err = node.historical.replica.addCollection(collectionMeta.ID, collectionMeta.Schema)
-	assert.NoError(t, err)
+	node.historical.replica.addCollection(collectionMeta.ID, collectionMeta.Schema)
 
 	collection, err := node.historical.replica.getCollectionByID(collectionID)
 	assert.NoError(t, err)

internal/querynode/segment_loader_test.go

@@ -270,8 +270,7 @@ func TestSegmentLoader_invalid(t *testing.T) {
 				}),
 			},
 		}
-		err = loader.historicalReplica.addCollection(defaultCollectionID, schema)
-		assert.NoError(t, err)
+		loader.historicalReplica.addCollection(defaultCollectionID, schema)
 
 		req := &querypb.LoadSegmentsRequest{
 			Base: &commonpb.MsgBase{