// 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" "encoding/json" "fmt" "math" "strconv" "testing" "time" "github.com/cockroachdb/errors" "github.com/samber/lo" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" "github.com/stretchr/testify/suite" "google.golang.org/grpc" "github.com/milvus-io/milvus-proto/go-api/v2/commonpb" "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" "github.com/milvus-io/milvus/internal/metastore/model" "github.com/milvus-io/milvus/internal/mocks" "github.com/milvus-io/milvus/internal/proto/internalpb" mockrootcoord "github.com/milvus-io/milvus/internal/rootcoord/mocks" "github.com/milvus-io/milvus/internal/util/proxyutil" interalratelimitutil "github.com/milvus-io/milvus/internal/util/ratelimitutil" "github.com/milvus-io/milvus/pkg/common" "github.com/milvus-io/milvus/pkg/metrics" "github.com/milvus-io/milvus/pkg/util/merr" "github.com/milvus-io/milvus/pkg/util/metricsinfo" "github.com/milvus-io/milvus/pkg/util/paramtable" "github.com/milvus-io/milvus/pkg/util/ratelimitutil" "github.com/milvus-io/milvus/pkg/util/testutils" "github.com/milvus-io/milvus/pkg/util/tsoutil" "github.com/milvus-io/milvus/pkg/util/typeutil" ) func TestQuotaCenter(t *testing.T) { paramtable.Init() ctx, cancel := context.WithCancel(context.Background()) defer cancel() core, err := NewCore(ctx, nil) assert.NoError(t, err) core.tsoAllocator = newMockTsoAllocator() pcm := proxyutil.NewMockProxyClientManager(t) pcm.EXPECT().GetProxyMetrics(mock.Anything).Return(nil, nil).Maybe() dc := mocks.NewMockDataCoordClient(t) dc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(nil, nil).Maybe() collectionIDToPartitionIDs := map[int64][]int64{ 1: {}, 2: {}, 3: {}, } collectionIDToDBID := typeutil.NewConcurrentMap[int64, int64]() collectionIDToDBID.Insert(1, 0) collectionIDToDBID.Insert(2, 0) collectionIDToDBID.Insert(3, 0) collectionIDToDBID.Insert(4, 1) t.Run("test QuotaCenter", func(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByIDWithMaxTs(mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) quotaCenter.Start() time.Sleep(10 * time.Millisecond) quotaCenter.stop() }) t.Run("test QuotaCenter stop", func(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) meta := mockrootcoord.NewIMetaTable(t) paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaCenterCollectInterval.Key, "1") defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaCenterCollectInterval.Key) qc.ExpectedCalls = nil // mock query coord stuck for at most 10s qc.EXPECT().GetMetrics(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, gmr *milvuspb.GetMetricsRequest, co ...grpc.CallOption) (*milvuspb.GetMetricsResponse, error) { counter := 0 for { select { case <-ctx.Done(): return nil, merr.ErrCollectionNotFound default: if counter < 10 { time.Sleep(1 * time.Second) counter++ } } } }) meta.EXPECT().GetCollectionByID(mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() meta.EXPECT().ListDatabases(mock.Anything, mock.Anything).Return([]*model.Database{ { Name: "default", ID: 1, }, }, nil).Maybe() quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) quotaCenter.Start() time.Sleep(3 * time.Second) // assert stop won't stuck more than 5s start := time.Now() quotaCenter.stop() assert.True(t, time.Since(start).Seconds() <= 5) }) t.Run("test collectMetrics", func(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByIDWithMaxTs(mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() meta.EXPECT().ListDatabases(mock.Anything, mock.Anything).Return([]*model.Database{ { Name: "default", ID: 1, }, }, nil).Maybe() qc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{Status: merr.Success()}, nil) quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) err = quotaCenter.collectMetrics() assert.Error(t, err) // for empty response quotaCenter = NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) err = quotaCenter.collectMetrics() assert.Error(t, err) dc.ExpectedCalls = nil dc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{ Status: merr.Status(errors.New("mock error")), }, nil) quotaCenter = NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) err = quotaCenter.collectMetrics() assert.Error(t, err) dc.ExpectedCalls = nil dc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(nil, errors.New("mock error")) qc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(nil, fmt.Errorf("mock err")) quotaCenter = NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) err = quotaCenter.collectMetrics() assert.Error(t, err) qc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{ Status: merr.Status(err), }, nil) quotaCenter = NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) err = quotaCenter.collectMetrics() assert.Error(t, err) }) t.Run("list database fail", func(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) dc2 := mocks.NewMockDataCoordClient(t) pcm2 := proxyutil.NewMockProxyClientManager(t) meta := mockrootcoord.NewIMetaTable(t) emptyQueryCoordTopology := &metricsinfo.QueryCoordTopology{} queryBytes, _ := json.Marshal(emptyQueryCoordTopology) qc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{ Status: merr.Success(), Response: string(queryBytes), }, nil).Once() emptyDataCoordTopology := &metricsinfo.DataCoordTopology{} dataBytes, _ := json.Marshal(emptyDataCoordTopology) dc2.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{ Status: merr.Success(), Response: string(dataBytes), }, nil).Once() pcm2.EXPECT().GetProxyMetrics(mock.Anything).Return([]*milvuspb.GetMetricsResponse{}, nil).Once() meta.EXPECT().ListDatabases(mock.Anything, mock.Anything).Return(nil, errors.New("mock error")).Once() quotaCenter := NewQuotaCenter(pcm2, qc, dc2, core.tsoAllocator, meta) err = quotaCenter.collectMetrics() assert.Error(t, err) }) t.Run("get collection by id fail, querynode", func(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) dc2 := mocks.NewMockDataCoordClient(t) pcm2 := proxyutil.NewMockProxyClientManager(t) meta := mockrootcoord.NewIMetaTable(t) emptyQueryCoordTopology := &metricsinfo.QueryCoordTopology{ Cluster: metricsinfo.QueryClusterTopology{ ConnectedNodes: []metricsinfo.QueryNodeInfos{ { QuotaMetrics: &metricsinfo.QueryNodeQuotaMetrics{ Effect: metricsinfo.NodeEffect{ CollectionIDs: []int64{1000}, }, }, CollectionMetrics: &metricsinfo.QueryNodeCollectionMetrics{ CollectionRows: map[int64]int64{ 1000: 100, }, }, }, }, }, } queryBytes, _ := json.Marshal(emptyQueryCoordTopology) qc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{ Status: merr.Success(), Response: string(queryBytes), }, nil).Once() emptyDataCoordTopology := &metricsinfo.DataCoordTopology{} dataBytes, _ := json.Marshal(emptyDataCoordTopology) dc2.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{ Status: merr.Success(), Response: string(dataBytes), }, nil).Once() pcm2.EXPECT().GetProxyMetrics(mock.Anything).Return([]*milvuspb.GetMetricsResponse{}, nil).Once() meta.EXPECT().ListDatabases(mock.Anything, mock.Anything).Return([]*model.Database{ { ID: 1, Name: "default", }, }, nil).Once() meta.EXPECT().GetCollectionByIDWithMaxTs(mock.Anything, mock.Anything).Return(nil, errors.New("mock err: get collection by id")).Once() quotaCenter := NewQuotaCenter(pcm2, qc, dc2, core.tsoAllocator, meta) err = quotaCenter.collectMetrics() assert.NoError(t, err) }) t.Run("get collection by id fail, datanode", func(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) dc2 := mocks.NewMockDataCoordClient(t) pcm2 := proxyutil.NewMockProxyClientManager(t) meta := mockrootcoord.NewIMetaTable(t) emptyQueryCoordTopology := &metricsinfo.QueryCoordTopology{} queryBytes, _ := json.Marshal(emptyQueryCoordTopology) qc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{ Status: merr.Success(), Response: string(queryBytes), }, nil).Once() emptyDataCoordTopology := &metricsinfo.DataCoordTopology{ Cluster: metricsinfo.DataClusterTopology{ ConnectedDataNodes: []metricsinfo.DataNodeInfos{ { QuotaMetrics: &metricsinfo.DataNodeQuotaMetrics{ Effect: metricsinfo.NodeEffect{ CollectionIDs: []int64{1000}, }, }, }, }, }, } dataBytes, _ := json.Marshal(emptyDataCoordTopology) dc2.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{ Status: merr.Success(), Response: string(dataBytes), }, nil).Once() pcm2.EXPECT().GetProxyMetrics(mock.Anything).Return([]*milvuspb.GetMetricsResponse{}, nil).Once() meta.EXPECT().ListDatabases(mock.Anything, mock.Anything).Return([]*model.Database{ { ID: 1, Name: "default", }, }, nil).Once() meta.EXPECT().GetCollectionByIDWithMaxTs(mock.Anything, mock.Anything).Return(nil, errors.New("mock err: get collection by id")).Once() quotaCenter := NewQuotaCenter(pcm2, qc, dc2, core.tsoAllocator, meta) err = quotaCenter.collectMetrics() assert.NoError(t, err) }) t.Run("test force deny reading collection", func(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByIDWithMaxTs(mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) quotaCenter.readableCollections = map[int64]map[int64][]int64{ 0: collectionIDToPartitionIDs, } err := quotaCenter.resetAllCurrentRates() assert.NoError(t, err) Params.Save(Params.QuotaConfig.ForceDenyReading.Key, "true") defer Params.Reset(Params.QuotaConfig.ForceDenyReading.Key) quotaCenter.calculateReadRates() for collectionID := range collectionIDToPartitionIDs { collectionLimiters := quotaCenter.rateLimiter.GetCollectionLimiters(0, collectionID) assert.NotNil(t, collectionLimiters) limiters := collectionLimiters.GetLimiters() assert.NotNil(t, limiters) for _, rt := range []internalpb.RateType{ internalpb.RateType_DQLSearch, internalpb.RateType_DQLQuery, } { ret, ok := limiters.Get(rt) assert.True(t, ok) assert.Equal(t, ret.Limit(), Limit(0)) } } }) t.Run("test force deny writing", func(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT(). GetCollectionByIDWithMaxTs(mock.Anything, mock.Anything). Return(nil, merr.ErrCollectionNotFound). Maybe() quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) quotaCenter.collectionIDToDBID = typeutil.NewConcurrentMap[int64, int64]() quotaCenter.collectionIDToDBID.Insert(1, 0) quotaCenter.collectionIDToDBID.Insert(2, 0) quotaCenter.collectionIDToDBID.Insert(3, 0) quotaCenter.writableCollections = map[int64]map[int64][]int64{ 0: collectionIDToPartitionIDs, } quotaCenter.writableCollections[0][1] = append(quotaCenter.writableCollections[0][1], 1000) err := quotaCenter.resetAllCurrentRates() assert.NoError(t, err) err = quotaCenter.forceDenyWriting(commonpb.ErrorCode_ForceDeny, false, nil, []int64{4}, nil) assert.Error(t, err) err = quotaCenter.forceDenyWriting(commonpb.ErrorCode_ForceDeny, false, nil, []int64{1, 2, 3}, map[int64][]int64{ 1: {1000}, }) assert.NoError(t, err) for collectionID := range collectionIDToPartitionIDs { collectionLimiters := quotaCenter.rateLimiter.GetCollectionLimiters(0, collectionID) assert.NotNil(t, collectionLimiters) limiters := collectionLimiters.GetLimiters() assert.NotNil(t, limiters) for _, rt := range []internalpb.RateType{ internalpb.RateType_DMLInsert, internalpb.RateType_DMLUpsert, internalpb.RateType_DMLDelete, internalpb.RateType_DMLBulkLoad, } { ret, ok := limiters.Get(rt) assert.True(t, ok) assert.Equal(t, ret.Limit(), Limit(0)) } } err = quotaCenter.forceDenyWriting(commonpb.ErrorCode_ForceDeny, false, []int64{0}, nil, nil) assert.NoError(t, err) dbLimiters := quotaCenter.rateLimiter.GetDatabaseLimiters(0) assert.NotNil(t, dbLimiters) limiters := dbLimiters.GetLimiters() assert.NotNil(t, limiters) for _, rt := range []internalpb.RateType{ internalpb.RateType_DMLInsert, internalpb.RateType_DMLUpsert, internalpb.RateType_DMLDelete, internalpb.RateType_DMLBulkLoad, } { ret, ok := limiters.Get(rt) assert.True(t, ok) assert.Equal(t, ret.Limit(), Limit(0)) } }) t.Run("test calculateRates", func(t *testing.T) { forceBak := Params.QuotaConfig.ForceDenyWriting.GetValue() paramtable.Get().Save(Params.QuotaConfig.ForceDenyWriting.Key, "false") defer func() { paramtable.Get().Save(Params.QuotaConfig.ForceDenyWriting.Key, forceBak) }() qc := mocks.NewMockQueryCoordClient(t) meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByIDWithMaxTs(mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) quotaCenter.clearMetrics() err = quotaCenter.calculateRates() assert.NoError(t, err) alloc := newMockTsoAllocator() alloc.GenerateTSOF = func(count uint32) (typeutil.Timestamp, error) { return 0, fmt.Errorf("mock tso err") } quotaCenter.tsoAllocator = alloc quotaCenter.clearMetrics() err = quotaCenter.calculateRates() assert.Error(t, err) }) t.Run("test getTimeTickDelayFactor factors", func(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByIDWithMaxTs(mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) type ttCase struct { maxTtDelay time.Duration curTt time.Time fgTt time.Time expectedFactor float64 } t0 := time.Now() ttCases := []ttCase{ {10 * time.Second, t0, t0.Add(1 * time.Second), 1}, {10 * time.Second, t0, t0, 1}, {10 * time.Second, t0.Add(1 * time.Second), t0, 0.9}, {10 * time.Second, t0.Add(2 * time.Second), t0, 0.8}, {10 * time.Second, t0.Add(5 * time.Second), t0, 0.5}, {10 * time.Second, t0.Add(7 * time.Second), t0, 0.3}, {10 * time.Second, t0.Add(9 * time.Second), t0, 0.1}, {10 * time.Second, t0.Add(10 * time.Second), t0, 0}, {10 * time.Second, t0.Add(100 * time.Second), t0, 0}, } backup := Params.QuotaConfig.MaxTimeTickDelay.GetValue() for _, c := range ttCases { paramtable.Get().Save(Params.QuotaConfig.MaxTimeTickDelay.Key, fmt.Sprintf("%f", c.maxTtDelay.Seconds())) fgTs := tsoutil.ComposeTSByTime(c.fgTt, 0) quotaCenter.queryNodeMetrics = map[UniqueID]*metricsinfo.QueryNodeQuotaMetrics{ 1: { Fgm: metricsinfo.FlowGraphMetric{ NumFlowGraph: 1, MinFlowGraphTt: fgTs, MinFlowGraphChannel: "dml", }, }, } curTs := tsoutil.ComposeTSByTime(c.curTt, 0) factors := quotaCenter.getTimeTickDelayFactor(curTs) for _, factor := range factors { assert.True(t, math.Abs(factor-c.expectedFactor) < 0.01) } } Params.Save(Params.QuotaConfig.MaxTimeTickDelay.Key, backup) }) t.Run("test TimeTickDelayFactor factors", func(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByIDWithMaxTs(mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() meta.EXPECT().GetDatabaseByID(mock.Anything, mock.Anything, mock.Anything).Return(nil, merr.ErrDatabaseNotFound).Maybe() quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) type ttCase struct { delay time.Duration expectedFactor float64 } ttCases := []ttCase{ {0 * time.Second, 1}, {1 * time.Second, 0.9}, {2 * time.Second, 0.8}, {5 * time.Second, 0.5}, {7 * time.Second, 0.3}, {9 * time.Second, 0.1}, {10 * time.Second, 0}, {100 * time.Second, 0}, } backup := Params.QuotaConfig.MaxTimeTickDelay.GetValue() paramtable.Get().Save(Params.QuotaConfig.DMLLimitEnabled.Key, "true") paramtable.Get().Save(Params.QuotaConfig.TtProtectionEnabled.Key, "true") paramtable.Get().Save(Params.QuotaConfig.MaxTimeTickDelay.Key, "10.0") paramtable.Get().Save(Params.QuotaConfig.DMLMaxInsertRatePerCollection.Key, "100.0") paramtable.Get().Save(Params.QuotaConfig.DMLMinInsertRatePerCollection.Key, "0.0") paramtable.Get().Save(Params.QuotaConfig.DMLMaxUpsertRatePerCollection.Key, "100.0") paramtable.Get().Save(Params.QuotaConfig.DMLMinUpsertRatePerCollection.Key, "0.0") paramtable.Get().Save(Params.QuotaConfig.DMLMaxDeleteRatePerCollection.Key, "100.0") paramtable.Get().Save(Params.QuotaConfig.DMLMinDeleteRatePerCollection.Key, "0.0") forceBak := Params.QuotaConfig.ForceDenyWriting.GetValue() paramtable.Get().Save(Params.QuotaConfig.ForceDenyWriting.Key, "false") defer func() { paramtable.Get().Save(Params.QuotaConfig.ForceDenyWriting.Key, forceBak) }() alloc := newMockTsoAllocator() quotaCenter.tsoAllocator = alloc for _, c := range ttCases { minTS := tsoutil.ComposeTSByTime(time.Now(), 0) hackCurTs := tsoutil.ComposeTSByTime(time.Now().Add(c.delay), 0) alloc.GenerateTSOF = func(count uint32) (typeutil.Timestamp, error) { return hackCurTs, nil } quotaCenter.queryNodeMetrics = map[UniqueID]*metricsinfo.QueryNodeQuotaMetrics{ 1: { Fgm: metricsinfo.FlowGraphMetric{ NumFlowGraph: 1, MinFlowGraphTt: minTS, MinFlowGraphChannel: "dml", }, Effect: metricsinfo.NodeEffect{ CollectionIDs: []int64{1, 2, 3}, }, }, } quotaCenter.dataNodeMetrics = map[UniqueID]*metricsinfo.DataNodeQuotaMetrics{ 11: { Fgm: metricsinfo.FlowGraphMetric{ NumFlowGraph: 1, MinFlowGraphTt: minTS, MinFlowGraphChannel: "dml", }, Effect: metricsinfo.NodeEffect{ CollectionIDs: []int64{1, 2, 3}, }, }, } quotaCenter.writableCollections = map[int64]map[int64][]int64{ 0: collectionIDToPartitionIDs, } quotaCenter.collectionIDToDBID = collectionIDToDBID err = quotaCenter.resetAllCurrentRates() assert.NoError(t, err) err = quotaCenter.calculateWriteRates() assert.NoError(t, err) limit, ok := quotaCenter.rateLimiter.GetCollectionLimiters(0, 1).GetLimiters().Get(internalpb.RateType_DMLDelete) assert.True(t, ok) assert.NotNil(t, limit) deleteFactor := float64(limit.Limit()) / Params.QuotaConfig.DMLMaxInsertRatePerCollection.GetAsFloat() assert.True(t, math.Abs(deleteFactor-c.expectedFactor) < 0.01) } Params.Save(Params.QuotaConfig.MaxTimeTickDelay.Key, backup) }) t.Run("test calculateReadRates", func(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByIDWithMaxTs(mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() meta.EXPECT().ListDatabases(mock.Anything, mock.Anything).Return([]*model.Database{ { ID: 0, Name: "default", }, { ID: 1, Name: "db1", }, }, nil).Maybe() meta.EXPECT().GetDatabaseByID(mock.Anything, mock.Anything, mock.Anything).Return(nil, merr.ErrDatabaseNotFound).Maybe() quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) quotaCenter.clearMetrics() quotaCenter.collectionIDToDBID = collectionIDToDBID quotaCenter.readableCollections = map[int64]map[int64][]int64{ 0: {1: {}}, 1: {2: {}}, } quotaCenter.dbs.Insert("default", 0) quotaCenter.dbs.Insert("db1", 1) paramtable.Get().Save(Params.QuotaConfig.ForceDenyReading.Key, "false") meta.EXPECT().GetDatabaseByID(mock.Anything, mock.Anything, mock.Anything).Unset() meta.EXPECT().GetDatabaseByID(mock.Anything, mock.Anything, mock.Anything). RunAndReturn(func(ctx context.Context, i int64, u uint64) (*model.Database, error) { if i == 1 { return &model.Database{ ID: 1, Name: "db1", Properties: []*commonpb.KeyValuePair{ { Key: common.DatabaseForceDenyReadingKey, Value: "true", }, }, }, nil } return nil, errors.New("mock error") }).Maybe() quotaCenter.resetAllCurrentRates() err = quotaCenter.calculateReadRates() assert.NoError(t, err) assert.NoError(t, err) rln := quotaCenter.rateLimiter.GetDatabaseLimiters(0) limiters := rln.GetLimiters() a, _ := limiters.Get(internalpb.RateType_DQLSearch) assert.NotEqual(t, Limit(0), a.Limit()) b, _ := limiters.Get(internalpb.RateType_DQLQuery) assert.NotEqual(t, Limit(0), b.Limit()) rln = quotaCenter.rateLimiter.GetDatabaseLimiters(1) limiters = rln.GetLimiters() a, _ = limiters.Get(internalpb.RateType_DQLSearch) assert.Equal(t, Limit(0), a.Limit()) b, _ = limiters.Get(internalpb.RateType_DQLQuery) assert.Equal(t, Limit(0), b.Limit()) }) t.Run("test calculateWriteRates", func(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByIDWithMaxTs(mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) err = quotaCenter.calculateWriteRates() assert.NoError(t, err) // force deny paramtable.Get().Save(Params.QuotaConfig.ForceDenyWriting.Key, "true") quotaCenter.writableCollections = map[int64]map[int64][]int64{ 0: collectionIDToPartitionIDs, 1: {4: {}}, } quotaCenter.collectionIDToDBID = collectionIDToDBID quotaCenter.collectionIDToDBID = collectionIDToDBID quotaCenter.resetAllCurrentRates() err = quotaCenter.calculateWriteRates() assert.NoError(t, err) limiters := quotaCenter.rateLimiter.GetRootLimiters().GetLimiters() a, _ := limiters.Get(internalpb.RateType_DMLInsert) assert.Equal(t, Limit(0), a.Limit()) b, _ := limiters.Get(internalpb.RateType_DMLUpsert) assert.Equal(t, Limit(0), b.Limit()) c, _ := limiters.Get(internalpb.RateType_DMLDelete) assert.Equal(t, Limit(0), c.Limit()) paramtable.Get().Reset(Params.QuotaConfig.ForceDenyWriting.Key) // force deny writing for databases meta.EXPECT().GetDatabaseByID(mock.Anything, mock.Anything, mock.Anything). RunAndReturn(func(ctx context.Context, i int64, u uint64) (*model.Database, error) { if i == 1 { return &model.Database{ ID: 1, Name: "db4", Properties: []*commonpb.KeyValuePair{ { Key: common.DatabaseForceDenyWritingKey, Value: "true", }, }, }, nil } return nil, errors.New("mock error") }).Maybe() quotaCenter.resetAllCurrentRates() err = quotaCenter.calculateWriteRates() assert.NoError(t, err) rln := quotaCenter.rateLimiter.GetDatabaseLimiters(0) limiters = rln.GetLimiters() a, _ = limiters.Get(internalpb.RateType_DMLInsert) assert.NotEqual(t, Limit(0), a.Limit()) b, _ = limiters.Get(internalpb.RateType_DMLUpsert) assert.NotEqual(t, Limit(0), b.Limit()) c, _ = limiters.Get(internalpb.RateType_DMLDelete) assert.NotEqual(t, Limit(0), c.Limit()) rln = quotaCenter.rateLimiter.GetDatabaseLimiters(1) limiters = rln.GetLimiters() a, _ = limiters.Get(internalpb.RateType_DMLInsert) assert.Equal(t, Limit(0), a.Limit()) b, _ = limiters.Get(internalpb.RateType_DMLUpsert) assert.Equal(t, Limit(0), b.Limit()) c, _ = limiters.Get(internalpb.RateType_DMLDelete) assert.Equal(t, Limit(0), c.Limit()) meta.EXPECT().GetDatabaseByID(mock.Anything, mock.Anything, mock.Anything).Unset() meta.EXPECT().GetDatabaseByID(mock.Anything, mock.Anything, mock.Anything).Return(nil, merr.ErrDatabaseNotFound).Maybe() // disable tt delay protection disableTtBak := Params.QuotaConfig.TtProtectionEnabled.GetValue() paramtable.Get().Save(Params.QuotaConfig.TtProtectionEnabled.Key, "false") quotaCenter.resetAllCurrentRates() quotaCenter.queryNodeMetrics = make(map[UniqueID]*metricsinfo.QueryNodeQuotaMetrics) quotaCenter.queryNodeMetrics[0] = &metricsinfo.QueryNodeQuotaMetrics{ Hms: metricsinfo.HardwareMetrics{ Memory: 100, MemoryUsage: 100, }, Effect: metricsinfo.NodeEffect{CollectionIDs: []int64{1, 2, 3}}, } err = quotaCenter.calculateWriteRates() assert.NoError(t, err) for db, collections := range quotaCenter.writableCollections { for collection := range collections { states := quotaCenter.rateLimiter.GetCollectionLimiters(db, collection).GetQuotaStates() code, _ := states.Get(milvuspb.QuotaState_DenyToWrite) if db == 0 { assert.Equal(t, commonpb.ErrorCode_MemoryQuotaExhausted, code) } else { assert.Equal(t, commonpb.ErrorCode_Success, code) } } } paramtable.Get().Save(Params.QuotaConfig.TtProtectionEnabled.Key, disableTtBak) }) t.Run("test MemoryFactor factors", func(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByIDWithMaxTs(mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) type memCase struct { lowWater float64 highWater float64 memUsage uint64 memTotal uint64 expectedFactor float64 } memCases := []memCase{ {0.8, 0.9, 10, 100, 1}, {0.8, 0.9, 80, 100, 1}, {0.8, 0.9, 82, 100, 0.8}, {0.8, 0.9, 85, 100, 0.5}, {0.8, 0.9, 88, 100, 0.2}, {0.8, 0.9, 90, 100, 0}, {0.85, 0.95, 25, 100, 1}, {0.85, 0.95, 85, 100, 1}, {0.85, 0.95, 87, 100, 0.8}, {0.85, 0.95, 90, 100, 0.5}, {0.85, 0.95, 93, 100, 0.2}, {0.85, 0.95, 95, 100, 0}, } quotaCenter.writableCollections = map[int64]map[int64][]int64{ 0: collectionIDToPartitionIDs, } for _, c := range memCases { paramtable.Get().Save(Params.QuotaConfig.QueryNodeMemoryLowWaterLevel.Key, fmt.Sprintf("%f", c.lowWater)) paramtable.Get().Save(Params.QuotaConfig.QueryNodeMemoryHighWaterLevel.Key, fmt.Sprintf("%f", c.highWater)) quotaCenter.queryNodeMetrics = map[UniqueID]*metricsinfo.QueryNodeQuotaMetrics{ 1: { Hms: metricsinfo.HardwareMetrics{ MemoryUsage: c.memUsage, Memory: c.memTotal, }, Effect: metricsinfo.NodeEffect{ NodeID: 1, CollectionIDs: []int64{1, 2, 3}, }, }, } factors := quotaCenter.getMemoryFactor() for _, factor := range factors { assert.True(t, math.Abs(factor-c.expectedFactor) < 0.01) } } paramtable.Get().Reset(Params.QuotaConfig.QueryNodeMemoryLowWaterLevel.Key) paramtable.Get().Reset(Params.QuotaConfig.QueryNodeMemoryHighWaterLevel.Key) }) t.Run("test GrowingSegmentsSize factors", func(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByIDWithMaxTs(mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) defaultRatio := Params.QuotaConfig.GrowingSegmentsSizeMinRateRatio.GetAsFloat() tests := []struct { low float64 high float64 growingSize int64 memTotal uint64 expectedFactor float64 }{ {0.8, 0.9, 10, 100, 1}, {0.8, 0.9, 80, 100, 1}, {0.8, 0.9, 82, 100, 0.8}, {0.8, 0.9, 85, 100, 0.5}, {0.8, 0.9, 88, 100, defaultRatio}, {0.8, 0.9, 90, 100, defaultRatio}, {0.85, 0.95, 25, 100, 1}, {0.85, 0.95, 85, 100, 1}, {0.85, 0.95, 87, 100, 0.8}, {0.85, 0.95, 90, 100, 0.5}, {0.85, 0.95, 93, 100, defaultRatio}, {0.85, 0.95, 95, 100, defaultRatio}, } quotaCenter.writableCollections = map[int64]map[int64][]int64{ 0: collectionIDToPartitionIDs, } paramtable.Get().Save(Params.QuotaConfig.GrowingSegmentsSizeProtectionEnabled.Key, "true") for _, test := range tests { paramtable.Get().Save(Params.QuotaConfig.GrowingSegmentsSizeLowWaterLevel.Key, fmt.Sprintf("%f", test.low)) paramtable.Get().Save(Params.QuotaConfig.GrowingSegmentsSizeHighWaterLevel.Key, fmt.Sprintf("%f", test.high)) quotaCenter.queryNodeMetrics = map[UniqueID]*metricsinfo.QueryNodeQuotaMetrics{ 1: { Hms: metricsinfo.HardwareMetrics{ Memory: test.memTotal, }, Effect: metricsinfo.NodeEffect{ NodeID: 1, CollectionIDs: []int64{1, 2, 3}, }, GrowingSegmentsSize: test.growingSize, }, } factors := quotaCenter.getGrowingSegmentsSizeFactor() for _, factor := range factors { assert.True(t, math.Abs(factor-test.expectedFactor) < 0.01) } } paramtable.Get().Reset(Params.QuotaConfig.GrowingSegmentsSizeLowWaterLevel.Key) paramtable.Get().Reset(Params.QuotaConfig.GrowingSegmentsSizeHighWaterLevel.Key) }) t.Run("test checkDiskQuota", func(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByIDWithMaxTs(mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() meta.EXPECT().GetDatabaseByID(mock.Anything, mock.Anything, mock.Anything).Return(nil, merr.ErrDatabaseNotFound).Maybe() quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) quotaCenter.checkDiskQuota(nil) checkLimiter := func(notEquals ...int64) { for db, collections := range quotaCenter.writableCollections { for collection := range collections { limiters := quotaCenter.rateLimiter.GetCollectionLimiters(db, collection).GetLimiters() if lo.Contains(notEquals, collection) { a, _ := limiters.Get(internalpb.RateType_DMLInsert) assert.NotEqual(t, Limit(0), a.Limit()) b, _ := limiters.Get(internalpb.RateType_DMLUpsert) assert.NotEqual(t, Limit(0), b.Limit()) c, _ := limiters.Get(internalpb.RateType_DMLDelete) assert.NotEqual(t, Limit(0), c.Limit()) } else { a, _ := limiters.Get(internalpb.RateType_DMLInsert) assert.Equal(t, Limit(0), a.Limit()) b, _ := limiters.Get(internalpb.RateType_DMLUpsert) assert.Equal(t, Limit(0), b.Limit()) c, _ := limiters.Get(internalpb.RateType_DMLDelete) assert.Equal(t, Limit(0), c.Limit()) } } } } // total DiskQuota exceeded paramtable.Get().Save(Params.QuotaConfig.DiskQuota.Key, "99") paramtable.Get().Save(Params.QuotaConfig.DiskQuotaPerCollection.Key, "90") quotaCenter.dataCoordMetrics = &metricsinfo.DataCoordQuotaMetrics{ TotalBinlogSize: 10 * 1024 * 1024, CollectionBinlogSize: map[int64]int64{ 1: 100 * 1024 * 1024, 2: 100 * 1024 * 1024, 3: 100 * 1024 * 1024, }, } quotaCenter.writableCollections = map[int64]map[int64][]int64{ 0: collectionIDToPartitionIDs, } quotaCenter.collectionIDToDBID = collectionIDToDBID quotaCenter.resetAllCurrentRates() quotaCenter.checkDiskQuota(nil) checkLimiter() paramtable.Get().Reset(Params.QuotaConfig.DiskQuota.Key) paramtable.Get().Reset(Params.QuotaConfig.DiskQuotaPerCollection.Key) // collection DiskQuota exceeded colQuotaBackup := Params.QuotaConfig.DiskQuotaPerCollection.GetValue() paramtable.Get().Save(Params.QuotaConfig.DiskQuotaPerCollection.Key, "30") quotaCenter.dataCoordMetrics = &metricsinfo.DataCoordQuotaMetrics{CollectionBinlogSize: map[int64]int64{ 1: 20 * 1024 * 1024, 2: 30 * 1024 * 1024, 3: 60 * 1024 * 1024, }} quotaCenter.writableCollections = map[int64]map[int64][]int64{ 0: collectionIDToPartitionIDs, } quotaCenter.resetAllCurrentRates() quotaCenter.checkDiskQuota(nil) checkLimiter(1) paramtable.Get().Save(Params.QuotaConfig.DiskQuotaPerCollection.Key, colQuotaBackup) }) t.Run("test setRates", func(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) pcm.EXPECT().GetProxyCount().Return(1) pcm.EXPECT().SetRates(mock.Anything, mock.Anything).Return(nil) meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByIDWithMaxTs(mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) quotaCenter.writableCollections = map[int64]map[int64][]int64{ 0: collectionIDToPartitionIDs, } quotaCenter.readableCollections = map[int64]map[int64][]int64{ 0: collectionIDToPartitionIDs, } quotaCenter.resetAllCurrentRates() collectionID := int64(1) limitNode := quotaCenter.rateLimiter.GetCollectionLimiters(0, collectionID) limitNode.GetLimiters().Insert(internalpb.RateType_DMLInsert, ratelimitutil.NewLimiter(100, 100)) limitNode.GetQuotaStates().Insert(milvuspb.QuotaState_DenyToWrite, commonpb.ErrorCode_MemoryQuotaExhausted) limitNode.GetQuotaStates().Insert(milvuspb.QuotaState_DenyToRead, commonpb.ErrorCode_ForceDeny) err = quotaCenter.sendRatesToProxy() assert.NoError(t, err) }) t.Run("test recordMetrics", func(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByIDWithMaxTs(mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) quotaCenter.writableCollections = map[int64]map[int64][]int64{ 0: collectionIDToPartitionIDs, } quotaCenter.readableCollections = map[int64]map[int64][]int64{ 0: collectionIDToPartitionIDs, } quotaCenter.resetAllCurrentRates() collectionID := int64(1) limitNode := quotaCenter.rateLimiter.GetCollectionLimiters(0, collectionID) limitNode.GetQuotaStates().Insert(milvuspb.QuotaState_DenyToWrite, commonpb.ErrorCode_MemoryQuotaExhausted) limitNode.GetQuotaStates().Insert(milvuspb.QuotaState_DenyToRead, commonpb.ErrorCode_ForceDeny) quotaCenter.recordMetrics() }) t.Run("test guaranteeMinRate", func(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByIDWithMaxTs(mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) quotaCenter.readableCollections = map[int64]map[int64][]int64{ 0: collectionIDToPartitionIDs, } quotaCenter.resetAllCurrentRates() minRate := Limit(100) collectionID := int64(1) limitNode := quotaCenter.rateLimiter.GetCollectionLimiters(0, collectionID) limitNode.GetLimiters().Insert(internalpb.RateType_DQLSearch, ratelimitutil.NewLimiter(50, 50)) quotaCenter.guaranteeMinRate(float64(minRate), internalpb.RateType_DQLSearch, limitNode) limiter, _ := limitNode.GetLimiters().Get(internalpb.RateType_DQLSearch) assert.EqualValues(t, minRate, limiter.Limit()) }) t.Run("test diskAllowance", func(t *testing.T) { tests := []struct { name string totalDiskQuota string collDiskQuota string totalDiskUsage int64 // in MB collDiskUsage int64 // in MB expectAllowance float64 // in bytes }{ {"test max", "-1", "-1", 100, 100, math.MaxFloat64}, {"test total quota exceeded", "100", "-1", 100, 100, 0}, {"test coll quota exceeded", "-1", "20", 100, 20, 0}, {"test not exceeded", "100", "20", 80, 10, 10 * 1024 * 1024}, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { collection := UniqueID(0) meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByIDWithMaxTs(mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() quotaCenter := NewQuotaCenter(pcm, nil, dc, core.tsoAllocator, meta) quotaCenter.resetAllCurrentRates() quotaBackup := Params.QuotaConfig.DiskQuota.GetValue() colQuotaBackup := Params.QuotaConfig.DiskQuotaPerCollection.GetValue() paramtable.Get().Save(Params.QuotaConfig.DiskQuota.Key, test.totalDiskQuota) paramtable.Get().Save(Params.QuotaConfig.DiskQuotaPerCollection.Key, test.collDiskQuota) quotaCenter.diskMu.Lock() quotaCenter.dataCoordMetrics = &metricsinfo.DataCoordQuotaMetrics{} quotaCenter.dataCoordMetrics.CollectionBinlogSize = map[int64]int64{collection: test.collDiskUsage * 1024 * 1024} quotaCenter.totalBinlogSize = test.totalDiskUsage * 1024 * 1024 quotaCenter.diskMu.Unlock() allowance := quotaCenter.diskAllowance(collection) assert.Equal(t, test.expectAllowance, allowance) paramtable.Get().Save(Params.QuotaConfig.DiskQuota.Key, quotaBackup) paramtable.Get().Save(Params.QuotaConfig.DiskQuotaPerCollection.Key, colQuotaBackup) }) } }) t.Run("test reset current rates", func(t *testing.T) { meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByIDWithMaxTs(mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe() quotaCenter := NewQuotaCenter(pcm, nil, dc, core.tsoAllocator, meta) quotaCenter.readableCollections = map[int64]map[int64][]int64{ 0: {1: {}}, } quotaCenter.writableCollections = map[int64]map[int64][]int64{ 0: {1: {}}, } quotaCenter.collectionIDToDBID = collectionIDToDBID quotaCenter.resetAllCurrentRates() limiters := quotaCenter.rateLimiter.GetCollectionLimiters(0, 1).GetLimiters() getRate := func(m *typeutil.ConcurrentMap[internalpb.RateType, *ratelimitutil.Limiter], key internalpb.RateType) float64 { v, _ := m.Get(key) return float64(v.Limit()) } assert.Equal(t, getRate(limiters, internalpb.RateType_DMLInsert), Params.QuotaConfig.DMLMaxInsertRatePerCollection.GetAsFloat()) assert.Equal(t, getRate(limiters, internalpb.RateType_DMLUpsert), Params.QuotaConfig.DMLMaxUpsertRatePerCollection.GetAsFloat()) assert.Equal(t, getRate(limiters, internalpb.RateType_DMLDelete), Params.QuotaConfig.DMLMaxDeleteRatePerCollection.GetAsFloat()) assert.Equal(t, getRate(limiters, internalpb.RateType_DMLBulkLoad), Params.QuotaConfig.DMLMaxBulkLoadRatePerCollection.GetAsFloat()) assert.Equal(t, getRate(limiters, internalpb.RateType_DQLSearch), Params.QuotaConfig.DQLMaxSearchRatePerCollection.GetAsFloat()) assert.Equal(t, getRate(limiters, internalpb.RateType_DQLQuery), Params.QuotaConfig.DQLMaxQueryRatePerCollection.GetAsFloat()) meta.ExpectedCalls = nil meta.EXPECT().GetCollectionByIDWithMaxTs(mock.Anything, mock.Anything).Return(&model.Collection{ Properties: []*commonpb.KeyValuePair{ { Key: common.CollectionInsertRateMaxKey, Value: "1", }, { Key: common.CollectionDeleteRateMaxKey, Value: "2", }, { Key: common.CollectionBulkLoadRateMaxKey, Value: "3", }, { Key: common.CollectionQueryRateMaxKey, Value: "4", }, { Key: common.CollectionSearchRateMaxKey, Value: "5", }, { Key: common.CollectionUpsertRateMaxKey, Value: "6", }, }, }, nil) quotaCenter.resetAllCurrentRates() limiters = quotaCenter.rateLimiter.GetCollectionLimiters(0, 1).GetLimiters() assert.Equal(t, getRate(limiters, internalpb.RateType_DMLInsert), float64(1*1024*1024)) assert.Equal(t, getRate(limiters, internalpb.RateType_DMLDelete), float64(2*1024*1024)) assert.Equal(t, getRate(limiters, internalpb.RateType_DMLBulkLoad), float64(3*1024*1024)) assert.Equal(t, getRate(limiters, internalpb.RateType_DQLQuery), float64(4)) assert.Equal(t, getRate(limiters, internalpb.RateType_DQLSearch), float64(5)) assert.Equal(t, getRate(limiters, internalpb.RateType_DMLUpsert), float64(6*1024*1024)) }) } type QuotaCenterSuite struct { testutils.PromMetricsSuite core *Core pcm *proxyutil.MockProxyClientManager dc *mocks.MockDataCoordClient qc *mocks.MockQueryCoordClient meta *mockrootcoord.IMetaTable } func (s *QuotaCenterSuite) SetupSuite() { paramtable.Init() var err error s.core, err = NewCore(context.Background(), nil) s.Require().NoError(err) } func (s *QuotaCenterSuite) SetupTest() { s.pcm = proxyutil.NewMockProxyClientManager(s.T()) s.dc = mocks.NewMockDataCoordClient(s.T()) s.qc = mocks.NewMockQueryCoordClient(s.T()) s.meta = mockrootcoord.NewIMetaTable(s.T()) } func (s *QuotaCenterSuite) getEmptyQCMetricsRsp() string { metrics := &metricsinfo.QueryCoordTopology{ Cluster: metricsinfo.QueryClusterTopology{}, } resp, err := metricsinfo.MarshalTopology(metrics) s.Require().NoError(err) return resp } func (s *QuotaCenterSuite) getEmptyDCMetricsRsp() string { metrics := &metricsinfo.DataCoordTopology{ Cluster: metricsinfo.DataClusterTopology{}, } resp, err := metricsinfo.MarshalTopology(metrics) s.Require().NoError(err) return resp } func (s *QuotaCenterSuite) TestSyncMetricsSuccess() { pcm := s.pcm dc := s.dc qc := s.qc meta := s.meta core := s.core call := meta.EXPECT().ListDatabases(mock.Anything, mock.Anything).Return([]*model.Database{ { ID: 1, Name: "default", }, }, nil) defer call.Unset() s.Run("querycoord_cluster", func() { pcm.EXPECT().GetProxyMetrics(mock.Anything).Return(nil, nil).Once() dc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{ Status: merr.Status(nil), Response: s.getEmptyDCMetricsRsp(), }, nil).Once() metrics := &metricsinfo.QueryCoordTopology{ Cluster: metricsinfo.QueryClusterTopology{ ConnectedNodes: []metricsinfo.QueryNodeInfos{ {BaseComponentInfos: metricsinfo.BaseComponentInfos{ID: 1}, QuotaMetrics: &metricsinfo.QueryNodeQuotaMetrics{Effect: metricsinfo.NodeEffect{NodeID: 1, CollectionIDs: []int64{100, 200}}}}, {BaseComponentInfos: metricsinfo.BaseComponentInfos{ID: 2}, QuotaMetrics: &metricsinfo.QueryNodeQuotaMetrics{Effect: metricsinfo.NodeEffect{NodeID: 2, CollectionIDs: []int64{200, 300}}}}, }, }, } resp, err := metricsinfo.MarshalTopology(metrics) s.Require().NoError(err) qc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{ Status: merr.Status(nil), Response: resp, }, nil).Once() meta.EXPECT().GetCollectionByIDWithMaxTs(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, i int64) (*model.Collection, error) { return &model.Collection{CollectionID: i, DBID: 1}, nil }).Times(3) quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) err = quotaCenter.collectMetrics() s.Require().NoError(err) s.ElementsMatch([]int64{100, 200, 300}, lo.Keys(quotaCenter.readableCollections[1])) nodes := lo.Keys(quotaCenter.queryNodeMetrics) s.ElementsMatch([]int64{1, 2}, nodes) }) s.Run("datacoord_cluster", func() { pcm.EXPECT().GetProxyMetrics(mock.Anything).Return(nil, nil).Once() qc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{ Status: merr.Status(nil), Response: s.getEmptyQCMetricsRsp(), }, nil).Once() metrics := &metricsinfo.DataCoordTopology{ Cluster: metricsinfo.DataClusterTopology{ ConnectedDataNodes: []metricsinfo.DataNodeInfos{ {BaseComponentInfos: metricsinfo.BaseComponentInfos{ID: 1}, QuotaMetrics: &metricsinfo.DataNodeQuotaMetrics{Effect: metricsinfo.NodeEffect{NodeID: 1, CollectionIDs: []int64{100, 200}}}}, {BaseComponentInfos: metricsinfo.BaseComponentInfos{ID: 2}, QuotaMetrics: &metricsinfo.DataNodeQuotaMetrics{Effect: metricsinfo.NodeEffect{NodeID: 2, CollectionIDs: []int64{200, 300}}}}, }, }, } resp, err := metricsinfo.MarshalTopology(metrics) s.Require().NoError(err) dc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{ Status: merr.Status(nil), Response: resp, }, nil).Once() meta.EXPECT().GetCollectionByIDWithMaxTs(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, i int64) (*model.Collection, error) { return &model.Collection{CollectionID: i, DBID: 1}, nil }).Times(3) quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) err = quotaCenter.collectMetrics() s.Require().NoError(err) s.ElementsMatch([]int64{100, 200, 300}, lo.Keys(quotaCenter.writableCollections[1])) nodes := lo.Keys(quotaCenter.dataNodeMetrics) s.ElementsMatch([]int64{1, 2}, nodes) }) } func (s *QuotaCenterSuite) TestSyncMetricsFailure() { pcm := s.pcm dc := s.dc qc := s.qc meta := s.meta core := s.core call := meta.EXPECT().ListDatabases(mock.Anything, mock.Anything).Return([]*model.Database{ { ID: 1, Name: "default", }, }, nil) defer call.Unset() s.Run("querycoord_failure", func() { pcm.EXPECT().GetProxyMetrics(mock.Anything).Return(nil, nil).Once() dc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{ Status: merr.Status(nil), Response: s.getEmptyDCMetricsRsp(), }, nil).Once() qc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(nil, errors.New("mock")).Once() quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) err := quotaCenter.collectMetrics() s.Error(err) }) s.Run("querycoord_bad_response", func() { pcm.EXPECT().GetProxyMetrics(mock.Anything).Return(nil, nil).Once() dc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{ Status: merr.Status(nil), Response: s.getEmptyDCMetricsRsp(), }, nil).Once() qc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{ Status: merr.Status(nil), Response: "abc", }, nil).Once() quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) err := quotaCenter.collectMetrics() s.Error(err) }) s.Run("datacoord_failure", func() { pcm.EXPECT().GetProxyMetrics(mock.Anything).Return(nil, nil).Once() qc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{ Status: merr.Status(nil), Response: s.getEmptyQCMetricsRsp(), }, nil).Once() dc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(nil, errors.New("mocked")).Once() quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) err := quotaCenter.collectMetrics() s.Error(err) }) s.Run("datacoord_bad_response", func() { pcm.EXPECT().GetProxyMetrics(mock.Anything).Return(nil, nil).Once() qc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{ Status: merr.Status(nil), Response: s.getEmptyQCMetricsRsp(), }, nil).Once() dc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{ Status: merr.Status(nil), Response: "abc", }, nil).Once() quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) err := quotaCenter.collectMetrics() s.Error(err) }) s.Run("proxy_manager_return_failure", func() { qc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{ Status: merr.Status(nil), Response: s.getEmptyQCMetricsRsp(), }, nil).Once() dc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{ Status: merr.Status(nil), Response: s.getEmptyDCMetricsRsp(), }, nil).Once() pcm.EXPECT().GetProxyMetrics(mock.Anything).Return(nil, errors.New("mocked")).Once() quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) err := quotaCenter.collectMetrics() s.Error(err) }) s.Run("proxy_manager_bad_response", func() { qc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{ Status: merr.Status(nil), Response: s.getEmptyQCMetricsRsp(), }, nil).Once() dc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{ Status: merr.Status(nil), Response: s.getEmptyDCMetricsRsp(), }, nil).Once() pcm.EXPECT().GetProxyMetrics(mock.Anything).Return([]*milvuspb.GetMetricsResponse{ { Status: merr.Status(nil), Response: "abc", }, }, nil).Once() quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) err := quotaCenter.collectMetrics() s.Error(err) }) } func (s *QuotaCenterSuite) TestNodeOffline() { pcm := s.pcm dc := s.dc qc := s.qc meta := s.meta core := s.core call := meta.EXPECT().GetCollectionByIDWithMaxTs(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, i int64) (*model.Collection, error) { return &model.Collection{CollectionID: i, DBID: 1}, nil }).Maybe() defer call.Unset() dbCall := meta.EXPECT().ListDatabases(mock.Anything, mock.Anything).Return([]*model.Database{ { ID: 1, Name: "default", }, }, nil) defer dbCall.Unset() metrics.RootCoordTtDelay.Reset() Params.Save(Params.QuotaConfig.TtProtectionEnabled.Key, "true") defer Params.Reset(Params.QuotaConfig.TtProtectionEnabled.Key) // proxy pcm.EXPECT().GetProxyMetrics(mock.Anything).Return(nil, nil) // qc first time qcMetrics := &metricsinfo.QueryCoordTopology{ Cluster: metricsinfo.QueryClusterTopology{ ConnectedNodes: []metricsinfo.QueryNodeInfos{ { BaseComponentInfos: metricsinfo.BaseComponentInfos{ID: 1}, QuotaMetrics: &metricsinfo.QueryNodeQuotaMetrics{ Fgm: metricsinfo.FlowGraphMetric{NumFlowGraph: 2, MinFlowGraphChannel: "dml_0"}, Effect: metricsinfo.NodeEffect{ NodeID: 1, CollectionIDs: []int64{100, 200}, }, }, }, { BaseComponentInfos: metricsinfo.BaseComponentInfos{ID: 2}, QuotaMetrics: &metricsinfo.QueryNodeQuotaMetrics{ Fgm: metricsinfo.FlowGraphMetric{NumFlowGraph: 2, MinFlowGraphChannel: "dml_0"}, Effect: metricsinfo.NodeEffect{ NodeID: 2, CollectionIDs: []int64{100, 200}, }, }, }, }, }, } resp, err := metricsinfo.MarshalTopology(qcMetrics) s.Require().NoError(err) qc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{ Status: merr.Status(nil), Response: resp, }, nil).Once() // dc first time dcMetrics := &metricsinfo.DataCoordTopology{ Cluster: metricsinfo.DataClusterTopology{ ConnectedDataNodes: []metricsinfo.DataNodeInfos{ { BaseComponentInfos: metricsinfo.BaseComponentInfos{ID: 3}, QuotaMetrics: &metricsinfo.DataNodeQuotaMetrics{ Fgm: metricsinfo.FlowGraphMetric{NumFlowGraph: 2, MinFlowGraphChannel: "dml_0"}, Effect: metricsinfo.NodeEffect{NodeID: 3, CollectionIDs: []int64{100, 200}}, }, }, { BaseComponentInfos: metricsinfo.BaseComponentInfos{ID: 4}, QuotaMetrics: &metricsinfo.DataNodeQuotaMetrics{ Fgm: metricsinfo.FlowGraphMetric{NumFlowGraph: 2, MinFlowGraphChannel: "dml_0"}, Effect: metricsinfo.NodeEffect{NodeID: 4, CollectionIDs: []int64{200, 300}}, }, }, }, }, } resp, err = metricsinfo.MarshalTopology(dcMetrics) s.Require().NoError(err) dc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{ Status: merr.Status(nil), Response: resp, }, nil).Once() quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) err = quotaCenter.collectMetrics() s.Require().NoError(err) quotaCenter.getTimeTickDelayFactor(tsoutil.ComposeTSByTime(time.Now(), 0)) s.CollectCntEqual(metrics.RootCoordTtDelay, 4) // qc second time qcMetrics = &metricsinfo.QueryCoordTopology{ Cluster: metricsinfo.QueryClusterTopology{ ConnectedNodes: []metricsinfo.QueryNodeInfos{ { BaseComponentInfos: metricsinfo.BaseComponentInfos{ID: 2}, QuotaMetrics: &metricsinfo.QueryNodeQuotaMetrics{ Fgm: metricsinfo.FlowGraphMetric{NumFlowGraph: 2, MinFlowGraphChannel: "dml_0"}, Effect: metricsinfo.NodeEffect{NodeID: 2, CollectionIDs: []int64{200, 300}}, }, }, }, }, } resp, err = metricsinfo.MarshalTopology(qcMetrics) s.Require().NoError(err) qc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{ Status: merr.Status(nil), Response: resp, }, nil).Once() // dc second time dcMetrics = &metricsinfo.DataCoordTopology{ Cluster: metricsinfo.DataClusterTopology{ ConnectedDataNodes: []metricsinfo.DataNodeInfos{ { BaseComponentInfos: metricsinfo.BaseComponentInfos{ID: 4}, QuotaMetrics: &metricsinfo.DataNodeQuotaMetrics{ Fgm: metricsinfo.FlowGraphMetric{NumFlowGraph: 2, MinFlowGraphChannel: "dml_0"}, Effect: metricsinfo.NodeEffect{NodeID: 2, CollectionIDs: []int64{200, 300}}, }, }, }, }, } resp, err = metricsinfo.MarshalTopology(dcMetrics) s.Require().NoError(err) dc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(&milvuspb.GetMetricsResponse{ Status: merr.Status(nil), Response: resp, }, nil).Once() err = quotaCenter.collectMetrics() s.Require().NoError(err) quotaCenter.getTimeTickDelayFactor(tsoutil.ComposeTSByTime(time.Now(), 0)) s.CollectCntEqual(metrics.RootCoordTtDelay, 2) } func TestQuotaCenterSuite(t *testing.T) { suite.Run(t, new(QuotaCenterSuite)) } func TestUpdateLimiter(t *testing.T) { t.Run("nil node", func(t *testing.T) { updateLimiter(nil, nil, internalpb.RateScope_Database, dql) }) t.Run("normal op", func(t *testing.T) { node := interalratelimitutil.NewRateLimiterNode(internalpb.RateScope_Collection) node.GetLimiters().Insert(internalpb.RateType_DQLSearch, ratelimitutil.NewLimiter(5, 5)) newLimit := ratelimitutil.NewLimiter(10, 10) updateLimiter(node, newLimit, internalpb.RateScope_Collection, dql) searchLimit, _ := node.GetLimiters().Get(internalpb.RateType_DQLSearch) assert.Equal(t, Limit(10), searchLimit.Limit()) }) } func TestGetRateType(t *testing.T) { t.Run("invalid rate type", func(t *testing.T) { assert.Panics(t, func() { getRateTypes(internalpb.RateScope(100), ddl) }) }) t.Run("ddl cluster scope", func(t *testing.T) { a := getRateTypes(internalpb.RateScope_Cluster, ddl) assert.Equal(t, 5, a.Len()) }) } func TestResetAllCurrentRates(t *testing.T) { paramtable.Init() ctx := context.Background() qc := mocks.NewMockQueryCoordClient(t) meta := mockrootcoord.NewIMetaTable(t) pcm := proxyutil.NewMockProxyClientManager(t) dc := mocks.NewMockDataCoordClient(t) core, _ := NewCore(ctx, nil) core.tsoAllocator = newMockTsoAllocator() meta.EXPECT().GetCollectionByIDWithMaxTs(mock.Anything, mock.Anything).Return(nil, errors.New("mock error")) quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) quotaCenter.readableCollections = map[int64]map[int64][]int64{ 1: {}, } quotaCenter.writableCollections = map[int64]map[int64][]int64{ 2: { 100: []int64{}, }, } err := quotaCenter.resetAllCurrentRates() assert.NoError(t, err) db1 := quotaCenter.rateLimiter.GetDatabaseLimiters(1) assert.NotNil(t, db1) db2 := quotaCenter.rateLimiter.GetDatabaseLimiters(2) assert.NotNil(t, db2) collection := quotaCenter.rateLimiter.GetCollectionLimiters(2, 100) assert.NotNil(t, collection) } func newQuotaCenterForTesting(t *testing.T, ctx context.Context, meta IMetaTable) *QuotaCenter { qc := mocks.NewMockQueryCoordClient(t) pcm := proxyutil.NewMockProxyClientManager(t) dc := mocks.NewMockDataCoordClient(t) core, _ := NewCore(ctx, nil) core.tsoAllocator = newMockTsoAllocator() quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) quotaCenter.rateLimiter.GetRootLimiters().GetLimiters().Insert(internalpb.RateType_DMLInsert, ratelimitutil.NewLimiter(500, 500)) quotaCenter.rateLimiter.GetOrCreatePartitionLimiters(1, 10, 100, newParamLimiterFunc(internalpb.RateScope_Database, allOps), newParamLimiterFunc(internalpb.RateScope_Collection, allOps), newParamLimiterFunc(internalpb.RateScope_Partition, allOps), ) quotaCenter.rateLimiter.GetOrCreatePartitionLimiters(1, 10, 101, newParamLimiterFunc(internalpb.RateScope_Database, allOps), newParamLimiterFunc(internalpb.RateScope_Collection, allOps), newParamLimiterFunc(internalpb.RateScope_Partition, allOps), ) quotaCenter.rateLimiter.GetOrCreatePartitionLimiters(2, 20, 200, newParamLimiterFunc(internalpb.RateScope_Database, allOps), newParamLimiterFunc(internalpb.RateScope_Collection, allOps), newParamLimiterFunc(internalpb.RateScope_Partition, allOps), ) quotaCenter.rateLimiter.GetOrCreatePartitionLimiters(2, 30, 300, newParamLimiterFunc(internalpb.RateScope_Database, allOps), newParamLimiterFunc(internalpb.RateScope_Collection, allOps), newParamLimiterFunc(internalpb.RateScope_Partition, allOps), ) quotaCenter.rateLimiter.GetOrCreatePartitionLimiters(4, 40, 400, newParamLimiterFunc(internalpb.RateScope_Database, allOps), newParamLimiterFunc(internalpb.RateScope_Collection, allOps), newParamLimiterFunc(internalpb.RateScope_Partition, allOps), ) quotaCenter.dataCoordMetrics = &metricsinfo.DataCoordQuotaMetrics{ TotalBinlogSize: 200 * 1024 * 1024, CollectionBinlogSize: map[int64]int64{ 10: 15 * 1024 * 1024, 20: 6 * 1024 * 1024, 30: 6 * 1024 * 1024, 40: 4 * 1024 * 1024, }, PartitionsBinlogSize: map[int64]map[int64]int64{ 10: { 100: 10 * 1024 * 1024, 101: 5 * 1024 * 1024, }, 20: { 200: 6 * 1024 * 1024, }, 30: { 300: 6 * 1024 * 1024, }, 40: { 400: 4 * 1024 * 1024, }, }, } quotaCenter.collectionIDToDBID = typeutil.NewConcurrentMap[int64, int64]() quotaCenter.collectionIDToDBID.Insert(10, 1) quotaCenter.collectionIDToDBID.Insert(20, 2) quotaCenter.collectionIDToDBID.Insert(30, 2) quotaCenter.collectionIDToDBID.Insert(40, 4) return quotaCenter } func TestCheckDiskQuota(t *testing.T) { paramtable.Init() ctx := context.Background() t.Run("disk quota check disable", func(t *testing.T) { qc := mocks.NewMockQueryCoordClient(t) meta := mockrootcoord.NewIMetaTable(t) pcm := proxyutil.NewMockProxyClientManager(t) dc := mocks.NewMockDataCoordClient(t) core, _ := NewCore(ctx, nil) core.tsoAllocator = newMockTsoAllocator() quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) Params.Save(Params.QuotaConfig.DiskProtectionEnabled.Key, "false") defer Params.Reset(Params.QuotaConfig.DiskProtectionEnabled.Key) err := quotaCenter.checkDiskQuota(nil) assert.NoError(t, err) }) t.Run("disk quota check enable", func(t *testing.T) { diskQuotaStr := "10" Params.Save(Params.QuotaConfig.DiskProtectionEnabled.Key, "true") defer Params.Reset(Params.QuotaConfig.DiskProtectionEnabled.Key) Params.Save(Params.QuotaConfig.DiskQuota.Key, "150") defer Params.Reset(Params.QuotaConfig.DiskQuota.Key) Params.Save(Params.QuotaConfig.DiskQuotaPerDB.Key, diskQuotaStr) defer Params.Reset(Params.QuotaConfig.DiskQuotaPerDB.Key) Params.Save(Params.QuotaConfig.DiskQuotaPerCollection.Key, diskQuotaStr) defer Params.Reset(Params.QuotaConfig.DiskQuotaPerCollection.Key) Params.Save(Params.QuotaConfig.DiskQuotaPerPartition.Key, diskQuotaStr) defer Params.Reset(Params.QuotaConfig.DiskQuotaPerPartition.Key) Params.Save(Params.QuotaConfig.DMLLimitEnabled.Key, "true") defer Params.Reset(Params.QuotaConfig.DMLLimitEnabled.Key) Params.Save(Params.QuotaConfig.DMLMaxInsertRate.Key, "10") defer Params.Reset(Params.QuotaConfig.DMLMaxInsertRate.Key) Params.Save(Params.QuotaConfig.DMLMaxInsertRatePerDB.Key, "10") defer Params.Reset(Params.QuotaConfig.DMLMaxInsertRatePerDB.Key) Params.Save(Params.QuotaConfig.DMLMaxInsertRatePerCollection.Key, "10") defer Params.Reset(Params.QuotaConfig.DMLMaxInsertRatePerCollection.Key) Params.Save(Params.QuotaConfig.DMLMaxInsertRatePerPartition.Key, "10") defer Params.Reset(Params.QuotaConfig.DMLMaxInsertRatePerPartition.Key) meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetCollectionByIDWithMaxTs(mock.Anything, mock.Anything).Return(nil, errors.New("mock error")) meta.EXPECT().GetDatabaseByID(mock.Anything, mock.Anything, mock.Anything). RunAndReturn(func(ctx context.Context, i int64, u uint64) (*model.Database, error) { if i == 4 { return &model.Database{ ID: 1, Name: "db4", Properties: []*commonpb.KeyValuePair{ { Key: common.DatabaseDiskQuotaKey, Value: "2", }, }, }, nil } return nil, errors.New("mock error") }).Maybe() quotaCenter := newQuotaCenterForTesting(t, ctx, meta) checkRate := func(rateNode *interalratelimitutil.RateLimiterNode, expectValue float64) { insertRate, ok := rateNode.GetLimiters().Get(internalpb.RateType_DMLInsert) assert.True(t, ok) assert.EqualValues(t, expectValue, insertRate.Limit()) } diskQuota, err := strconv.ParseFloat(diskQuotaStr, 64) assert.NoError(t, err) configQuotaValue := 1024 * 1024 * diskQuota { err := quotaCenter.checkDiskQuota(nil) assert.NoError(t, err) checkRate(quotaCenter.rateLimiter.GetRootLimiters(), 0) } { Params.Save(Params.QuotaConfig.DiskQuota.Key, "999") err := quotaCenter.checkDiskQuota(nil) assert.NoError(t, err) checkRate(quotaCenter.rateLimiter.GetDatabaseLimiters(1), 0) checkRate(quotaCenter.rateLimiter.GetDatabaseLimiters(2), 0) checkRate(quotaCenter.rateLimiter.GetDatabaseLimiters(4), 0) checkRate(quotaCenter.rateLimiter.GetCollectionLimiters(1, 10), 0) checkRate(quotaCenter.rateLimiter.GetCollectionLimiters(2, 20), configQuotaValue) checkRate(quotaCenter.rateLimiter.GetCollectionLimiters(2, 30), configQuotaValue) checkRate(quotaCenter.rateLimiter.GetCollectionLimiters(4, 40), configQuotaValue) checkRate(quotaCenter.rateLimiter.GetPartitionLimiters(1, 10, 100), 0) checkRate(quotaCenter.rateLimiter.GetPartitionLimiters(1, 10, 101), configQuotaValue) checkRate(quotaCenter.rateLimiter.GetPartitionLimiters(2, 20, 200), configQuotaValue) checkRate(quotaCenter.rateLimiter.GetPartitionLimiters(2, 30, 300), configQuotaValue) checkRate(quotaCenter.rateLimiter.GetPartitionLimiters(4, 40, 400), configQuotaValue) } }) } func TestTORequestLimiter(t *testing.T) { ctx := context.Background() qc := mocks.NewMockQueryCoordClient(t) meta := mockrootcoord.NewIMetaTable(t) pcm := proxyutil.NewMockProxyClientManager(t) dc := mocks.NewMockDataCoordClient(t) core, _ := NewCore(ctx, nil) core.tsoAllocator = newMockTsoAllocator() quotaCenter := NewQuotaCenter(pcm, qc, dc, core.tsoAllocator, meta) pcm.EXPECT().GetProxyCount().Return(2) limitNode := interalratelimitutil.NewRateLimiterNode(internalpb.RateScope_Cluster) a := ratelimitutil.NewLimiter(500, 500) a.SetLimit(200) b := ratelimitutil.NewLimiter(100, 100) limitNode.GetLimiters().Insert(internalpb.RateType_DMLInsert, a) limitNode.GetLimiters().Insert(internalpb.RateType_DMLDelete, b) limitNode.GetLimiters().Insert(internalpb.RateType_DMLBulkLoad, GetInfLimiter(internalpb.RateType_DMLBulkLoad)) limitNode.GetQuotaStates().Insert(milvuspb.QuotaState_DenyToRead, commonpb.ErrorCode_ForceDeny) quotaCenter.rateAllocateStrategy = Average proxyLimit := quotaCenter.toRequestLimiter(limitNode) assert.Equal(t, 1, len(proxyLimit.Rates)) assert.Equal(t, internalpb.RateType_DMLInsert, proxyLimit.Rates[0].Rt) assert.Equal(t, float64(100), proxyLimit.Rates[0].R) assert.Equal(t, 1, len(proxyLimit.States)) assert.Equal(t, milvuspb.QuotaState_DenyToRead, proxyLimit.States[0]) assert.Equal(t, 1, len(proxyLimit.Codes)) assert.Equal(t, commonpb.ErrorCode_ForceDeny, proxyLimit.Codes[0]) }