Skip to content

Commit 05c08d0

Browse files
yanavlasovyanavlasov
authored
Fix quota result when all limits were exceeded (#1059)
* Fix quota result when all limits were exceeded Signed-off-by: yavlasov <yavlasov@google.com> * Address comments Signed-off-by: yavlasov <yavlasov@google.com> * Fix comment Signed-off-by: yavlasov <yavlasov@google.com> * fix comment Signed-off-by: yavlasov <yavlasov@google.com> * Remove log into a file Signed-off-by: yavlasov <yavlasov@google.com> --------- Signed-off-by: yavlasov <yavlasov@google.com>
1 parent 4fe3db4 commit 05c08d0

4 files changed

Lines changed: 246 additions & 24 deletions

File tree

‎integration-test/scripts/token-quota.sh‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,8 @@ fi
1818
# Quota is debited from service_2 bucket on the response path so only the 4th request should be rejected
1919
response=$(curl -f -s -H "request-no: 4" http://envoy-proxy:8888/tokenquota)
2020

21-
if [ $? -ne 0 ]; then
22-
echo "Quota mode does not deny requests yet"
21+
if [ $? -eq 0 ]; then
22+
echo "Quota limiting should fail the request"
2323
exit 1
2424
fi
2525

‎src/redis/fixed_cache_impl.go‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -214,7 +214,6 @@ func (this *fixedRateLimitCacheImpl) DoLimit(
214214

215215
responseDescriptorStatuses[i] = this.baseRateLimiter.GetResponseDescriptorStatus(cacheKey.Key,
216216
limitInfo, isOverLimitWithLocalCache[i], hitsAddends[i])
217-
218217
}
219218

220219
return responseDescriptorStatuses

‎src/service/ratelimit.go‎

Lines changed: 17 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -201,14 +201,16 @@ func (this *service) shouldRateLimitWorker(
201201

202202
response := &pb.RateLimitResponse{}
203203
response.Statuses = make([]*pb.RateLimitResponse_DescriptorStatus, len(request.Descriptors))
204-
finalCode := pb.RateLimitResponse_OK
205204

206205
// Keep track of the descriptor which is closest to hit the ratelimit
207206
minLimitRemaining := MaxUint32
208207
var minimumDescriptor *pb.RateLimitResponse_DescriptorStatus = nil
209208

210209
// Track quota mode violations for metadata
211210
var quotaModeViolations []int
211+
failedRateLimitDescriptors := 0
212+
failedQuotaDescriptors := 0
213+
totalQuotaDescriptors := 0
212214

213215
for i, descriptorStatus := range responseDescriptorStatuses {
214216
// Keep track of the descriptor closest to hit the ratelimit
@@ -226,28 +228,31 @@ func (this *service) shouldRateLimitWorker(
226228
}
227229
} else {
228230
response.Statuses[i] = descriptorStatus
231+
isQuotaMode := globalQuotaMode || (limitsToCheck[i] != nil && limitsToCheck[i].QuotaMode)
229232
if descriptorStatus.Code == pb.RateLimitResponse_OVER_LIMIT {
230-
// Check if this limit is in quota mode (individual or global)
231-
isQuotaMode := globalQuotaMode || (limitsToCheck[i] != nil && limitsToCheck[i].QuotaMode)
232-
233233
if isQuotaMode {
234234
// In quota mode: track the violation for metadata but keep response as OK
235235
quotaModeViolations = append(quotaModeViolations, i)
236-
response.Statuses[i] = &pb.RateLimitResponse_DescriptorStatus{
237-
Code: pb.RateLimitResponse_OK,
238-
CurrentLimit: descriptorStatus.CurrentLimit,
239-
LimitRemaining: descriptorStatus.LimitRemaining,
240-
}
236+
failedQuotaDescriptors += 1
241237
} else {
242-
// Normal rate limit: set final code to OVER_LIMIT
243-
finalCode = descriptorStatus.Code
238+
failedRateLimitDescriptors += 1
244239
minimumDescriptor = descriptorStatus
245240
minLimitRemaining = 0
246241
}
247242
}
243+
if isQuotaMode {
244+
totalQuotaDescriptors += 1
245+
}
248246
}
249247
}
250248

249+
finalCode := pb.RateLimitResponse_OK
250+
// The final code is OVER_LIMIT iff at least one rate limit descriptor is over the limit
251+
// or all quota descriptors are over the limit.
252+
if failedRateLimitDescriptors > 0 || (totalQuotaDescriptors > 0 && totalQuotaDescriptors == failedQuotaDescriptors) {
253+
finalCode = pb.RateLimitResponse_OVER_LIMIT
254+
}
255+
251256
// Add Headers if requested
252257
if this.customHeadersEnabled && minimumDescriptor != nil {
253258
response.ResponseHeadersToAdd = []*core.HeaderValue{
@@ -380,6 +385,7 @@ func (this *service) ShouldRateLimit(
380385
ctx context.Context,
381386
request *pb.RateLimitRequest,
382387
) (finalResponse *pb.RateLimitResponse, finalError error) {
388+
logger.Debugf("ShouldRateLimit: %+v", request)
383389
// Generate trace
384390
_, span := tracer.Start(ctx, "ShouldRateLimit Execution",
385391
trace.WithAttributes(

‎test/service/ratelimit_test.go‎

Lines changed: 227 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -705,14 +705,14 @@ func TestServiceGlobalQuotaMode(test *testing.T) {
705705
})
706706
response, err := service.ShouldRateLimit(context.Background(), request)
707707

708-
// OK overall code even if limit response was OVER_LIMIT, because global quota mode is enabled
708+
// OVER_LIMIT overall code since all quota limits were OVER_LIMIT
709709
common.AssertProtoEqual(
710710
t.assert,
711711
&pb.RateLimitResponse{
712-
OverallCode: pb.RateLimitResponse_OK,
712+
OverallCode: pb.RateLimitResponse_OVER_LIMIT,
713713
Statuses: []*pb.RateLimitResponse_DescriptorStatus{
714-
{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 0},
715-
{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[1].Limit, LimitRemaining: 0},
714+
{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0},
715+
{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[1].Limit, LimitRemaining: 0},
716716
},
717717
},
718718
response)
@@ -809,11 +809,11 @@ func TestServicePerDescriptorQuotaMode(test *testing.T) {
809809
t.cache.EXPECT().DoLimit(context.Background(), request, limits).Return(
810810
[]*pb.RateLimitResponse_DescriptorStatus{
811811
{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0},
812-
{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[1].Limit, LimitRemaining: 0},
812+
{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[1].Limit, LimitRemaining: 0},
813813
})
814814
response, err := service.ShouldRateLimit(context.Background(), request)
815815

816-
// Regular limit should cause OVER_LIMIT overall, but quota mode limit should be converted to OK
816+
// Regular limit should cause OVER_LIMIT overall, even though quota mode is under the limit
817817
common.AssertProtoEqual(
818818
t.assert,
819819
&pb.RateLimitResponse{
@@ -827,7 +827,111 @@ func TestServicePerDescriptorQuotaMode(test *testing.T) {
827827
t.assert.Nil(err)
828828
}
829829

830-
func TestServiceQuotaModeOnly(test *testing.T) {
830+
func TestServiceMixedPerDescriptorModes(test *testing.T) {
831+
t := commonSetup(test)
832+
defer t.controller.Finish()
833+
834+
// No Global Quota mode
835+
service := t.setupBasicService()
836+
837+
request := common.NewRateLimitRequest(
838+
"quota-domain", [][][2]string{{{"regular", "limit"}}, {{"quota", "limit"}}}, 1)
839+
840+
// Create limits with one having quota mode enabled per-descriptor
841+
// In this configuration the limits will be evaluated as rate limits.
842+
limits := []*config.RateLimit{
843+
// Regular limit
844+
{
845+
FullKey: "regular_limit",
846+
Limit: &pb.RateLimitResponse_RateLimit{RequestsPerUnit: 5, Unit: pb.RateLimitResponse_RateLimit_MINUTE},
847+
QuotaMode: false,
848+
ShadowMode: false,
849+
},
850+
// Quota mode limit
851+
{
852+
FullKey: "quota_limit",
853+
Limit: &pb.RateLimitResponse_RateLimit{RequestsPerUnit: 3, Unit: pb.RateLimitResponse_RateLimit_MINUTE},
854+
QuotaMode: true,
855+
ShadowMode: false,
856+
},
857+
}
858+
859+
t.config.EXPECT().GetLimit(context.Background(), "quota-domain", request.Descriptors[0]).Return(limits[0])
860+
t.config.EXPECT().GetLimit(context.Background(), "quota-domain", request.Descriptors[1]).Return(limits[1])
861+
t.cache.EXPECT().DoLimit(context.Background(), request, limits).Return(
862+
[]*pb.RateLimitResponse_DescriptorStatus{
863+
{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 0},
864+
{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[1].Limit, LimitRemaining: 0},
865+
})
866+
response, err := service.ShouldRateLimit(context.Background(), request)
867+
868+
// Overall result is OVER_LIMIT, since all quota limits were exceeded
869+
common.AssertProtoEqual(
870+
t.assert,
871+
&pb.RateLimitResponse{
872+
OverallCode: pb.RateLimitResponse_OVER_LIMIT,
873+
Statuses: []*pb.RateLimitResponse_DescriptorStatus{
874+
{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 0},
875+
{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[1].Limit, LimitRemaining: 0},
876+
},
877+
},
878+
response)
879+
t.assert.Nil(err)
880+
}
881+
882+
func TestServiceMixedPerDescriptorModesUnderLimit(test *testing.T) {
883+
t := commonSetup(test)
884+
defer t.controller.Finish()
885+
886+
// No Global Quota mode
887+
service := t.setupBasicService()
888+
889+
request := common.NewRateLimitRequest(
890+
"quota-domain", [][][2]string{{{"regular", "limit"}}, {{"quota", "limit"}}}, 1)
891+
892+
// Create limits with one having quota mode enabled per-descriptor
893+
// In this configuration the limits will be evaluated as rate limits.
894+
limits := []*config.RateLimit{
895+
// Regular limit
896+
{
897+
FullKey: "regular_limit",
898+
Limit: &pb.RateLimitResponse_RateLimit{RequestsPerUnit: 5, Unit: pb.RateLimitResponse_RateLimit_MINUTE},
899+
QuotaMode: false,
900+
ShadowMode: false,
901+
},
902+
// Quota mode limit
903+
{
904+
FullKey: "quota_limit",
905+
Limit: &pb.RateLimitResponse_RateLimit{RequestsPerUnit: 3, Unit: pb.RateLimitResponse_RateLimit_MINUTE},
906+
QuotaMode: true,
907+
ShadowMode: false,
908+
},
909+
}
910+
911+
t.config.EXPECT().GetLimit(context.Background(), "quota-domain", request.Descriptors[0]).Return(limits[0])
912+
t.config.EXPECT().GetLimit(context.Background(), "quota-domain", request.Descriptors[1]).Return(limits[1])
913+
t.cache.EXPECT().DoLimit(context.Background(), request, limits).Return(
914+
[]*pb.RateLimitResponse_DescriptorStatus{
915+
{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 0},
916+
{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[1].Limit, LimitRemaining: 0},
917+
})
918+
response, err := service.ShouldRateLimit(context.Background(), request)
919+
920+
// Overall result is OVER_LIMIT, since all quota limits were exceeded
921+
common.AssertProtoEqual(
922+
t.assert,
923+
&pb.RateLimitResponse{
924+
OverallCode: pb.RateLimitResponse_OK,
925+
Statuses: []*pb.RateLimitResponse_DescriptorStatus{
926+
{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 0},
927+
{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[1].Limit, LimitRemaining: 0},
928+
},
929+
},
930+
response)
931+
t.assert.Nil(err)
932+
}
933+
934+
func TestServiceQuotaModeOnlyAllOverTheLimit(test *testing.T) {
831935
t := commonSetup(test)
832936
defer t.controller.Finish()
833937

@@ -861,13 +965,61 @@ func TestServiceQuotaModeOnly(test *testing.T) {
861965
})
862966
response, err := service.ShouldRateLimit(context.Background(), request)
863967

864-
// All quota mode limits should result in OK overall code
968+
// Since quota limits were exceeded overall result in OVER_LIMIT
969+
common.AssertProtoEqual(
970+
t.assert,
971+
&pb.RateLimitResponse{
972+
OverallCode: pb.RateLimitResponse_OVER_LIMIT,
973+
Statuses: []*pb.RateLimitResponse_DescriptorStatus{
974+
{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0},
975+
{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[1].Limit, LimitRemaining: 0},
976+
},
977+
},
978+
response)
979+
t.assert.Nil(err)
980+
}
981+
982+
func TestServiceQuotaModeOnlySomeOverTheLimit(test *testing.T) {
983+
t := commonSetup(test)
984+
defer t.controller.Finish()
985+
986+
service := t.setupBasicService()
987+
988+
request := common.NewRateLimitRequest(
989+
"quota-domain", [][][2]string{{{"quota1", "limit"}}, {{"quota2", "limit"}}}, 1)
990+
991+
// Both limits are in quota mode
992+
limits := []*config.RateLimit{
993+
{
994+
FullKey: "quota_limit_1",
995+
Limit: &pb.RateLimitResponse_RateLimit{RequestsPerUnit: 5, Unit: pb.RateLimitResponse_RateLimit_MINUTE},
996+
QuotaMode: true,
997+
ShadowMode: false,
998+
},
999+
{
1000+
FullKey: "quota_limit_2",
1001+
Limit: &pb.RateLimitResponse_RateLimit{RequestsPerUnit: 3, Unit: pb.RateLimitResponse_RateLimit_MINUTE},
1002+
QuotaMode: true,
1003+
ShadowMode: false,
1004+
},
1005+
}
1006+
1007+
t.config.EXPECT().GetLimit(context.Background(), "quota-domain", request.Descriptors[0]).Return(limits[0])
1008+
t.config.EXPECT().GetLimit(context.Background(), "quota-domain", request.Descriptors[1]).Return(limits[1])
1009+
t.cache.EXPECT().DoLimit(context.Background(), request, limits).Return(
1010+
[]*pb.RateLimitResponse_DescriptorStatus{
1011+
{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0},
1012+
{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[1].Limit, LimitRemaining: 0},
1013+
})
1014+
response, err := service.ShouldRateLimit(context.Background(), request)
1015+
1016+
// Since only some quota limits were exceeded overall result is OK
8651017
common.AssertProtoEqual(
8661018
t.assert,
8671019
&pb.RateLimitResponse{
8681020
OverallCode: pb.RateLimitResponse_OK,
8691021
Statuses: []*pb.RateLimitResponse_DescriptorStatus{
870-
{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 0},
1022+
{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0},
8711023
{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[1].Limit, LimitRemaining: 0},
8721024
},
8731025
},
@@ -895,6 +1047,71 @@ func TestServiceQuotaModeWithShadowMode(test *testing.T) {
8951047
t.configUpdateEventChan <- t.configUpdateEvent
8961048
barrier.wait()
8971049

1050+
request := common.NewRateLimitRequest(
1051+
"quota-domain", [][][2]string{{{"regular", "limit"}}, {{"quota", "limit"}}}, 1)
1052+
1053+
// Mix of regular and quota mode limits with global shadow mode
1054+
limits := []*config.RateLimit{
1055+
{
1056+
FullKey: "regular_limit",
1057+
Limit: &pb.RateLimitResponse_RateLimit{RequestsPerUnit: 5, Unit: pb.RateLimitResponse_RateLimit_MINUTE},
1058+
QuotaMode: true,
1059+
ShadowMode: false,
1060+
},
1061+
{
1062+
FullKey: "quota_limit",
1063+
Limit: &pb.RateLimitResponse_RateLimit{RequestsPerUnit: 3, Unit: pb.RateLimitResponse_RateLimit_MINUTE},
1064+
QuotaMode: true,
1065+
ShadowMode: false,
1066+
},
1067+
}
1068+
1069+
t.config.EXPECT().GetLimit(context.Background(), "quota-domain", request.Descriptors[0]).Return(limits[0])
1070+
t.config.EXPECT().GetLimit(context.Background(), "quota-domain", request.Descriptors[1]).Return(limits[1])
1071+
t.cache.EXPECT().DoLimit(context.Background(), request, limits).Return(
1072+
[]*pb.RateLimitResponse_DescriptorStatus{
1073+
{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0},
1074+
{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[1].Limit, LimitRemaining: 0},
1075+
})
1076+
response, err := service.ShouldRateLimit(context.Background(), request)
1077+
1078+
// Global shadow mode should override everything and result in OK
1079+
common.AssertProtoEqual(
1080+
t.assert,
1081+
&pb.RateLimitResponse{
1082+
OverallCode: pb.RateLimitResponse_OK,
1083+
Statuses: []*pb.RateLimitResponse_DescriptorStatus{
1084+
{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0},
1085+
{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[1].Limit, LimitRemaining: 0},
1086+
},
1087+
},
1088+
response)
1089+
t.assert.Nil(err)
1090+
1091+
// Verify global shadow mode counter is incremented
1092+
t.assert.EqualValues(1, t.statStore.NewCounter("global_shadow_mode").Value())
1093+
}
1094+
1095+
func TestServiceMixedModeWithShadowMode(test *testing.T) {
1096+
os.Setenv("SHADOW_MODE", "true")
1097+
defer func() {
1098+
os.Unsetenv("SHADOW_MODE")
1099+
}()
1100+
1101+
t := commonSetup(test)
1102+
defer t.controller.Finish()
1103+
1104+
service := t.setupBasicService()
1105+
1106+
// Force a config reload to pick up environment variables.
1107+
barrier := newBarrier()
1108+
t.configUpdateEvent.EXPECT().GetConfig().DoAndReturn(func() (config.RateLimitConfig, any) {
1109+
barrier.signal()
1110+
return t.config, nil
1111+
})
1112+
t.configUpdateEventChan <- t.configUpdateEvent
1113+
barrier.wait()
1114+
8981115
request := common.NewRateLimitRequest(
8991116
"quota-domain", [][][2]string{{{"regular", "limit"}}, {{"quota", "limit"}}}, 1)
9001117

@@ -930,7 +1147,7 @@ func TestServiceQuotaModeWithShadowMode(test *testing.T) {
9301147
OverallCode: pb.RateLimitResponse_OK,
9311148
Statuses: []*pb.RateLimitResponse_DescriptorStatus{
9321149
{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0},
933-
{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[1].Limit, LimitRemaining: 0},
1150+
{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[1].Limit, LimitRemaining: 0},
9341151
},
9351152
},
9361153
response)

0 commit comments

Comments
 (0)