Skip to content

Commit 292c602

Browse files
aabchooaabchoo
committed
shadow mode with rate limit request header
Signed-off-by: achoo30 <achoo30@bloomberg.net>
1 parent 3fb7025 commit 292c602

File tree

6 files changed

+149
-26
lines changed

6 files changed

+149
-26
lines changed

‎src/limiter/base_limiter.go‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -131,11 +131,8 @@ func (this *BaseRateLimiter) GetResponseDescriptorStatus(key string, limitInfo *
131131
}
132132
}
133133

134-
// If the limit is in ShadowMode, it should be always return OK
135134
if isOverLimit && limitInfo.limit.ShadowMode {
136135
logger.Debugf("Limit with key %s, is in shadow_mode", limitInfo.limit.FullKey)
137-
responseDescriptorStatus.Code = pb.RateLimitResponse_OK
138-
// Increase shadow mode stats if the limit was actually over the limit
139136
this.increaseShadowModeStats(isOverLimitWithLocalCache, limitInfo, hitsAddend)
140137
}
141138

‎src/service/ratelimit.go‎

Lines changed: 25 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -51,9 +51,10 @@ type service struct {
5151
customHeaderRemainingHeader string
5252
customHeaderResetHeader string
5353
customHeaderClock utils.TimeSource
54-
globalShadowMode bool
55-
globalQuotaMode bool
56-
responseDynamicMetadataEnabled bool
54+
globalShadowMode bool
55+
globalQuotaMode bool
56+
responseDynamicMetadataEnabled bool
57+
shadowModeExceededHeaderEnabled bool
5758
}
5859

