Release before remove when releasing flowgraphs (#27191)

GetAndRemove removes the fg from manager immediately,
while the flowgraph is still releasing. This PR will remove
the fg from flowgraphManager AFTER flowgraphs released.

- Add Remove for ConcurrentMap
- Move collections() into flowgraph manager

Signed-off-by: yangxuan <xuan.yang@zilliz.com>

internal/datanode/flow_graph_manager.go

@@ -137,11 +137,13 @@ func (fm *flowgraphManager) addAndStart(dn *DataNode, vchan *datapb.VchannelInfo
 }
 
 func (fm *flowgraphManager) release(vchanName string) {
-	if fg, loaded := fm.flowgraphs.GetAndRemove(vchanName); loaded {
+	if fg, loaded := fm.flowgraphs.Get(vchanName); loaded {
 		fg.close()
+		fm.flowgraphs.Remove(vchanName)
+
 		metrics.DataNodeNumFlowGraphs.WithLabelValues(fmt.Sprint(paramtable.GetNodeID())).Dec()
-	}
 		rateCol.removeFlowGraphChannel(vchanName)
+	}
 }
 
 func (fm *flowgraphManager) getFlushCh(segID UniqueID) (chan<- flushMsg, error) {
@@ -227,3 +229,13 @@ func (fm *flowgraphManager) dropAll() {
 		return true
 	})
 }
+
+func (fm *flowgraphManager) collections() []int64 {
+	collectionSet := typeutil.UniqueSet{}
+	fm.flowgraphs.Range(func(key string, value *dataSyncService) bool {
+		collectionSet.Insert(value.channel.getCollectionID())
+		return true
+	})
+
+	return collectionSet.Collect()
+}

internal/datanode/flow_graph_manager_test.go

@@ -55,6 +55,7 @@ func TestFlowGraphManager(t *testing.T) {
 	defer func() {
 		fm.dropAll()
 	}()
+
 	t.Run("Test addAndStart", func(t *testing.T) {
 		vchanName := "by-dev-rootcoord-dml-test-flowgraphmanager-addAndStart"
 		vchan := &datapb.VchannelInfo{

internal/datanode/metrics_info.go

@@ -49,15 +49,6 @@ func (node *DataNode) getQuotaMetrics() (*metricsinfo.DataNodeQuotaMetrics, erro
 		return nil, err
 	}
 
-	getAllCollections := func() []int64 {
-		collectionSet := typeutil.UniqueSet{}
-		node.flowgraphManager.flowgraphs.Range(func(key string, fg *dataSyncService) bool {
-			collectionSet.Insert(fg.channel.getCollectionID())
-			return true
-		})
-
-		return collectionSet.Collect()
-	}
 	minFGChannel, minFGTt := rateCol.getMinFlowGraphTt()
 	return &metricsinfo.DataNodeQuotaMetrics{
 		Hms: metricsinfo.HardwareMetrics{},
@@ -69,7 +60,7 @@ func (node *DataNode) getQuotaMetrics() (*metricsinfo.DataNodeQuotaMetrics, erro
 		},
 		Effect: metricsinfo.NodeEffect{
 			NodeID:        node.GetSession().ServerID,
-			CollectionIDs: getAllCollections(),
+			CollectionIDs: node.flowgraphManager.collections(),
 		},
 	}, nil
 }

pkg/util/typeutil/map.go

@@ -101,6 +101,14 @@ func (m *ConcurrentMap[K, V]) GetAndRemove(key K) (V, bool) {
 	return value.(V), true
 }
 
+// Remove removes the `key`, `value` set if `key` is in the map,
+// does nothing if `key` not in the map.
+func (m *ConcurrentMap[K, V]) Remove(key K) {
+	if _, loaded := m.inner.LoadAndDelete(key); loaded {
+		m.len.Dec()
+	}
+}
+
 func (m *ConcurrentMap[K, V]) Len() int {
 	return int(m.len.Load())
 }

pkg/util/typeutil/map_test.go

@@ -127,6 +127,23 @@ func (suite *MapUtilSuite) TestConcurrentMap() {
 		suite.FailNow("empty map range")
 		return false
 	})
+
+	suite.Run("TestRemove", func() {
+		currMap := NewConcurrentMap[int64, string]()
+		suite.Equal(0, currMap.Len())
+
+		currMap.Remove(100)
+		suite.Equal(0, currMap.Len())
+
+		suite.Equal(currMap.Len(), 0)
+		v, loaded := currMap.GetOrInsert(100, "v-100")
+		suite.Equal("v-100", v)
+		suite.Equal(false, loaded)
+		suite.Equal(1, currMap.Len())
+
+		currMap.Remove(100)
+		suite.Equal(0, currMap.Len())
+	})
 }
 
 func TestMapUtil(t *testing.T) {