Improve rate limiter interceptor (#24389)

Signed-off-by: bigsheeper <yihao.dai@zilliz.com>
This commit is contained in:
yihao.dai 2023-05-25 14:55:32 +08:00 committed by GitHub
parent 10c779d3d3
commit 6e17aa2608
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -66,41 +66,31 @@ func getRequestInfo(req interface{}) (int64, internalpb.RateType, int, error) {
return collectionID, internalpb.RateType_DQLSearch, int(r.GetNq()), nil return collectionID, internalpb.RateType_DQLSearch, int(r.GetNq()), nil
case *milvuspb.QueryRequest: case *milvuspb.QueryRequest:
collectionID, _ := globalMetaCache.GetCollectionID(context.TODO(), r.GetCollectionName()) collectionID, _ := globalMetaCache.GetCollectionID(context.TODO(), r.GetCollectionName())
return collectionID, internalpb.RateType_DQLQuery, 1, nil // think of the query request's nq as 1 return collectionID, internalpb.RateType_DQLQuery, 1, nil // we regard the nq of query as equivalent to 1.
case *milvuspb.CreateCollectionRequest: case *milvuspb.CreateCollectionRequest:
collectionID, _ := globalMetaCache.GetCollectionID(context.TODO(), r.GetCollectionName()) return 0, internalpb.RateType_DDLCollection, 1, nil
return collectionID, internalpb.RateType_DDLCollection, 1, nil
case *milvuspb.DropCollectionRequest: case *milvuspb.DropCollectionRequest:
collectionID, _ := globalMetaCache.GetCollectionID(context.TODO(), r.GetCollectionName()) return 0, internalpb.RateType_DDLCollection, 1, nil
return collectionID, internalpb.RateType_DDLCollection, 1, nil
case *milvuspb.LoadCollectionRequest: case *milvuspb.LoadCollectionRequest:
collectionID, _ := globalMetaCache.GetCollectionID(context.TODO(), r.GetCollectionName()) return 0, internalpb.RateType_DDLCollection, 1, nil
return collectionID, internalpb.RateType_DDLCollection, 1, nil
case *milvuspb.ReleaseCollectionRequest: case *milvuspb.ReleaseCollectionRequest:
collectionID, _ := globalMetaCache.GetCollectionID(context.TODO(), r.GetCollectionName()) return 0, internalpb.RateType_DDLCollection, 1, nil
return collectionID, internalpb.RateType_DDLCollection, 1, nil
case *milvuspb.CreatePartitionRequest: case *milvuspb.CreatePartitionRequest:
collectionID, _ := globalMetaCache.GetCollectionID(context.TODO(), r.GetCollectionName()) return 0, internalpb.RateType_DDLPartition, 1, nil
return collectionID, internalpb.RateType_DDLPartition, 1, nil
case *milvuspb.DropPartitionRequest: case *milvuspb.DropPartitionRequest:
collectionID, _ := globalMetaCache.GetCollectionID(context.TODO(), r.GetCollectionName()) return 0, internalpb.RateType_DDLPartition, 1, nil
return collectionID, internalpb.RateType_DDLPartition, 1, nil
case *milvuspb.LoadPartitionsRequest: case *milvuspb.LoadPartitionsRequest:
collectionID, _ := globalMetaCache.GetCollectionID(context.TODO(), r.GetCollectionName()) return 0, internalpb.RateType_DDLPartition, 1, nil
return collectionID, internalpb.RateType_DDLPartition, 1, nil
case *milvuspb.ReleasePartitionsRequest: case *milvuspb.ReleasePartitionsRequest:
collectionID, _ := globalMetaCache.GetCollectionID(context.TODO(), r.GetCollectionName()) return 0, internalpb.RateType_DDLPartition, 1, nil
return collectionID, internalpb.RateType_DDLPartition, 1, nil
case *milvuspb.CreateIndexRequest: case *milvuspb.CreateIndexRequest:
collectionID, _ := globalMetaCache.GetCollectionID(context.TODO(), r.GetCollectionName()) return 0, internalpb.RateType_DDLIndex, 1, nil
return collectionID, internalpb.RateType_DDLIndex, 1, nil
case *milvuspb.DropIndexRequest: case *milvuspb.DropIndexRequest:
collectionID, _ := globalMetaCache.GetCollectionID(context.TODO(), r.GetCollectionName()) return 0, internalpb.RateType_DDLIndex, 1, nil
return collectionID, internalpb.RateType_DDLIndex, 1, nil
case *milvuspb.FlushRequest: case *milvuspb.FlushRequest:
return 0, internalpb.RateType_DDLFlush, 1, nil return 0, internalpb.RateType_DDLFlush, 1, nil
case *milvuspb.ManualCompactionRequest: case *milvuspb.ManualCompactionRequest:
return r.GetCollectionID(), internalpb.RateType_DDLCompaction, 1, nil return 0, internalpb.RateType_DDLCompaction, 1, nil
// TODO: support more request // TODO: support more request
default: default:
if req == nil { if req == nil {