Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 21 additions & 7 deletions src/service/ratelimit.go
Original file line number Diff line number Diff line change
Expand Up @@ -214,6 +214,8 @@ func (this *service) shouldRateLimitWorker(
response := &pb.RateLimitResponse{}
response.Statuses = make([]*pb.RateLimitResponse_DescriptorStatus, len(request.Descriptors))

overLimitSeen := false

// Keep track of the descriptor which is closest to hit the ratelimit
minLimitRemaining := MaxUint32
var minimumDescriptor *pb.RateLimitResponse_DescriptorStatus = nil
Expand Down Expand Up @@ -242,6 +244,7 @@ func (this *service) shouldRateLimitWorker(
response.Statuses[i] = descriptorStatus
isQuotaMode := globalQuotaMode || (limitsToCheck[i] != nil && limitsToCheck[i].QuotaMode)
if descriptorStatus.Code == pb.RateLimitResponse_OVER_LIMIT {
overLimitSeen = true
if isQuotaMode {
failedQuotaDescriptors += 1
} else {
Expand Down Expand Up @@ -290,25 +293,28 @@ func (this *service) shouldRateLimitWorker(
this.stats.GlobalShadowMode.Inc()
}

// If response dynamic data enabled, set dynamic data on response.
if this.responseDynamicMetadataEnabled {
response.DynamicMetadata = ratelimitToMetadata(request, passedDescriptors, limitsToCheck)
if this.responseDynamicMetadataEnabled && overLimitSeen {
response.DynamicMetadata = ratelimitToMetadata(request, response.Statuses, passedDescriptors, limitsToCheck)
}

response.OverallCode = finalCode
return response
}

func ratelimitToMetadata(req *pb.RateLimitRequest, passedDescriptors []int, limitsToCheck []*config.RateLimit) *structpb.Struct {
func ratelimitToMetadata(req *pb.RateLimitRequest, statuses []*pb.RateLimitResponse_DescriptorStatus, passedDescriptors []int, limitsToCheck []*config.RateLimit) *structpb.Struct {
fields := make(map[string]*structpb.Value)

// Domain
fields["domain"] = structpb.NewStringValue(req.Domain)

// Descriptors
descriptorsValues := make([]*structpb.Value, 0, len(req.Descriptors))
for _, descriptor := range req.Descriptors {
s := descriptorToStruct(descriptor)
for i, descriptor := range req.Descriptors {
var status *pb.RateLimitResponse_DescriptorStatus
if i < len(statuses) {
status = statuses[i]
}
s := descriptorToStruct(descriptor, status)
if s == nil {
continue
}
Expand Down Expand Up @@ -340,7 +346,7 @@ func ratelimitToMetadata(req *pb.RateLimitRequest, passedDescriptors []int, limi
return &structpb.Struct{Fields: fields}
}

func descriptorToStruct(descriptor *ratelimitv3.RateLimitDescriptor) *structpb.Struct {
func descriptorToStruct(descriptor *ratelimitv3.RateLimitDescriptor, status *pb.RateLimitResponse_DescriptorStatus) *structpb.Struct {
if descriptor == nil {
return nil
}
Expand All @@ -367,6 +373,14 @@ func descriptorToStruct(descriptor *ratelimitv3.RateLimitDescriptor) *structpb.S
fields["hitsAddend"] = structpb.NewNumberValue(float64(hitsAddend.GetValue()))
}

// Per-descriptor shadow_mode rewrites code to OK, so shadowed over-limit is not visible.
if status != nil {
fields["code"] = structpb.NewStringValue(status.GetCode().String())
if name := status.GetCurrentLimit().GetName(); name != "" {
fields["limitName"] = structpb.NewStringValue(name)
}
}

return &structpb.Struct{Fields: fields}
}

Expand Down
127 changes: 126 additions & 1 deletion src/service/ratelimit_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ func TestRatelimitToMetadata(t *testing.T) {
cases := []struct {
name string
req *pb.RateLimitRequest
statuses []*pb.RateLimitResponse_DescriptorStatus
passedDescriptors []int
limitsToCheck []*config.RateLimit
expected string
Expand All @@ -43,6 +44,7 @@ func TestRatelimitToMetadata(t *testing.T) {
},
},
},
statuses: nil,
passedDescriptors: nil,
limitsToCheck: []*config.RateLimit{nil},
expected: `{
Expand Down Expand Up @@ -71,6 +73,7 @@ func TestRatelimitToMetadata(t *testing.T) {
},
},
},
statuses: nil,
passedDescriptors: []int{0},
limitsToCheck: []*config.RateLimit{
{
Expand Down Expand Up @@ -127,6 +130,7 @@ func TestRatelimitToMetadata(t *testing.T) {
},
},
},
statuses: nil,
passedDescriptors: []int{1, 2},
limitsToCheck: []*config.RateLimit{
{
Expand Down Expand Up @@ -194,6 +198,7 @@ func TestRatelimitToMetadata(t *testing.T) {
},
},
},
statuses: nil,
passedDescriptors: []int{0},
limitsToCheck: []*config.RateLimit{
{
Expand All @@ -210,13 +215,133 @@ func TestRatelimitToMetadata(t *testing.T) {
],
"domain": "addend-domain",
"hitsAddend": 5
}`,
},
{
name: "Statuses carry outcome and matched limit name",
req: &pb.RateLimitRequest{
Domain: "fake-domain",
Descriptors: []*ratelimitv3.RateLimitDescriptor{
{
Entries: []*ratelimitv3.RateLimitDescriptor_Entry{
{Key: "key1", Value: "val1"},
},
},
{
Entries: []*ratelimitv3.RateLimitDescriptor_Entry{
{Key: "key2", Value: "val2"},
},
},
},
},
statuses: []*pb.RateLimitResponse_DescriptorStatus{
{
Code: pb.RateLimitResponse_OK,
CurrentLimit: &pb.RateLimitResponse_RateLimit{Name: "within-limit-rule"},
},
{
Code: pb.RateLimitResponse_OVER_LIMIT,
CurrentLimit: &pb.RateLimitResponse_RateLimit{Name: "over-limit-rule"},
},
},
passedDescriptors: nil,
limitsToCheck: nil,
expected: `{
"descriptors": [
{
"entries": [
"key1=val1"
],
"code": "OK",
"limitName": "within-limit-rule"
},
{
"entries": [
"key2=val2"
],
"code": "OVER_LIMIT",
"limitName": "over-limit-rule"
}
],
"domain": "fake-domain"
}`,
},
{
name: "Unnamed limit omits limitName",
req: &pb.RateLimitRequest{
Domain: "fake-domain",
Descriptors: []*ratelimitv3.RateLimitDescriptor{
{
Entries: []*ratelimitv3.RateLimitDescriptor_Entry{
{Key: "key1", Value: "val1"},
},
},
},
},
statuses: []*pb.RateLimitResponse_DescriptorStatus{
{
Code: pb.RateLimitResponse_OVER_LIMIT,
CurrentLimit: &pb.RateLimitResponse_RateLimit{RequestsPerUnit: 10},
},
},
passedDescriptors: nil,
limitsToCheck: nil,
expected: `{
"descriptors": [
{
"entries": [
"key1=val1"
],
"code": "OVER_LIMIT"
}
],
"domain": "fake-domain"
}`,
},
{
name: "Fewer statuses than descriptors omits outcome",
req: &pb.RateLimitRequest{
Domain: "fake-domain",
Descriptors: []*ratelimitv3.RateLimitDescriptor{
{
Entries: []*ratelimitv3.RateLimitDescriptor_Entry{
{Key: "key1", Value: "val1"},
},
},
{
Entries: []*ratelimitv3.RateLimitDescriptor_Entry{
{Key: "key2", Value: "val2"},
},
},
},
},
statuses: []*pb.RateLimitResponse_DescriptorStatus{
{Code: pb.RateLimitResponse_OVER_LIMIT},
},
passedDescriptors: nil,
limitsToCheck: nil,
expected: `{
"descriptors": [
{
"entries": [
"key1=val1"
],
"code": "OVER_LIMIT"
},
{
"entries": [
"key2=val2"
]
}
],
"domain": "fake-domain"
}`,
},
}

for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
got := ratelimitToMetadata(tc.req, tc.passedDescriptors, tc.limitsToCheck)
got := ratelimitToMetadata(tc.req, tc.statuses, tc.passedDescriptors, tc.limitsToCheck)
expected := &structpb.Struct{}
err := protojson.Unmarshal([]byte(tc.expected), expected)
require.NoError(t, err)
Expand Down
83 changes: 71 additions & 12 deletions test/service/ratelimit_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -1084,7 +1084,7 @@ func TestMetadataReturnedForPassedDescriptors(test *testing.T) {
t.assert.Equal("service_1", nameVal.GetStringValue())
}

func TestMetadataReturnedForAllPassedDescriptors(test *testing.T) {
func TestMetadataNotReturnedWhenAllDescriptorsPass(test *testing.T) {
os.Setenv("QUOTA_MODE", "true")
os.Setenv("RESPONSE_DYNAMIC_METADATA", "true")
defer func() {
Expand Down Expand Up @@ -1125,27 +1125,86 @@ func TestMetadataReturnedForAllPassedDescriptors(test *testing.T) {
{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[1].Limit, LimitRemaining: 6},
})
response, err := service.ShouldRateLimit(context.Background(), request)

t.assert.Nil(err)
t.assert.Equal(pb.RateLimitResponse_OK, response.OverallCode)
// Dynamic metadata is not built when no descriptor is over the limit
t.assert.Nil(response.DynamicMetadata)
}

func TestMetadataReturnedWithCodeAndLimitName(test *testing.T) {
os.Setenv("QUOTA_MODE", "true")
os.Setenv("RESPONSE_DYNAMIC_METADATA", "true")
defer func() {
os.Unsetenv("QUOTA_MODE")
os.Unsetenv("RESPONSE_DYNAMIC_METADATA")
}()

t := commonSetup(test)
defer t.controller.Finish()

service := t.setupBasicService()

// Force a config reload to pick up environment variables.
barrier := newBarrier()
t.configUpdateEvent.EXPECT().GetConfig().DoAndReturn(func() (config.RateLimitConfig, any) {
barrier.signal()
return t.config, nil
})
t.configUpdateEventChan <- t.configUpdateEvent
barrier.wait()

// Make a request.
request := common.NewRateLimitRequest(
"quota-domain", [][][2]string{{{"regular", "limit"}}, {{"quota", "limit"}}}, 1)

limits := []*config.RateLimit{
config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_MINUTE, t.statsManager.NewStats("key"), false, false, true, "first-limit", nil, false),
config.NewRateLimit(5, pb.RateLimitResponse_RateLimit_MINUTE, t.statsManager.NewStats("key2"), false, false, true, "second-limit", nil, false),
}
limits[0].Metadata = &structpb.Struct{Fields: map[string]*structpb.Value{"name": structpb.NewStringValue("service_1")}}
limits[1].Metadata = &structpb.Struct{Fields: map[string]*structpb.Value{"some_other_name": structpb.NewStringValue("service_2")}}

t.config.EXPECT().GetLimit(context.Background(), "quota-domain", request.Descriptors[0]).Return(limits[0])
t.config.EXPECT().GetLimit(context.Background(), "quota-domain", request.Descriptors[1]).Return(limits[1])
t.cache.EXPECT().DoLimit(context.Background(), request, limits).Return(
[]*pb.RateLimitResponse_DescriptorStatus{
{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 5},
{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[1].Limit, LimitRemaining: 0},
})
response, err := service.ShouldRateLimit(context.Background(), request)
test.Logf("DynamicMetadata: %+v", response.DynamicMetadata)

// Verify response includes metadata about quota violations
t.assert.Nil(err)
// Quota mode: OVER_LIMIT on quota descriptor alone doesn't cause overall OVER_LIMIT
t.assert.Equal(pb.RateLimitResponse_OK, response.OverallCode)
t.assert.NotNil(response.DynamicMetadata)

// Verify metadata for passed limits
// Verify descriptors include code and limitName
descriptorsVal, ok := response.DynamicMetadata.GetFields()["descriptors"]
t.assert.True(ok)
descriptors := descriptorsVal.GetListValue().GetValues()
t.assert.Equal(2, len(descriptors))

// First descriptor: OK
desc0 := descriptors[0].GetStructValue().GetFields()
t.assert.Equal("OK", desc0["code"].GetStringValue())
t.assert.Equal("first-limit", desc0["limitName"].GetStringValue())

// Second descriptor: OVER_LIMIT
desc1 := descriptors[1].GetStructValue().GetFields()
t.assert.Equal("OVER_LIMIT", desc1["code"].GetStringValue())
t.assert.Equal("second-limit", desc1["limitName"].GetStringValue())

// Verify passed metadata still works
passedMetadataVal, ok := response.DynamicMetadata.GetFields()["metadata"]
t.assert.True(ok)
passedMetadata := passedMetadataVal.GetStructValue()
t.assert.NotNil(passedMetadata)

fields := passedMetadata.GetFields()
nameVal, ok := fields["name"]
t.assert.True(ok)
// Both descriptors have passed metadata should contain values from both descriptors
t.assert.Equal("service_1", nameVal.GetStringValue())
nameVal, ok = fields["some_other_name"]
t.assert.True(ok)
t.assert.Equal("service_2", nameVal.GetStringValue())
}

func TestOverlappingMetadataReturnsTheFirstValue(test *testing.T) {
Expand Down Expand Up @@ -1183,20 +1242,20 @@ func TestOverlappingMetadataReturnsTheFirstValue(test *testing.T) {

t.config.EXPECT().GetLimit(context.Background(), "quota-domain", request.Descriptors[0]).Return(limits[0])
t.config.EXPECT().GetLimit(context.Background(), "quota-domain", request.Descriptors[1]).Return(limits[1])
// Need at least one OVER_LIMIT to trigger metadata building
t.cache.EXPECT().DoLimit(context.Background(), request, limits).Return(
[]*pb.RateLimitResponse_DescriptorStatus{
{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 5},
{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[1].Limit, LimitRemaining: 6},
{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[1].Limit, LimitRemaining: 0},
})
response, err := service.ShouldRateLimit(context.Background(), request)
test.Logf("DynamicMetadata: %+v", response.DynamicMetadata)

// Verify response includes metadata about quota violations
t.assert.Nil(err)
t.assert.Equal(pb.RateLimitResponse_OK, response.OverallCode)
t.assert.NotNil(response.DynamicMetadata)

// Verify metadata for passed limits
// Verify metadata for passed limits - first descriptor's metadata takes precedence
passedMetadataVal, ok := response.DynamicMetadata.GetFields()["metadata"]
t.assert.True(ok)
passedMetadata := passedMetadataVal.GetStructValue()
Expand All @@ -1205,7 +1264,7 @@ func TestOverlappingMetadataReturnsTheFirstValue(test *testing.T) {
fields := passedMetadata.GetFields()
nameVal, ok := fields["name"]
t.assert.True(ok)
// Metadata from the first descriptor takes precendence
// Metadata from the first descriptor takes precedence
t.assert.Equal("service_1", nameVal.GetStringValue())
}

Expand Down