milvus-io
/
milvus
mirror of https://github.com/milvus-io/milvus.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

support collection level rate limit (#22767) 

Signed-off-by: Wei Liu <wei.liu@zilliz.com>
pull/23695/head
wei liu 2023-04-25 15:54:35 +08:00 committed by GitHub
parent 6653e2c3b0
commit 4fb8919a97
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
18 changed files with 5741 additions and 690 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
internal/datanode/metrics_info.go
View File

@ -48,6 +48,17 @@ func (node *DataNode) getQuotaMetrics() (*metricsinfo.DataNodeQuotaMetrics, erro
if err != nil {
return nil, err
}
getAllCollections := func() []int64 {
collectionSet := typeutil.UniqueSet{}
node.flowgraphManager.flowgraphs.Range(func(key, value any) bool {
fg := value.(*dataSyncService)
collectionSet.Insert(fg.channel.getCollectionID())
return true
})
return collectionSet.Collect()
}
minFGChannel, minFGTt := rateCol.getMinFlowGraphTt()
return &metricsinfo.DataNodeQuotaMetrics{
Hms: metricsinfo.HardwareMetrics{},
@ -57,6 +68,10 @@ func (node *DataNode) getQuotaMetrics() (*metricsinfo.DataNodeQuotaMetrics, erro
MinFlowGraphTt: minFGTt,
NumFlowGraph: node.flowgraphManager.getFlowGraphNum(),
},
Effect: metricsinfo.NodeEffect{
NodeID: node.GetSession().ServerID,
CollectionIDs: getAllCollections(),
},
}, nil
}

9
internal/proto/proxy.proto
View File

@ -50,9 +50,14 @@ message RefreshPolicyInfoCacheRequest {
string opKey = 3;
}
message SetRatesRequest {
common.MsgBase base = 1;
message CollectionRate {
int64 collection = 1;
repeated internal.Rate rates = 2;
repeated milvus.QuotaState states = 3;
repeated common.ErrorCode codes = 4;
}
message SetRatesRequest {
common.MsgBase base = 1;
repeated CollectionRate rates = 2;
}

166
internal/proto/proxypb/proxy.pb.go
View File

@ -251,8 +251,8 @@ func (m *RefreshPolicyInfoCacheRequest) GetOpKey() string {
return ""
}
type SetRatesRequest struct {
Base *commonpb.MsgBase `protobuf:"bytes,1,opt,name=base,proto3" json:"base,omitempty"`
type CollectionRate struct {
Collection int64 `protobuf:"varint,1,opt,name=collection,proto3" json:"collection,omitempty"`
Rates []*internalpb.Rate `protobuf:"bytes,2,rep,name=rates,proto3" json:"rates,omitempty"`
States []milvuspb.QuotaState `protobuf:"varint,3,rep,packed,name=states,proto3,enum=milvus.proto.milvus.QuotaState" json:"states,omitempty"`
Codes []commonpb.ErrorCode `protobuf:"varint,4,rep,packed,name=codes,proto3,enum=milvus.proto.common.ErrorCode" json:"codes,omitempty"`
@ -261,11 +261,72 @@ type SetRatesRequest struct {
XXX_sizecache int32 `json:"-"`
}
func (m *CollectionRate) Reset() { *m = CollectionRate{} }
func (m *CollectionRate) String() string { return proto.CompactTextString(m) }
func (*CollectionRate) ProtoMessage() {}
func (*CollectionRate) Descriptor() ([]byte, []int) {
return fileDescriptor_700b50b08ed8dbaf, []int{4}
}
func (m *CollectionRate) XXX_Unmarshal(b []byte) error {
return xxx_messageInfo_CollectionRate.Unmarshal(m, b)
}
func (m *CollectionRate) XXX_Marshal(b []byte, deterministic bool) ([]byte, error) {
return xxx_messageInfo_CollectionRate.Marshal(b, m, deterministic)
}
func (m *CollectionRate) XXX_Merge(src proto.Message) {
xxx_messageInfo_CollectionRate.Merge(m, src)
}
func (m *CollectionRate) XXX_Size() int {
return xxx_messageInfo_CollectionRate.Size(m)
}
func (m *CollectionRate) XXX_DiscardUnknown() {
xxx_messageInfo_CollectionRate.DiscardUnknown(m)
}
var xxx_messageInfo_CollectionRate proto.InternalMessageInfo
func (m *CollectionRate) GetCollection() int64 {
if m != nil {
return m.Collection
}
return 0
}
func (m *CollectionRate) GetRates() []*internalpb.Rate {
if m != nil {
return m.Rates
}
return nil
}
func (m *CollectionRate) GetStates() []milvuspb.QuotaState {
if m != nil {
return m.States
}
return nil
}
func (m *CollectionRate) GetCodes() []commonpb.ErrorCode {
if m != nil {
return m.Codes
}
return nil
}
type SetRatesRequest struct {
Base *commonpb.MsgBase `protobuf:"bytes,1,opt,name=base,proto3" json:"base,omitempty"`
Rates []*CollectionRate `protobuf:"bytes,2,rep,name=rates,proto3" json:"rates,omitempty"`
XXX_NoUnkeyedLiteral struct{} `json:"-"`
XXX_unrecognized []byte `json:"-"`
XXX_sizecache int32 `json:"-"`
}
func (m *SetRatesRequest) Reset() { *m = SetRatesRequest{} }
func (m *SetRatesRequest) String() string { return proto.CompactTextString(m) }
func (*SetRatesRequest) ProtoMessage() {}
func (*SetRatesRequest) Descriptor() ([]byte, []int) {
return fileDescriptor_700b50b08ed8dbaf, []int{4}
return fileDescriptor_700b50b08ed8dbaf, []int{5}
}
func (m *SetRatesRequest) XXX_Unmarshal(b []byte) error {
@ -293,78 +354,67 @@ func (m *SetRatesRequest) GetBase() *commonpb.MsgBase {
return nil
}
func (m *SetRatesRequest) GetRates() []*internalpb.Rate {
func (m *SetRatesRequest) GetRates() []*CollectionRate {
if m != nil {
return m.Rates
}
return nil
}
func (m *SetRatesRequest) GetStates() []milvuspb.QuotaState {
if m != nil {
return m.States
}
return nil
}
func (m *SetRatesRequest) GetCodes() []commonpb.ErrorCode {
if m != nil {
return m.Codes
}
return nil
}
func init() {
proto.RegisterType((*InvalidateCollMetaCacheRequest)(nil), "milvus.proto.proxy.InvalidateCollMetaCacheRequest")
proto.RegisterType((*InvalidateCredCacheRequest)(nil), "milvus.proto.proxy.InvalidateCredCacheRequest")
proto.RegisterType((*UpdateCredCacheRequest)(nil), "milvus.proto.proxy.UpdateCredCacheRequest")
proto.RegisterType((*RefreshPolicyInfoCacheRequest)(nil), "milvus.proto.proxy.RefreshPolicyInfoCacheRequest")
proto.RegisterType((*CollectionRate)(nil), "milvus.proto.proxy.CollectionRate")
proto.RegisterType((*SetRatesRequest)(nil), "milvus.proto.proxy.SetRatesRequest")
}
func init() { proto.RegisterFile("proxy.proto", fileDescriptor_700b50b08ed8dbaf) }
var fileDescriptor_700b50b08ed8dbaf = []byte{
// 625 bytes of a gzipped FileDescriptorProto
0x1f, 0x8b, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0xff, 0xb4, 0x54, 0x51, 0x4f, 0x13, 0x4d,
0x14, 0x65, 0x29, 0x2d, 0x7c, 0x97, 0x06, 0x92, 0x09, 0x1f, 0xd6, 0x22, 0xd8, 0x2c, 0x46, 0x1a,
0x12, 0x5b, 0xa9, 0x24, 0xbe, 0x53, 0x4c, 0x43, 0x0c, 0x04, 0xb7, 0xfa, 0xe2, 0x8b, 0x99, 0xdd,
0xbd, 0xd0, 0x21, 0xdb, 0x99, 0x65, 0x66, 0x16, 0xed, 0x93, 0x89, 0xff, 0xc8, 0x37, 0xff, 0x90,
0xff, 0xc3, 0xec, 0xce, 0x76, 0x61, 0xeb, 0x96, 0x8d, 0x12, 0xdf, 0x7a, 0x66, 0xcf, 0x9d, 0x73,
0xee, 0xed, 0xdc, 0x03, 0xab, 0xa1, 0x14, 0x5f, 0x26, 0x9d, 0x50, 0x0a, 0x2d, 0x08, 0x19, 0xb3,
0xe0, 0x26, 0x52, 0x06, 0x75, 0x92, 0x2f, 0xcd, 0xba, 0x27, 0xc6, 0x63, 0xc1, 0xcd, 0x59, 0x73,
0x8d, 0x71, 0x8d, 0x92, 0xd3, 0x20, 0xc5, 0xf5, 0xbb, 0x15, 0xf6, 0x0f, 0x0b, 0x76, 0x4e, 0xf8,
0x0d, 0x0d, 0x98, 0x4f, 0x35, 0xf6, 0x45, 0x10, 0x9c, 0xa2, 0xa6, 0x7d, 0xea, 0x8d, 0xd0, 0xc1,
0xeb, 0x08, 0x95, 0x26, 0x2f, 0x61, 0xc9, 0xa5, 0x0a, 0x1b, 0x56, 0xcb, 0x6a, 0xaf, 0xf6, 0x9e,
0x74, 0x72, 0x8a, 0xa9, 0xd4, 0xa9, 0xba, 0x3c, 0xa2, 0x0a, 0x9d, 0x84, 0x49, 0x1e, 0xc1, 0xb2,
0xef, 0x7e, 0xe2, 0x74, 0x8c, 0x8d, 0xc5, 0x96, 0xd5, 0xfe, 0xcf, 0xa9, 0xf9, 0xee, 0x19, 0x1d,
0x23, 0xd9, 0x83, 0x75, 0x4f, 0x04, 0x01, 0x7a, 0x9a, 0x09, 0x6e, 0x08, 0x95, 0x84, 0xb0, 0x76,
0x7b, 0x9c, 0x10, 0x6d, 0xa8, 0xdf, 0x9e, 0x9c, 0x1c, 0x37, 0x96, 0x5a, 0x56, 0xbb, 0xe2, 0xe4,
0xce, 0xec, 0x2b, 0x68, 0xde, 0x71, 0x2e, 0xd1, 0x7f, 0xa0, 0xeb, 0x26, 0xac, 0x44, 0x2a, 0x9e,
0x54, 0x66, 0x3b, 0xc3, 0xf6, 0x37, 0x0b, 0x36, 0x3f, 0x84, 0xff, 0x5e, 0x28, 0xfe, 0x16, 0x52,
0xa5, 0x3e, 0x0b, 0xe9, 0xa7, 0xa3, 0xc9, 0xb0, 0xfd, 0x15, 0xb6, 0x1d, 0xbc, 0x90, 0xa8, 0x46,
0xe7, 0x22, 0x60, 0xde, 0xe4, 0x84, 0x5f, 0x88, 0x07, 0x5a, 0xd9, 0x84, 0x9a, 0x08, 0xdf, 0x4f,
0x42, 0x63, 0xa4, 0xea, 0xa4, 0x88, 0x6c, 0x40, 0x55, 0x84, 0x6f, 0x71, 0x92, 0x7a, 0x30, 0xc0,
0xfe, 0x69, 0xc1, 0xfa, 0x10, 0xb5, 0x43, 0x35, 0xaa, 0xbf, 0xd7, 0x3c, 0x80, 0xaa, 0x8c, 0x6f,
0x68, 0x2c, 0xb6, 0x2a, 0xed, 0xd5, 0xde, 0x56, 0xbe, 0x24, 0x7b, 0xad, 0xb1, 0x8a, 0x63, 0x98,
0xe4, 0x35, 0xd4, 0x94, 0x4e, 0x6a, 0x2a, 0xad, 0x4a, 0x7b, 0xad, 0xf7, 0x34, 0x5f, 0x93, 0x82,
0x77, 0x91, 0xd0, 0x74, 0x18, 0xf3, 0x9c, 0x94, 0x4e, 0x0e, 0xa1, 0xea, 0x09, 0x1f, 0x55, 0x63,
0x29, 0xa9, 0xdb, 0x29, 0xb4, 0xf7, 0x46, 0x4a, 0x21, 0xfb, 0xc2, 0x47, 0xc7, 0x90, 0x7b, 0xdf,
0x97, 0xa1, 0x7a, 0x1e, 0xaf, 0x12, 0x09, 0x80, 0x0c, 0x50, 0xf7, 0xc5, 0x38, 0x14, 0x1c, 0xb9,
0x1e, 0x9a, 0x5b, 0x3b, 0x85, 0xf2, 0xbf, 0x13, 0xd3, 0x19, 0x35, 0x9f, 0x15, 0xf2, 0x67, 0xc8,
0xf6, 0x02, 0xb9, 0x86, 0x8d, 0x01, 0x26, 0x90, 0x29, 0xcd, 0x3c, 0xd5, 0x1f, 0x51, 0xce, 0x31,
0x20, 0xbd, 0x39, 0x23, 0x2a, 0x22, 0x4f, 0x35, 0x77, 0x0b, 0x35, 0x87, 0x5a, 0x32, 0x7e, 0xe9,
0xa0, 0x0a, 0x05, 0x57, 0x68, 0x2f, 0x10, 0x09, 0xdb, 0xf9, 0xf5, 0x37, 0xeb, 0x95, 0x85, 0xc0,
0xac, 0xb6, 0xc9, 0x9e, 0xfb, 0x13, 0xa3, 0xb9, 0x55, 0x38, 0xe6, 0xd8, 0x6a, 0x14, 0xb7, 0x49,
0xa1, 0x3e, 0x40, 0x7d, 0xec, 0x4f, 0xdb, 0xdb, 0x9f, 0xdf, 0x5e, 0x46, 0xfa, 0xc3, 0xb6, 0xae,
0xe0, 0x71, 0x3e, 0x1b, 0x90, 0x6b, 0x46, 0x03, 0xd3, 0x52, 0xa7, 0xa4, 0xa5, 0x99, 0x0d, 0x2f,
0x6b, 0xc7, 0x85, 0xff, 0x6f, 0xa3, 0xe1, 0xae, 0xce, 0x7e, 0x91, 0x4e, 0x71, 0x8a, 0x94, 0x69,
0x5c, 0xc1, 0x66, 0xf1, 0xea, 0x93, 0x83, 0x22, 0x91, 0x7b, 0x63, 0xa2, 0x4c, 0xcb, 0x87, 0xf5,
0x01, 0xea, 0xe4, 0xfd, 0x9f, 0xa2, 0x96, 0xcc, 0x53, 0xe4, 0xf9, 0xbc, 0x07, 0x9f, 0x12, 0xa6,
0x37, 0xef, 0x95, 0xf2, 0xb2, 0x7f, 0xe8, 0x0c, 0x56, 0xa6, 0x51, 0x42, 0x76, 0x8b, 0x7a, 0x98,
0x09, 0x9a, 0x12, 0xd7, 0x47, 0x87, 0x1f, 0x7b, 0x97, 0x4c, 0x8f, 0x22, 0x37, 0xfe, 0xd2, 0x35,
0xd4, 0x17, 0x4c, 0xa4, 0xbf, 0xba, 0xd3, 0x47, 0xd5, 0x4d, 0xaa, 0xbb, 0x89, 0x44, 0xe8, 0xba,
0xb5, 0x04, 0xbe, 0xfa, 0x15, 0x00, 0x00, 0xff, 0xff, 0x47, 0x77, 0x96, 0xe8, 0x54, 0x07, 0x00,
// 657 bytes of a gzipped FileDescriptorProto
0x1f, 0x8b, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0xff, 0xb4, 0x54, 0xd1, 0x4e, 0x13, 0x41,
0x14, 0x65, 0x29, 0x2d, 0x78, 0x69, 0x4a, 0x32, 0x41, 0xac, 0x45, 0xb0, 0x59, 0x8c, 0x34, 0x24,
0x6e, 0xa5, 0x92, 0xe8, 0x33, 0xc5, 0x34, 0xc4, 0x40, 0x70, 0xab, 0x2f, 0xbe, 0x98, 0xd9, 0xdd,
0x0b, 0x5d, 0xb2, 0x9d, 0x59, 0x66, 0xa6, 0x68, 0x1f, 0x8c, 0x89, 0x7f, 0xe4, 0x9b, 0xdf, 0xe1,
0x17, 0x99, 0xdd, 0xd9, 0x6e, 0xbb, 0x65, 0x61, 0xa3, 0xc4, 0xb7, 0xde, 0x99, 0x73, 0x7b, 0xce,
0xb9, 0x77, 0xf6, 0xc0, 0x6a, 0x28, 0xf8, 0xd7, 0xb1, 0x15, 0x0a, 0xae, 0x38, 0x21, 0x43, 0x3f,
0xb8, 0x1e, 0x49, 0x5d, 0x59, 0xf1, 0x4d, 0xa3, 0xea, 0xf2, 0xe1, 0x90, 0x33, 0x7d, 0xd6, 0xa8,
0xf9, 0x4c, 0xa1, 0x60, 0x34, 0x48, 0xea, 0xea, 0x6c, 0x87, 0xf9, 0xcb, 0x80, 0xed, 0x63, 0x76,
0x4d, 0x03, 0xdf, 0xa3, 0x0a, 0xbb, 0x3c, 0x08, 0x4e, 0x50, 0xd1, 0x2e, 0x75, 0x07, 0x68, 0xe3,
0xd5, 0x08, 0xa5, 0x22, 0x2f, 0x61, 0xc9, 0xa1, 0x12, 0xeb, 0x46, 0xd3, 0x68, 0xad, 0x76, 0x9e,
0x58, 0x19, 0xc6, 0x84, 0xea, 0x44, 0x5e, 0x1c, 0x52, 0x89, 0x76, 0x8c, 0x24, 0x8f, 0x60, 0xd9,
0x73, 0x3e, 0x33, 0x3a, 0xc4, 0xfa, 0x62, 0xd3, 0x68, 0x3d, 0xb0, 0x2b, 0x9e, 0x73, 0x4a, 0x87,
0x48, 0x76, 0x61, 0xcd, 0xe5, 0x41, 0x80, 0xae, 0xf2, 0x39, 0xd3, 0x80, 0x52, 0x0c, 0xa8, 0x4d,
0x8f, 0x63, 0xa0, 0x09, 0xd5, 0xe9, 0xc9, 0xf1, 0x51, 0x7d, 0xa9, 0x69, 0xb4, 0x4a, 0x76, 0xe6,
0xcc, 0xbc, 0x84, 0xc6, 0x8c, 0x72, 0x81, 0xde, 0x3d, 0x55, 0x37, 0x60, 0x65, 0x24, 0xa3, 0x49,
0xa5, 0xb2, 0xd3, 0xda, 0xfc, 0x61, 0xc0, 0xc6, 0xc7, 0xf0, 0xff, 0x13, 0x45, 0x77, 0x21, 0x95,
0xf2, 0x0b, 0x17, 0x5e, 0x32, 0x9a, 0xb4, 0x36, 0xbf, 0xc3, 0x96, 0x8d, 0xe7, 0x02, 0xe5, 0xe0,
0x8c, 0x07, 0xbe, 0x3b, 0x3e, 0x66, 0xe7, 0xfc, 0x9e, 0x52, 0x36, 0xa0, 0xc2, 0xc3, 0x0f, 0xe3,
0x50, 0x0b, 0x29, 0xdb, 0x49, 0x45, 0xd6, 0xa1, 0xcc, 0xc3, 0x77, 0x38, 0x4e, 0x34, 0xe8, 0xc2,
0xfc, 0x6d, 0x40, 0xad, 0x9b, 0xae, 0xc0, 0xa6, 0x0a, 0xc9, 0x36, 0xc0, 0x74, 0x29, 0x31, 0x71,
0xc9, 0x9e, 0x39, 0x21, 0xfb, 0x50, 0x16, 0x54, 0xa1, 0xac, 0x2f, 0x36, 0x4b, 0xad, 0xd5, 0xce,
0x66, 0x56, 0x53, 0xfa, 0x34, 0xa3, 0xff, 0xb2, 0x35, 0x92, 0xbc, 0x86, 0x8a, 0x54, 0x71, 0x4f,
0xa9, 0x59, 0x6a, 0xd5, 0x3a, 0x4f, 0xb3, 0x3d, 0x49, 0xf1, 0x7e, 0xc4, 0x15, 0xed, 0x47, 0x38,
0x3b, 0x81, 0x93, 0x03, 0x28, 0xbb, 0xdc, 0x43, 0x59, 0x5f, 0x8a, 0xfb, 0xb6, 0x73, 0xfd, 0xbf,
0x15, 0x82, 0x8b, 0x2e, 0xf7, 0xd0, 0xd6, 0x60, 0xf3, 0x1b, 0xac, 0xf5, 0x51, 0x45, 0x02, 0xe4,
0xbf, 0xcf, 0xf1, 0x4d, 0xd6, 0xa6, 0x69, 0xdd, 0xfc, 0x2c, 0xad, 0xec, 0xe4, 0x12, 0xb7, 0x9d,
0x9f, 0xcb, 0x50, 0x3e, 0x8b, 0xee, 0x49, 0x00, 0xa4, 0x87, 0xaa, 0xcb, 0x87, 0x21, 0x67, 0xc8,
0x54, 0x5f, 0x9b, 0xb2, 0x72, 0xdd, 0xdf, 0x04, 0x26, 0xda, 0x1b, 0xcf, 0x72, 0xf1, 0x73, 0x60,
0x73, 0x81, 0x5c, 0xc1, 0x7a, 0x0f, 0xe3, 0xd2, 0x97, 0xca, 0x77, 0x65, 0x77, 0x40, 0x19, 0xc3,
0x80, 0x74, 0x6e, 0xd9, 0x50, 0x1e, 0x78, 0xc2, 0xb9, 0x93, 0xcb, 0xd9, 0x57, 0xc2, 0x67, 0x17,
0x36, 0xca, 0x90, 0x33, 0x89, 0xe6, 0x02, 0x11, 0xb0, 0x95, 0x8d, 0x1a, 0x3d, 0x8d, 0x34, 0x70,
0xe6, 0xb9, 0xf5, 0xd8, 0xee, 0x4e, 0xa7, 0xc6, 0x66, 0xee, 0x76, 0x22, 0xa9, 0xa3, 0xc8, 0x26,
0x85, 0x6a, 0x0f, 0xd5, 0x91, 0x37, 0xb1, 0xb7, 0x77, 0xbb, 0xbd, 0x14, 0xf4, 0x97, 0xb6, 0x2e,
0xe1, 0x71, 0x36, 0x87, 0x90, 0x29, 0x9f, 0x06, 0xda, 0x92, 0x55, 0x60, 0x69, 0x2e, 0x4d, 0x8a,
0xec, 0x38, 0xf0, 0x70, 0x1a, 0x43, 0xb3, 0x3c, 0x7b, 0x79, 0x3c, 0xf9, 0x89, 0x55, 0xc4, 0x71,
0x09, 0x1b, 0xf9, 0x31, 0x43, 0xf6, 0xf3, 0x48, 0xee, 0x8c, 0xa4, 0x22, 0x2e, 0x0f, 0xd6, 0x7a,
0xa8, 0xe2, 0xf7, 0x7f, 0x82, 0x4a, 0xf8, 0xae, 0x24, 0xcf, 0x6f, 0x7b, 0xf0, 0x09, 0x60, 0xf2,
0xcf, 0xbb, 0x85, 0xb8, 0x74, 0x43, 0xa7, 0xb0, 0x32, 0xf9, 0xc4, 0xc9, 0x4e, 0x9e, 0x87, 0xb9,
0x00, 0x28, 0x50, 0x7d, 0x78, 0xf0, 0xa9, 0x73, 0xe1, 0xab, 0xc1, 0xc8, 0x89, 0x6e, 0xda, 0x1a,
0xfa, 0xc2, 0xe7, 0xc9, 0xaf, 0xf6, 0xe4, 0x51, 0xb5, 0xe3, 0xee, 0x76, 0x4c, 0x11, 0x3a, 0x4e,
0x25, 0x2e, 0x5f, 0xfd, 0x09, 0x00, 0x00, 0xff, 0xff, 0x9a, 0xf6, 0x3d, 0x98, 0xc0, 0x07, 0x00,
0x00,
}

17
internal/proxy/impl.go
View File

@ -26,7 +26,6 @@ import (
"github.com/cockroachdb/errors"
"github.com/golang/protobuf/proto"
"github.com/samber/lo"
"go.opentelemetry.io/otel"
"go.uber.org/zap"
"golang.org/x/sync/errgroup"
@ -53,7 +52,6 @@ import (
"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/timerecord"
"github.com/milvus-io/milvus/pkg/util/tsoutil"
"github.com/milvus-io/milvus/pkg/util/typeutil"
@ -4453,25 +4451,12 @@ func (node *Proxy) SetRates(ctx context.Context, request *proxypb.SetRatesReques
return resp, nil
}
err := node.multiRateLimiter.globalRateLimiter.setRates(request.GetRates())
err := node.multiRateLimiter.SetRates(request.GetRates())
// TODO: set multiple rate limiter rates
if err != nil {
resp.Reason = err.Error()
return resp, nil
}
node.multiRateLimiter.SetQuotaStates(request.GetStates(), request.GetCodes())
rateStrs := lo.FilterMap(request.GetRates(), func(r *internalpb.Rate, _ int) (string, bool) {
return fmt.Sprintf("rateType:%s, rate:%s", r.GetRt().String(), ratelimitutil.Limit(r.GetR()).String()),
ratelimitutil.Limit(r.GetR()) != ratelimitutil.Inf
})
if len(rateStrs) > 0 {
log.RatedInfo(30, "current rates in proxy", zap.Int64("proxyNodeID", paramtable.GetNodeID()), zap.Strings("rates", rateStrs))
}
if len(request.GetStates()) != 0 {
for i := range request.GetStates() {
log.Warn("Proxy set quota states", zap.String("state", request.GetStates()[i].String()), zap.String("reason", request.GetCodes()[i].String()))
}
}
resp.ErrorCode = commonpb.ErrorCode_Success
return resp, nil
}

8
internal/proxy/impl_test.go
View File

@ -141,7 +141,13 @@ func TestProxy_CheckHealth(t *testing.T) {
states := []milvuspb.QuotaState{milvuspb.QuotaState_DenyToWrite, milvuspb.QuotaState_DenyToRead}
codes := []commonpb.ErrorCode{commonpb.ErrorCode_MemoryQuotaExhausted, commonpb.ErrorCode_ForceDeny}
node.multiRateLimiter.SetQuotaStates(states, codes)
node.multiRateLimiter.SetRates([]*proxypb.CollectionRate{
{
Collection: 1,
States: states,
Codes: codes,
},
})
resp, err = node.CheckHealth(context.Background(), &milvuspb.CheckHealthRequest{})
assert.NoError(t, err)
assert.Equal(t, true, resp.IsHealthy)

151
internal/proxy/multi_rate_limiter.go
View File

@ -17,6 +17,7 @@
package proxy
import (
"context"
"fmt"
"strconv"
"sync"
@ -27,6 +28,7 @@ import (
"github.com/milvus-io/milvus-proto/go-api/commonpb"
"github.com/milvus-io/milvus-proto/go-api/milvuspb"
"github.com/milvus-io/milvus/internal/proto/internalpb"
"github.com/milvus-io/milvus/internal/proto/proxypb"
"github.com/milvus-io/milvus/pkg/config"
"github.com/milvus-io/milvus/pkg/log"
"github.com/milvus-io/milvus/pkg/metrics"
@ -49,27 +51,33 @@ func GetQuotaErrorString(errCode commonpb.ErrorCode) string {
// MultiRateLimiter includes multilevel rate limiters, such as global rateLimiter,
// collection level rateLimiter and so on. It also implements Limiter interface.
type MultiRateLimiter struct {
globalRateLimiter *rateLimiter
// TODO: add collection level rateLimiter
quotaStatesMu sync.RWMutex
quotaStates map[milvuspb.QuotaState]commonpb.ErrorCode
quotaStatesMu sync.RWMutex
collectionLimiters map[int64]*rateLimiter
}
// NewMultiRateLimiter returns a new MultiRateLimiter.
func NewMultiRateLimiter() *MultiRateLimiter {
m := &MultiRateLimiter{}
m.globalRateLimiter = newRateLimiter()
m := &MultiRateLimiter{
collectionLimiters: make(map[int64]*rateLimiter, 0),
}
return m
}
// Check checks if request would be limited or denied.
func (m *MultiRateLimiter) Check(rt internalpb.RateType, n int) commonpb.ErrorCode {
func (m *MultiRateLimiter) Check(collectionID int64, rt internalpb.RateType, n int) commonpb.ErrorCode {
if !Params.QuotaConfig.QuotaAndLimitsEnabled.GetAsBool() {
return commonpb.ErrorCode_Success
}
limit, rate := m.globalRateLimiter.limit(rt, n)
limiter := m.collectionLimiters[collectionID]
if limiter == nil {
return commonpb.ErrorCode_Success
}
limit, rate := limiter.limit(rt, n)
if rate == 0 {
return m.GetErrorCode(rt)
return limiter.getErrorCode(rt)
}
if limit {
return commonpb.ErrorCode_RateLimit
@ -77,52 +85,73 @@ func (m *MultiRateLimiter) Check(rt internalpb.RateType, n int) commonpb.ErrorCo
return commonpb.ErrorCode_Success
}
func (m *MultiRateLimiter) GetErrorCode(rt internalpb.RateType) commonpb.ErrorCode {
switch rt {
case internalpb.RateType_DMLInsert, internalpb.RateType_DMLDelete, internalpb.RateType_DMLBulkLoad:
m.quotaStatesMu.RLock()
defer m.quotaStatesMu.RUnlock()
return m.quotaStates[milvuspb.QuotaState_DenyToWrite]
case internalpb.RateType_DQLSearch, internalpb.RateType_DQLQuery:
m.quotaStatesMu.RLock()
defer m.quotaStatesMu.RUnlock()
return m.quotaStates[milvuspb.QuotaState_DenyToRead]
}
return commonpb.ErrorCode_Success
}
// GetQuotaStates returns quota states.
func (m *MultiRateLimiter) GetQuotaStates() ([]milvuspb.QuotaState, []string) {
m.quotaStatesMu.RLock()
defer m.quotaStatesMu.RUnlock()
states := make([]milvuspb.QuotaState, 0, len(m.quotaStates))
reasons := make([]string, 0, len(m.quotaStates))
for k, v := range m.quotaStates {
states = append(states, k)
reasons = append(reasons, GetQuotaErrorString(v))
serviceStates := make(map[milvuspb.QuotaState]typeutil.Set[commonpb.ErrorCode])
// deduplicate same (state, code) pair from different collection
for _, limiter := range m.collectionLimiters {
limiter.quotaStates.Range(func(state milvuspb.QuotaState, errCode commonpb.ErrorCode) bool {
if serviceStates[state] == nil {
serviceStates[state] = typeutil.NewSet[commonpb.ErrorCode]()
}
serviceStates[state].Insert(errCode)
return true
})
}
states := make([]milvuspb.QuotaState, 0)
reasons := make([]string, 0)
for state, errCodes := range serviceStates {
for errCode := range errCodes {
states = append(states, state)
reasons = append(reasons, GetQuotaErrorString(errCode))
}
}
return states, reasons
}
// SetQuotaStates sets quota states for MultiRateLimiter.
func (m *MultiRateLimiter) SetQuotaStates(states []milvuspb.QuotaState, codes []commonpb.ErrorCode) {
func (m *MultiRateLimiter) SetRates(rates []*proxypb.CollectionRate) error {
m.quotaStatesMu.Lock()
defer m.quotaStatesMu.Unlock()
m.quotaStates = make(map[milvuspb.QuotaState]commonpb.ErrorCode, len(states))
for i := 0; i < len(states); i++ {
m.quotaStates[states[i]] = codes[i]
collectionSet := typeutil.NewUniqueSet()
for _, collectionRates := range rates {
collectionSet.Insert(collectionRates.Collection)
rateLimiter, ok := m.collectionLimiters[collectionRates.GetCollection()]
if !ok {
rateLimiter = newRateLimiter()
}
err := rateLimiter.setRates(collectionRates)
if err != nil {
return err
}
m.collectionLimiters[collectionRates.GetCollection()] = rateLimiter
}
// remove dropped collection's rate limiter
for collectionID := range m.collectionLimiters {
if !collectionSet.Contain(collectionID) {
delete(m.collectionLimiters, collectionID)
}
}
return nil
}
// rateLimiter implements Limiter.
type rateLimiter struct {
limiters *typeutil.ConcurrentMap[internalpb.RateType, *ratelimitutil.Limiter]
limiters *typeutil.ConcurrentMap[internalpb.RateType, *ratelimitutil.Limiter]
quotaStates *typeutil.ConcurrentMap[milvuspb.QuotaState, commonpb.ErrorCode]
}
// newRateLimiter returns a new RateLimiter.
func newRateLimiter() *rateLimiter {
rl := &rateLimiter{
limiters: typeutil.NewConcurrentMap[internalpb.RateType, *ratelimitutil.Limiter](),
limiters: typeutil.NewConcurrentMap[internalpb.RateType, *ratelimitutil.Limiter](),
quotaStates: typeutil.NewConcurrentMap[milvuspb.QuotaState, commonpb.ErrorCode](),
}
rl.registerLimiters()
return rl
@ -138,34 +167,67 @@ func (rl *rateLimiter) limit(rt internalpb.RateType, n int) (bool, float64) {
return !limit.AllowN(time.Now(), n), float64(limit.Limit())
}
// setRates sets new rates for the limiters.
func (rl *rateLimiter) setRates(rates []*internalpb.Rate) error {
for _, r := range rates {
func (rl *rateLimiter) setRates(collectionRate *proxypb.CollectionRate) error {
log := log.Ctx(context.TODO()).WithRateGroup("proxy.rateLimiter", 1.0, 60.0).With(
zap.Int64("proxyNodeID", paramtable.GetNodeID()),
zap.Int64("CollectionID", collectionRate.Collection),
)
for _, r := range collectionRate.GetRates() {
if limit, ok := rl.limiters.Get(r.GetRt()); ok {
limit.SetLimit(ratelimitutil.Limit(r.GetR()))
setRateGaugeByRateType(r.GetRt(), paramtable.GetNodeID(), r.GetR())
setRateGaugeByRateType(r.GetRt(), paramtable.GetNodeID(), collectionRate.Collection, r.GetR())
} else {
return fmt.Errorf("unregister rateLimiter for rateType %s", r.GetRt().String())
}
log.RatedInfo(30, "current collection rates in proxy",
zap.String("rateType", r.Rt.String()),
zap.String("rateLimit", ratelimitutil.Limit(r.GetR()).String()),
)
}
// clear old quota states
rl.quotaStates = typeutil.NewConcurrentMap[milvuspb.QuotaState, commonpb.ErrorCode]()
for i := 0; i < len(collectionRate.GetStates()); i++ {
rl.quotaStates.Insert(collectionRate.States[i], collectionRate.Codes[i])
log.RatedWarn(30, "Proxy set collection quota states",
zap.String("state", collectionRate.GetStates()[i].String()),
zap.String("reason", collectionRate.GetCodes()[i].String()),
)
}
return nil
}
func (rl *rateLimiter) getErrorCode(rt internalpb.RateType) commonpb.ErrorCode {
switch rt {
case internalpb.RateType_DMLInsert, internalpb.RateType_DMLDelete, internalpb.RateType_DMLBulkLoad:
if errCode, ok := rl.quotaStates.Get(milvuspb.QuotaState_DenyToWrite); ok {
return errCode
}
case internalpb.RateType_DQLSearch, internalpb.RateType_DQLQuery:
if errCode, ok := rl.quotaStates.Get(milvuspb.QuotaState_DenyToRead); ok {
return errCode
}
}
return commonpb.ErrorCode_Success
}
// setRateGaugeByRateType sets ProxyLimiterRate metrics.
func setRateGaugeByRateType(rateType internalpb.RateType, nodeID int64, rate float64) {
func setRateGaugeByRateType(rateType internalpb.RateType, nodeID int64, collectionID int64, rate float64) {
if ratelimitutil.Limit(rate) == ratelimitutil.Inf {
return
}
nodeIDStr := strconv.FormatInt(nodeID, 10)
collectionIDStr := strconv.FormatInt(collectionID, 10)
switch rateType {
case internalpb.RateType_DMLInsert:
metrics.ProxyLimiterRate.WithLabelValues(nodeIDStr, metrics.InsertLabel).Set(rate)
metrics.ProxyLimiterRate.WithLabelValues(nodeIDStr, collectionIDStr, metrics.InsertLabel).Set(rate)
case internalpb.RateType_DMLDelete:
metrics.ProxyLimiterRate.WithLabelValues(nodeIDStr, metrics.DeleteLabel).Set(rate)
metrics.ProxyLimiterRate.WithLabelValues(nodeIDStr, collectionIDStr, metrics.DeleteLabel).Set(rate)
case internalpb.RateType_DQLSearch:
metrics.ProxyLimiterRate.WithLabelValues(nodeIDStr, metrics.SearchLabel).Set(rate)
metrics.ProxyLimiterRate.WithLabelValues(nodeIDStr, collectionIDStr, metrics.SearchLabel).Set(rate)
case internalpb.RateType_DQLQuery:
metrics.ProxyLimiterRate.WithLabelValues(nodeIDStr, metrics.QueryLabel).Set(rate)
metrics.ProxyLimiterRate.WithLabelValues(nodeIDStr, collectionIDStr, metrics.QueryLabel).Set(rate)
}
}
@ -181,6 +243,7 @@ func (rl *rateLimiter) printRates(rates []*internalpb.Rate) {
// registerLimiters register limiter for all rate types.
func (rl *rateLimiter) registerLimiters() {
log := log.Ctx(context.TODO()).WithRateGroup("proxy.rateLimiter", 1.0, 60.0)
quotaConfig := &Params.QuotaConfig
for rt := range internalpb.RateType_name {
var r *paramtable.ParamItem
@ -228,9 +291,9 @@ func (rl *rateLimiter) registerLimiters() {
}
}(internalpb.RateType(rt))
paramtable.Get().Watch(r.Key, config.NewHandler(fmt.Sprintf("rateLimiter-%d", rt), onEvent))
log.Info("RateLimiter register for rateType",
log.RatedInfo(30, "RateLimiter register for rateType",
zap.String("rateType", internalpb.RateType_name[rt]),
zap.String("rate", ratelimitutil.Limit(r.GetAsFloat()).String()),
zap.String("rateLimit", ratelimitutil.Limit(r.GetAsFloat()).String()),
zap.String("burst", fmt.Sprintf("%v", burst)))
}
}

129
internal/proxy/multi_rate_limiter_test.go
View File

@ -25,7 +25,9 @@ import (
"time"
"github.com/milvus-io/milvus-proto/go-api/commonpb"
"github.com/milvus-io/milvus-proto/go-api/milvuspb"
"github.com/milvus-io/milvus/internal/proto/internalpb"
"github.com/milvus-io/milvus/internal/proto/proxypb"
"github.com/milvus-io/milvus/pkg/util/etcd"
"github.com/milvus-io/milvus/pkg/util/paramtable"
"github.com/milvus-io/milvus/pkg/util/ratelimitutil"
@ -33,19 +35,21 @@ import (
)
func TestMultiRateLimiter(t *testing.T) {
collectionID := int64(1)
t.Run("test multiRateLimiter", func(t *testing.T) {
bak := Params.QuotaConfig.QuotaAndLimitsEnabled
paramtable.Get().Save(Params.QuotaConfig.QuotaAndLimitsEnabled.Key, "true")
multiLimiter := NewMultiRateLimiter()
multiLimiter.collectionLimiters[collectionID] = newRateLimiter()
for _, rt := range internalpb.RateType_value {
multiLimiter.globalRateLimiter.limiters.Insert(internalpb.RateType(rt), ratelimitutil.NewLimiter(ratelimitutil.Limit(1000), 1))
multiLimiter.collectionLimiters[collectionID].limiters.Insert(internalpb.RateType(rt), ratelimitutil.NewLimiter(ratelimitutil.Limit(1000), 1))
}
for _, rt := range internalpb.RateType_value {
errCode := multiLimiter.Check(internalpb.RateType(rt), 1)
errCode := multiLimiter.Check(collectionID, internalpb.RateType(rt), 1)
assert.Equal(t, commonpb.ErrorCode_Success, errCode)
errCode = multiLimiter.Check(internalpb.RateType(rt), math.MaxInt)
errCode = multiLimiter.Check(collectionID, internalpb.RateType(rt), math.MaxInt)
assert.Equal(t, commonpb.ErrorCode_Success, errCode)
errCode = multiLimiter.Check(internalpb.RateType(rt), math.MaxInt)
errCode = multiLimiter.Check(collectionID, internalpb.RateType(rt), math.MaxInt)
assert.Equal(t, commonpb.ErrorCode_RateLimit, errCode)
}
Params.QuotaConfig.QuotaAndLimitsEnabled = bak
@ -53,10 +57,11 @@ func TestMultiRateLimiter(t *testing.T) {
t.Run("not enable quotaAndLimit", func(t *testing.T) {
multiLimiter := NewMultiRateLimiter()
multiLimiter.collectionLimiters[collectionID] = newRateLimiter()
bak := Params.QuotaConfig.QuotaAndLimitsEnabled
paramtable.Get().Save(Params.QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
for _, rt := range internalpb.RateType_value {
errCode := multiLimiter.Check(internalpb.RateType(rt), 1)
errCode := multiLimiter.Check(collectionID, internalpb.RateType(rt), 1)
assert.Equal(t, commonpb.ErrorCode_Success, errCode)
}
Params.QuotaConfig.QuotaAndLimitsEnabled = bak
@ -69,7 +74,7 @@ func TestMultiRateLimiter(t *testing.T) {
multiLimiter := NewMultiRateLimiter()
bak := Params.QuotaConfig.QuotaAndLimitsEnabled
paramtable.Get().Save(Params.QuotaConfig.QuotaAndLimitsEnabled.Key, "true")
errCode := multiLimiter.Check(internalpb.RateType_DMLInsert, 1*1024*1024)
errCode := multiLimiter.Check(collectionID, internalpb.RateType_DMLInsert, 1*1024*1024)
assert.Equal(t, commonpb.ErrorCode_Success, errCode)
Params.QuotaConfig.QuotaAndLimitsEnabled = bak
Params.QuotaConfig.DMLMaxInsertRate = bakInsertRate
@ -80,6 +85,69 @@ func TestMultiRateLimiter(t *testing.T) {
run(math.MaxFloat64 / 3)
run(math.MaxFloat64 / 10000)
})
t.Run("test set rates", func(t *testing.T) {
multiLimiter := NewMultiRateLimiter()
zeroRates := make([]*internalpb.Rate, 0, len(internalpb.RateType_value))
for _, rt := range internalpb.RateType_value {
zeroRates = append(zeroRates, &internalpb.Rate{
Rt: internalpb.RateType(rt), R: 0,
})
}
err := multiLimiter.SetRates([]*proxypb.CollectionRate{
{
Collection: 1,
Rates: zeroRates,
},
{
Collection: 2,
Rates: zeroRates,
},
})
assert.NoError(t, err)
})
t.Run("test quota states", func(t *testing.T) {
multiLimiter := NewMultiRateLimiter()
zeroRates := make([]*internalpb.Rate, 0, len(internalpb.RateType_value))
for _, rt := range internalpb.RateType_value {
zeroRates = append(zeroRates, &internalpb.Rate{
Rt: internalpb.RateType(rt), R: 0,
})
}
err := multiLimiter.SetRates([]*proxypb.CollectionRate{
{
Collection: 1,
Rates: zeroRates,
States: []milvuspb.QuotaState{
milvuspb.QuotaState_DenyToWrite,
},
Codes: []commonpb.ErrorCode{
commonpb.ErrorCode_DiskQuotaExhausted,
},
},
{
Collection: 2,
Rates: zeroRates,
States: []milvuspb.QuotaState{
milvuspb.QuotaState_DenyToRead,
},
Codes: []commonpb.ErrorCode{
commonpb.ErrorCode_ForceDeny,
},
},
})
assert.NoError(t, err)
states, codes := multiLimiter.GetQuotaStates()
assert.Len(t, states, 2)
assert.Len(t, codes, 2)
assert.Contains(t, codes, GetQuotaErrorString(commonpb.ErrorCode_DiskQuotaExhausted))
assert.Contains(t, codes, GetQuotaErrorString(commonpb.ErrorCode_ForceDeny))
})
}
func TestRateLimiter(t *testing.T) {
@ -110,7 +178,10 @@ func TestRateLimiter(t *testing.T) {
Rt: internalpb.RateType(rt), R: 0,
})
}
err := limiter.setRates(zeroRates)
err := limiter.setRates(&proxypb.CollectionRate{
Collection: 1,
Rates: zeroRates,
})
assert.NoError(t, err)
for _, rt := range internalpb.RateType_value {
for i := 0; i < 100; i++ {
@ -118,6 +189,50 @@ func TestRateLimiter(t *testing.T) {
assert.True(t, ok)
}
}
err = limiter.setRates(&proxypb.CollectionRate{
Collection: 1,
States: []milvuspb.QuotaState{milvuspb.QuotaState_DenyToRead, milvuspb.QuotaState_DenyToWrite},
Codes: []commonpb.ErrorCode{commonpb.ErrorCode_DiskQuotaExhausted, commonpb.ErrorCode_DiskQuotaExhausted},
})
assert.NoError(t, err)
assert.Equal(t, limiter.quotaStates.Len(), 2)
err = limiter.setRates(&proxypb.CollectionRate{
Collection: 1,
States: []milvuspb.QuotaState{},
})
assert.NoError(t, err)
assert.Equal(t, limiter.quotaStates.Len(), 0)
})
t.Run("test get error code", func(t *testing.T) {
limiter := newRateLimiter()
for _, rt := range internalpb.RateType_value {
limiter.limiters.Insert(internalpb.RateType(rt), ratelimitutil.NewLimiter(ratelimitutil.Limit(1000), 1))
}
zeroRates := make([]*internalpb.Rate, 0, len(internalpb.RateType_value))
for _, rt := range internalpb.RateType_value {
zeroRates = append(zeroRates, &internalpb.Rate{
Rt: internalpb.RateType(rt), R: 0,
})
}
err := limiter.setRates(&proxypb.CollectionRate{
Collection: 1,
Rates: zeroRates,
States: []milvuspb.QuotaState{
milvuspb.QuotaState_DenyToWrite,
milvuspb.QuotaState_DenyToRead,
},
Codes: []commonpb.ErrorCode{
commonpb.ErrorCode_DiskQuotaExhausted,
commonpb.ErrorCode_ForceDeny,
},
})
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_ForceDeny, limiter.getErrorCode(internalpb.RateType_DQLQuery))
assert.Equal(t, commonpb.ErrorCode_DiskQuotaExhausted, limiter.getErrorCode(internalpb.RateType_DMLInsert))
})
t.Run("tests refresh rate by config", func(t *testing.T) {

7
internal/proxy/proxy_rpc_test.go
View File

@ -63,7 +63,12 @@ func TestProxyRpcLimit(t *testing.T) {
commonpbutil.WithMsgID(int64(0)),
commonpbutil.WithTimeStamp(0),
),
Rates: rates,
Rates: []*proxypb.CollectionRate{
{
Collection: 1,
Rates: rates,
},
},
}
_, err = client.SetRates(ctx, req)
// should be limited because of the rpc limit

107
internal/proxy/rate_limit_interceptor.go
View File

@ -22,65 +22,116 @@ import (
"reflect"
"github.com/golang/protobuf/proto"
"go.uber.org/zap"
"google.golang.org/grpc"
"github.com/milvus-io/milvus-proto/go-api/commonpb"
"github.com/milvus-io/milvus-proto/go-api/milvuspb"
"github.com/milvus-io/milvus/internal/proto/internalpb"
"github.com/milvus-io/milvus/internal/types"
"github.com/milvus-io/milvus/pkg/log"
)
// RateLimitInterceptor returns a new unary server interceptors that performs request rate limiting.
func RateLimitInterceptor(limiter types.Limiter) grpc.UnaryServerInterceptor {
return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
rt, n, err := getRequestInfo(req)
collectionID, rt, n, err := getRequestInfo(req)
if err != nil {
return handler(ctx, req)
}
code := limiter.Check(rt, n)
if code != commonpb.ErrorCode_Success {
rsp := getFailedResponse(req, rt, code, info.FullMethod)
if rsp != nil {
return rsp, nil
switch request := req.(type) {
case *milvuspb.FlushRequest:
collectionNames := request.GetCollectionNames()
for _, collectionName := range collectionNames {
collectionID, err = globalMetaCache.GetCollectionID(context.TODO(), collectionName)
if err != nil {
log.Info("skip limit check",
zap.String("method", info.FullMethod),
zap.Error(err),
)
return handler(ctx, req)
}
code := limiter.Check(collectionID, rt, n)
if code != commonpb.ErrorCode_Success {
rsp := getFailedResponse(req, rt, code, info.FullMethod)
if rsp != nil {
return rsp, nil
}
}
}
return handler(ctx, req)
default:
code := limiter.Check(collectionID, rt, n)
if code != commonpb.ErrorCode_Success {
rsp := getFailedResponse(req, rt, code, info.FullMethod)
if rsp != nil {
return rsp, nil
}
}
return handler(ctx, req)
}
return handler(ctx, req)
}
}
// getRequestInfo returns rateType of request and return tokens needed.
func getRequestInfo(req interface{}) (internalpb.RateType, int, error) {
// getRequestInfo returns collection name and rateType of request and return tokens needed.
func getRequestInfo(req interface{}) (int64, internalpb.RateType, int, error) {
switch r := req.(type) {
case *milvuspb.InsertRequest:
return internalpb.RateType_DMLInsert, proto.Size(r), nil
collectionID, _ := globalMetaCache.GetCollectionID(context.TODO(), r.GetCollectionName())
return collectionID, internalpb.RateType_DMLInsert, proto.Size(r), nil
case *milvuspb.DeleteRequest:
return internalpb.RateType_DMLDelete, proto.Size(r), nil
collectionID, _ := globalMetaCache.GetCollectionID(context.TODO(), r.GetCollectionName())
return collectionID, internalpb.RateType_DMLDelete, proto.Size(r), nil
case *milvuspb.ImportRequest:
return internalpb.RateType_DMLBulkLoad, proto.Size(r), nil
collectionID, _ := globalMetaCache.GetCollectionID(context.TODO(), r.GetCollectionName())
return collectionID, internalpb.RateType_DMLBulkLoad, proto.Size(r), nil
case *milvuspb.SearchRequest:
return internalpb.RateType_DQLSearch, int(r.GetNq()), nil
collectionID, _ := globalMetaCache.GetCollectionID(context.TODO(), r.GetCollectionName())
return collectionID, internalpb.RateType_DQLSearch, int(r.GetNq()), nil
case *milvuspb.QueryRequest:
return internalpb.RateType_DQLQuery, 1, nil // think of the query request's nq as 1
case *milvuspb.CreateCollectionRequest, *milvuspb.DropCollectionRequest:
return internalpb.RateType_DDLCollection, 1, nil
case *milvuspb.LoadCollectionRequest, *milvuspb.ReleaseCollectionRequest:
return internalpb.RateType_DDLCollection, 1, nil
case *milvuspb.CreatePartitionRequest, *milvuspb.DropPartitionRequest:
return internalpb.RateType_DDLPartition, 1, nil
case *milvuspb.LoadPartitionsRequest, *milvuspb.ReleasePartitionsRequest:
return internalpb.RateType_DDLPartition, 1, nil
case *milvuspb.CreateIndexRequest, *milvuspb.DropIndexRequest:
return internalpb.RateType_DDLIndex, 1, nil
collectionID, _ := globalMetaCache.GetCollectionID(context.TODO(), r.GetCollectionName())
return collectionID, internalpb.RateType_DQLQuery, 1, nil // think of the query request's nq as 1
case *milvuspb.CreateCollectionRequest:
collectionID, _ := globalMetaCache.GetCollectionID(context.TODO(), r.GetCollectionName())
return collectionID, internalpb.RateType_DDLCollection, 1, nil
case *milvuspb.DropCollectionRequest:
collectionID, _ := globalMetaCache.GetCollectionID(context.TODO(), r.GetCollectionName())
return collectionID, internalpb.RateType_DDLCollection, 1, nil
case *milvuspb.LoadCollectionRequest:
collectionID, _ := globalMetaCache.GetCollectionID(context.TODO(), r.GetCollectionName())
return collectionID, internalpb.RateType_DDLCollection, 1, nil
case *milvuspb.ReleaseCollectionRequest:
collectionID, _ := globalMetaCache.GetCollectionID(context.TODO(), r.GetCollectionName())
return collectionID, internalpb.RateType_DDLCollection, 1, nil
case *milvuspb.CreatePartitionRequest:
collectionID, _ := globalMetaCache.GetCollectionID(context.TODO(), r.GetCollectionName())
return collectionID, internalpb.RateType_DDLPartition, 1, nil
case *milvuspb.DropPartitionRequest:
collectionID, _ := globalMetaCache.GetCollectionID(context.TODO(), r.GetCollectionName())
return collectionID, internalpb.RateType_DDLPartition, 1, nil
case *milvuspb.LoadPartitionsRequest:
collectionID, _ := globalMetaCache.GetCollectionID(context.TODO(), r.GetCollectionName())
return collectionID, internalpb.RateType_DDLPartition, 1, nil
case *milvuspb.ReleasePartitionsRequest:
collectionID, _ := globalMetaCache.GetCollectionID(context.TODO(), r.GetCollectionName())
return collectionID, internalpb.RateType_DDLPartition, 1, nil
case *milvuspb.CreateIndexRequest:
collectionID, _ := globalMetaCache.GetCollectionID(context.TODO(), r.GetCollectionName())
return collectionID, internalpb.RateType_DDLIndex, 1, nil
case *milvuspb.DropIndexRequest:
collectionID, _ := globalMetaCache.GetCollectionID(context.TODO(), r.GetCollectionName())
return collectionID, internalpb.RateType_DDLIndex, 1, nil
case *milvuspb.FlushRequest:
return internalpb.RateType_DDLFlush, 1, nil
return 0, internalpb.RateType_DDLFlush, 1, nil
case *milvuspb.ManualCompactionRequest:
return internalpb.RateType_DDLCompaction, 1, nil
return r.GetCollectionID(), internalpb.RateType_DDLCompaction, 1, nil
// TODO: support more request
default:
if req == nil {
return 0, 0, fmt.Errorf("null request")
return 0, 0, 0, fmt.Errorf("null request")
}
return 0, 0, fmt.Errorf("unsupported request type %s", reflect.TypeOf(req).Name())
return 0, 0, 0, fmt.Errorf("unsupported request type %s", reflect.TypeOf(req).Name())
}
}

76
internal/proxy/rate_limit_interceptor_test.go
View File

@ -36,7 +36,7 @@ type limiterMock struct {
quotaStateReasons []commonpb.ErrorCode
}
func (l *limiterMock) Check(rt internalpb.RateType, n int) commonpb.ErrorCode {
func (l *limiterMock) Check(collection int64, rt internalpb.RateType, n int) commonpb.ErrorCode {
if l.rate == 0 {
return commonpb.ErrorCode_ForceDeny
}
@ -48,55 +48,107 @@ func (l *limiterMock) Check(rt internalpb.RateType, n int) commonpb.ErrorCode {
func TestRateLimitInterceptor(t *testing.T) {
t.Run("test getRequestInfo", func(t *testing.T) {
rt, size, err := getRequestInfo(&milvuspb.InsertRequest{})
globalMetaCache = newMockCache()
collection, rt, size, err := getRequestInfo(&milvuspb.InsertRequest{})
assert.NoError(t, err)
assert.Equal(t, proto.Size(&milvuspb.InsertRequest{}), size)
assert.Equal(t, internalpb.RateType_DMLInsert, rt)
assert.Equal(t, collection, int64(0))
rt, size, err = getRequestInfo(&milvuspb.DeleteRequest{})
collection, rt, size, err = getRequestInfo(&milvuspb.DeleteRequest{})
assert.NoError(t, err)
assert.Equal(t, proto.Size(&milvuspb.DeleteRequest{}), size)
assert.Equal(t, internalpb.RateType_DMLDelete, rt)
assert.Equal(t, collection, int64(0))
rt, size, err = getRequestInfo(&milvuspb.ImportRequest{})
collection, rt, size, err = getRequestInfo(&milvuspb.ImportRequest{})
assert.NoError(t, err)
assert.Equal(t, proto.Size(&milvuspb.ImportRequest{}), size)
assert.Equal(t, internalpb.RateType_DMLBulkLoad, rt)
assert.Equal(t, collection, int64(0))
rt, size, err = getRequestInfo(&milvuspb.SearchRequest{})
collection, rt, size, err = getRequestInfo(&milvuspb.SearchRequest{Nq: 5})
assert.NoError(t, err)
assert.Equal(t, proto.Size(&milvuspb.SearchRequest{}), size)
assert.Equal(t, 5, size)
assert.Equal(t, internalpb.RateType_DQLSearch, rt)
assert.Equal(t, collection, int64(0))
rt, size, err = getRequestInfo(&milvuspb.QueryRequest{})
collection, rt, size, err = getRequestInfo(&milvuspb.QueryRequest{})
assert.NoError(t, err)
assert.Equal(t, 1, size)
assert.Equal(t, internalpb.RateType_DQLQuery, rt)
assert.Equal(t, collection, int64(0))
rt, size, err = getRequestInfo(&milvuspb.CreateCollectionRequest{})
collection, rt, size, err = getRequestInfo(&milvuspb.CreateCollectionRequest{})
assert.NoError(t, err)
assert.Equal(t, 1, size)
assert.Equal(t, internalpb.RateType_DDLCollection, rt)
assert.Equal(t, collection, int64(0))
rt, size, err = getRequestInfo(&milvuspb.CreatePartitionRequest{})
collection, rt, size, err = getRequestInfo(&milvuspb.LoadCollectionRequest{})
assert.NoError(t, err)
assert.Equal(t, 1, size)
assert.Equal(t, internalpb.RateType_DDLCollection, rt)
assert.Equal(t, collection, int64(0))
collection, rt, size, err = getRequestInfo(&milvuspb.ReleaseCollectionRequest{})
assert.NoError(t, err)
assert.Equal(t, 1, size)
assert.Equal(t, internalpb.RateType_DDLCollection, rt)
assert.Equal(t, collection, int64(0))
collection, rt, size, err = getRequestInfo(&milvuspb.DropCollectionRequest{})
assert.NoError(t, err)
assert.Equal(t, 1, size)
assert.Equal(t, internalpb.RateType_DDLCollection, rt)
assert.Equal(t, collection, int64(0))
collection, rt, size, err = getRequestInfo(&milvuspb.CreatePartitionRequest{})
assert.NoError(t, err)
assert.Equal(t, 1, size)
assert.Equal(t, internalpb.RateType_DDLPartition, rt)
assert.Equal(t, collection, int64(0))
rt, size, err = getRequestInfo(&milvuspb.CreateIndexRequest{})
collection, rt, size, err = getRequestInfo(&milvuspb.LoadPartitionsRequest{})
assert.NoError(t, err)
assert.Equal(t, 1, size)
assert.Equal(t, internalpb.RateType_DDLPartition, rt)
assert.Equal(t, collection, int64(0))
collection, rt, size, err = getRequestInfo(&milvuspb.ReleasePartitionsRequest{})
assert.NoError(t, err)
assert.Equal(t, 1, size)
assert.Equal(t, internalpb.RateType_DDLPartition, rt)
assert.Equal(t, collection, int64(0))
collection, rt, size, err = getRequestInfo(&milvuspb.DropPartitionRequest{})
assert.NoError(t, err)
assert.Equal(t, 1, size)
assert.Equal(t, internalpb.RateType_DDLPartition, rt)
assert.Equal(t, collection, int64(0))
collection, rt, size, err = getRequestInfo(&milvuspb.CreateIndexRequest{})
assert.NoError(t, err)
assert.Equal(t, 1, size)
assert.Equal(t, internalpb.RateType_DDLIndex, rt)
assert.Equal(t, collection, int64(0))
rt, size, err = getRequestInfo(&milvuspb.FlushRequest{})
collection, rt, size, err = getRequestInfo(&milvuspb.DropIndexRequest{})
assert.NoError(t, err)
assert.Equal(t, 1, size)
assert.Equal(t, internalpb.RateType_DDLIndex, rt)
assert.Equal(t, collection, int64(0))
_, rt, size, err = getRequestInfo(&milvuspb.FlushRequest{})
assert.NoError(t, err)
assert.Equal(t, 1, size)
assert.Equal(t, internalpb.RateType_DDLFlush, rt)
rt, size, err = getRequestInfo(&milvuspb.ManualCompactionRequest{})
collection, rt, size, err = getRequestInfo(&milvuspb.ManualCompactionRequest{})
assert.NoError(t, err)
assert.Equal(t, 1, size)
assert.Equal(t, internalpb.RateType_DDLCompaction, rt)
assert.Equal(t, collection, int64(0))
})
t.Run("test getFailedResponse", func(t *testing.T) {

619
internal/rootcoord/quota_center.go
View File

@ -39,6 +39,7 @@ import (
"github.com/milvus-io/milvus/pkg/util/ratelimitutil"
"github.com/milvus-io/milvus/pkg/util/tsoutil"
"github.com/milvus-io/milvus/pkg/util/typeutil"
"github.com/samber/lo"
)
const (
@ -59,6 +60,10 @@ const Inf = ratelimitutil.Inf
type Limit = ratelimitutil.Limit
type collectionRates = map[internalpb.RateType]Limit
type collectionStates = map[milvuspb.QuotaState]commonpb.ErrorCode
// QuotaCenter manages the quota and limitations of the whole cluster,
// it receives metrics info from DataNodes, QueryNodes and Proxies, and
// notifies Proxies to limit rate of requests from clients or reject
@ -83,13 +88,15 @@ type QuotaCenter struct {
dataCoord types.DataCoord
// metrics
queryNodeMetrics map[UniqueID]*metricsinfo.QueryNodeQuotaMetrics
dataNodeMetrics map[UniqueID]*metricsinfo.DataNodeQuotaMetrics
proxyMetrics map[UniqueID]*metricsinfo.ProxyQuotaMetrics
dataCoordMetrics *metricsinfo.DataCoordQuotaMetrics
queryNodeMetrics map[UniqueID]*metricsinfo.QueryNodeQuotaMetrics
dataNodeMetrics map[UniqueID]*metricsinfo.DataNodeQuotaMetrics
proxyMetrics map[UniqueID]*metricsinfo.ProxyQuotaMetrics
dataCoordMetrics *metricsinfo.DataCoordQuotaMetrics
readableCollections []int64
writableCollections []int64
currentRates map[internalpb.RateType]Limit
quotaStates map[milvuspb.QuotaState]commonpb.ErrorCode
currentRates map[int64]collectionRates
quotaStates map[int64]collectionStates
tsoAllocator tso.Allocator
rateAllocateStrategy RateAllocateStrategy
@ -101,12 +108,14 @@ type QuotaCenter struct {
// NewQuotaCenter returns a new QuotaCenter.
func NewQuotaCenter(proxies *proxyClientManager, queryCoord types.QueryCoord, dataCoord types.DataCoord, tsoAllocator tso.Allocator) *QuotaCenter {
return &QuotaCenter{
proxies: proxies,
queryCoord: queryCoord,
dataCoord: dataCoord,
currentRates: make(map[internalpb.RateType]Limit),
quotaStates: make(map[milvuspb.QuotaState]commonpb.ErrorCode),
tsoAllocator: tsoAllocator,
proxies: proxies,
queryCoord: queryCoord,
dataCoord: dataCoord,
currentRates: make(map[int64]map[internalpb.RateType]Limit),
quotaStates: make(map[int64]map[milvuspb.QuotaState]commonpb.ErrorCode),
tsoAllocator: tsoAllocator,
readableCollections: make([]int64, 0),
writableCollections: make([]int64, 0),
rateAllocateStrategy: DefaultRateAllocateStrategy,
stopChan: make(chan struct{}),
@ -183,11 +192,15 @@ func (q *QuotaCenter) syncMetrics() error {
if err != nil {
return err
}
collections := typeutil.NewUniqueSet()
for _, queryNodeMetric := range queryCoordTopology.Cluster.ConnectedNodes {
if queryNodeMetric.QuotaMetrics != nil {
q.queryNodeMetrics[queryNodeMetric.ID] = queryNodeMetric.QuotaMetrics
collections.Insert(queryNodeMetric.QuotaMetrics.Effect.CollectionIDs...)
}
}
q.readableCollections = collections.Collect()
return nil
})
// get Data cluster metrics
@ -204,11 +217,15 @@ func (q *QuotaCenter) syncMetrics() error {
if err != nil {
return err
}
collections := typeutil.NewUniqueSet()
for _, dataNodeMetric := range dataCoordTopology.Cluster.ConnectedDataNodes {
if dataNodeMetric.QuotaMetrics != nil {
q.dataNodeMetrics[dataNodeMetric.ID] = dataNodeMetric.QuotaMetrics
collections.Insert(dataNodeMetric.QuotaMetrics.Effect.CollectionIDs...)
}
}
q.writableCollections = collections.Collect()
if dataCoordTopology.Cluster.Self.QuotaMetrics != nil {
q.dataCoordMetrics = dataCoordTopology.Cluster.Self.QuotaMetrics
}
@ -246,20 +263,36 @@ func (q *QuotaCenter) syncMetrics() error {
}
// forceDenyWriting sets dml rates to 0 to reject all dml requests.
func (q *QuotaCenter) forceDenyWriting(errorCode commonpb.ErrorCode) {
q.currentRates[internalpb.RateType_DMLInsert] = 0
q.currentRates[internalpb.RateType_DMLDelete] = 0
q.currentRates[internalpb.RateType_DMLBulkLoad] = 0
log.Warn("QuotaCenter force to deny writing", zap.String("reason", errorCode.String()))
q.quotaStates[milvuspb.QuotaState_DenyToWrite] = errorCode
func (q *QuotaCenter) forceDenyWriting(errorCode commonpb.ErrorCode, collections ...int64) {
if len(collections) == 0 && len(q.writableCollections) != 0 {
// default to all writable collections
collections = q.writableCollections
}
for _, collection := range collections {
q.currentRates[collection][internalpb.RateType_DMLInsert] = 0
q.currentRates[collection][internalpb.RateType_DMLDelete] = 0
q.currentRates[collection][internalpb.RateType_DMLBulkLoad] = 0
q.quotaStates[collection][milvuspb.QuotaState_DenyToWrite] = errorCode
}
log.Warn("QuotaCenter force to deny writing",
zap.Int64s("collectionIDs", collections),
zap.String("reason", errorCode.String()))
}
// forceDenyWriting sets dql rates to 0 to reject all dql requests.
func (q *QuotaCenter) forceDenyReading(errorCode commonpb.ErrorCode) {
q.currentRates[internalpb.RateType_DQLSearch] = 0
q.currentRates[internalpb.RateType_DQLQuery] = 0
log.Warn("QuotaCenter force to deny reading", zap.String("reason", errorCode.String()))
q.quotaStates[milvuspb.QuotaState_DenyToRead] = errorCode
// forceDenyReading sets dql rates to 0 to reject all dql requests.
func (q *QuotaCenter) forceDenyReading(errorCode commonpb.ErrorCode, collections ...int64) {
if len(collections) == 0 {
// default to all readable collections
collections = q.readableCollections
}
for _, collection := range collections {
q.currentRates[collection][internalpb.RateType_DQLSearch] = 0
q.currentRates[collection][internalpb.RateType_DQLQuery] = 0
q.quotaStates[collection][milvuspb.QuotaState_DenyToRead] = errorCode
}
log.Warn("QuotaCenter force to deny reading",
zap.Int64s("collectionIDs", collections),
zap.String("reason", errorCode.String()))
}
// getRealTimeRate return real time rate in Proxy.
@ -276,9 +309,11 @@ func (q *QuotaCenter) getRealTimeRate(rateType internalpb.RateType) float64 {
}
// guaranteeMinRate make sure the rate will not be less than the min rate.
func (q *QuotaCenter) guaranteeMinRate(minRate float64, rateType internalpb.RateType) {
if minRate > 0 && q.currentRates[rateType] < Limit(minRate) {
q.currentRates[rateType] = Limit(minRate)
func (q *QuotaCenter) guaranteeMinRate(minRate float64, rateType internalpb.RateType, collections ...int64) {
for _, collection := range collections {
if minRate > 0 && q.currentRates[collection][rateType] < Limit(minRate) {
q.currentRates[collection][rateType] = Limit(minRate)
}
}
}
@ -289,42 +324,87 @@ func (q *QuotaCenter) calculateReadRates() {
return
}
enableQueueProtection := Params.QuotaConfig.QueueProtectionEnabled.GetAsBool()
if !enableQueueProtection {
return
}
limitCollectionSet := typeutil.NewUniqueSet()
// query latency
queueLatencyThreshold := Params.QuotaConfig.QueueLatencyThreshold.GetAsDuration(time.Second)
// queueLatencyThreshold >= 0 means enable queue latency protection
if queueLatencyThreshold >= 0 {
for _, metric := range q.queryNodeMetrics {
searchLatency := metric.SearchQueue.AvgQueueDuration
queryLatency := metric.QueryQueue.AvgQueueDuration
if searchLatency >= queueLatencyThreshold || queryLatency >= queueLatencyThreshold {
limitCollectionSet.Insert(metric.Effect.CollectionIDs...)
}
}
}
// queue length
nqInQueueThreshold := Params.QuotaConfig.NQInQueueThreshold.GetAsInt64()
if nqInQueueThreshold >= 0 {
// >= 0 means enable queue length protection
sum := func(ri metricsinfo.ReadInfoInQueue) int64 {
return ri.UnsolvedQueue + ri.ReadyQueue + ri.ReceiveChan + ri.ExecuteChan
}
for _, metric := range q.queryNodeMetrics {
searchNQSum := sum(metric.SearchQueue)
queryTasksSum := sum(metric.QueryQueue)
nqInQueue := searchNQSum + queryTasksSum // We think of the NQ of query request as 1.
if nqInQueue >= nqInQueueThreshold {
limitCollectionSet.Insert(metric.Effect.CollectionIDs...)
}
}
}
// read result
enableResultProtection := Params.QuotaConfig.ResultProtectionEnabled.GetAsBool()
if enableResultProtection {
maxRate := Params.QuotaConfig.MaxReadResultRate.GetAsFloat()
rateCount := float64(0)
for _, metric := range q.proxyMetrics {
for _, rm := range metric.Rms {
if rm.Label == metricsinfo.ReadResultThroughput {
rateCount += rm.Rate
}
}
}
if rateCount >= maxRate {
limitCollectionSet.Insert(q.readableCollections...)
}
}
coolOffSpeed := Params.QuotaConfig.CoolOffSpeed.GetAsFloat()
coolOff := func(realTimeSearchRate float64, realTimeQueryRate float64) {
if q.currentRates[internalpb.RateType_DQLSearch] != Inf && realTimeSearchRate > 0 {
q.currentRates[internalpb.RateType_DQLSearch] = Limit(realTimeSearchRate * coolOffSpeed)
coolOff := func(realTimeSearchRate float64, realTimeQueryRate float64, collections ...int64) {
for _, collection := range collections {
if q.currentRates[collection][internalpb.RateType_DQLSearch] != Inf && realTimeSearchRate > 0 {
q.currentRates[collection][internalpb.RateType_DQLSearch] = Limit(realTimeSearchRate * coolOffSpeed)
log.Warn("QuotaCenter cool read rates off done",
zap.Int64("collectionID", collection),
zap.Any("searchRate", q.currentRates[collection][internalpb.RateType_DQLSearch]))
}
if q.currentRates[collection][internalpb.RateType_DQLQuery] != Inf && realTimeQueryRate > 0 {
q.currentRates[collection][internalpb.RateType_DQLQuery] = Limit(realTimeQueryRate * coolOffSpeed)
log.Warn("QuotaCenter cool read rates off done",
zap.Int64("collectionID", collection),
zap.Any("queryRate", q.currentRates[collection][internalpb.RateType_DQLQuery]))
}
}
if q.currentRates[internalpb.RateType_DQLQuery] != Inf && realTimeSearchRate > 0 {
q.currentRates[internalpb.RateType_DQLQuery] = Limit(realTimeQueryRate * coolOffSpeed)
}
q.guaranteeMinRate(Params.QuotaConfig.DQLMinSearchRate.GetAsFloat(), internalpb.RateType_DQLSearch)
q.guaranteeMinRate(Params.QuotaConfig.DQLMinQueryRate.GetAsFloat(), internalpb.RateType_DQLQuery)
log.Warn("QuotaCenter cool read rates off done",
zap.Any("searchRate", q.currentRates[internalpb.RateType_DQLSearch]),
zap.Any("queryRate", q.currentRates[internalpb.RateType_DQLQuery]))
log.Info("QueryNodeMetrics when cool-off", zap.Any("metrics", q.queryNodeMetrics))
q.guaranteeMinRate(Params.QuotaConfig.DQLMinSearchRate.GetAsFloat(), internalpb.RateType_DQLSearch, collections...)
q.guaranteeMinRate(Params.QuotaConfig.DQLMinQueryRate.GetAsFloat(), internalpb.RateType_DQLQuery, collections...)
log.Info("QueryNodeMetrics when cool-off",
zap.Any("metrics", q.queryNodeMetrics))
}
// TODO: unify search and query?
realTimeSearchRate := q.getRealTimeRate(internalpb.RateType_DQLSearch)
realTimeQueryRate := q.getRealTimeRate(internalpb.RateType_DQLQuery)
queueLatencyFactor := q.getQueryLatencyFactor()
if Limit(queueLatencyFactor) == Limit(coolOffSpeed) {
coolOff(realTimeSearchRate, realTimeQueryRate)
return
}
queueLengthFactor := q.getNQInQueryFactor()
if Limit(queueLengthFactor) == Limit(coolOffSpeed) {
coolOff(realTimeSearchRate, realTimeQueryRate)
return
}
resultRateFactor := q.getReadResultFactor()
if Limit(resultRateFactor) == Limit(coolOffSpeed) {
coolOff(realTimeSearchRate, realTimeQueryRate)
}
coolOff(realTimeSearchRate, realTimeQueryRate, limitCollectionSet.Collect()...)
}
// calculateWriteRates calculates and sets dml rates.
@ -340,40 +420,128 @@ func (q *QuotaCenter) calculateWriteRates() error {
return nil
}
collectionFactor := map[int64]float64{}
for _, collection := range q.writableCollections {
collectionFactor[collection] = 1.0
}
ts, err := q.tsoAllocator.GenerateTSO(1)
if err != nil {
return err
}
ttFactor := q.getTimeTickDelayFactor(ts)
if ttFactor <= 0 {
q.forceDenyWriting(commonpb.ErrorCode_TimeTickLongDelay) // tt protection
return nil
maxDelay := Params.QuotaConfig.MaxTimeTickDelay.GetAsDuration(time.Second)
curTs, _ := tsoutil.ParseTS(ts)
processTtDelay := func(role string, nodeID int64, vchannel string, minTs time.Time, delay time.Duration, effectedCollections []int64) {
log := log.With(zap.String("role", role),
zap.Int64("nodeID", nodeID),
zap.String("vchannel", vchannel),
zap.Time("curTs", curTs),
zap.Time("minTs", minTs),
zap.Duration("delay", delay),
zap.Duration("MaxDelay", maxDelay),
zap.Int64s("effectedCollections", effectedCollections),
)
factor := float64(maxDelay.Nanoseconds()-delay.Nanoseconds()) / float64(maxDelay.Nanoseconds())
if factor <= 0 {
log.Warn("QuotaCenter: force deny writing due to long timeTick delay")
q.forceDenyWriting(commonpb.ErrorCode_TimeTickLongDelay, effectedCollections...)
}
for _, collection := range effectedCollections {
if factor < collectionFactor[collection] {
collectionFactor[collection] = factor
}
}
}
memFactor := q.getMemoryFactor()
if memFactor <= 0 {
q.forceDenyWriting(commonpb.ErrorCode_MemoryQuotaExhausted) // memory protection
return nil
processMemUsage := func(role string, nodeID int64, hms metricsinfo.HardwareMetrics, memLowLevel float64, memHighLevel float64, effectedCollections []int64) {
memoryWaterLevel := float64(hms.MemoryUsage) / float64(hms.Memory)
log := log.With(
zap.String("role", role),
zap.Int64("nodeID", nodeID),
zap.Uint64("usedMem", hms.MemoryUsage),
zap.Uint64("totalMem", hms.Memory),
zap.Float64("memoryWaterLevel", memoryWaterLevel),
zap.Float64("memoryHighWaterLevel", memHighLevel),
zap.Int64s("effectedCollections", effectedCollections),
)
factor := (memHighLevel - memoryWaterLevel) / (memHighLevel - memLowLevel)
if factor <= 0 {
log.Warn("QuotaCenter: force deny writing due to memory reach high water")
q.forceDenyWriting(commonpb.ErrorCode_MemoryQuotaExhausted, effectedCollections...)
}
if factor < 0.9 {
log.Warn("QuotaCenter: Node memory to low water level, limit writing rate",
zap.Float64("factor", factor))
}
for _, collection := range effectedCollections {
if factor < collectionFactor[collection] {
collectionFactor[collection] = factor
}
}
}
if memFactor < ttFactor {
ttFactor = memFactor
enableTtProtection := Params.QuotaConfig.TtProtectionEnabled.GetAsBool()
enableMemProtection := Params.QuotaConfig.MemProtectionEnabled.GetAsBool()
dataNodeMemoryLowWaterLevel := Params.QuotaConfig.DataNodeMemoryLowWaterLevel.GetAsFloat()
dataNodeMemoryHighWaterLevel := Params.QuotaConfig.DataNodeMemoryHighWaterLevel.GetAsFloat()
queryNodeMemoryLowWaterLevel := Params.QuotaConfig.QueryNodeMemoryLowWaterLevel.GetAsFloat()
queryNodeMemoryHighWaterLevel := Params.QuotaConfig.QueryNodeMemoryHighWaterLevel.GetAsFloat()
for nodeID, metric := range q.queryNodeMetrics {
if enableTtProtection && maxDelay >= 0 && metric.Fgm.NumFlowGraph > 0 && metric.Fgm.MinFlowGraphChannel != "" {
minTs, _ := tsoutil.ParseTS(metric.Fgm.MinFlowGraphTt)
delay := curTs.Sub(minTs)
processTtDelay(typeutil.QueryNodeRole, nodeID, metric.Fgm.MinFlowGraphChannel, minTs, delay, metric.Effect.CollectionIDs)
metrics.RootCoordTtDelay.WithLabelValues(typeutil.QueryNodeRole, strconv.FormatInt(nodeID, 10)).Set(float64(delay.Milliseconds()))
}
if enableMemProtection {
processMemUsage(typeutil.QueryNodeRole, nodeID, metric.Hms, queryNodeMemoryLowWaterLevel, queryNodeMemoryHighWaterLevel, metric.Effect.CollectionIDs)
}
}
if q.currentRates[internalpb.RateType_DMLInsert] != Inf {
q.currentRates[internalpb.RateType_DMLInsert] *= Limit(ttFactor)
for nodeID, metric := range q.dataNodeMetrics {
if enableTtProtection && maxDelay >= 0 && metric.Fgm.NumFlowGraph > 0 && metric.Fgm.MinFlowGraphChannel != "" {
minTs, _ := tsoutil.ParseTS(metric.Fgm.MinFlowGraphTt)
delay := curTs.Sub(minTs)
processTtDelay(typeutil.DataNodeRole, nodeID, metric.Fgm.MinFlowGraphChannel, minTs, delay, metric.Effect.CollectionIDs)
metrics.RootCoordTtDelay.WithLabelValues(typeutil.DataNodeRole, strconv.FormatInt(nodeID, 10)).Set(float64(delay.Milliseconds()))
}
if enableMemProtection {
processMemUsage(typeutil.QueryNodeRole, nodeID, metric.Hms, dataNodeMemoryLowWaterLevel, dataNodeMemoryHighWaterLevel, metric.Effect.CollectionIDs)
}
}
if q.currentRates[internalpb.RateType_DMLDelete] != Inf {
q.currentRates[internalpb.RateType_DMLDelete] *= Limit(ttFactor)
for collection, cf := range collectionFactor {
if q.currentRates[collection][internalpb.RateType_DMLInsert] != Inf && cf > 0 {
q.currentRates[collection][internalpb.RateType_DMLInsert] *= Limit(cf)
}
if q.currentRates[collection][internalpb.RateType_DMLDelete] != Inf && cf > 0 {
q.currentRates[collection][internalpb.RateType_DMLDelete] *= Limit(cf)
}
q.guaranteeMinRate(Params.QuotaConfig.DMLMinInsertRate.GetAsFloat(), internalpb.RateType_DMLInsert)
q.guaranteeMinRate(Params.QuotaConfig.DMLMinDeleteRate.GetAsFloat(), internalpb.RateType_DMLDelete)
log.Warn("QuotaCenter cool write rates off done",
zap.Int64("collectionID", collection),
zap.Float64("factor", cf))
}
q.guaranteeMinRate(Params.QuotaConfig.DMLMinInsertRate.GetAsFloat(), internalpb.RateType_DMLInsert)
q.guaranteeMinRate(Params.QuotaConfig.DMLMinDeleteRate.GetAsFloat(), internalpb.RateType_DMLDelete)
return nil
}
// calculateRates calculates target rates by different strategies.
func (q *QuotaCenter) calculateRates() error {
q.resetCurrentRates()
q.resetAllCurrentRates()
err := q.calculateWriteRates()
if err != nil {
@ -385,236 +553,45 @@ func (q *QuotaCenter) calculateRates() error {
return nil
}
func (q *QuotaCenter) resetAllCurrentRates() {
q.quotaStates = make(map[int64]map[milvuspb.QuotaState]commonpb.ErrorCode)
q.currentRates = map[int64]map[internalpb.RateType]ratelimitutil.Limit{}
for _, collection := range q.writableCollections {
q.resetCurrentRate(internalpb.RateType_DMLInsert, collection)
q.resetCurrentRate(internalpb.RateType_DMLDelete, collection)
q.resetCurrentRate(internalpb.RateType_DMLBulkLoad, collection)
}
for _, collection := range q.readableCollections {
q.resetCurrentRate(internalpb.RateType_DQLSearch, collection)
q.resetCurrentRate(internalpb.RateType_DQLQuery, collection)
}
}
// resetCurrentRates resets all current rates to configured rates.
func (q *QuotaCenter) resetCurrentRates() {
for _, rateType := range internalpb.RateType_value {
rt := internalpb.RateType(rateType)
switch rt {
case internalpb.RateType_DMLInsert:
q.currentRates[rt] = Limit(Params.QuotaConfig.DMLMaxInsertRate.GetAsFloat())
case internalpb.RateType_DMLDelete:
q.currentRates[rt] = Limit(Params.QuotaConfig.DMLMaxDeleteRate.GetAsFloat())
case internalpb.RateType_DMLBulkLoad:
q.currentRates[rt] = Limit(Params.QuotaConfig.DMLMaxBulkLoadRate.GetAsFloat())
case internalpb.RateType_DQLSearch:
q.currentRates[rt] = Limit(Params.QuotaConfig.DQLMaxSearchRate.GetAsFloat())
case internalpb.RateType_DQLQuery:
q.currentRates[rt] = Limit(Params.QuotaConfig.DQLMaxQueryRate.GetAsFloat())
}
if q.currentRates[rt] < 0 {
q.currentRates[rt] = Inf // no limit
}
}
q.quotaStates = make(map[milvuspb.QuotaState]commonpb.ErrorCode)
}
// getTimeTickDelayFactor gets time tick delay of DataNodes and QueryNodes,
// and return the factor according to max tolerable time tick delay.
func (q *QuotaCenter) getTimeTickDelayFactor(ts Timestamp) float64 {
var curMaxDelay time.Duration
var role, vchannel string
var minTt time.Time
t1, _ := tsoutil.ParseTS(ts)
for nodeID, metric := range q.queryNodeMetrics {
if metric.Fgm.NumFlowGraph > 0 && metric.Fgm.MinFlowGraphChannel != "" {
t2, _ := tsoutil.ParseTS(metric.Fgm.MinFlowGraphTt)
delay := t1.Sub(t2)
if delay.Nanoseconds() > curMaxDelay.Nanoseconds() {
curMaxDelay = delay
role = fmt.Sprintf("%s-%d", typeutil.QueryNodeRole, nodeID)
vchannel = metric.Fgm.MinFlowGraphChannel
minTt = t2
}
metrics.RootCoordTtDelay.WithLabelValues(typeutil.QueryNodeRole, strconv.FormatInt(nodeID, 10)).Set(float64(curMaxDelay.Milliseconds()))
}
}
for nodeID, metric := range q.dataNodeMetrics {
if metric.Fgm.NumFlowGraph > 0 && metric.Fgm.MinFlowGraphChannel != "" {
t2, _ := tsoutil.ParseTS(metric.Fgm.MinFlowGraphTt)
delay := t1.Sub(t2)
if delay.Nanoseconds() > curMaxDelay.Nanoseconds() {
curMaxDelay = delay
role = fmt.Sprintf("%s-%d", typeutil.DataNodeRole, nodeID)
vchannel = metric.Fgm.MinFlowGraphChannel
minTt = t2
}
metrics.RootCoordTtDelay.WithLabelValues(typeutil.DataNodeRole, strconv.FormatInt(nodeID, 10)).Set(float64(curMaxDelay.Milliseconds()))
}
func (q *QuotaCenter) resetCurrentRate(rt internalpb.RateType, collection int64) {
if q.currentRates[collection] == nil {
q.currentRates[collection] = make(map[internalpb.RateType]ratelimitutil.Limit)
}
if !Params.QuotaConfig.TtProtectionEnabled.GetAsBool() {
return 1
if q.quotaStates[collection] == nil {
q.quotaStates[collection] = make(map[milvuspb.QuotaState]commonpb.ErrorCode)
}
maxDelay := Params.QuotaConfig.MaxTimeTickDelay.GetAsDuration(time.Second)
if maxDelay < 0 {
// < 0 means disable tt protection
return 1
switch rt {
case internalpb.RateType_DMLInsert:
q.currentRates[collection][rt] = Limit(Params.QuotaConfig.DMLMaxInsertRate.GetAsFloat())
case internalpb.RateType_DMLDelete:
q.currentRates[collection][rt] = Limit(Params.QuotaConfig.DMLMaxDeleteRate.GetAsFloat())
case internalpb.RateType_DMLBulkLoad:
q.currentRates[collection][rt] = Limit(Params.QuotaConfig.DMLMaxBulkLoadRate.GetAsFloat())
case internalpb.RateType_DQLSearch:
q.currentRates[collection][rt] = Limit(Params.QuotaConfig.DQLMaxSearchRate.GetAsFloat())
case internalpb.RateType_DQLQuery:
q.currentRates[collection][rt] = Limit(Params.QuotaConfig.DQLMaxQueryRate.GetAsFloat())
}
if curMaxDelay.Nanoseconds() >= maxDelay.Nanoseconds() {
log.Warn("QuotaCenter force deny writing due to long timeTick delay",
zap.String("node", role),
zap.String("vchannel", vchannel),
zap.Time("curTs", t1),
zap.Time("minTs", minTt),
zap.Duration("delay", curMaxDelay),
zap.Duration("MaxDelay", maxDelay))
log.Info("DataNode and QueryNode Metrics",
zap.Any("QueryNodeMetrics", q.queryNodeMetrics),
zap.Any("DataNodeMetrics", q.dataNodeMetrics))
return 0
if q.currentRates[collection][rt] < 0 {
q.currentRates[collection][rt] = Inf // no limit
}
factor := float64(maxDelay.Nanoseconds()-curMaxDelay.Nanoseconds()) / float64(maxDelay.Nanoseconds())
if factor <= 0.9 {
log.Warn("QuotaCenter: limit writing due to long timeTick delay",
zap.String("node", role),
zap.String("vchannel", vchannel),
zap.Time("curTs", t1),
zap.Time("minTs", minTt),
zap.Duration("delay", curMaxDelay),
zap.Duration("MaxDelay", maxDelay),
zap.Float64("factor", factor))
}
return factor
}
// getNQInQueryFactor checks search&query nq in QueryNode,
// and return the factor according to NQInQueueThreshold.
func (q *QuotaCenter) getNQInQueryFactor() float64 {
if !Params.QuotaConfig.QueueProtectionEnabled.GetAsBool() {
return 1
}
sum := func(ri metricsinfo.ReadInfoInQueue) int64 {
return ri.UnsolvedQueue + ri.ReadyQueue + ri.ReceiveChan + ri.ExecuteChan
}
nqInQueueThreshold := Params.QuotaConfig.NQInQueueThreshold.GetAsInt64()
if nqInQueueThreshold < 0 {
// < 0 means disable queue length protection
return 1
}
for _, metric := range q.queryNodeMetrics {
searchNQSum := sum(metric.SearchQueue)
queryTasksSum := sum(metric.QueryQueue)
nqInQueue := searchNQSum + queryTasksSum // We think of the NQ of query request as 1.
if nqInQueue >= nqInQueueThreshold {
return Params.QuotaConfig.CoolOffSpeed.GetAsFloat()
}
}
return 1
}
// getQueryLatencyFactor checks queueing latency in QueryNode for search&query requests,
// and return the factor according to QueueLatencyThreshold.
func (q *QuotaCenter) getQueryLatencyFactor() float64 {
if !Params.QuotaConfig.QueueProtectionEnabled.GetAsBool() {
return 1
}
queueLatencyThreshold := Params.QuotaConfig.QueueLatencyThreshold.GetAsDuration(time.Second)
if queueLatencyThreshold < 0 {
// < 0 means disable queue latency protection
return 1
}
for _, metric := range q.queryNodeMetrics {
searchLatency := metric.SearchQueue.AvgQueueDuration
queryLatency := metric.QueryQueue.AvgQueueDuration
if searchLatency >= queueLatencyThreshold || queryLatency >= queueLatencyThreshold {
return Params.QuotaConfig.CoolOffSpeed.GetAsFloat()
}
}
return 1
}
// getReadResultFactor checks search result rate in Proxy,
// and return the factor according to MaxReadResultRate.
func (q *QuotaCenter) getReadResultFactor() float64 {
if !Params.QuotaConfig.ResultProtectionEnabled.GetAsBool() {
return 1
}
maxRate := Params.QuotaConfig.MaxReadResultRate.GetAsFloat()
rateCount := float64(0)
for _, metric := range q.proxyMetrics {
for _, rm := range metric.Rms {
if rm.Label == metricsinfo.ReadResultThroughput {
rateCount += rm.Rate
}
}
}
if rateCount >= maxRate {
return Params.QuotaConfig.CoolOffSpeed.GetAsFloat()
}
return 1
}
// getMemoryFactor checks whether any node has memory resource issue,
// and return the factor according to max memory water level.
func (q *QuotaCenter) getMemoryFactor() float64 {
factor := float64(1)
if !Params.QuotaConfig.MemProtectionEnabled.GetAsBool() {
return 1
}
dataNodeMemoryLowWaterLevel := Params.QuotaConfig.DataNodeMemoryLowWaterLevel.GetAsFloat()
dataNodeMemoryHighWaterLevel := Params.QuotaConfig.DataNodeMemoryHighWaterLevel.GetAsFloat()
queryNodeMemoryLowWaterLevel := Params.QuotaConfig.QueryNodeMemoryLowWaterLevel.GetAsFloat()
queryNodeMemoryHighWaterLevel := Params.QuotaConfig.QueryNodeMemoryHighWaterLevel.GetAsFloat()
for nodeID, metric := range q.queryNodeMetrics {
memoryWaterLevel := float64(metric.Hms.MemoryUsage) / float64(metric.Hms.Memory)
if memoryWaterLevel <= queryNodeMemoryLowWaterLevel {
continue
}
if memoryWaterLevel >= queryNodeMemoryHighWaterLevel {
log.Warn("QuotaCenter: QueryNode memory to high water level",
zap.String("Node", fmt.Sprintf("%s-%d", typeutil.QueryNodeRole, nodeID)),
zap.Uint64("UsedMem", metric.Hms.MemoryUsage),
zap.Uint64("TotalMem", metric.Hms.Memory),
zap.Float64("memoryWaterLevel", memoryWaterLevel),
zap.Float64("memoryHighWaterLevel", queryNodeMemoryHighWaterLevel))
return 0
}
p := (queryNodeMemoryHighWaterLevel - memoryWaterLevel) / (queryNodeMemoryHighWaterLevel - queryNodeMemoryLowWaterLevel)
if p < factor {
log.Warn("QuotaCenter: QueryNode memory to low water level, limit writing rate",
zap.String("Node", fmt.Sprintf("%s-%d", typeutil.QueryNodeRole, nodeID)),
zap.Uint64("UsedMem", metric.Hms.MemoryUsage),
zap.Uint64("TotalMem", metric.Hms.Memory),
zap.Float64("memoryWaterLevel", memoryWaterLevel),
zap.Float64("memoryLowWaterLevel", queryNodeMemoryLowWaterLevel))
factor = p
}
}
for nodeID, metric := range q.dataNodeMetrics {
memoryWaterLevel := float64(metric.Hms.MemoryUsage) / float64(metric.Hms.Memory)
if memoryWaterLevel <= dataNodeMemoryLowWaterLevel {
continue
}
if memoryWaterLevel >= dataNodeMemoryHighWaterLevel {
log.Warn("QuotaCenter: DataNode memory to high water level",
zap.String("Node", fmt.Sprintf("%s-%d", typeutil.DataNodeRole, nodeID)),
zap.Uint64("UsedMem", metric.Hms.MemoryUsage),
zap.Uint64("TotalMem", metric.Hms.Memory),
zap.Float64("memoryWaterLevel", memoryWaterLevel),
zap.Float64("memoryHighWaterLevel", dataNodeMemoryHighWaterLevel))
return 0
}
p := (dataNodeMemoryHighWaterLevel - memoryWaterLevel) / (dataNodeMemoryHighWaterLevel - dataNodeMemoryLowWaterLevel)
if p < factor {
log.Warn("QuotaCenter: DataNode memory to low water level, limit writing rate",
zap.String("Node", fmt.Sprintf("%s-%d", typeutil.DataNodeRole, nodeID)),
zap.Uint64("UsedMem", metric.Hms.MemoryUsage),
zap.Uint64("TotalMem", metric.Hms.Memory),
zap.Float64("memoryWaterLevel", memoryWaterLevel),
zap.Float64("memoryLowWaterLevel", dataNodeMemoryLowWaterLevel))
factor = p
}
}
return factor
}
// ifDiskQuotaExceeded checks if disk quota exceeded.
@ -642,32 +619,38 @@ func (q *QuotaCenter) ifDiskQuotaExceeded() bool {
func (q *QuotaCenter) setRates() error {
ctx, cancel := context.WithTimeout(context.Background(), SetRatesTimeout)
defer cancel()
var map2List func() []*internalpb.Rate
switch q.rateAllocateStrategy {
case Average:
map2List = func() []*internalpb.Rate {
toCollectionRate := func(collection int64, currentRates map[internalpb.RateType]ratelimitutil.Limit) *proxypb.CollectionRate {
rates := make([]*internalpb.Rate, 0, len(q.currentRates))
switch q.rateAllocateStrategy {
case Average:
proxyNum := q.proxies.GetProxyCount()
if proxyNum == 0 {
return nil
}
rates := make([]*internalpb.Rate, 0, len(q.currentRates))
for rt, r := range q.currentRates {
for rt, r := range currentRates {
if r == Inf {
rates = append(rates, &internalpb.Rate{Rt: rt, R: float64(r)})
} else {
rates = append(rates, &internalpb.Rate{Rt: rt, R: float64(r) / float64(proxyNum)})
}
}
return rates
case ByRateWeight:
// TODO: support ByRateWeight
}
return &proxypb.CollectionRate{
Collection: collection,
Rates: rates,
States: lo.Keys(q.quotaStates[collection]),
Codes: lo.Values(q.quotaStates[collection]),
}
case ByRateWeight:
// TODO: support ByRateWeight
}
states := make([]milvuspb.QuotaState, 0, len(q.quotaStates))
codes := make([]commonpb.ErrorCode, 0, len(q.quotaStates))
for k, v := range q.quotaStates {
states = append(states, k)
codes = append(codes, v)
collectionRates := make([]*proxypb.CollectionRate, 0)
for collection, rates := range q.currentRates {
collectionRates = append(collectionRates, toCollectionRate(collection, rates))
}
timestamp := tsoutil.ComposeTSByTime(time.Now(), 0)
req := &proxypb.SetRatesRequest{
@ -675,9 +658,7 @@ func (q *QuotaCenter) setRates() error {
commonpbutil.WithMsgID(int64(timestamp)),
commonpbutil.WithTimeStamp(timestamp),
),
Rates: map2List(),
States: states,
Codes: codes,
Rates: collectionRates,
}
return q.proxies.SetRates(ctx, req)
}
@ -685,11 +666,13 @@ func (q *QuotaCenter) setRates() error {
// recordMetrics records metrics of quota states.
func (q *QuotaCenter) recordMetrics() {
record := func(errorCode commonpb.ErrorCode) {
for _, v := range q.quotaStates {
if v == errorCode {
metrics.RootCoordQuotaStates.WithLabelValues(errorCode.String()).Set(1)
return
for _, states := range q.quotaStates {
for _, state := range states {
if state == errorCode {
metrics.RootCoordQuotaStates.WithLabelValues(errorCode.String()).Set(1)
}
}
return
}
metrics.RootCoordQuotaStates.WithLabelValues(errorCode.String()).Set(0)
}

333
internal/rootcoord/quota_center_test.go
View File

@ -19,7 +19,6 @@ package rootcoord
import (
"context"
"fmt"
"math"
"testing"
"time"
@ -30,8 +29,10 @@ import (
"github.com/milvus-io/milvus-proto/go-api/milvuspb"
"github.com/milvus-io/milvus/internal/proto/internalpb"
"github.com/milvus-io/milvus/internal/types"
"github.com/milvus-io/milvus/pkg/log"
"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/tsoutil"
"github.com/milvus-io/milvus/pkg/util/typeutil"
)
@ -105,12 +106,21 @@ func TestQuotaCenter(t *testing.T) {
t.Run("test forceDeny", func(t *testing.T) {
qc := types.NewMockQueryCoord(t)
quotaCenter := NewQuotaCenter(pcm, qc, &dataCoordMockForQuota{}, core.tsoAllocator)
quotaCenter.readableCollections = []int64{1, 2, 3}
quotaCenter.resetAllCurrentRates()
quotaCenter.forceDenyReading(commonpb.ErrorCode_ForceDeny)
assert.Equal(t, Limit(0), quotaCenter.currentRates[internalpb.RateType_DQLQuery])
assert.Equal(t, Limit(0), quotaCenter.currentRates[internalpb.RateType_DQLQuery])
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_DQLQuery])
}
quotaCenter.writableCollections = []int64{1, 2, 3, 4}
quotaCenter.resetAllCurrentRates()
quotaCenter.forceDenyWriting(commonpb.ErrorCode_ForceDeny)
assert.Equal(t, Limit(0), quotaCenter.currentRates[internalpb.RateType_DMLInsert])
assert.Equal(t, Limit(0), quotaCenter.currentRates[internalpb.RateType_DMLDelete])
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])
}
})
t.Run("test calculateRates", func(t *testing.T) {
@ -127,186 +137,82 @@ func TestQuotaCenter(t *testing.T) {
assert.Error(t, err)
})
t.Run("test getTimeTickDelayFactor", func(t *testing.T) {
qc := types.NewMockQueryCoord(t)
// test MaxTimestamp
quotaCenter := NewQuotaCenter(pcm, qc, &dataCoordMockForQuota{}, core.tsoAllocator)
factor := quotaCenter.getTimeTickDelayFactor(0)
assert.Equal(t, float64(1), factor)
now := time.Now()
paramtable.Get().Save(Params.QuotaConfig.TtProtectionEnabled.Key, "true")
paramtable.Get().Save(Params.QuotaConfig.MaxTimeTickDelay.Key, "3")
// test force deny writing
alloc := newMockTsoAllocator()
alloc.GenerateTSOF = func(count uint32) (typeutil.Timestamp, error) {
added := now.Add(Params.QuotaConfig.MaxTimeTickDelay.GetAsDuration(time.Second))
ts := tsoutil.ComposeTSByTime(added, 0)
return ts, nil
}
quotaCenter.tsoAllocator = alloc
quotaCenter.queryNodeMetrics = map[UniqueID]*metricsinfo.QueryNodeQuotaMetrics{
1: {Fgm: metricsinfo.FlowGraphMetric{
MinFlowGraphTt: tsoutil.ComposeTSByTime(now, 0),
NumFlowGraph: 1,
MinFlowGraphChannel: "dml",
}}}
ts, err := quotaCenter.tsoAllocator.GenerateTSO(1)
assert.NoError(t, err)
factor = quotaCenter.getTimeTickDelayFactor(ts)
assert.Equal(t, float64(0), factor)
// test one-third time tick delay
alloc.GenerateTSOF = func(count uint32) (typeutil.Timestamp, error) {
oneThirdDelay := Params.QuotaConfig.MaxTimeTickDelay.GetAsDuration(time.Second) / 3
added := now.Add(oneThirdDelay)
oneThirdTs := tsoutil.ComposeTSByTime(added, 0)
return oneThirdTs, nil
}
quotaCenter.queryNodeMetrics = map[UniqueID]*metricsinfo.QueryNodeQuotaMetrics{
1: {Fgm: metricsinfo.FlowGraphMetric{
MinFlowGraphTt: tsoutil.ComposeTSByTime(now, 0),
NumFlowGraph: 1,
MinFlowGraphChannel: "dml",
}}}
ts, err = quotaCenter.tsoAllocator.GenerateTSO(1)
assert.NoError(t, err)
factor = quotaCenter.getTimeTickDelayFactor(ts)
ok := math.Abs(factor-2.0/3.0) < 0.0001
assert.True(t, ok)
})
t.Run("test getTimeTickDelayFactor factors", func(t *testing.T) {
t.Run("test TimeTickDelayFactor factors", func(t *testing.T) {
qc := types.NewMockQueryCoord(t)
quotaCenter := NewQuotaCenter(pcm, qc, &dataCoordMockForQuota{}, core.tsoAllocator)
type ttCase struct {
maxTtDelay time.Duration
curTt time.Time
fgTt time.Time
delay time.Duration
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},
{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
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.DMLMaxInsertRate.Key, "100.0")
paramtable.Get().Save(Params.QuotaConfig.DMLMinInsertRate.Key, "0.0")
paramtable.Get().Save(Params.QuotaConfig.DMLMaxDeleteRate.Key, "100.0")
paramtable.Get().Save(Params.QuotaConfig.DMLMinDeleteRate.Key, "0.0")
for i, c := range ttCases {
paramtable.Get().Save(Params.QuotaConfig.MaxTimeTickDelay.Key, fmt.Sprintf("%f", c.maxTtDelay.Seconds()))
fgTs := tsoutil.ComposeTSByTime(c.fgTt, 0)
log.Info(paramtable.Get().Get(Params.QuotaConfig.DMLMaxInsertRate.Key))
log.Info(Params.QuotaConfig.DMLMaxInsertRate.GetValue())
quotaCenter.writableCollections = []int64{1, 2, 3}
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: fgTs,
MinFlowGraphTt: minTS,
MinFlowGraphChannel: "dml",
},
Effect: metricsinfo.NodeEffect{
CollectionIDs: []int64{1, 2, 3},
},
},
}
curTs := tsoutil.ComposeTSByTime(c.curTt, 0)
factor := quotaCenter.getTimeTickDelayFactor(curTs)
if math.Abs(factor-c.expectedFactor) > 0.000001 {
t.Errorf("case %d failed: curTs[%v], fgTs[%v], expectedFactor: %f, actualFactor: %f",
i, c.curTt, c.fgTt, c.expectedFactor, factor)
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.resetAllCurrentRates()
quotaCenter.calculateWriteRates()
deleteFactor := float64(quotaCenter.currentRates[1][internalpb.RateType_DMLDelete]) / Params.QuotaConfig.DMLMaxInsertRate.GetAsFloat()
assert.Equal(t, c.expectedFactor, deleteFactor)
}
Params.QuotaConfig.MaxTimeTickDelay = backup
})
t.Run("test getNQInQueryFactor", func(t *testing.T) {
qc := types.NewMockQueryCoord(t)
quotaCenter := NewQuotaCenter(pcm, qc, &dataCoordMockForQuota{}, core.tsoAllocator)
factor := quotaCenter.getNQInQueryFactor()
assert.Equal(t, float64(1), factor)
// test cool off
paramtable.Get().Save(Params.QuotaConfig.QueueProtectionEnabled.Key, "true")
paramtable.Get().Save(Params.QuotaConfig.NQInQueueThreshold.Key, "100")
quotaCenter.queryNodeMetrics = map[UniqueID]*metricsinfo.QueryNodeQuotaMetrics{
1: {SearchQueue: metricsinfo.ReadInfoInQueue{
UnsolvedQueue: Params.QuotaConfig.NQInQueueThreshold.GetAsInt64(),
}}}
factor = quotaCenter.getNQInQueryFactor()
assert.Equal(t, Params.QuotaConfig.CoolOffSpeed.GetAsFloat(), factor)
// test no cool off
quotaCenter.queryNodeMetrics = map[UniqueID]*metricsinfo.QueryNodeQuotaMetrics{
1: {SearchQueue: metricsinfo.ReadInfoInQueue{
UnsolvedQueue: Params.QuotaConfig.NQInQueueThreshold.GetAsInt64() - 1,
}}}
factor = quotaCenter.getNQInQueryFactor()
assert.Equal(t, 1.0, factor)
//ok := math.Abs(factor-1.0) < 0.0001
//assert.True(t, ok)
})
t.Run("test getQueryLatencyFactor", func(t *testing.T) {
qc := types.NewMockQueryCoord(t)
quotaCenter := NewQuotaCenter(pcm, qc, &dataCoordMockForQuota{}, core.tsoAllocator)
factor := quotaCenter.getQueryLatencyFactor()
assert.Equal(t, float64(1), factor)
// test cool off
paramtable.Get().Save(Params.QuotaConfig.QueueProtectionEnabled.Key, "true")
paramtable.Get().Save(Params.QuotaConfig.QueueLatencyThreshold.Key, "3")
quotaCenter.queryNodeMetrics = map[UniqueID]*metricsinfo.QueryNodeQuotaMetrics{
1: {SearchQueue: metricsinfo.ReadInfoInQueue{
AvgQueueDuration: Params.QuotaConfig.QueueLatencyThreshold.GetAsDuration(time.Second),
}}}
factor = quotaCenter.getQueryLatencyFactor()
assert.Equal(t, Params.QuotaConfig.CoolOffSpeed.GetAsFloat(), factor)
// test no cool off
quotaCenter.queryNodeMetrics = map[UniqueID]*metricsinfo.QueryNodeQuotaMetrics{
1: {SearchQueue: metricsinfo.ReadInfoInQueue{
AvgQueueDuration: 1 * time.Second,
}}}
factor = quotaCenter.getQueryLatencyFactor()
assert.Equal(t, 1.0, factor)
})
t.Run("test checkReadResult", func(t *testing.T) {
qc := types.NewMockQueryCoord(t)
quotaCenter := NewQuotaCenter(pcm, qc, &dataCoordMockForQuota{}, core.tsoAllocator)
factor := quotaCenter.getReadResultFactor()
assert.Equal(t, float64(1), factor)
// test cool off
paramtable.Get().Save(Params.QuotaConfig.ResultProtectionEnabled.Key, "true")
paramtable.Get().Save(Params.QuotaConfig.MaxReadResultRate.Key, fmt.Sprintf("%f", 1.0/1024/1024))
quotaCenter.proxyMetrics = map[UniqueID]*metricsinfo.ProxyQuotaMetrics{
1: {Rms: []metricsinfo.RateMetric{
{Label: metricsinfo.ReadResultThroughput, Rate: 1.2},
}}}
factor = quotaCenter.getReadResultFactor()
assert.Equal(t, Params.QuotaConfig.CoolOffSpeed.GetAsFloat(), factor)
// test no cool off
quotaCenter.proxyMetrics = map[UniqueID]*metricsinfo.ProxyQuotaMetrics{
1: {Rms: []metricsinfo.RateMetric{
{Label: metricsinfo.ReadResultThroughput, Rate: 0.8},
}}}
factor = quotaCenter.getReadResultFactor()
assert.Equal(t, 1.0, factor)
})
t.Run("test calculateReadRates", func(t *testing.T) {
qc := types.NewMockQueryCoord(t)
quotaCenter := NewQuotaCenter(pcm, qc, &dataCoordMockForQuota{}, core.tsoAllocator)
quotaCenter.readableCollections = []int64{1, 2, 3}
quotaCenter.proxyMetrics = map[UniqueID]*metricsinfo.ProxyQuotaMetrics{
1: {Rms: []metricsinfo.RateMetric{
{Label: internalpb.RateType_DQLSearch.String(), Rate: 100},
@ -316,13 +222,22 @@ func TestQuotaCenter(t *testing.T) {
paramtable.Get().Save(Params.QuotaConfig.ForceDenyReading.Key, "false")
paramtable.Get().Save(Params.QuotaConfig.QueueProtectionEnabled.Key, "true")
paramtable.Get().Save(Params.QuotaConfig.QueueLatencyThreshold.Key, "100")
paramtable.Get().Save(Params.QuotaConfig.DQLLimitEnabled.Key, "true")
paramtable.Get().Save(Params.QuotaConfig.DQLMaxQueryRate.Key, "500")
paramtable.Get().Save(Params.QuotaConfig.DQLMaxSearchRate.Key, "500")
quotaCenter.resetAllCurrentRates()
quotaCenter.queryNodeMetrics = map[UniqueID]*metricsinfo.QueryNodeQuotaMetrics{
1: {SearchQueue: metricsinfo.ReadInfoInQueue{
AvgQueueDuration: Params.QuotaConfig.QueueLatencyThreshold.GetAsDuration(time.Second),
}, Effect: metricsinfo.NodeEffect{
NodeID: 1,
CollectionIDs: []int64{1, 2, 3},
}}}
quotaCenter.calculateReadRates()
assert.Equal(t, Limit(100.0*0.9), quotaCenter.currentRates[internalpb.RateType_DQLSearch])
assert.Equal(t, Limit(100.0*0.9), quotaCenter.currentRates[internalpb.RateType_DQLQuery])
for _, collection := range quotaCenter.readableCollections {
assert.Equal(t, Limit(100.0*0.9), quotaCenter.currentRates[collection][internalpb.RateType_DQLSearch])
assert.Equal(t, Limit(100.0*0.9), quotaCenter.currentRates[collection][internalpb.RateType_DQLQuery])
}
paramtable.Get().Save(Params.QuotaConfig.NQInQueueThreshold.Key, "100")
quotaCenter.queryNodeMetrics = map[UniqueID]*metricsinfo.QueryNodeQuotaMetrics{
@ -330,8 +245,10 @@ func TestQuotaCenter(t *testing.T) {
UnsolvedQueue: Params.QuotaConfig.NQInQueueThreshold.GetAsInt64(),
}}}
quotaCenter.calculateReadRates()
assert.Equal(t, Limit(100.0*0.9), quotaCenter.currentRates[internalpb.RateType_DQLSearch])
assert.Equal(t, Limit(100.0*0.9), quotaCenter.currentRates[internalpb.RateType_DQLQuery])
for _, collection := range quotaCenter.readableCollections {
assert.Equal(t, Limit(100.0*0.9), quotaCenter.currentRates[collection][internalpb.RateType_DQLSearch])
assert.Equal(t, Limit(100.0*0.9), quotaCenter.currentRates[collection][internalpb.RateType_DQLQuery])
}
paramtable.Get().Save(Params.QuotaConfig.ResultProtectionEnabled.Key, "true")
paramtable.Get().Save(Params.QuotaConfig.MaxReadResultRate.Key, "1")
@ -343,8 +260,10 @@ func TestQuotaCenter(t *testing.T) {
}}}
quotaCenter.queryNodeMetrics = map[UniqueID]*metricsinfo.QueryNodeQuotaMetrics{1: {SearchQueue: metricsinfo.ReadInfoInQueue{}}}
quotaCenter.calculateReadRates()
assert.Equal(t, Limit(100.0*0.9), quotaCenter.currentRates[internalpb.RateType_DQLSearch])
assert.Equal(t, Limit(100.0*0.9), quotaCenter.currentRates[internalpb.RateType_DQLQuery])
for _, collection := range quotaCenter.readableCollections {
assert.Equal(t, Limit(100.0*0.9), quotaCenter.currentRates[collection][internalpb.RateType_DQLSearch])
assert.Equal(t, Limit(100.0*0.9), quotaCenter.currentRates[collection][internalpb.RateType_DQLQuery])
}
})
t.Run("test calculateWriteRates", func(t *testing.T) {
@ -356,37 +275,32 @@ func TestQuotaCenter(t *testing.T) {
// DiskQuota exceeded
quotaBackup := Params.QuotaConfig.DiskQuota
paramtable.Get().Save(Params.QuotaConfig.DiskQuota.Key, "99")
quotaCenter.dataCoordMetrics = &metricsinfo.DataCoordQuotaMetrics{TotalBinlogSize: 100}
quotaCenter.dataCoordMetrics = &metricsinfo.DataCoordQuotaMetrics{TotalBinlogSize: 100 * 1024 * 1024}
quotaCenter.writableCollections = []int64{1, 2, 3}
quotaCenter.resetAllCurrentRates()
err = quotaCenter.calculateWriteRates()
assert.NoError(t, err)
assert.Equal(t, Limit(0), quotaCenter.currentRates[internalpb.RateType_DMLInsert])
assert.Equal(t, Limit(0), quotaCenter.currentRates[internalpb.RateType_DMLDelete])
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])
}
Params.QuotaConfig.DiskQuota = quotaBackup
// force deny
forceBak := Params.QuotaConfig.ForceDenyWriting
paramtable.Get().Save(Params.QuotaConfig.ForceDenyWriting.Key, "true")
quotaCenter.writableCollections = []int64{1, 2, 3}
quotaCenter.resetAllCurrentRates()
err = quotaCenter.calculateWriteRates()
assert.NoError(t, err)
assert.Equal(t, Limit(0), quotaCenter.currentRates[internalpb.RateType_DMLInsert])
assert.Equal(t, Limit(0), quotaCenter.currentRates[internalpb.RateType_DMLDelete])
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])
}
Params.QuotaConfig.ForceDenyWriting = forceBak
})
t.Run("test getMemoryFactor basic", func(t *testing.T) {
qc := types.NewMockQueryCoord(t)
quotaCenter := NewQuotaCenter(pcm, qc, &dataCoordMockForQuota{}, core.tsoAllocator)
factor := quotaCenter.getMemoryFactor()
assert.Equal(t, float64(1), factor)
quotaCenter.dataNodeMetrics = map[UniqueID]*metricsinfo.DataNodeQuotaMetrics{1: {Hms: metricsinfo.HardwareMetrics{MemoryUsage: 100, Memory: 100}}}
factor = quotaCenter.getMemoryFactor()
assert.Equal(t, float64(0), factor)
quotaCenter.queryNodeMetrics = map[UniqueID]*metricsinfo.QueryNodeQuotaMetrics{1: {Hms: metricsinfo.HardwareMetrics{MemoryUsage: 100, Memory: 100}}}
factor = quotaCenter.getMemoryFactor()
assert.Equal(t, float64(0), factor)
})
t.Run("test getMemoryFactor factors", func(t *testing.T) {
t.Run("test MemoryFactor factors", func(t *testing.T) {
qc := types.NewMockQueryCoord(t)
quotaCenter := NewQuotaCenter(pcm, qc, &dataCoordMockForQuota{}, core.tsoAllocator)
type memCase struct {
@ -412,15 +326,24 @@ func TestQuotaCenter(t *testing.T) {
{0.85, 0.95, 95, 100, 0},
}
for i, c := range memCases {
quotaCenter.writableCollections = append(quotaCenter.writableCollections, 1, 2, 3)
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}}}
factor := quotaCenter.getMemoryFactor()
if math.Abs(factor-c.expectedFactor) > 0.000001 {
t.Errorf("case %d failed: waterLever[low:%v, high:%v], memMetric[used:%d, total:%d], expectedFactor: %f, actualFactor: %f",
i, c.lowWater, c.highWater, c.memUsage, c.memTotal, c.expectedFactor, factor)
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},
},
},
}
quotaCenter.resetAllCurrentRates()
quotaCenter.calculateWriteRates()
}
paramtable.Get().Reset(Params.QuotaConfig.QueryNodeMemoryLowWaterLevel.Key)
@ -451,10 +374,19 @@ func TestQuotaCenter(t *testing.T) {
t.Run("test setRates", func(t *testing.T) {
qc := types.NewMockQueryCoord(t)
p1 := types.NewMockProxy(t)
p1.EXPECT().SetRates(mock.Anything, mock.Anything).Return(nil, nil)
pcm := &proxyClientManager{proxyClient: map[int64]types.Proxy{
TestProxyID: p1,
}}
quotaCenter := NewQuotaCenter(pcm, qc, &dataCoordMockForQuota{}, core.tsoAllocator)
quotaCenter.currentRates[internalpb.RateType_DMLInsert] = 100
quotaCenter.quotaStates[milvuspb.QuotaState_DenyToWrite] = commonpb.ErrorCode_MemoryQuotaExhausted
quotaCenter.quotaStates[milvuspb.QuotaState_DenyToRead] = commonpb.ErrorCode_ForceDeny
quotaCenter.resetAllCurrentRates()
collectionID := int64(1)
quotaCenter.currentRates[collectionID] = make(map[internalpb.RateType]ratelimitutil.Limit)
quotaCenter.quotaStates[collectionID] = make(map[milvuspb.QuotaState]commonpb.ErrorCode)
quotaCenter.currentRates[collectionID][internalpb.RateType_DMLInsert] = 100
quotaCenter.quotaStates[collectionID][milvuspb.QuotaState_DenyToWrite] = commonpb.ErrorCode_MemoryQuotaExhausted
quotaCenter.quotaStates[collectionID][milvuspb.QuotaState_DenyToRead] = commonpb.ErrorCode_ForceDeny
err = quotaCenter.setRates()
assert.NoError(t, err)
})
@ -462,17 +394,22 @@ func TestQuotaCenter(t *testing.T) {
t.Run("test recordMetrics", func(t *testing.T) {
qc := types.NewMockQueryCoord(t)
quotaCenter := NewQuotaCenter(pcm, qc, &dataCoordMockForQuota{}, core.tsoAllocator)
quotaCenter.quotaStates[milvuspb.QuotaState_DenyToWrite] = commonpb.ErrorCode_MemoryQuotaExhausted
quotaCenter.quotaStates[milvuspb.QuotaState_DenyToRead] = commonpb.ErrorCode_ForceDeny
collectionID := int64(1)
quotaCenter.quotaStates[collectionID] = make(map[milvuspb.QuotaState]commonpb.ErrorCode)
quotaCenter.quotaStates[collectionID][milvuspb.QuotaState_DenyToWrite] = commonpb.ErrorCode_MemoryQuotaExhausted
quotaCenter.quotaStates[collectionID][milvuspb.QuotaState_DenyToRead] = commonpb.ErrorCode_ForceDeny
quotaCenter.recordMetrics()
})
t.Run("test guaranteeMinRate", func(t *testing.T) {
qc := types.NewMockQueryCoord(t)
quotaCenter := NewQuotaCenter(pcm, qc, &dataCoordMockForQuota{}, core.tsoAllocator)
quotaCenter.resetAllCurrentRates()
minRate := Limit(100)
quotaCenter.currentRates[internalpb.RateType_DQLSearch] = Limit(50)
quotaCenter.guaranteeMinRate(float64(minRate), internalpb.RateType_DQLSearch)
assert.Equal(t, minRate, quotaCenter.currentRates[internalpb.RateType_DQLSearch])
collectionID := int64(1)
quotaCenter.currentRates[collectionID] = make(map[internalpb.RateType]ratelimitutil.Limit)
quotaCenter.currentRates[collectionID][internalpb.RateType_DQLSearch] = Limit(50)
quotaCenter.guaranteeMinRate(float64(minRate), internalpb.RateType_DQLSearch, 1)
assert.Equal(t, minRate, quotaCenter.currentRates[collectionID][internalpb.RateType_DQLSearch])
})
}

4768
internal/types/mock_proxy.go Normal file
View File

File diff suppressed because it is too large Load Diff

2
internal/types/types.go
View File

@ -41,7 +41,7 @@ type TimeTickProvider interface {
// If Limit function return true, the request will be rejected.
// Otherwise, the request will pass. Limit also returns limit of limiter.
type Limiter interface {
Check(rt internalpb.RateType, n int) commonpb.ErrorCode
Check(collectionID int64, rt internalpb.RateType, n int) commonpb.ErrorCode
}
// Component is the interface all services implement

2
pkg/metrics/proxy_metrics.go
View File

@ -222,7 +222,7 @@ var (
Subsystem: typeutil.ProxyRole,
Name: "limiter_rate",
Help: "",
}, []string{nodeIDLabelName, msgTypeLabelName})
}, []string{nodeIDLabelName, collectionIDLabelName, msgTypeLabelName})
ProxyHookFunc = prometheus.NewCounterVec(
prometheus.CounterOpts{

14
pkg/util/metricsinfo/quota_metric.go
View File

@ -67,6 +67,12 @@ type ReadInfoInQueue struct {
AvgQueueDuration time.Duration
}
// NodeEffect contains the a node and its effected collection info.
type NodeEffect struct {
NodeID int64
CollectionIDs []int64
}
// QueryNodeQuotaMetrics are metrics of QueryNode.
type QueryNodeQuotaMetrics struct {
Hms HardwareMetrics
@ -74,6 +80,7 @@ type QueryNodeQuotaMetrics struct {
Fgm FlowGraphMetric
SearchQueue ReadInfoInQueue
QueryQueue ReadInfoInQueue
Effect NodeEffect
}
type DataCoordQuotaMetrics struct {
@ -82,9 +89,10 @@ type DataCoordQuotaMetrics struct {
// DataNodeQuotaMetrics are metrics of DataNode.
type DataNodeQuotaMetrics struct {
Hms HardwareMetrics
Rms []RateMetric
Fgm FlowGraphMetric
Hms HardwareMetrics
Rms []RateMetric
Fgm FlowGraphMetric
Effect NodeEffect
}
// ProxyQuotaMetrics are metrics of Proxy.

4
pkg/util/typeutil/map.go
View File

@ -95,3 +95,7 @@ func (m *ConcurrentMap[K, V]) GetAndRemove(key K) (V, bool) {
m.len.Dec()
return value.(V), true
}
func (m *ConcurrentMap[K, V]) Len() int {
return int(m.len.Load())
}

4
pkg/util/typeutil/map_test.go
View File

@ -57,6 +57,8 @@ func (suite *MapUtilSuite) TestGetMapKeys() {
func (suite *MapUtilSuite) TestConcurrentMap() {
currMap := NewConcurrentMap[int64, string]()
suite.Equal(currMap.Len(), 0)
v, loaded := currMap.GetOrInsert(100, "v-100")
suite.Equal("v-100", v)
suite.Equal(false, loaded)
@ -67,6 +69,8 @@ func (suite *MapUtilSuite) TestConcurrentMap() {
suite.Equal("v-100", v)
suite.Equal(true, loaded)
suite.Equal(currMap.Len(), 1)
currMap.GetOrInsert(100, "v-100-new")
currMap.GetOrInsert(200, "v-200")
currMap.GetOrInsert(300, "v-300")