enhance: reduce quota and SetRates allocations

internal/datacoord/meta.go

@@ -874,9 +874,13 @@ func (m *meta) GetQuotaInfo() *metricsinfo.DataCoordQuotaMetrics {
 
 	segments := m.segments.GetSegments()
 	var total int64
-	storedBinlogSize := make(map[string]map[string]int64) // map[collectionID]map[segment_state]size
-	binlogFileCount := make(map[string]int64)             // map[collectionID]count
+	// In aggregate mode these are keyed directly by database / all, rather
+	// than building and discarding collection-level display metrics first.
+	storedBinlogSize := make(map[string]map[string]int64)
+	binlogFileCount := make(map[string]int64)
 	coll2DbName := make(map[string]string)
+	storedRowsByDB := make(map[string]map[commonpb.SegmentState]int64)
+	l0DeleteEntriesByDB := make(map[string]int64)
 
 	for _, segment := range segments {
 		segmentSize := segment.getSegmentSize()
@@ -893,22 +897,34 @@ func (m *meta) GetQuotaInfo() *metricsinfo.DataCoordQuotaMetrics {
 
 			coll, ok := m.collections.Get(segment.GetCollectionID())
 			if ok {
-				collIDStr := strconv.FormatInt(segment.GetCollectionID(), 10)
-				coll2DbName[collIDStr] = coll.DatabaseName
-				if _, ok := storedBinlogSize[collIDStr]; !ok {
-					storedBinlogSize[collIDStr] = make(map[string]int64)
+				metricKey, fileCountKey := coll.DatabaseName, metrics.AllLabel
+				if !aggregateCollectionMetrics {
+					metricKey = strconv.FormatInt(segment.GetCollectionID(), 10)
+					fileCountKey = metricKey
+					coll2DbName[metricKey] = coll.DatabaseName
+				}
+				if _, ok := storedBinlogSize[metricKey]; !ok {
+					storedBinlogSize[metricKey] = make(map[string]int64)
+				}
+				storedBinlogSize[metricKey][segment.GetState().String()] += segmentSize
+				binlogFileCount[fileCountKey] += int64(getBinlogFileCount(segment.SegmentInfo))
+				if aggregateCollectionMetrics {
+					if _, ok := storedRowsByDB[coll.DatabaseName]; !ok {
+						storedRowsByDB[coll.DatabaseName] = make(map[commonpb.SegmentState]int64)
+					}
+					storedRowsByDB[coll.DatabaseName][segment.GetState()] += segment.GetNumOfRows()
+					if segment.GetLevel() == datapb.SegmentLevel_L0 {
+						l0DeleteEntriesByDB[coll.DatabaseName] += segment.getDeltaCount()
+					}
 				}
-
-				storedBinlogSize[collIDStr][segment.GetState().String()] += segmentSize
-				binlogFileCount[collIDStr] += int64(getBinlogFileCount(segment.SegmentInfo))
-				// } else {
-				// log.Ctx(context.TODO()).Warn("not found database name", zap.Int64("collectionID", segment.GetCollectionID()))
 			}
 
-			if _, ok := collectionRowsNum[segment.GetCollectionID()]; !ok {
-				collectionRowsNum[segment.GetCollectionID()] = make(map[commonpb.SegmentState]int64)
+			if !aggregateCollectionMetrics {
+				if _, ok := collectionRowsNum[segment.GetCollectionID()]; !ok {
+					collectionRowsNum[segment.GetCollectionID()] = make(map[commonpb.SegmentState]int64)
+				}
+				collectionRowsNum[segment.GetCollectionID()][segment.GetState()] += segment.GetNumOfRows()
 			}
-			collectionRowsNum[segment.GetCollectionID()][segment.GetState()] += segment.GetNumOfRows()
 
 			if segment.GetLevel() == datapb.SegmentLevel_L0 {
 				collectionL0RowCounts[segment.GetCollectionID()] += segment.getDeltaCount()
@@ -919,17 +935,7 @@ func (m *meta) GetQuotaInfo() *metricsinfo.DataCoordQuotaMetrics {
 	// Reset to remove dropped collection
 	metrics.DataCoordStoredBinlogSize.Reset()
 	if aggregateCollectionMetrics {
-		storedBinlogSizeByDB := make(map[string]map[string]int64)
-		for collectionID, state2Size := range storedBinlogSize {
-			dbName := coll2DbName[collectionID]
-			if _, ok := storedBinlogSizeByDB[dbName]; !ok {
-				storedBinlogSizeByDB[dbName] = make(map[string]int64)
-			}
-			for state, size := range state2Size {
-				storedBinlogSizeByDB[dbName][state] += size
-			}
-		}
-		for dbName, state2Size := range storedBinlogSizeByDB {
+		for dbName, state2Size := range storedBinlogSize {
 			for state, size := range state2Size {
 				metrics.DataCoordStoredBinlogSize.WithLabelValues(dbName, metrics.AllLabel, state).Set(float64(size))
 			}
@@ -944,11 +950,7 @@ func (m *meta) GetQuotaInfo() *metricsinfo.DataCoordQuotaMetrics {
 	// Reset to remove dropped collection
 	metrics.DataCoordSegmentBinLogFileCount.Reset()
 	if aggregateCollectionMetrics {
-		var totalBinlogFileCount int64
-		for _, size := range binlogFileCount {
-			totalBinlogFileCount += size
-		}
-		metrics.DataCoordSegmentBinLogFileCount.WithLabelValues(metrics.AllLabel).Set(float64(totalBinlogFileCount))
+		metrics.DataCoordSegmentBinLogFileCount.WithLabelValues(metrics.AllLabel).Set(float64(binlogFileCount[metrics.AllLabel]))
 	} else {
 		for collectionID, size := range binlogFileCount {
 			metrics.DataCoordSegmentBinLogFileCount.WithLabelValues(collectionID).Set(float64(size))
@@ -957,19 +959,6 @@ func (m *meta) GetQuotaInfo() *metricsinfo.DataCoordQuotaMetrics {
 
 	metrics.DataCoordNumStoredRows.Reset()
 	if aggregateCollectionMetrics {
-		storedRowsByDB := make(map[string]map[commonpb.SegmentState]int64)
-		for collectionID, statesRows := range collectionRowsNum {
-			coll, ok := m.collections.Get(collectionID)
-			if !ok {
-				continue
-			}
-			if _, ok := storedRowsByDB[coll.DatabaseName]; !ok {
-				storedRowsByDB[coll.DatabaseName] = make(map[commonpb.SegmentState]int64)
-			}
-			for state, rows := range statesRows {
-				storedRowsByDB[coll.DatabaseName][state] += rows
-			}
-		}
 		for dbName, statesRows := range storedRowsByDB {
 			for state, rows := range statesRows {
 				metrics.DataCoordNumStoredRows.WithLabelValues(
@@ -989,13 +978,6 @@ func (m *meta) GetQuotaInfo() *metricsinfo.DataCoordQuotaMetrics {
 
 	metrics.DataCoordL0DeleteEntriesNum.Reset()
 	if aggregateCollectionMetrics {
-		l0DeleteEntriesByDB := make(map[string]int64)
-		for collectionID, entriesNum := range collectionL0RowCounts {
-			coll, ok := m.collections.Get(collectionID)
-			if ok {
-				l0DeleteEntriesByDB[coll.DatabaseName] += entriesNum
-			}
-		}
 		for dbName, entriesNum := range l0DeleteEntriesByDB {
 			metrics.DataCoordL0DeleteEntriesNum.WithLabelValues(dbName, metrics.AllLabel).Set(float64(entriesNum))
 		}

internal/datacoord/meta_test.go

@@ -3894,7 +3894,7 @@ func TestGetQuotaInfoAggregatesCollectionMetrics(t *testing.T) {
 	segment1 := buildSegment(1, 10, 100, "channel-1")
 	segment1.NumOfRows = 10
 	segment1.Level = datapb.SegmentLevel_L0
-	segment1.Stats = &datapb.Statistics{InsertBinlogSize: 100, DeleteNumRows: 3}
+	segment1.Stats = &datapb.Statistics{InsertBinlogSize: 100, InsertBinlogCount: 1, DeleteNumRows: 3}
 	segment1.Binlogs = []*datapb.FieldBinlog{{
 		FieldID: 1,
 		Binlogs: []*datapb.Binlog{{LogID: 1}},
@@ -3904,7 +3904,7 @@ func TestGetQuotaInfoAggregatesCollectionMetrics(t *testing.T) {
 	segment2 := buildSegment(2, 20, 200, "channel-2")
 	segment2.NumOfRows = 20
 	segment2.Level = datapb.SegmentLevel_L0
-	segment2.Stats = &datapb.Statistics{InsertBinlogSize: 200, DeleteNumRows: 4}
+	segment2.Stats = &datapb.Statistics{InsertBinlogSize: 200, InsertBinlogCount: 2, DeleteNumRows: 4}
 	segment2.Binlogs = []*datapb.FieldBinlog{{
 		FieldID: 1,
 		Binlogs: []*datapb.Binlog{{LogID: 2}, {LogID: 3}},

internal/datacoord/quota_aggregation_test.go

@@ -0,0 +1,166 @@
+// 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 datacoord
+
+import (
+	"strconv"
+	"testing"
+
+	"github.com/prometheus/client_golang/prometheus/testutil"
+	"github.com/stretchr/testify/require"
+
+	"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
+	"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
+	"github.com/milvus-io/milvus/pkg/v3/metrics"
+	"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
+	"github.com/milvus-io/milvus/pkg/v3/util/metricsinfo"
+	"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
+)
+
+func resetQuotaDisplayMetrics() {
+	metrics.DataCoordStoredBinlogSize.Reset()
+	metrics.DataCoordSegmentBinLogFileCount.Reset()
+	metrics.DataCoordNumStoredRows.Reset()
+	metrics.DataCoordL0DeleteEntriesNum.Reset()
+}
+
+func TestQuotaInfoAggregatePreservesControlData(t *testing.T) {
+	previous := metrics.CollectionLevelMetricsMode()
+	t.Cleanup(func() {
+		metrics.SetCollectionLevelMetricsMode(previous)
+		resetQuotaDisplayMetrics()
+	})
+	m := &meta{
+		collections: typeutil.NewConcurrentMap[int64, *collectionInfo](),
+		segments:    NewCachedSegmentsInfo(),
+	}
+	for id, db := range map[int64]string{1: "db-a", 2: "db-a", 3: "db-b"} {
+		m.collections.Insert(id, &collectionInfo{
+			ID: id, DatabaseName: db,
+			Schema: &schemapb.CollectionSchema{Name: strconv.FormatInt(id, 10)},
+		})
+	}
+	for _, fixture := range []struct {
+		id, collection, partition, rows, size, files, deletes int64
+		state                                                 commonpb.SegmentState
+		level                                                 datapb.SegmentLevel
+		importing                                             bool
+	}{
+		{10, 1, 11, 10, 100, 1, 3, commonpb.SegmentState_Growing, datapb.SegmentLevel_L0, false},
+		{11, 1, 12, 5, 50, 2, 0, commonpb.SegmentState_Flushed, datapb.SegmentLevel_L1, false},
+		{20, 2, 21, 20, 200, 2, 4, commonpb.SegmentState_Growing, datapb.SegmentLevel_L0, false},
+		{30, 3, 31, 7, 70, 1, 0, commonpb.SegmentState_Growing, datapb.SegmentLevel_L1, false},
+		// Quota must still count healthy segments whose collection is not cached.
+		{40, 999, 41, 9, 90, 1, 5, commonpb.SegmentState_Growing, datapb.SegmentLevel_L0, false},
+		{50, 1, 11, 100, 1000, 10, 0, commonpb.SegmentState_Growing, datapb.SegmentLevel_L1, true},
+		{60, 1, 11, 200, 2000, 20, 0, commonpb.SegmentState_Dropped, datapb.SegmentLevel_L1, false},
+	} {
+		segment := buildSegment(fixture.collection, fixture.partition, fixture.id, "channel")
+		segment.NumOfRows = fixture.rows
+		segment.State = fixture.state
+		segment.Level = fixture.level
+		segment.IsImporting = fixture.importing
+		segment.Stats = &datapb.Statistics{
+			InsertBinlogSize:  fixture.size,
+			InsertBinlogCount: fixture.files, DeleteNumRows: fixture.deletes,
+		}
+		m.segments.SetSegment(fixture.id, segment, 1)
+	}
+	expected := &metricsinfo.DataCoordQuotaMetrics{
+		TotalBinlogSize:      510,
+		CollectionBinlogSize: map[int64]int64{1: 150, 2: 200, 3: 70, 999: 90},
+		PartitionsBinlogSize: map[int64]map[int64]int64{
+			1: {11: 100, 12: 50}, 2: {21: 200}, 3: {31: 70}, 999: {41: 90},
+		},
+		CollectionL0RowCount: map[int64]int64{1: 3, 2: 4, 999: 5},
+	}
+	// Repeat mode changes: stale collection-level series must not leak into all.
+	for _, mode := range []string{
+		metrics.CollectionLevelMetricsModeFull,
+		metrics.CollectionLevelMetricsModeAggregate, metrics.CollectionLevelMetricsModeFull,
+		metrics.CollectionLevelMetricsModeAggregate,
+	} {
+		metrics.SetCollectionLevelMetricsMode(mode)
+		require.Equal(t, expected, m.GetQuotaInfo())
+		if mode == metrics.CollectionLevelMetricsModeAggregate {
+			require.Equal(t, 3, testutil.CollectAndCount(metrics.DataCoordStoredBinlogSize))
+			require.Equal(t, 1, testutil.CollectAndCount(metrics.DataCoordSegmentBinLogFileCount))
+			require.Equal(t, 3, testutil.CollectAndCount(metrics.DataCoordNumStoredRows))
+			require.Equal(t, 1, testutil.CollectAndCount(metrics.DataCoordL0DeleteEntriesNum))
+			require.Equal(t, float64(300), testutil.ToFloat64(metrics.DataCoordStoredBinlogSize.WithLabelValues(
+				"db-a", metrics.AllLabel, commonpb.SegmentState_Growing.String())))
+			require.Equal(t, float64(6), testutil.ToFloat64(metrics.DataCoordSegmentBinLogFileCount.WithLabelValues(metrics.AllLabel)))
+			require.Equal(t, float64(30), testutil.ToFloat64(metrics.DataCoordNumStoredRows.WithLabelValues(
+				"db-a", metrics.AllLabel, metrics.AllLabel, commonpb.SegmentState_Growing.String())))
+			require.Equal(t, float64(7), testutil.ToFloat64(metrics.DataCoordL0DeleteEntriesNum.WithLabelValues("db-a", metrics.AllLabel)))
+		} else {
+			require.Equal(t, 4, testutil.CollectAndCount(metrics.DataCoordStoredBinlogSize))
+			require.Equal(t, 3, testutil.CollectAndCount(metrics.DataCoordSegmentBinLogFileCount))
+			require.Equal(t, 4, testutil.CollectAndCount(metrics.DataCoordNumStoredRows))
+			require.Equal(t, 2, testutil.CollectAndCount(metrics.DataCoordL0DeleteEntriesNum))
+			require.Equal(t, float64(100), testutil.ToFloat64(metrics.DataCoordStoredBinlogSize.WithLabelValues(
+				"db-a", "1", commonpb.SegmentState_Growing.String())))
+		}
+	}
+
+	// Metadata removal clears display series, but never changes disk admission data.
+	for _, id := range []int64{1, 2, 3} {
+		m.collections.Remove(id)
+	}
+	require.Equal(t, expected, m.GetQuotaInfo())
+	require.Zero(t, testutil.CollectAndCount(metrics.DataCoordStoredBinlogSize))
+	require.Zero(t, testutil.CollectAndCount(metrics.DataCoordNumStoredRows))
+	require.Zero(t, testutil.CollectAndCount(metrics.DataCoordL0DeleteEntriesNum))
+	require.Equal(t, float64(0), testutil.ToFloat64(metrics.DataCoordSegmentBinLogFileCount.WithLabelValues(metrics.AllLabel)))
+	for _, segment := range m.segments.GetSegments() {
+		m.segments.DropSegment(segment.GetID(), 2)
+	}
+	empty := m.GetQuotaInfo()
+	require.Zero(t, empty.TotalBinlogSize)
+	require.Empty(t, empty.CollectionBinlogSize)
+	require.Empty(t, empty.PartitionsBinlogSize)
+	require.Empty(t, empty.CollectionL0RowCount)
+}
+
+func BenchmarkQuotaInfoAggregate(b *testing.B) {
+	previous := metrics.CollectionLevelMetricsMode()
+	metrics.SetCollectionLevelMetricsMode(metrics.CollectionLevelMetricsModeAggregate)
+	b.Cleanup(func() {
+		metrics.SetCollectionLevelMetricsMode(previous)
+		resetQuotaDisplayMetrics()
+	})
+	for _, collections := range []int{10000, 100000} {
+		m := &meta{
+			collections: typeutil.NewConcurrentMap[int64, *collectionInfo](),
+			segments:    NewCachedSegmentsInfo(),
+		}
+		for i := range collections {
+			id := int64(i + 1)
+			m.collections.Insert(id, &collectionInfo{ID: id, DatabaseName: "quota-db"})
+			segment := buildSegment(id, id, id, "channel")
+			segment.NumOfRows = 100
+			segment.Stats = &datapb.Statistics{InsertBinlogSize: 1024, InsertBinlogCount: 1}
+			m.segments.SetSegment(id, segment, 1)
+		}
+		b.Run(strconv.Itoa(collections), func(b *testing.B) {
+			b.ReportAllocs()
+			for b.Loop() {
+				m.GetQuotaInfo()
+			}
+		})
+	}
+}

internal/distributed/proxy/client/client.go

@@ -17,6 +17,7 @@
 package grpcproxyclient
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 
@@ -181,7 +182,7 @@ func (c *Client) GetProxyMetrics(ctx context.Context, req *milvuspb.GetMetricsRe
 
 // SetRates notifies Proxy to limit rates of requests.
 func (c *Client) SetRates(ctx context.Context, req *proxypb.SetRatesRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
-	req = typeutil.Clone(req)
+	req = cloneSetRatesEnvelope(req)
 	commonpbutil.UpdateMsgBase(
 		req.GetBase(),
 		commonpbutil.FillMsgBaseFromClient(paramtable.GetNodeID(), commonpbutil.WithTargetID(c.grpcClient.GetNodeID())),
@@ -191,6 +192,22 @@ func (c *Client) SetRates(ctx context.Context, req *proxypb.SetRatesRequest, opt
 	})
 }
 
+// cloneSetRatesEnvelope isolates the per-client routing header. The quota
+// snapshot is read-only for the entire fan-out, including retries and encoding.
+// Do not copy the generated struct: it contains protobuf runtime state.
+func cloneSetRatesEnvelope(req *proxypb.SetRatesRequest) *proxypb.SetRatesRequest {
+	if req == nil {
+		return nil
+	}
+	cloned := &proxypb.SetRatesRequest{
+		Base:        typeutil.Clone(req.GetBase()),
+		Rates:       req.GetRates(),
+		RootLimiter: req.GetRootLimiter(),
+	}
+	cloned.ProtoReflect().SetUnknown(bytes.Clone(req.ProtoReflect().GetUnknown()))
+	return cloned
+}
+
 func (c *Client) ListClientInfos(ctx context.Context, req *proxypb.ListClientInfosRequest, opts ...grpc.CallOption) (*proxypb.ListClientInfosResponse, error) {
 	req = typeutil.Clone(req)
 	commonpbutil.UpdateMsgBase(