5960
func (this *service) SetConfig(updateEvent provider.ConfigUpdateEvent, healthyWithAtLeastOneConfigLoad bool) {
@@ -90,6 +91,7 @@ func (this *service) SetConfig(updateEvent provider.ConfigUpdateEvent, healthyWi
9091
this.globalShadowMode = rlSettings.GlobalShadowMode
9192
this.globalQuotaMode = rlSettings.GlobalQuotaMode
9293
this.responseDynamicMetadataEnabled = rlSettings.ResponseDynamicMetadata
94+
this.shadowModeExceededHeaderEnabled = rlSettings.ShadowModeExceededHeaderEnabled
9395

9496
if rlSettings.RateLimitResponseHeadersEnabled {
9597
this.customHeadersEnabled = true
@@ -207,6 +209,7 @@ func (this *service) shouldRateLimitWorker(
207209

208210
// Track quota mode violations for metadata
209211
var quotaModeViolations []int
212+
shadowModeExceeded := false
210213

211214
for i, descriptorStatus := range responseDescriptorStatuses {
212215
// Keep track of the descriptor closest to hit the ratelimit
@@ -225,10 +228,18 @@ func (this *service) shouldRateLimitWorker(
225228
} else {
226229
response.Statuses[i] = descriptorStatus
227230
if descriptorStatus.Code == pb.RateLimitResponse_OVER_LIMIT {
231+
isShadowMode := limitsToCheck[i] != nil && limitsToCheck[i].ShadowMode
228232
// Check if this limit is in quota mode (individual or global)
229233
isQuotaMode := globalQuotaMode || (limitsToCheck[i] != nil && limitsToCheck[i].QuotaMode)
230234

231-
if isQuotaMode {
235+
if isShadowMode {
236+
shadowModeExceeded = true
237+
response.Statuses[i] = &pb.RateLimitResponse_DescriptorStatus{
238+
Code: pb.RateLimitResponse_OK,
239+
CurrentLimit: descriptorStatus.CurrentLimit,
240+
LimitRemaining: descriptorStatus.LimitRemaining,
241+
}
242+
} else if isQuotaMode {
232243
// In quota mode: track the violation for metadata but keep response as OK
233244
quotaModeViolations = append(quotaModeViolations, i)
234245
response.Statuses[i] = &pb.RateLimitResponse_DescriptorStatus{
@@ -258,9 +269,19 @@ func (this *service) shouldRateLimitWorker(
258269
// If there is a global shadow_mode, it should always return OK
259270
if finalCode == pb.RateLimitResponse_OVER_LIMIT && globalShadowMode {
260271
finalCode = pb.RateLimitResponse_OK
272+
shadowModeExceeded = true
261273
this.stats.GlobalShadowMode.Inc()
262274
}
263275

276+
if this.shadowModeExceededHeaderEnabled {
277+
response.RequestHeadersToAdd = []*core.HeaderValue{
278+
{
279+
Key: "x-ratelimit-exceeded-shadow-mode",
280+
Value: strconv.FormatBool(shadowModeExceeded),
281+
},
282+
}
283+
}
284+
264285
// If response dynamic data enabled, set dynamic data on response.
265286
if this.responseDynamicMetadataEnabled {
266287
response.DynamicMetadata = ratelimitToMetadata(request, quotaModeViolations, limitsToCheck)

‎src/settings/settings.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -208,7 +208,8 @@ type Settings struct {
208208
MemcacheTlsSkipHostnameVerification bool `envconfig:"MEMCACHE_TLS_SKIP_HOSTNAME_VERIFICATION" default:"false"`
209209

210210
// Should the ratelimiting be running in Global shadow-mode, ie. never report a ratelimit status, unless a rate was provided from envoy as an override
211-
GlobalShadowMode bool `envconfig:"SHADOW_MODE" default:"false"`
211+
GlobalShadowMode bool `envconfig:"SHADOW_MODE" default:"false"`
212+
ShadowModeExceededHeaderEnabled bool `envconfig:"SHADOW_MODE_EXCEEDED_HEADER_ENABLED" default:"false"`
212213

213214
// Should the ratelimiting be running in Global quota-mode, ie. set metadata but never report OVER_LIMIT status when quota limits are exceeded
214215
GlobalQuotaMode bool `envconfig:"QUOTA_MODE" default:"false"`

‎test/limiter/base_limiter_test.go‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -210,8 +210,8 @@ func TestGetResponseStatusOverLimitWithLocalCacheShadowMode(t *testing.T) {
210210
limitInfo := limiter.NewRateLimitInfo(limits[0], 2, 6, 4, 5)
211211
// As `isOverLimitWithLocalCache` is passed as `true`, immediate response is returned with no checks of the limits.
212212
responseStatus := baseRateLimit.GetResponseDescriptorStatus("key", limitInfo, true, 2)
213-
// Limit is reached, but response is still OK due to ShadowMode
214-
assert.Equal(pb.RateLimitResponse_OK, responseStatus.GetCode())
213+
// Limit is reached, shadow mode code conversion now happens in service layer
214+
assert.Equal(pb.RateLimitResponse_OVER_LIMIT, responseStatus.GetCode())
215215
assert.Equal(uint32(0), responseStatus.GetLimitRemaining())
216216
assert.Equal(limits[0].Limit, responseStatus.GetCurrentLimit())
217217
assert.Equal(uint64(2), limits[0].Stats.OverLimit.Value())
@@ -259,7 +259,8 @@ func TestGetResponseStatusOverLimitShadowMode(t *testing.T) {
259259
limits := []*config.RateLimit{config.NewRateLimit(5, pb.RateLimitResponse_RateLimit_SECOND, sm.NewStats("key_value"), false, true, false, "", nil, false)}
260260
limitInfo := limiter.NewRateLimitInfo(limits[0], 2, 7, 4, 5)
261261
responseStatus := baseRateLimit.GetResponseDescriptorStatus("key", limitInfo, false, 1)
262-
assert.Equal(pb.RateLimitResponse_OK, responseStatus.GetCode())
262+
// Shadow mode code conversion now happens in service layer
263+
assert.Equal(pb.RateLimitResponse_OVER_LIMIT, responseStatus.GetCode())
263264
assert.Equal(uint32(0), responseStatus.GetLimitRemaining())
264265
assert.Equal(limits[0].Limit, responseStatus.GetCurrentLimit())
265266
result, _ := localCache.Get([]byte("key"))

‎test/redis/fixed_cache_impl_test.go‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -564,10 +564,10 @@ func TestOverLimitWithLocalCacheShadowRule(t *testing.T) {
564564
"EXPIRE", "domain_key4_value4_997200", int64(3600)).DoAndReturn(pipeAppend)
565565
client.EXPECT().PipeDo(gomock.Any()).Return(nil)
566566

567-
// The result should be OK since limit is in ShadowMode
567+
// Shadow mode code conversion now happens in service layer, cache returns OVER_LIMIT
568568
assert.Equal(
569569
[]*pb.RateLimitResponse_DescriptorStatus{
570-
{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 0, DurationUntilReset: utils.CalculateReset(&limits[0].Limit.Unit, timeSource)},
570+
{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0, DurationUntilReset: utils.CalculateReset(&limits[0].Limit.Unit, timeSource)},
571571
},
572572
cache.DoLimit(context.Background(), request, limits))
573573
assert.Equal(uint64(3), limits[0].Stats.TotalHits.Value())
@@ -586,10 +586,10 @@ func TestOverLimitWithLocalCacheShadowRule(t *testing.T) {
586586
client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(),
587587
"EXPIRE", "domain_key4_value4_997200", int64(3600)).Times(0)
588588

589-
// The result should be OK since limit is in ShadowMode
589+
// Shadow mode code conversion now happens in service layer, cache returns OVER_LIMIT
590590
assert.Equal(
591591
[]*pb.RateLimitResponse_DescriptorStatus{
592-
{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 0, DurationUntilReset: utils.CalculateReset(&limits[0].Limit.Unit, timeSource)},
592+
{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0, DurationUntilReset: utils.CalculateReset(&limits[0].Limit.Unit, timeSource)},
593593
},
594594
cache.DoLimit(context.Background(), request, limits))
595595

‎test/service/ratelimit_test.go‎

Lines changed: 114 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -288,16 +288,16 @@ func TestRuleShadowMode(test *testing.T) {
288288
t.config.EXPECT().GetLimit(context.Background(), "different-domain", request.Descriptors[1]).Return(limits[1])
289289
t.cache.EXPECT().DoLimit(context.Background(), request, limits).Return(
290290
[]*pb.RateLimitResponse_DescriptorStatus{
291-
{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 0},
292-
{Code: pb.RateLimitResponse_OK, CurrentLimit: nil, LimitRemaining: 0},
291+
{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0},
292+
{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[1].Limit, LimitRemaining: 0},
293293
})
294294
response, err := service.ShouldRateLimit(context.Background(), request)
295295
t.assert.Equal(
296296
&pb.RateLimitResponse{
297297
OverallCode: pb.RateLimitResponse_OK,
298298
Statuses: []*pb.RateLimitResponse_DescriptorStatus{
299299
{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 0},
300-
{Code: pb.RateLimitResponse_OK, CurrentLimit: nil, LimitRemaining: 0},
300+
{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[1].Limit, LimitRemaining: 0},
301301
},
302302
},
303303
response)
@@ -319,16 +319,10 @@ func TestMixedRuleShadowMode(test *testing.T) {
319319
}
320320
t.config.EXPECT().GetLimit(context.Background(), "different-domain", request.Descriptors[0]).Return(limits[0])
321321
t.config.EXPECT().GetLimit(context.Background(), "different-domain", request.Descriptors[1]).Return(limits[1])
322-
testResults := []pb.RateLimitResponse_Code{pb.RateLimitResponse_OVER_LIMIT, pb.RateLimitResponse_OVER_LIMIT}
323-
for i := 0; i < len(limits); i++ {
324-
if limits[i].ShadowMode {
325-
testResults[i] = pb.RateLimitResponse_OK
326-
}
327-
}
328322
t.cache.EXPECT().DoLimit(context.Background(), request, limits).Return(
329323
[]*pb.RateLimitResponse_DescriptorStatus{
330-
{Code: testResults[0], CurrentLimit: limits[0].Limit, LimitRemaining: 0},
331-
{Code: testResults[1], CurrentLimit: nil, LimitRemaining: 0},
324+
{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0},
325+
{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: nil, LimitRemaining: 0},
332326
})
333327
response, err := service.ShouldRateLimit(context.Background(), request)
334328
t.assert.Equal(
@@ -345,6 +339,115 @@ func TestMixedRuleShadowMode(test *testing.T) {
345339
t.assert.EqualValues(0, t.statStore.NewCounter("global_shadow_mode").Value())
346340
}
347341

342+
func TestShadowModeExceededHeader(test *testing.T) {
343+
os.Setenv("SHADOW_MODE_EXCEEDED_HEADER_ENABLED", "true")
344+
defer os.Unsetenv("SHADOW_MODE_EXCEEDED_HEADER_ENABLED")
345+
346+
t := commonSetup(test)
347+
defer t.controller.Finish()
348+
service := t.setupBasicService()
349+
350+
// Test 1: Shadow mode descriptor exceeded - header should be true
351+
request := common.NewRateLimitRequest(
352+
"different-domain", [][][2]string{{{"foo", "bar"}}}, 1)
353+
limits := []*config.RateLimit{
354+
config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_MINUTE, t.statsManager.NewStats("key"), false, true, false, "", nil, false),
355+
}
356+
t.config.EXPECT().GetLimit(context.Background(), "different-domain", request.Descriptors[0]).Return(limits[0])
357+
t.cache.EXPECT().DoLimit(context.Background(), request, limits).Return(
358+
[]*pb.RateLimitResponse_DescriptorStatus{
359+
{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0},
360+
})
361+
response, err := service.ShouldRateLimit(context.Background(), request)
362+
t.assert.Nil(err)
363+
t.assert.Equal(pb.RateLimitResponse_OK, response.OverallCode)
364+
t.assert.Equal(pb.RateLimitResponse_OK, response.Statuses[0].Code)
365+
t.assert.Equal(1, len(response.RequestHeadersToAdd))
366+
t.assert.Equal("x-ratelimit-exceeded-shadow-mode", response.RequestHeadersToAdd[0].Key)
367+
t.assert.Equal("true", response.RequestHeadersToAdd[0].Value)
368+
}
369+
370+
func TestShadowModeExceededHeaderNotExceeded(test *testing.T) {
371+
os.Setenv("SHADOW_MODE_EXCEEDED_HEADER_ENABLED", "true")
372+
defer os.Unsetenv("SHADOW_MODE_EXCEEDED_HEADER_ENABLED")
373+
374+
t := commonSetup(test)
375+
defer t.controller.Finish()
376+
service := t.setupBasicService()
377+
378+
// Test: Shadow mode descriptor NOT exceeded - header should be false
379+
request := common.NewRateLimitRequest(
380+
"different-domain", [][][2]string{{{"foo", "bar"}}}, 1)
381+
limits := []*config.RateLimit{
382+
config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_MINUTE, t.statsManager.NewStats("key"), false, true, false, "", nil, false),
383+
}
384+
t.config.EXPECT().GetLimit(context.Background(), "different-domain", request.Descriptors[0]).Return(limits[0])
385+
t.cache.EXPECT().DoLimit(context.Background(), request, limits).Return(
386+
[]*pb.RateLimitResponse_DescriptorStatus{
387+
{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 5},
388+
})
389+
response, err := service.ShouldRateLimit(context.Background(), request)
390+
t.assert.Nil(err)
391+
t.assert.Equal(pb.RateLimitResponse_OK, response.OverallCode)
392+
t.assert.Equal(1, len(response.RequestHeadersToAdd))
393+
t.assert.Equal("x-ratelimit-exceeded-shadow-mode", response.RequestHeadersToAdd[0].Key)
394+
t.assert.Equal("false", response.RequestHeadersToAdd[0].Value)
395+
}
396+
397+
func TestShadowModeExceededHeaderDisabled(test *testing.T) {
398+
// Feature flag not set, so no header should be added
399+
t := commonSetup(test)
400+
defer t.controller.Finish()
401+
service := t.setupBasicService()
402+
403+
request := common.NewRateLimitRequest(
404+
"different-domain", [][][2]string{{{"foo", "bar"}}}, 1)
405+
limits := []*config.RateLimit{
406+
config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_MINUTE, t.statsManager.NewStats("key"), false, true, false, "", nil, false),
407+
}
408+
t.config.EXPECT().GetLimit(context.Background(), "different-domain", request.Descriptors[0]).Return(limits[0])
409+
t.cache.EXPECT().DoLimit(context.Background(), request, limits).Return(
410+
[]*pb.RateLimitResponse_DescriptorStatus{
411+
{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0},
412+
})
413+
response, err := service.ShouldRateLimit(context.Background(), request)
414+
t.assert.Nil(err)
415+
t.assert.Equal(pb.RateLimitResponse_OK, response.OverallCode)
416+
t.assert.Nil(response.RequestHeadersToAdd)
417+
}
418+
419+
func TestShadowModeExceededHeaderGlobalShadowMode(test *testing.T) {
420+
os.Setenv("SHADOW_MODE", "true")
421+
os.Setenv("SHADOW_MODE_EXCEEDED_HEADER_ENABLED", "true")
422+
defer func() {
423+
os.Unsetenv("SHADOW_MODE")
424+
os.Unsetenv("SHADOW_MODE_EXCEEDED_HEADER_ENABLED")
425+
}()
426+
427+
t := commonSetup(test)
428+
defer t.controller.Finish()
429+
service := t.setupBasicService()
430+
431+
// Non-shadow-mode descriptor that is over limit, but global shadow mode converts it
432+
request := common.NewRateLimitRequest(
433+
"different-domain", [][][2]string{{{"foo", "bar"}}}, 1)
434+
limits := []*config.RateLimit{
435+
config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_MINUTE, t.statsManager.NewStats("key"), false, false, false, "", nil, false),
436+
}
437+
t.config.EXPECT().GetLimit(context.Background(), "different-domain", request.Descriptors[0]).Return(limits[0])
438+
t.cache.EXPECT().DoLimit(context.Background(), request, limits).Return(
439+
[]*pb.RateLimitResponse_DescriptorStatus{
440+
{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0},
441+
})
442+
response, err := service.ShouldRateLimit(context.Background(), request)
443+
t.assert.Nil(err)
444+
t.assert.Equal(pb.RateLimitResponse_OK, response.OverallCode)
445+
t.assert.Equal(1, len(response.RequestHeadersToAdd))
446+
t.assert.Equal("x-ratelimit-exceeded-shadow-mode", response.RequestHeadersToAdd[0].Key)
447+
t.assert.Equal("true", response.RequestHeadersToAdd[0].Value)
448+
t.assert.EqualValues(1, t.statStore.NewCounter("global_shadow_mode").Value())
449+
}
450+
348451
func TestServiceWithCustomRatelimitHeaders(test *testing.T) {
349452
os.Setenv("LIMIT_RESPONSE_HEADERS_ENABLED", "true")
350453
os.Setenv("LIMIT_LIMIT_HEADER", "A-Ratelimit-Limit")

0 commit comments

Comments
 (0)