implement quota mode for soft rate limit check (#1045)

* implement quota mode for rate limit check

Signed-off-by: Dan Sun <dsun20@bloomberg.net>

* fix format

Signed-off-by: Dan Sun <dsun20@bloomberg.net>

* fix format

Signed-off-by: Dan Sun <dsun20@bloomberg.net>

* fix tests

Signed-off-by: Dan Sun <dsun20@bloomberg.net>

* fix quota mode flag

Signed-off-by: Dan Sun <dsun20@bloomberg.net>

---------

Signed-off-by: Dan Sun <dsun20@bloomberg.net>

Author: yuzisun

src/config/config.go

@@ -22,6 +22,7 @@ type RateLimit struct {
 	Limit          *pb.RateLimitResponse_RateLimit
 	Unlimited      bool
 	ShadowMode     bool
+	QuotaMode      bool
 	Name           string
 	Replaces       []string
 	DetailedMetric bool

src/config/config_impl.go

@@ -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"`
@@ -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,
@@ -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(),
@@ -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,
@@ -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()
@@ -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 == "" {
@@ -336,6 +339,7 @@ func (this *rateLimitConfigImpl) GetLimit(
 			this.statsManager.NewStats(rateLimitKey),
 			false,
 			false,
+			false,
 			"",
 			[]string{},
 			false,
@@ -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,
@@ -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
 			}
 
@@ -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
 		}
 	}

src/service/ratelimit.go

@@ -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)
 }
 
@@ -52,6 +52,7 @@ type service struct {
 	customHeaderResetHeader        string
 	customHeaderClock              utils.TimeSource
 	globalShadowMode               bool
+	globalQuotaMode                bool
 	responseDynamicMetadataEnabled bool
 }
 
@@ -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 {
@@ -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))
@@ -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 &&
@@ -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
+				}
 			}
 		}
 	}
@@ -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
@@ -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}
 }
 
@@ -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,
@@ -396,6 +435,7 @@ func NewService(cache limiter.RateLimitCache, configProvider provider.RateLimitC
 		stats:             statsManager.NewServiceStats(),
 		health:            health,
 		globalShadowMode:  shadowMode,
+		globalQuotaMode:   false,
 		customHeaderClock: clock,
 	}