Skip to content

implement quota mode for soft rate limit check #1045

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
collin-lee merged 5 commits into from
Jan 15, 2026
Merged

implement quota mode for soft rate limit check #1045

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
14f65e1
implement quota mode for rate limit check
yuzisun Jan 10, 2026
8e2fe71
fix format
yuzisun Jan 10, 2026
ea64227
fix format
yuzisun Jan 10, 2026
98be9cb
fix tests
yuzisun Jan 11, 2026
63a43e0
fix quota mode flag
yuzisun Jan 11, 2026
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
1 change: 1 addition & 0 deletions src/config/config.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ type RateLimit struct {
Limit *pb.RateLimitResponse_RateLimit
Unlimited bool
ShadowMode bool
QuotaMode bool
Name string
Replaces []string
DetailedMetric bool
Expand Down
21 changes: 13 additions & 8 deletions src/config/config_impl.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ type YamlDescriptor struct {
RateLimit *YamlRateLimit `yaml:"rate_limit"`
Descriptors []YamlDescriptor
ShadowMode bool `yaml:"shadow_mode"`
QuotaMode bool `yaml:"quota_mode"`
DetailedMetric bool `yaml:"detailed_metric"`
ValueToMetric bool `yaml:"value_to_metric"`
ShareThreshold bool `yaml:"share_threshold"`
Expand Down Expand Up @@ -70,6 +71,7 @@ var validKeys = map[string]bool{
"requests_per_unit": true,
"unlimited": true,
"shadow_mode": true,
"quota_mode": true,
"name": true,
"replaces": true,
"detailed_metric": true,
Expand All @@ -84,7 +86,7 @@ var validKeys = map[string]bool{
// @param unlimited supplies whether the rate limit is unlimited
// @return the new config entry.
func NewRateLimit(requestsPerUnit uint32, unit pb.RateLimitResponse_RateLimit_Unit, rlStats stats.RateLimitStats,
unlimited bool, shadowMode bool, name string, replaces []string, detailedMetric bool,
unlimited bool, shadowMode bool, quotaMode bool, name string, replaces []string, detailedMetric bool,
) *RateLimit {
return &RateLimit{
FullKey: rlStats.GetKey(),
Expand All @@ -96,6 +98,7 @@ func NewRateLimit(requestsPerUnit uint32, unit pb.RateLimitResponse_RateLimit_Un
},
Unlimited: unlimited,
ShadowMode: shadowMode,
QuotaMode: quotaMode,
Name: name,
Replaces: replaces,
DetailedMetric: detailedMetric,
Expand All @@ -108,8 +111,8 @@ func (this *rateLimitDescriptor) dump() string {
ret := ""
if this.limit != nil {
ret += fmt.Sprintf(
"%s: unit=%s requests_per_unit=%d, shadow_mode: %t\n", this.limit.FullKey,
this.limit.Limit.Unit.String(), this.limit.Limit.RequestsPerUnit, this.limit.ShadowMode)
"%s: unit=%s requests_per_unit=%d, shadow_mode: %t, quota_mode: %t\n", this.limit.FullKey,
this.limit.Limit.Unit.String(), this.limit.Limit.RequestsPerUnit, this.limit.ShadowMode, this.limit.QuotaMode)
}
for _, descriptor := range this.descriptors {
ret += descriptor.dump()
Expand Down Expand Up @@ -174,12 +177,12 @@ func (this *rateLimitDescriptor) loadDescriptors(config RateLimitConfigToLoad, p

rateLimit = NewRateLimit(
descriptorConfig.RateLimit.RequestsPerUnit, pb.RateLimitResponse_RateLimit_Unit(value),
statsManager.NewStats(newParentKey), unlimited, descriptorConfig.ShadowMode,
statsManager.NewStats(newParentKey), unlimited, descriptorConfig.ShadowMode, descriptorConfig.QuotaMode,
descriptorConfig.RateLimit.Name, replaces, descriptorConfig.DetailedMetric,
)
rateLimitDebugString = fmt.Sprintf(
" ratelimit={requests_per_unit=%d, unit=%s, unlimited=%t, shadow_mode=%t}", rateLimit.Limit.RequestsPerUnit,
rateLimit.Limit.Unit.String(), rateLimit.Unlimited, rateLimit.ShadowMode)
" ratelimit={requests_per_unit=%d, unit=%s, unlimited=%t, shadow_mode=%t, quota_mode=%t}", rateLimit.Limit.RequestsPerUnit,
rateLimit.Limit.Unit.String(), rateLimit.Unlimited, rateLimit.ShadowMode, rateLimit.QuotaMode)

for _, replaces := range descriptorConfig.RateLimit.Replaces {
if replaces.Name == "" {
Expand Down Expand Up @@ -336,6 +339,7 @@ func (this *rateLimitConfigImpl) GetLimit(
this.statsManager.NewStats(rateLimitKey),
false,
false,
false,
"",
[]string{},
false,
Expand Down Expand Up @@ -452,6 +456,7 @@ func (this *rateLimitConfigImpl) GetLimit(
Limit: originalLimit.Limit,
Unlimited: originalLimit.Unlimited,
ShadowMode: originalLimit.ShadowMode,
QuotaMode: originalLimit.QuotaMode,
Name: originalLimit.Name,
Replaces: originalLimit.Replaces,
DetailedMetric: originalLimit.DetailedMetric,
Expand Down Expand Up @@ -481,7 +486,7 @@ func (this *rateLimitConfigImpl) GetLimit(
if rateLimit != nil && rateLimit.DetailedMetric {
// Preserve ShareThresholdKeyPattern when recreating rate limit
originalShareThresholdKeyPattern := rateLimit.ShareThresholdKeyPattern
rateLimit = NewRateLimit(rateLimit.Limit.RequestsPerUnit, rateLimit.Limit.Unit, this.statsManager.NewStats(rateLimit.FullKey), rateLimit.Unlimited, rateLimit.ShadowMode, rateLimit.Name, rateLimit.Replaces, rateLimit.DetailedMetric)
rateLimit = NewRateLimit(rateLimit.Limit.RequestsPerUnit, rateLimit.Limit.Unit, this.statsManager.NewStats(rateLimit.FullKey), rateLimit.Unlimited, rateLimit.ShadowMode, rateLimit.QuotaMode, rateLimit.Name, rateLimit.Replaces, rateLimit.DetailedMetric)
rateLimit.ShareThresholdKeyPattern = originalShareThresholdKeyPattern
}

Expand Down Expand Up @@ -531,7 +536,7 @@ func (this *rateLimitConfigImpl) GetLimit(
if enhancedKey != rateLimit.FullKey {
// Recreate to ensure a clean stats struct, then set to enhanced stats
originalShareThresholdKeyPattern := rateLimit.ShareThresholdKeyPattern
rateLimit = NewRateLimit(rateLimit.Limit.RequestsPerUnit, rateLimit.Limit.Unit, this.statsManager.NewStats(enhancedKey), rateLimit.Unlimited, rateLimit.ShadowMode, rateLimit.Name, rateLimit.Replaces, rateLimit.DetailedMetric)
rateLimit = NewRateLimit(rateLimit.Limit.RequestsPerUnit, rateLimit.Limit.Unit, this.statsManager.NewStats(enhancedKey), rateLimit.Unlimited, rateLimit.ShadowMode, rateLimit.QuotaMode, rateLimit.Name, rateLimit.Replaces, rateLimit.DetailedMetric)
rateLimit.ShareThresholdKeyPattern = originalShareThresholdKeyPattern
}
}
Expand Down
60 changes: 50 additions & 10 deletions src/service/ratelimit.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ var tracer = otel.Tracer("ratelimit")

type RateLimitServiceServer interface {
pb.RateLimitServiceServer
GetCurrentConfig() (config.RateLimitConfig, bool)
GetCurrentConfig() (config.RateLimitConfig, bool, bool)
SetConfig(updateEvent provider.ConfigUpdateEvent, healthyWithAtLeastOneConfigLoad bool)
}

Expand All @@ -52,6 +52,7 @@ type service struct {
customHeaderResetHeader string
customHeaderClock utils.TimeSource
globalShadowMode bool
globalQuotaMode bool
responseDynamicMetadataEnabled bool
}

Expand Down Expand Up @@ -87,6 +88,7 @@ func (this *service) SetConfig(updateEvent provider.ConfigUpdateEvent, healthyWi

rlSettings := settings.NewSettings()
this.globalShadowMode = rlSettings.GlobalShadowMode
this.globalQuotaMode = rlSettings.GlobalQuotaMode
this.responseDynamicMetadataEnabled = rlSettings.ResponseDynamicMetadata

if rlSettings.RateLimitResponseHeadersEnabled {
Expand Down Expand Up @@ -186,7 +188,7 @@ func (this *service) shouldRateLimitWorker(
checkServiceErr(request.Domain != "", "rate limit domain must not be empty")
checkServiceErr(len(request.Descriptors) != 0, "rate limit descriptor list must not be empty")

snappedConfig, globalShadowMode := this.GetCurrentConfig()
snappedConfig, globalShadowMode, globalQuotaMode := this.GetCurrentConfig()
limitsToCheck, isUnlimited := this.constructLimitsToCheck(request, ctx, snappedConfig)

assert.Assert(len(limitsToCheck) == len(isUnlimited))
Expand All @@ -203,6 +205,9 @@ func (this *service) shouldRateLimitWorker(
minLimitRemaining := MaxUint32
var minimumDescriptor *pb.RateLimitResponse_DescriptorStatus = nil

// Track quota mode violations for metadata
var quotaModeViolations []int

for i, descriptorStatus := range responseDescriptorStatuses {
// Keep track of the descriptor closest to hit the ratelimit
if this.customHeadersEnabled &&
Expand All @@ -220,10 +225,23 @@ func (this *service) shouldRateLimitWorker(
} else {
response.Statuses[i] = descriptorStatus
if descriptorStatus.Code == pb.RateLimitResponse_OVER_LIMIT {
finalCode = descriptorStatus.Code

minimumDescriptor = descriptorStatus
minLimitRemaining = 0
// Check if this limit is in quota mode (individual or global)
isQuotaMode := globalQuotaMode || (limitsToCheck[i] != nil && limitsToCheck[i].QuotaMode)

if isQuotaMode {
// In quota mode: track the violation for metadata but keep response as OK
quotaModeViolations = append(quotaModeViolations, i)
response.Statuses[i] = &pb.RateLimitResponse_DescriptorStatus{
Code: pb.RateLimitResponse_OK,
CurrentLimit: descriptorStatus.CurrentLimit,
LimitRemaining: descriptorStatus.LimitRemaining,
}
} else {
// Normal rate limit: set final code to OVER_LIMIT
finalCode = descriptorStatus.Code
minimumDescriptor = descriptorStatus
minLimitRemaining = 0
}
}
}
}
Expand All @@ -245,14 +263,14 @@ func (this *service) shouldRateLimitWorker(

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

response.OverallCode = finalCode
return response
}

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

// Domain
Expand All @@ -276,6 +294,27 @@ func ratelimitToMetadata(req *pb.RateLimitRequest) *structpb.Struct {
fields["hitsAddend"] = structpb.NewNumberValue(float64(hitsAddend))
}

// Quota mode information
if len(quotaModeViolations) > 0 {
violationValues := make([]*structpb.Value, len(quotaModeViolations))
for i, violationIndex := range quotaModeViolations {
violationValues[i] = structpb.NewNumberValue(float64(violationIndex))
}
fields["quotaModeViolations"] = structpb.NewListValue(&structpb.ListValue{
Values: violationValues,
})
}

// Check if any limits have quota mode enabled
quotaModeEnabled := false
for _, limit := range limitsToCheck {
if limit != nil && limit.QuotaMode {
quotaModeEnabled = true
break
}
}
fields["quotaModeEnabled"] = structpb.NewBoolValue(quotaModeEnabled)

return &structpb.Struct{Fields: fields}
}

Expand Down Expand Up @@ -379,10 +418,10 @@ func (this *service) ShouldRateLimit(
return response, nil
}

func (this *service) GetCurrentConfig() (config.RateLimitConfig, bool) {
func (this *service) GetCurrentConfig() (config.RateLimitConfig, bool, bool) {
this.configLock.RLock()
defer this.configLock.RUnlock()
return this.config, this.globalShadowMode
return this.config, this.globalShadowMode, this.globalQuotaMode
}

func NewService(cache limiter.RateLimitCache, configProvider provider.RateLimitConfigProvider, statsManager stats.Manager,
Expand All @@ -396,6 +435,7 @@ func NewService(cache limiter.RateLimitCache, configProvider provider.RateLimitC
stats: statsManager.NewServiceStats(),
health: health,
globalShadowMode: shadowMode,
globalQuotaMode: false,
customHeaderClock: clock,
}

Expand Down
Loading
Loading