test/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2024-11-30 02:48:45 +08:00
Code Issues Actions 11 Packages Projects Releases Wiki Activity

Fix panic due to collection not exist in quota effect (#24320)

Browse Source 
Signed-off-by: bigsheeper <yihao.dai@zilliz.com>
This commit is contained in:
yihao.dai 2023-05-24 20:53:30 +08:00 committed by GitHub
parent 4ea0830c25
commit 1471da846d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 18 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
internal/rootcoord/quota_center.go
View File

@ -276,6 +276,10 @@ func (q *QuotaCenter) forceDenyWriting(errorCode commonpb.ErrorCode, collections
collections = q.writableCollections
}
for _, collection := range collections {
if _, ok := q.currentRates[collection]; !ok {
q.currentRates[collection] = make(map[internalpb.RateType]Limit)
q.quotaStates[collection] = make(map[milvuspb.QuotaState]commonpb.ErrorCode)
}
q.currentRates[collection][internalpb.RateType_DMLInsert] = 0
q.currentRates[collection][internalpb.RateType_DMLDelete] = 0
q.currentRates[collection][internalpb.RateType_DMLBulkLoad] = 0
@ -293,6 +297,10 @@ func (q *QuotaCenter) forceDenyReading(errorCode commonpb.ErrorCode, collections
collections = q.readableCollections
}
for _, collection := range collections {
if _, ok := q.currentRates[collection]; !ok {
q.currentRates[collection] = make(map[internalpb.RateType]Limit)
q.quotaStates[collection] = make(map[milvuspb.QuotaState]commonpb.ErrorCode)
}
q.currentRates[collection][internalpb.RateType_DQLSearch] = 0
q.currentRates[collection][internalpb.RateType_DQLQuery] = 0
q.quotaStates[collection][milvuspb.QuotaState_DenyToRead] = errorCode

14
internal/rootcoord/quota_center_test.go
View File

@ -108,19 +108,25 @@ func TestQuotaCenter(t *testing.T) {
quotaCenter := NewQuotaCenter(pcm, qc, &dataCoordMockForQuota{}, core.tsoAllocator)
quotaCenter.readableCollections = []int64{1, 2, 3}
quotaCenter.resetAllCurrentRates()
quotaCenter.forceDenyReading(commonpb.ErrorCode_ForceDeny)
quotaCenter.forceDenyReading(commonpb.ErrorCode_ForceDeny, 1, 2, 3, 4)
for _, collection := range quotaCenter.readableCollections {
assert.Equal(t, Limit(0), quotaCenter.currentRates[collection][internalpb.RateType_DQLQuery])
assert.Equal(t, Limit(0), quotaCenter.currentRates[collection][internalpb.RateType_DQLSearch])
assert.Equal(t, Limit(0), quotaCenter.currentRates[collection][internalpb.RateType_DQLQuery])
}
assert.Equal(t, Limit(0), quotaCenter.currentRates[4][internalpb.RateType_DQLSearch])
assert.Equal(t, Limit(0), quotaCenter.currentRates[4][internalpb.RateType_DQLQuery])
quotaCenter.writableCollections = []int64{1, 2, 3, 4}
quotaCenter.writableCollections = []int64{1, 2, 3}
quotaCenter.resetAllCurrentRates()
quotaCenter.forceDenyWriting(commonpb.ErrorCode_ForceDeny)
quotaCenter.forceDenyWriting(commonpb.ErrorCode_ForceDeny, 1, 2, 3, 4)
for _, collection := range quotaCenter.writableCollections {
assert.Equal(t, Limit(0), quotaCenter.currentRates[collection][internalpb.RateType_DMLInsert])
assert.Equal(t, Limit(0), quotaCenter.currentRates[collection][internalpb.RateType_DMLDelete])
assert.Equal(t, Limit(0), quotaCenter.currentRates[collection][internalpb.RateType_DMLBulkLoad])
}
assert.Equal(t, Limit(0), quotaCenter.currentRates[4][internalpb.RateType_DMLInsert])
assert.Equal(t, Limit(0), quotaCenter.currentRates[4][internalpb.RateType_DMLDelete])
assert.Equal(t, Limit(0), quotaCenter.currentRates[4][internalpb.RateType_DMLBulkLoad])
})
t.Run("test calculateRates", func(t *testing.T) {