diff --git a/internal/datacoord/meta.go b/internal/datacoord/meta.go index f6522e55760..f9f6860264e 100644 --- a/internal/datacoord/meta.go +++ b/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)) } diff --git a/internal/datacoord/meta_test.go b/internal/datacoord/meta_test.go index a4dd04d4cd3..bd61b799739 100644 --- a/internal/datacoord/meta_test.go +++ b/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}}, diff --git a/internal/datacoord/quota_aggregation_test.go b/internal/datacoord/quota_aggregation_test.go new file mode 100644 index 00000000000..b89385a57b1 --- /dev/null +++ b/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() + } + }) + } +} diff --git a/internal/distributed/proxy/client/client.go b/internal/distributed/proxy/client/client.go index 723586155f3..55b9e163e2d 100644 --- a/internal/distributed/proxy/client/client.go +++ b/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( diff --git a/internal/distributed/proxy/client/set_rates_snapshot_test.go b/internal/distributed/proxy/client/set_rates_snapshot_test.go new file mode 100644 index 00000000000..cf65808b61d --- /dev/null +++ b/internal/distributed/proxy/client/set_rates_snapshot_test.go @@ -0,0 +1,149 @@ +// 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 grpcproxyclient + +import ( + "context" + "strconv" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" + "google.golang.org/grpc" + "google.golang.org/grpc/encoding" + "google.golang.org/protobuf/proto" + + "github.com/milvus-io/milvus-proto/go-api/v3/commonpb" + "github.com/milvus-io/milvus/internal/mocks" + "github.com/milvus-io/milvus/pkg/v3/proto/internalpb" + "github.com/milvus-io/milvus/pkg/v3/proto/proxypb" + "github.com/milvus-io/milvus/pkg/v3/util/merr" + "github.com/milvus-io/milvus/pkg/v3/util/paramtable" +) + +func setRatesSnapshotForTest(collections int) *proxypb.SetRatesRequest { + children := make(map[int64]*proxypb.LimiterNode, collections) + for i := range collections { + children[int64(i)] = &proxypb.LimiterNode{Limiter: &proxypb.Limiter{ + Rates: []*internalpb.Rate{{Rt: internalpb.RateType_DQLSearch, R: 100}}, + }} + } + return &proxypb.SetRatesRequest{ + Base: &commonpb.MsgBase{MsgID: 42, Timestamp: 100, SourceID: 7, TargetID: 8}, + Rates: []*proxypb.CollectionRate{{Collection: 1}}, + RootLimiter: &proxypb.LimiterNode{Children: map[int64]*proxypb.LimiterNode{ + 1: {Children: children}, + }}, + } +} + +func TestCloneSetRatesEnvelope(t *testing.T) { + require.Nil(t, cloneSetRatesEnvelope(nil)) + require.Nil(t, cloneSetRatesEnvelope(&proxypb.SetRatesRequest{}).GetBase()) + request := setRatesSnapshotForTest(2) + // Unknown fields must survive old clients without sharing mutable header bytes. + request.ProtoReflect().SetUnknown([]byte{0xa0, 0x06, 0x01}) + request.Base.ProtoReflect().SetUnknown([]byte{0xa0, 0x06, 0x02}) + before := proto.Clone(request) + cloned := cloneSetRatesEnvelope(request) + require.True(t, proto.Equal(request, cloned)) + require.NotSame(t, request, cloned) + require.NotSame(t, request.Base, cloned.Base) + require.Same(t, request.RootLimiter, cloned.RootLimiter) + require.Same(t, request.Rates[0], cloned.Rates[0]) + cloned.Base.TargetID = 99 + cloned.ProtoReflect().GetUnknown()[2] = 3 + cloned.Base.ProtoReflect().GetUnknown()[2] = 4 + require.True(t, proto.Equal(before, request)) +} + +func TestSetRatesConcurrentSnapshotAndRetries(t *testing.T) { + paramtable.Init() + request := setRatesSnapshotForTest(128) + request.ProtoReflect().SetUnknown([]byte{0xa0, 0x06, 0x01}) + original := proto.Clone(request) + codec := encoding.GetCodecV2("proto") + require.NotNil(t, codec) + + const clients = 8 + var wg sync.WaitGroup + for i := range clients { + targetID := int64(i + 100) + proxy := mocks.NewMockProxyClient(t) + transport := mocks.NewMockGrpcClient[proxypb.ProxyClient](t) + transport.EXPECT().GetNodeID().Return(targetID).Once() + calls := 0 + proxy.EXPECT().SetRates(mock.Anything, mock.Anything).RunAndReturn( + func(_ context.Context, sent *proxypb.SetRatesRequest, _ ...grpc.CallOption) (*commonpb.Status, error) { + calls++ + assert.Same(t, request.RootLimiter, sent.RootLimiter) + assert.Equal(t, targetID, sent.GetBase().GetTargetID()) + assert.Equal(t, request.GetBase().GetSourceID(), sent.GetBase().GetSourceID()) + // Use the actual registered codec, including its size-cache path. + data, err := codec.Marshal(sent) + if !assert.NoError(t, err) { + return nil, err + } + defer data.Free() + decoded := &proxypb.SetRatesRequest{} + assert.NoError(t, codec.Unmarshal(data, decoded)) + assert.True(t, proto.Equal(sent, decoded)) + if calls == 1 { + return nil, context.DeadlineExceeded + } + return merr.Success(), nil + }).Twice() + transport.EXPECT().ReCall(mock.Anything, mock.Anything).RunAndReturn( + func(_ context.Context, call func(proxypb.ProxyClient) (interface{}, error)) (interface{}, error) { + _, err := call(proxy) + assert.ErrorIs(t, err, context.DeadlineExceeded) + return call(proxy) + }).Once() + client := &Client{grpcClient: transport} + wg.Add(1) + go func() { + defer wg.Done() + status, err := client.SetRates(context.Background(), request) + assert.NoError(t, err) + assert.NoError(t, merr.Error(status)) + }() + } + wg.Wait() + require.True(t, proto.Equal(original, request)) +} + +var benchmarkSetRatesRequest *proxypb.SetRatesRequest + +func BenchmarkSetRatesClone(b *testing.B) { + for _, collections := range []int{10000, 100000} { + request := setRatesSnapshotForTest(collections) + b.Run(strconv.Itoa(collections)+"/deep", func(b *testing.B) { + b.ReportAllocs() + for b.Loop() { + benchmarkSetRatesRequest = proto.Clone(request).(*proxypb.SetRatesRequest) + } + }) + b.Run(strconv.Itoa(collections)+"/envelope", func(b *testing.B) { + b.ReportAllocs() + for b.Loop() { + benchmarkSetRatesRequest = cloneSetRatesEnvelope(request) + } + }) + } +} diff --git a/internal/rootcoord/meta_table.go b/internal/rootcoord/meta_table.go index d91a1ae7257..d3a9cf56d54 100644 --- a/internal/rootcoord/meta_table.go +++ b/internal/rootcoord/meta_table.go @@ -1189,6 +1189,12 @@ func (mt *MetaTable) ListAllAvailCollections(ctx context.Context) map[int64][]in } func (mt *MetaTable) ListAllAvailPartitions(ctx context.Context) map[int64]map[int64][]int64 { + return mt.ListQuotaPartitions(ctx, true) +} + +// ListQuotaPartitions returns the same collection membership as +// ListAllAvailPartitions, without materializing partition IDs when not needed. +func (mt *MetaTable) ListQuotaPartitions(ctx context.Context, includePartitions bool) map[int64]map[int64][]int64 { mt.ddLock.RLock() defer mt.ddLock.RUnlock() @@ -1208,11 +1214,34 @@ func (mt *MetaTable) ListAllAvailPartitions(ctx context.Context) map[int64]map[i if _, ok := ret[dbID]; !ok { ret[dbID] = make(map[int64][]int64, 64) } - ret[dbID][collMeta.CollectionID] = lo.Map(collMeta.Partitions, func(part *model.Partition, _ int) int64 { return part.PartitionID }) + var partitionIDs []int64 + if includePartitions { + partitionIDs = lo.Map(collMeta.Partitions, func(part *model.Partition, _ int) int64 { return part.PartitionID }) + } + ret[dbID][collMeta.CollectionID] = partitionIDs } return ret } +// GetQuotaCollectionProperties copies only the collection properties for quota +// calculation, rather than cloning the collection and filtering its partitions. +func (mt *MetaTable) GetQuotaCollectionProperties(ctx context.Context, collectionID int64) (map[string]string, error) { + mt.ddLock.RLock() + defer mt.ddLock.RUnlock() + coll := mt.collID2Meta[collectionID] + if coll == nil || !coll.Available() { + return nil, merr.WrapErrCollectionNotFound(collectionID) + } + if len(coll.Properties) == 0 { + return nil, nil + } + properties := make(map[string]string, len(coll.Properties)) + for _, pair := range coll.Properties { + properties[pair.GetKey()] = pair.GetValue() + } + return properties, nil +} + func (mt *MetaTable) ListCollections(ctx context.Context, dbName string, ts Timestamp, onlyAvail bool) ([]*model.Collection, error) { mt.ddLock.RLock() defer mt.ddLock.RUnlock() diff --git a/internal/rootcoord/quota_allocation_test.go b/internal/rootcoord/quota_allocation_test.go new file mode 100644 index 00000000000..4546a6ae04e --- /dev/null +++ b/internal/rootcoord/quota_allocation_test.go @@ -0,0 +1,394 @@ +// 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 rootcoord + +import ( + "context" + "strconv" + "testing" + "time" + + "github.com/stretchr/testify/require" + "google.golang.org/protobuf/proto" + + "github.com/milvus-io/milvus-proto/go-api/v3/commonpb" + "github.com/milvus-io/milvus-proto/go-api/v3/milvuspb" + "github.com/milvus-io/milvus/internal/metastore/model" + "github.com/milvus-io/milvus/internal/proxy" + "github.com/milvus-io/milvus/internal/util/proxyutil" + "github.com/milvus-io/milvus/internal/util/quota" + rlinternal "github.com/milvus-io/milvus/internal/util/ratelimitutil" + "github.com/milvus-io/milvus/pkg/v3/common" + "github.com/milvus-io/milvus/pkg/v3/proto/etcdpb" + "github.com/milvus-io/milvus/pkg/v3/proto/internalpb" + "github.com/milvus-io/milvus/pkg/v3/proto/proxypb" + "github.com/milvus-io/milvus/pkg/v3/util" + "github.com/milvus-io/milvus/pkg/v3/util/merr" + "github.com/milvus-io/milvus/pkg/v3/util/paramtable" + "github.com/milvus-io/milvus/pkg/v3/util/ratelimitutil" +) + +func setQuotaTestParam(t testing.TB, param *paramtable.ParamItem, value string) { + t.Helper() + previous, err := paramtable.GetBaseTable().Load(param.Key) + require.NoError(t, Params.Save(param.Key, value)) + t.Cleanup(func() { + if err != nil { + require.NoError(t, Params.Reset(param.Key)) + } else { + require.NoError(t, Params.Save(param.Key, previous)) + } + }) +} + +func quotaMetaForTest(collections int) *MetaTable { + mt := &MetaTable{ + dbName2Meta: map[string]*model.Database{ + "quota": {ID: 10, Name: "quota"}, + "empty": {ID: 20, Name: "empty"}, + }, + collID2Meta: make(map[int64]*model.Collection, collections), + } + for i := range collections { + id := int64(i + 100) + mt.collID2Meta[id] = &model.Collection{ + DBID: 10, CollectionID: id, State: etcdpb.CollectionState_CollectionCreated, + Partitions: []*model.Partition{ + {PartitionID: id * 10, State: etcdpb.PartitionState_PartitionCreated}, + {PartitionID: id*10 + 1, State: etcdpb.PartitionState_PartitionDropping}, + }, + } + } + return mt +} + +type quotaRateSnapshot struct { + Limit Limit + Updated bool +} + +type quotaNodeSnapshot struct { + Level internalpb.RateScope + ID int64 + Rates map[internalpb.RateType]quotaRateSnapshot + States map[milvuspb.QuotaState]rlinternal.QuotaStateInfo +} + +func quotaTreeSnapshot(root *rlinternal.RateLimiterNode) map[string]quotaNodeSnapshot { + result := make(map[string]quotaNodeSnapshot) + var visit func(string, *rlinternal.RateLimiterNode) + visit = func(path string, node *rlinternal.RateLimiterNode) { + snapshot := quotaNodeSnapshot{ + Level: node.Level(), ID: node.GetID(), + Rates: make(map[internalpb.RateType]quotaRateSnapshot), + States: make(map[milvuspb.QuotaState]rlinternal.QuotaStateInfo), + } + node.GetLimiters().Range(func(rt internalpb.RateType, limiter *ratelimitutil.Limiter) bool { + snapshot.Rates[rt] = quotaRateSnapshot{limiter.Limit(), limiter.HasUpdated()} + return true + }) + node.GetQuotaStates().Range(func(state milvuspb.QuotaState, info *rlinternal.QuotaStateInfo) bool { + snapshot.States[state] = *info + return true + }) + result[path] = snapshot + node.GetChildren().Range(func(id int64, child *rlinternal.RateLimiterNode) bool { + visit(path+"/"+strconv.FormatInt(id, 10), child) + return true + }) + } + visit("", root) + return result +} + +func TestQuotaTreeReuseMatchesFreshCalculation(t *testing.T) { + setQuotaTestParam(t, &Params.QuotaConfig.DQLLimitEnabled, "true") + setQuotaTestParam(t, &Params.QuotaConfig.DMLLimitEnabled, "false") + setQuotaTestParam(t, &Params.QuotaConfig.DQLMaxSearchRate, "1000") + setQuotaTestParam(t, &Params.QuotaConfig.DQLMaxSearchRatePerDB, "800") + setQuotaTestParam(t, &Params.QuotaConfig.DQLMaxSearchRatePerCollection, "600") + setQuotaTestParam(t, &Params.QuotaConfig.DQLMaxSearchRatePerPartition, "400") + setQuotaTestParam(t, &Params.QuotaConfig.DQLMaxQueryRatePerPartition, strconv.FormatFloat(float64(Inf), 'g', -1, 64)) + + mt := quotaMetaForTest(2) + reused := NewQuotaCenter(nil, nil, nil, mt) + fresh := NewQuotaCenter(nil, nil, nil, mt) + t.Cleanup(reused.cancel) + t.Cleanup(fresh.cancel) + compareRound := func() { + t.Helper() + reused.clearMetrics() + fresh.clearMetrics() + require.NoError(t, reused.resetAllCurrentRates()) + require.NoError(t, resetQuotaFreshForTest(fresh)) + require.Equal(t, quotaTreeSnapshot(fresh.rateLimiter.GetRootLimiters()), + quotaTreeSnapshot(reused.rateLimiter.GetRootLimiters())) + } + compareRound() + root := reused.rateLimiter.GetRootLimiters() + collection := reused.rateLimiter.GetCollectionLimiters(10, 100) + partition := reused.rateLimiter.GetPartitionLimiters(10, 100, 1000) + search, ok := collection.GetLimiters().Get(internalpb.RateType_DQLSearch) + require.True(t, ok) + + // Simulate the previous calculation reducing rates and denying requests. + rlinternal.TraverseRateLimiterTree(root, func(_ internalpb.RateType, limiter *ratelimitutil.Limiter) bool { + limiter.SetLimit(0) + return true + }, nil) + for _, node := range []*rlinternal.RateLimiterNode{root, reused.rateLimiter.GetDatabaseLimiters(10), collection, partition} { + for _, state := range []milvuspb.QuotaState{milvuspb.QuotaState_DenyToWrite, milvuspb.QuotaState_DenyToRead, milvuspb.QuotaState_DenyToDDL} { + node.GetQuotaStates().Insert(state, + &rlinternal.QuotaStateInfo{ErrorCode: commonpb.ErrorCode_ForceDeny, Reason: "previous round"}) + } + } + collection.GetLimiters().Insert(internalpb.RateType(999), ratelimitutil.NewLimiter(0, 0)) + compareRound() + require.Same(t, root, reused.rateLimiter.GetRootLimiters()) + require.Same(t, collection, reused.rateLimiter.GetCollectionLimiters(10, 100)) + require.Same(t, partition, reused.rateLimiter.GetPartitionLimiters(10, 100, 1000)) + nextSearch, _ := collection.GetLimiters().Get(internalpb.RateType_DQLSearch) + require.Same(t, search, nextSearch) + + // Refresh overrides, drop membership, add a database and retain empty DBs. + mt.collID2Meta[100].Properties = []*commonpb.KeyValuePair{{Key: common.CollectionSearchRateMaxKey, Value: "123"}} + mt.collID2Meta[100].Partitions = mt.collID2Meta[100].Partitions[:1] + delete(mt.collID2Meta, 101) + mt.collID2Meta[200] = &model.Collection{DBID: 30, CollectionID: 200, State: etcdpb.CollectionState_CollectionCreated} + compareRound() + require.Nil(t, reused.rateLimiter.GetCollectionLimiters(10, 101)) + require.Nil(t, reused.rateLimiter.GetPartitionLimiters(10, 100, 1001)) + require.Equal(t, Limit(123), search.Limit()) + + // Turning partition limits off must remove old partition nodes; a finite + // collection property reverting to infinity must not retain HasUpdated. + mt.collID2Meta[100].Properties = nil + setQuotaTestParam(t, &Params.QuotaConfig.DQLMaxSearchRatePerCollection, strconv.FormatFloat(float64(Inf), 'g', -1, 64)) + setQuotaTestParam(t, &Params.QuotaConfig.DQLMaxSearchRatePerPartition, strconv.FormatFloat(float64(Inf), 'g', -1, 64)) + compareRound() + require.Zero(t, collection.GetChildren().Len()) + require.Equal(t, Inf, search.Limit()) + require.False(t, search.HasUpdated()) + + delete(mt.collID2Meta, 100) + delete(mt.dbName2Meta, "quota") + compareRound() + require.Nil(t, reused.rateLimiter.GetDatabaseLimiters(10)) + require.NotNil(t, reused.rateLimiter.GetDatabaseLimiters(20)) + // A later re-created collection and re-enabled partition limit are included. + mt.collID2Meta[100] = quotaMetaForTest(1).collID2Meta[100] + setQuotaTestParam(t, &Params.QuotaConfig.DQLLimitEnabled, "true") + setQuotaTestParam(t, &Params.QuotaConfig.DQLMaxSearchRatePerPartition, "200") + compareRound() + require.NotNil(t, reused.rateLimiter.GetPartitionLimiters(10, 100, 1000)) +} + +func TestQuotaMetadataProjection(t *testing.T) { + mt := quotaMetaForTest(3) + mt.collID2Meta[101].State = etcdpb.CollectionState_CollectionDropping + mt.collID2Meta[102].DBID = util.NonDBID + mt.collID2Meta[100].Properties = []*commonpb.KeyValuePair{ + {Key: "key", Value: "first"}, {Key: "key", Value: "last"}, + } + ctx := context.Background() + full := mt.ListAllAvailPartitions(ctx) + require.Equal(t, []int64{1000, 1001}, full[10][100]) + require.Contains(t, full[util.DefaultDBID], int64(102)) + require.NotContains(t, full[10], int64(101)) + require.Empty(t, full[20]) + projected := mt.ListQuotaPartitions(ctx, false) + for dbID, collections := range full { + require.Len(t, projected[dbID], len(collections)) + for collectionID := range collections { + require.Contains(t, projected[dbID], collectionID) + require.Nil(t, projected[dbID][collectionID]) + } + } + props, err := mt.GetQuotaCollectionProperties(ctx, 100) + require.NoError(t, err) + require.Equal(t, map[string]string{"key": "last"}, props) + props["key"] = "changed" + require.Equal(t, "last", mt.collID2Meta[100].Properties[1].Value) + empty, err := mt.GetQuotaCollectionProperties(ctx, 102) + require.NoError(t, err) + require.Nil(t, empty) + for _, id := range []int64{101, 999} { + _, err = mt.GetQuotaCollectionProperties(ctx, id) + require.ErrorIs(t, err, merr.ErrCollectionNotFound) + } +} + +func TestQuotaSnapshotDenialRecoveryAtProxy(t *testing.T) { + setQuotaTestParam(t, &Params.QuotaConfig.QuotaAndLimitsEnabled, "true") + setQuotaTestParam(t, &Params.QuotaConfig.DMLLimitEnabled, "false") + setQuotaTestParam(t, &Params.QuotaConfig.DQLLimitEnabled, "true") + setQuotaTestParam(t, &Params.QuotaConfig.DQLMaxSearchRatePerPartition, "100") + for _, scope := range []string{"cluster", "database", "collection", "partition"} { + t.Run(scope, func(t *testing.T) { + manager := proxyutil.NewMockProxyClientManager(t) + // Exactly one proxy-count snapshot for each full request, not per node. + manager.EXPECT().GetProxyCount().Return(2).Once() + manager.EXPECT().GetProxyCount().Return(3).Once() + q := NewQuotaCenter(manager, nil, nil, quotaMetaForTest(1)) + t.Cleanup(q.cancel) + require.NoError(t, q.resetAllCurrentRates()) + nodes := map[string]*rlinternal.RateLimiterNode{ + "cluster": q.rateLimiter.GetRootLimiters(), + "database": q.rateLimiter.GetDatabaseLimiters(10), + "collection": q.rateLimiter.GetCollectionLimiters(10, 100), + "partition": q.rateLimiter.GetPartitionLimiters(10, 100, 1000), + } + limiter, ok := nodes[scope].GetLimiters().Get(internalpb.RateType_DMLInsert) + require.True(t, ok) + limiter.SetLimit(0) + nodes[scope].GetQuotaStates().Insert(milvuspb.QuotaState_DenyToWrite, + &rlinternal.QuotaStateInfo{ErrorCode: commonpb.ErrorCode_DiskQuotaExhausted, Reason: "quota snapshot test"}) + assertPartitionSearchRate := func(request *proxypb.SetRatesRequest, want float64) { + t.Helper() + for _, rate := range request.RootLimiter.Children[10].Children[100].Children[1000].Limiter.Rates { + if rate.Rt == internalpb.RateType_DQLSearch { + require.InDelta(t, want, rate.R, 0.000001) + return + } + } + t.Fatal("partition search rate missing") + } + published := q.toRatesRequest() + assertPartitionSearchRate(published, 100.0/2) + original := proto.Clone(published) + admission := proxy.NewSimpleLimiter(time.Millisecond, 1) + apply := func(request *proxypb.SetRatesRequest) { + data, err := proto.Marshal(request) + require.NoError(t, err) + decoded := &proxypb.SetRatesRequest{} + require.NoError(t, proto.Unmarshal(data, decoded)) + require.NoError(t, admission.SetRates(decoded.RootLimiter)) + } + apply(published) + err := admission.Check(10, map[int64][]int64{100: {1000}}, internalpb.RateType_DMLInsert, 1) + require.ErrorContains(t, err, "quota snapshot test") + require.NoError(t, q.resetAllCurrentRates()) + require.True(t, proto.Equal(original, published), "later calculation changed a published snapshot") + recovered := q.toRatesRequest() + assertPartitionSearchRate(recovered, 100.0/3) + apply(recovered) + require.NoError(t, admission.Check(10, map[int64][]int64{100: {1000}}, internalpb.RateType_DMLInsert, 1)) + }) + } +} + +func BenchmarkQuotaReset(b *testing.B) { + setQuotaTestParam(b, &Params.QuotaConfig.DMLLimitEnabled, "false") + setQuotaTestParam(b, &Params.QuotaConfig.DQLLimitEnabled, "false") + for _, collections := range []int{10000, 100000} { + mt := quotaMetaForTest(collections) + for _, fresh := range []bool{true, false} { + name := "reuse" + if fresh { + name = "fresh" + } + b.Run(strconv.Itoa(collections)+"/"+name, func(b *testing.B) { + q := NewQuotaCenter(nil, nil, nil, mt) + defer q.cancel() + require.NoError(b, q.resetAllCurrentRates()) + b.ReportAllocs() + for b.Loop() { + clear(q.collectionProps) + if fresh { + require.NoError(b, resetQuotaFreshForTest(q)) + } else { + require.NoError(b, q.resetAllCurrentRates()) + } + } + }) + } + } +} + +// Pre-reuse calculation retained as a differential oracle and allocation +// baseline. Both paths use the same property reader; metadata savings are +// intentionally excluded from this benchmark. +func resetQuotaFreshForTest(q *QuotaCenter) error { + clusterLimiter := newParamLimiterFunc(internalpb.RateScope_Cluster, allOps)() + q.rateLimiter = rlinternal.NewRateLimiterTree(clusterLimiter) + + enablePartitionRateLimit := false + for rt := range getRateTypes(internalpb.RateScope_Partition, allOps) { + r := quota.GetQuotaValue(internalpb.RateScope_Partition, rt, Params) + if Limit(r) != Inf { + enablePartitionRateLimit = true + } + } + + // updateLimiterHasUpdated checks all limiters in a RateLimiterNode and sets hasUpdated to true + // for those with non-Inf values + updateLimiterHasUpdated := func(node *rlinternal.RateLimiterNode) { + if node == nil { + return + } + node.GetLimiters().Range(func(rateType internalpb.RateType, limiter *ratelimitutil.Limiter) bool { + if limiter.Limit() != Inf { + limiter.SetHasUpdated(true) + } + return true + }) + } + + collectionRateTypes := getRateTypes(internalpb.RateScope_Collection, allOps) + initLimiters := func(sourceCollections map[int64]map[int64][]int64) { + for dbID, collections := range sourceCollections { + for collectionID, partitionIDs := range collections { + collectionLimitVals := make(map[internalpb.RateType]Limit, collectionRateTypes.Len()) + collectionRateTypes.Range(func(rt internalpb.RateType) bool { + limitVal, err := q.getCollectionMaxLimit(rt, collectionID) + if err != nil { + limitVal = Limit(quota.GetQuotaValue(internalpb.RateScope_Collection, rt, Params)) + } + collectionLimitVals[rt] = limitVal + return true + }) + + getCollectionLimitVal := func(rateType internalpb.RateType) Limit { + return collectionLimitVals[rateType] + } + + collectionLimiter := q.rateLimiter.GetOrCreateCollectionLimiters(dbID, collectionID, + newParamLimiterFunc(internalpb.RateScope_Database, allOps), + newParamLimiterFuncWithLimitFunc(internalpb.RateScope_Collection, allOps, getCollectionLimitVal)) + updateLimiterHasUpdated(collectionLimiter) + + if !enablePartitionRateLimit { + continue + } + for _, partitionID := range partitionIDs { + partitionLimiter := q.rateLimiter.GetOrCreatePartitionLimiters(dbID, collectionID, partitionID, + newParamLimiterFunc(internalpb.RateScope_Database, allOps), + newParamLimiterFuncWithLimitFunc(internalpb.RateScope_Collection, allOps, getCollectionLimitVal), + newParamLimiterFunc(internalpb.RateScope_Partition, allOps)) + updateLimiterHasUpdated(partitionLimiter) + } + } + if len(collections) == 0 { + dbLimiter := q.rateLimiter.GetOrCreateDatabaseLimiters(dbID, newParamLimiterFunc(internalpb.RateScope_Database, allOps)) + updateLimiterHasUpdated(dbLimiter) + } + } + } + partitions := q.meta.ListAllAvailPartitions(q.ctx) + initLimiters(partitions) + return nil +} diff --git a/internal/rootcoord/quota_center.go b/internal/rootcoord/quota_center.go index 3a81267b0e4..bcd21203868 100644 --- a/internal/rootcoord/quota_center.go +++ b/internal/rootcoord/quota_center.go @@ -1391,76 +1391,115 @@ func (q *QuotaCenter) calculateEzStates() error { } func (q *QuotaCenter) resetAllCurrentRates() error { - clusterLimiter := newParamLimiterFunc(internalpb.RateScope_Cluster, allOps)() - q.rateLimiter = rlinternal.NewRateLimiterTree(clusterLimiter) - + clusterRates := quotaRateValues(internalpb.RateScope_Cluster) + databaseRates := quotaRateValues(internalpb.RateScope_Database) + collectionRates := quotaRateValues(internalpb.RateScope_Collection) + partitionRates := quotaRateValues(internalpb.RateScope_Partition) enablePartitionRateLimit := false - for rt := range getRateTypes(internalpb.RateScope_Partition, allOps) { - r := quota.GetQuotaValue(internalpb.RateScope_Partition, rt, Params) - if Limit(r) != Inf { - enablePartitionRateLimit = true - } + for _, value := range partitionRates { + enablePartitionRateLimit = enablePartitionRateLimit || value != Inf } - // updateLimiterHasUpdated checks all limiters in a RateLimiterNode and sets hasUpdated to true - // for those with non-Inf values - updateLimiterHasUpdated := func(node *rlinternal.RateLimiterNode) { - if node == nil { - return + // Only the quota loop mutates this calculation tree. Wire snapshots contain + // copied values, never aliases to the mutable limiters or quota states. + root := q.rateLimiter.GetRootLimiters() + resetQuotaLimiter(root, clusterRates, false) + partitions := q.quotaPartitionSnapshot(enablePartitionRateLimit) + root.GetChildren().Range(func(dbID int64, _ *rlinternal.RateLimiterNode) bool { + if _, exists := partitions[dbID]; !exists { + root.GetChildren().Remove(dbID) } - node.GetLimiters().Range(func(rateType internalpb.RateType, limiter *ratelimitutil.Limiter) bool { - if limiter.Limit() != Inf { - limiter.SetHasUpdated(true) + return true + }) + + // Reuse scratch maps across collections; no per-collection rate map/set. + livePartitions := make(map[int64]struct{}) + newDatabaseLimiter := newParamLimiterFunc(internalpb.RateScope_Database, allOps) + newCollectionLimiter := newParamLimiterFunc(internalpb.RateScope_Collection, allOps) + newPartitionLimiter := newParamLimiterFunc(internalpb.RateScope_Partition, allOps) + for dbID, collections := range partitions { + dbLimiter := q.rateLimiter.GetOrCreateDatabaseLimiters(dbID, newDatabaseLimiter) + // Preserve the existing wire-update flags, including empty databases. + resetQuotaLimiter(dbLimiter, databaseRates, len(collections) == 0) + dbLimiter.GetChildren().Range(func(collectionID int64, _ *rlinternal.RateLimiterNode) bool { + if _, exists := collections[collectionID]; !exists { + dbLimiter.GetChildren().Remove(collectionID) } return true }) - } - - collectionRateTypes := getRateTypes(internalpb.RateScope_Collection, allOps) - initLimiters := func(sourceCollections map[int64]map[int64][]int64) { - for dbID, collections := range sourceCollections { - for collectionID, partitionIDs := range collections { - collectionLimitVals := make(map[internalpb.RateType]Limit, collectionRateTypes.Len()) - collectionRateTypes.Range(func(rt internalpb.RateType) bool { - limitVal, err := q.getCollectionMaxLimit(rt, collectionID) - if err != nil { - limitVal = Limit(quota.GetQuotaValue(internalpb.RateScope_Collection, rt, Params)) - } - collectionLimitVals[rt] = limitVal - return true - }) - - getCollectionLimitVal := func(rateType internalpb.RateType) Limit { - return collectionLimitVals[rateType] - } - - collectionLimiter := q.rateLimiter.GetOrCreateCollectionLimiters(dbID, collectionID, - newParamLimiterFunc(internalpb.RateScope_Database, allOps), - newParamLimiterFuncWithLimitFunc(internalpb.RateScope_Collection, allOps, getCollectionLimitVal)) - updateLimiterHasUpdated(collectionLimiter) - - if !enablePartitionRateLimit { + for collectionID, partitionIDs := range collections { + for rt := range collectionRates { + // Flush has no collection-property override. Do not manufacture an + // unsupported-rate error for every collection just to use its default. + if rt == internalpb.RateType_DDLFlush { continue } + value, err := q.getCollectionMaxLimit(rt, collectionID) + if err != nil { + value = Limit(quota.GetQuotaValue(internalpb.RateScope_Collection, rt, Params)) + } + collectionRates[rt] = value + } + collectionLimiter := q.rateLimiter.GetOrCreateCollectionLimiters(dbID, collectionID, + newDatabaseLimiter, newCollectionLimiter) + resetQuotaLimiter(collectionLimiter, collectionRates, true) + + clear(livePartitions) + if enablePartitionRateLimit { for _, partitionID := range partitionIDs { + livePartitions[partitionID] = struct{}{} partitionLimiter := q.rateLimiter.GetOrCreatePartitionLimiters(dbID, collectionID, partitionID, - newParamLimiterFunc(internalpb.RateScope_Database, allOps), - newParamLimiterFuncWithLimitFunc(internalpb.RateScope_Collection, allOps, getCollectionLimitVal), - newParamLimiterFunc(internalpb.RateScope_Partition, allOps)) - updateLimiterHasUpdated(partitionLimiter) + newDatabaseLimiter, newCollectionLimiter, newPartitionLimiter) + resetQuotaLimiter(partitionLimiter, partitionRates, true) } } - if len(collections) == 0 { - dbLimiter := q.rateLimiter.GetOrCreateDatabaseLimiters(dbID, newParamLimiterFunc(internalpb.RateScope_Database, allOps)) - updateLimiterHasUpdated(dbLimiter) - } + collectionLimiter.GetChildren().Range(func(partitionID int64, _ *rlinternal.RateLimiterNode) bool { + if _, exists := livePartitions[partitionID]; !exists { + collectionLimiter.GetChildren().Remove(partitionID) + } + return true + }) } } - partitions := q.meta.ListAllAvailPartitions(q.ctx) - initLimiters(partitions) return nil } +// quotaRateValues reads each scope's defaults once per calculation round. +func quotaRateValues(scope internalpb.RateScope) map[internalpb.RateType]Limit { + rates := make(map[internalpb.RateType]Limit) + for rt := range getRateTypes(scope, allOps) { + rates[rt] = Limit(quota.GetQuotaValue(scope, rt, Params)) + } + return rates +} + +// resetQuotaLimiter restores the calculation baseline without replacing nodes. +// This is not a reset of a Proxy's live admission/token-bucket state. +func resetQuotaLimiter(node *rlinternal.RateLimiterNode, rates map[internalpb.RateType]Limit, markFinite bool) { + limiters := node.GetLimiters() + limiters.Range(func(rt internalpb.RateType, _ *ratelimitutil.Limiter) bool { + if _, exists := rates[rt]; !exists { + limiters.Remove(rt) + } + return true + }) + for rt, value := range rates { + limiter, exists := limiters.Get(rt) + if !exists { + limiter = ratelimitutil.NewLimiter(value, 0) + limiters.Insert(rt, limiter) + } else if limiter.Limit() != value { + limiter.SetLimit(value) + } + limiter.SetHasUpdated(markFinite && value != Inf) + } + states := node.GetQuotaStates() + states.Range(func(state milvuspb.QuotaState, _ *rlinternal.QuotaStateInfo) bool { + states.Remove(state) + return true + }) +} + // getCollectionMaxLimit get limit value from collection's properties. func (q *QuotaCenter) getCollectionMaxLimit(rt internalpb.RateType, collectionID int64) (ratelimitutil.Limit, error) { collectionProps := q.getCollectionLimitProperties(collectionID) @@ -1485,17 +1524,26 @@ func (q *QuotaCenter) getCollectionLimitProperties(collection int64) map[string] return props } - collectionInfo, err := q.meta.GetCollectionByIDWithMaxTs(context.TODO(), collection) + var properties map[string]string + var err error + if reader, ok := q.meta.(quotaMetadataReader); ok { + properties, err = reader.GetQuotaCollectionProperties(q.ctx, collection) + } else { + // Preserve compatibility with alternative metadata implementations. + var collectionInfo *model.Collection + collectionInfo, err = q.meta.GetCollectionByIDWithMaxTs(q.ctx, collection) + if err == nil && len(collectionInfo.Properties) > 0 { + properties = make(map[string]string, len(collectionInfo.Properties)) + for _, pair := range collectionInfo.Properties { + properties[pair.GetKey()] = pair.GetValue() + } + } + } if err != nil { mlog.RatedWarn(q.ctx, rate.Limit(10), "failed to get rate limit properties from collection meta", mlog.FieldCollectionID(collection), mlog.Err(err)) - return make(map[string]string) - } - - properties := make(map[string]string) - for _, pair := range collectionInfo.Properties { - properties[pair.GetKey()] = pair.GetValue() + return nil } q.collectionProps[collection] = properties @@ -1503,6 +1551,21 @@ func (q *QuotaCenter) getCollectionLimitProperties(collection int64) map[string] return properties } +// quotaMetadataReader avoids cloning collections and enumerating unused partition +// IDs. IMetaTable implementations without this optional projection retain the +// existing read path and semantics. +type quotaMetadataReader interface { + GetQuotaCollectionProperties(context.Context, int64) (map[string]string, error) + ListQuotaPartitions(context.Context, bool) map[int64]map[int64][]int64 +} + +func (q *QuotaCenter) quotaPartitionSnapshot(includePartitions bool) map[int64]map[int64][]int64 { + if reader, ok := q.meta.(quotaMetadataReader); ok { + return reader.ListQuotaPartitions(q.ctx, includePartitions) + } + return q.meta.ListAllAvailPartitions(q.ctx) +} + // checkDiskQuota checks if disk quota exceeded. func (q *QuotaCenter) checkDiskQuota(denyWritingDBs map[int64]struct{}) error { q.diskMu.Lock() @@ -1629,10 +1692,13 @@ func (q *QuotaCenter) checkDBDiskQuota(dbSizeInfo map[int64]int64) []int64 { } func (q *QuotaCenter) toRequestLimiter(limiter *rlinternal.RateLimiterNode) *proxypb.Limiter { + return q.toRequestLimiterForProxyCount(limiter, q.proxies.GetProxyCount()) +} + +func (q *QuotaCenter) toRequestLimiterForProxyCount(limiter *rlinternal.RateLimiterNode, proxyNum int) *proxypb.Limiter { var rates []*internalpb.Rate switch q.rateAllocateStrategy { case Average: - proxyNum := q.proxies.GetProxyCount() if proxyNum == 0 { return nil } @@ -1672,23 +1738,29 @@ func (q *QuotaCenter) toRequestLimiter(limiter *rlinternal.RateLimiterNode) *pro func (q *QuotaCenter) toRatesRequest() *proxypb.SetRatesRequest { clusterRateLimiter := q.rateLimiter.GetRootLimiters() + proxyNum := q.proxies.GetProxyCount() + toLimiter := func(node *rlinternal.RateLimiterNode) *proxypb.Limiter { + return q.toRequestLimiterForProxyCount(node, proxyNum) + } // collect db rate limit if clusterRateLimiter has database limiter children dbLimiters := make(map[int64]*proxypb.LimiterNode, clusterRateLimiter.GetChildren().Len()) clusterRateLimiter.GetChildren().Range(func(dbID int64, dbRateLimiters *rlinternal.RateLimiterNode) bool { - dbLimiter := q.toRequestLimiter(dbRateLimiters) + dbLimiter := toLimiter(dbRateLimiters) // collect collection rate limit if dbRateLimiters has collection limiter children collectionLimiters := make(map[int64]*proxypb.LimiterNode, dbRateLimiters.GetChildren().Len()) dbRateLimiters.GetChildren().Range(func(collectionID int64, collectionRateLimiters *rlinternal.RateLimiterNode) bool { - collectionLimiter := q.toRequestLimiter(collectionRateLimiters) + collectionLimiter := toLimiter(collectionRateLimiters) // collect partitions rate limit if collectionRateLimiters has partition limiter children - partitionLimiters := make(map[int64]*proxypb.LimiterNode, collectionRateLimiters.GetChildren().Len()) + var partitionLimiters map[int64]*proxypb.LimiterNode + if count := collectionRateLimiters.GetChildren().Len(); count > 0 { + partitionLimiters = make(map[int64]*proxypb.LimiterNode, count) + } collectionRateLimiters.GetChildren().Range(func(partitionID int64, partitionRateLimiters *rlinternal.RateLimiterNode) bool { partitionLimiters[partitionID] = &proxypb.LimiterNode{ - Limiter: q.toRequestLimiter(partitionRateLimiters), - Children: make(map[int64]*proxypb.LimiterNode, 0), + Limiter: toLimiter(partitionRateLimiters), } return true }) @@ -1709,7 +1781,7 @@ func (q *QuotaCenter) toRatesRequest() *proxypb.SetRatesRequest { }) clusterLimiter := &proxypb.LimiterNode{ - Limiter: q.toRequestLimiter(clusterRateLimiter), + Limiter: toLimiter(clusterRateLimiter), Children: dbLimiters, } diff --git a/internal/util/proxyutil/proxy_client_manager.go b/internal/util/proxyutil/proxy_client_manager.go index dc71c03ee02..5d52c878296 100644 --- a/internal/util/proxyutil/proxy_client_manager.go +++ b/internal/util/proxyutil/proxy_client_manager.go @@ -342,6 +342,8 @@ func (p *ProxyClientManager) GetProxyMetrics(ctx context.Context) ([]*milvuspb.G } // SetRates notifies Proxy to limit rates of requests. +// The request is an immutable snapshot until all clients, including retries, +// finish. Each client copies its routing header and shares the read-only payload. func (p *ProxyClientManager) SetRates(ctx context.Context, request *proxypb.SetRatesRequest) error { if p.proxyClient.Len() == 0 { mlog.Warn(ctx, "proxy client is empty, SetRates will not send to any client")