enhance: Set loaded collection/partition number to metrics (#38271)

Related to #36456 Previous PR: #38471 #38233 Signed-off-by: Congqi Xia <[email protected]>
milvus-io · Dec 6, 2024 · 36946cc · 36946cc
1 parent f16232f
commit 36946cc
Showing 1 changed file with 8 additions and 5 deletions.
diff --git a/internal/querycoordv2/meta/collection_manager.go b/internal/querycoordv2/meta/collection_manager.go
@@ -543,6 +543,11 @@ func (m *CollectionManager) putPartition(ctx context.Context, partitions []*Part
 	return nil
 }
 
+func (m *CollectionManager) updateLoadMetrics() {
+	metrics.QueryCoordNumCollections.WithLabelValues().Set(float64(len(lo.Filter(lo.Values(m.collections), func(coll *Collection, _ int) bool { return coll.LoadPercentage == 100 }))))
+	metrics.QueryCoordNumPartitions.WithLabelValues().Set(float64(len(lo.Filter(lo.Values(m.partitions), func(part *Partition, _ int) bool { return part.LoadPercentage == 100 }))))
+}
+
 func (m *CollectionManager) UpdateLoadPercent(ctx context.Context, partitionID int64, loadPercent int32) (int32, error) {
 	m.rwmutex.Lock()
 	defer m.rwmutex.Unlock()
@@ -590,8 +595,7 @@ func (m *CollectionManager) UpdateLoadPercent(ctx context.Context, partitionID i
 		// if collection becomes loaded, clear it's recoverTimes in load info
 		newCollection.RecoverTimes = 0
 
-		metrics.QueryCoordNumCollections.WithLabelValues().Set(float64(len(lo.Values(m.collections))))
-		metrics.QueryCoordNumPartitions.WithLabelValues().Set(float64(len(m.partitions)))
+		defer m.updateLoadMetrics()
 		elapsed := time.Since(newCollection.CreatedAt)
 		metrics.QueryCoordLoadLatency.WithLabelValues().Observe(float64(elapsed.Milliseconds()))
 		eventlog.Record(eventlog.NewRawEvt(eventlog.Level_Info, fmt.Sprintf("Collection %d loaded", newCollection.CollectionID)))
@@ -617,8 +621,7 @@ func (m *CollectionManager) RemoveCollection(ctx context.Context, collectionID t
 		delete(m.collectionPartitions, collectionID)
 	}
 	metrics.CleanQueryCoordMetricsWithCollectionID(collectionID)
-	metrics.QueryCoordNumPartitions.WithLabelValues().Set(float64(len(m.partitions)))
-	metrics.QueryCoordNumCollections.WithLabelValues().Set(float64(len(lo.Values(m.collections))))
+	m.updateLoadMetrics()
 	return nil
 }
 
@@ -631,7 +634,6 @@ func (m *CollectionManager) RemovePartition(ctx context.Context, collectionID ty
 	defer m.rwmutex.Unlock()
 
 	err := m.removePartition(ctx, collectionID, partitionIDs...)
-	metrics.QueryCoordNumPartitions.WithLabelValues().Set(float64(len(m.partitions)))
 	return err
 }
 
@@ -645,6 +647,7 @@ func (m *CollectionManager) removePartition(ctx context.Context, collectionID ty
 		delete(m.partitions, id)
 		delete(partitions, id)
 	}
+	m.updateLoadMetrics()
 
 	return nil
 }