Allow dynamic window multiplier via a new extension (#833)

Author: lahsivjar

processor/ratelimitprocessor/README.md

@@ -37,8 +37,9 @@ You can override one or more of the following fields:
 | Field                  | Description                                                                                                                                                                                                       | Required | Default    |
 |------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|----------|------------|
 | `enabled`              | Enables the dynamic rate limiting feature.                                                                                                                                                                        | No       | `false`    |
-| `window_multiplier`    | The factor by which the previous window rate is multiplied to get the dynamic limit.                                                                                                                             | No       | `1.3`      |
 | `window_duration`      | The time window duration for calculating traffic rates.                                                                                                                                                           | No       | `2m`       |
+| `default_window_multiplier` | The default factor by which the previous window rate is multiplied to get the dynamic limit. Can be overridden by providing a `window_configurator` extension.                                               | No       | `1.3`      |
+| `window_configurator`  | An optional extension to calculate window multiplier dynamically based on unique keys.                                                                                                                            | No       |            |
 
 ### Dynamic Rate Limiting Deep Dive
 
@@ -59,7 +60,14 @@ Where:
 
 * `previous_rate`: The rate of traffic in the previous window (normalized per second).
 * `static_rate`: The configured `rate` in the main configuration.
-* `window_multiplier`: A factor applied to the previous window rate to determine the dynamic limit.
+* `window_multiplier`: A factor applied to the previous window rate to determine the dynamic limit. It is derived based on the following formula:
+
+  ```text
+  window_multiplier = default_window_multiplier
+  if window_configurator is defined {
+    window_multiplier = window_configurator::Multiplier(...)
+  }
+  ```
 
 **Important Notes:**
 
@@ -72,7 +80,7 @@ Let's walk through a few examples to illustrate the behavior of the dynamic rate
 Assume the following configuration:
 
 * `window_duration`: 2m
-* `window_multiplier`: 1.5
+* `window_multiplier`: 1.5 (with no `window_configurator` provided)
 * `rate`: 1000 requests/second (this is our `static_rate`)
 
 #### Scenario 1: Initial Traffic
@@ -296,3 +304,18 @@ Telemetry and metrics:
 
   * `ratelimit.resolver.failures` — total number of resolver failures
   * `ratelimit.dynamic.escalations` — number of times dynamic rate was peeked (attributes: `class`, `source_kind`, `success`)
+
+### Throttling rate based on custom logic using `window_configurator`
+
+The `window_configurator` option configures a custom OpenTelemetry Collector extension to dynamically choose a window multiplier for the current period. The value is the extension ID (the extension's configured name). The extension MUST implement the `WindowConfigurator` interface. The multiplier can be used to scale up the rate limit from previous window (by returning a multiplier greater than `1`) or scale down the rate limit from previous window (by returning a multiplier less than `1`, greater than `0`). If the extension returns a negative value then the `default_window_multiplier` will be used. Note that the dynamic limit MUST be at least the configured `static_rate`, ensuring a minimum level of throughput. An example configuration including the window configurator:
+
+```yaml
+processors:
+  ratelimiter:
+    dynamic_limits:
+      enabled: true
+      window_duration: 2m
+      default_window_multiplier: 1.5
+      # This is an example window configurator. It doesn't exist.
+      window_configurator: kafkalagwindowconfiguratorextension
+```

processor/ratelimitprocessor/config.go

@@ -71,12 +71,22 @@ type Config struct {
 type DynamicRateLimiting struct {
 	// Enabled tells the processor to use dynamic rate limiting.
 	Enabled bool `mapstructure:"enabled"`
-	// WindowMultiplier is the factor by which the previous window rate is
-	// multiplied to get the dynamic part of the limit. Defaults to 1.3.
-	WindowMultiplier float64 `mapstructure:"window_multiplier"`
+
 	// WindowDuration defines the time window for which the dynamic rate limit
-	// is calculated on.
+	// is calculated on. Defaults to 2 minutes.
 	WindowDuration time.Duration `mapstructure:"window_duration"`
+
+	// DefaultWindowMultiplier is the factor by which the previous window rate is
+	// multiplied to get the dynamic part of the limit. Defaults to 1.3.
+	DefaultWindowMultiplier float64 `mapstructure:"default_window_multiplier"`
+
+	// WindowConfigurator is the component ID of the extension to dynamically
+	// determine the window multiplier. The extension is expected to implement
+	// the `WindowConfigurator` interface. The window configurator is used in
+	// the hot path so it should respond fast. The effective rate cannot go
+	// below the configured static rate limit settings. If the configurator
+	// returns a negative multiplier then the default multiplier will be used.
+	WindowConfigurator component.ID `mapstructure:"window_configurator"`
 }
 
 // Class defines a named rate limit class for class-based dynamic rate limiting.
@@ -98,8 +108,8 @@ func (d *DynamicRateLimiting) Validate() error {
 		return nil
 	}
 	var errs []error
-	if d.WindowMultiplier < 1 {
-		errs = append(errs, errors.New("window_multiplier must be greater than or equal to 1"))
+	if d.DefaultWindowMultiplier < 1 {
+		errs = append(errs, errors.New("default_window_multiplier must be greater than or equal to 1"))
 	}
 	if d.WindowDuration <= 0 {
 		errs = append(errs, errors.New("window_duration must be greater than zero"))
@@ -244,8 +254,8 @@ func createDefaultConfig() component.Config {
 			ThrottleInterval: DefaultThrottleInterval,
 		},
 		DynamicRateLimiting: DynamicRateLimiting{
-			WindowMultiplier: 1.3,
-			WindowDuration:   2 * time.Minute,
+			DefaultWindowMultiplier: 1.3,
+			WindowDuration:          2 * time.Minute,
 		},
 		Classes:      nil,
 		DefaultClass: "",

processor/ratelimitprocessor/config_test.go

@@ -32,8 +32,8 @@ import (
 )
 
 var defaultDynamicRateLimiting = DynamicRateLimiting{
-	WindowMultiplier: 1.3,
-	WindowDuration:   2 * time.Minute,
+	DefaultWindowMultiplier: 1.3,
+	WindowDuration:          2 * time.Minute,
 }
 
 func TestLoadConfig(t *testing.T) {
@@ -211,9 +211,9 @@ func TestLoadConfig(t *testing.T) {
 					ThrottleInterval: 1 * time.Second,
 				},
 				DynamicRateLimiting: DynamicRateLimiting{
-					Enabled:          true,
-					WindowMultiplier: 1.5,
-					WindowDuration:   time.Minute,
+					Enabled:                 true,
+					DefaultWindowMultiplier: 1.5,
+					WindowDuration:          time.Minute,
 				},
 			},
 		},

processor/ratelimitprocessor/gubernator.go

@@ -51,12 +51,29 @@ type ClassResolver interface {
 	ResolveClass(ctx context.Context, key string) (string, error)
 }
 
+// WindowConfigurator allows adjusting the rates dynamically by configuring
+// the multiplier for the next calculation window.
+//
+// NOTE(lahsivjar): We may want to make the duration configurable too.
+type WindowConfigurator interface {
+	// Multiplier returns the calculated multiplier for the next window.
+	Multiplier(ctx context.Context, window time.Duration, key string) float64
+}
+
 type noopResolver struct{}
 
 func (noopResolver) ResolveClass(context.Context, string) (string, error) {
 	return "", nil
 }
 
+type defaultWindowConfigurator struct {
+	multiplier float64
+}
+
+func (d defaultWindowConfigurator) Multiplier(context.Context, time.Duration, string) float64 {
+	return d.multiplier
+}
+
 type gubernatorRateLimiter struct {
 	cfg      *Config
 	logger   *zap.Logger
@@ -67,8 +84,9 @@ type gubernatorRateLimiter struct {
 	client     gubernator.V1Client
 	clientConn *grpc.ClientConn
 	// Class resolver for class-based rate limiting
-	classResolver    ClassResolver
-	telemetryBuilder *metadata.TelemetryBuilder
+	classResolver      ClassResolver
+	windowConfigurator WindowConfigurator
+	telemetryBuilder   *metadata.TelemetryBuilder
 }
 
 func newGubernatorDaemonConfig(logger *zap.Logger) (gubernator.DaemonConfig, error) {
@@ -102,12 +120,13 @@ func newGubernatorRateLimiter(cfg *Config, logger *zap.Logger, telemetryBuilder
 	}
 
 	return &gubernatorRateLimiter{
-		cfg:              cfg,
-		logger:           logger,
-		behavior:         gubernator.Behavior_BATCHING,
-		daemonCfg:        daemonCfg,
-		telemetryBuilder: telemetryBuilder,
-		classResolver:    noopResolver{},
+		cfg:                cfg,
+		logger:             logger,
+		behavior:           gubernator.Behavior_BATCHING,
+		daemonCfg:          daemonCfg,
+		telemetryBuilder:   telemetryBuilder,
+		classResolver:      noopResolver{},
+		windowConfigurator: defaultWindowConfigurator{multiplier: cfg.DynamicRateLimiting.DefaultWindowMultiplier},
 	}, nil
 }
 
@@ -123,6 +142,17 @@ func (r *gubernatorRateLimiter) Start(ctx context.Context, host component.Host)
 		r.classResolver = cr.(ClassResolver)
 	}
 
+	if wCon := r.cfg.DynamicRateLimiting.WindowConfigurator; wCon.String() != "" {
+		wc, ok := host.GetExtensions()[wCon]
+		if !ok {
+			return fmt.Errorf("window configurator %s not found", wCon)
+		}
+		if err := wc.Start(ctx, host); err != nil {
+			return fmt.Errorf("failed to start window configurator %s: %w", wCon, err)
+		}
+		r.windowConfigurator = wc.(WindowConfigurator)
+	}
+
 	r.daemon, err = gubernator.SpawnDaemon(ctx, r.daemonCfg)
 	if err != nil {
 		return fmt.Errorf("failed to spawn gubernator daemon: %w", err)
@@ -154,6 +184,11 @@ func (r *gubernatorRateLimiter) Shutdown(ctx context.Context) error {
 			return fmt.Errorf("failed to shutdown class resolver: %w", err)
 		}
 	}
+	if w, ok := r.windowConfigurator.(component.Component); ok {
+		if err := w.Shutdown(ctx); err != nil {
+			return fmt.Errorf("failed to shutdown window configurator: %w", err)
+		}
+	}
 	return nil
 }
 
@@ -356,20 +391,35 @@ func (r *gubernatorRateLimiter) getDynamicLimit(ctx context.Context,
 	// reqs/events/bytes per Throttle interval may not be 1s.
 	staticRate := float64(cfg.Rate) / r.cfg.ThrottleInterval.Seconds()
 	drc := newDynamicRateContext(uniqueKey, now, r.cfg.DynamicRateLimiting)
-	// Get current and previous window rates
-	current, previous, err := r.peekRates(ctx, int64(hits), drc)
+	// Get current and previous window rates, the current rates are without
+	// accounting for the new hits.
+	current, previous, err := r.peekRates(ctx, drc)
 	if err != nil {
 		return -1, err
 	}
+	windowMultiplier := r.windowConfigurator.Multiplier(
+		ctx,
+		drc.WindowDuration,
+		uniqueKey,
+	)
+	if windowMultiplier < 0 {
+		windowMultiplier = drc.DefaultWindowMultiplier
+	}
 	// Only record the incoming hits when the current rate is within the allowed
 	// range, otherwise, do not record the hits and return the calculated rate.
-	// The idea is to continuously increase the rate limit. MaxAllowed sets a
-	// ceiling on it with the window duration.
+	// MaxAllowed sets a ceiling on the rate with the window duration.
+	//
 	// NOTE(marclop) We may want to add a follow-up static ceiling to avoid
 	// unbounded growth.
-	maxAllowed := math.Max(staticRate, previous*drc.WindowMultiplier)
+	maxAllowed := math.Max(staticRate, previous*windowMultiplier)
+	// Normalise the current rate assuming no more events will occur during the
+	// rest of the window. This will ensure that we record hits based on the
+	// currently observed hits and NOT based on extrapolated data.
+	current = current * drc.elapsed.Seconds() / drc.WindowDuration.Seconds()
 	if current <= maxAllowed {
-		if err := r.recordHits(ctx, drc, hits); err != nil {
+		// Deduce how many hits to record to reach to the max allowed number
+		remainingHits := int((maxAllowed - current) * drc.WindowDuration.Seconds())
+		if err := r.recordHits(ctx, drc, min(hits, remainingHits)); err != nil {
 			return -1, err
 		}
 	}
@@ -397,9 +447,10 @@ func (r *gubernatorRateLimiter) newDynamicRequest(
 	}
 }
 
-// peekRates retrieves the current (including incoming hits) and previous rates
-// from Gubernator. All Rates are normalized per second.
-func (r *gubernatorRateLimiter) peekRates(ctx context.Context, hits int64,
+// peekRates retrieves the current and previous rates from Gubernator. All
+// Rates are normalized per second.
+func (r *gubernatorRateLimiter) peekRates(
+	ctx context.Context,
 	drc dynamicRateContext,
 ) (float64, float64, error) {
 	// ----------------------- PEEK PHASE -----------------------
@@ -417,9 +468,8 @@ func (r *gubernatorRateLimiter) peekRates(ctx context.Context, hits int64,
 		return -1, -1, err
 	}
 	// Normalize the current rate based on the elapsed time (since the window
-	// hasn't fully elapsed). Then add the current hits to it.
+	// hasn't fully elapsed).
 	currentRate := rateFromResponse(peekResponses[0], drc.elapsed)
-	currentRate += float64(hits) / drc.elapsed.Seconds()
 	// Normalize the PREVIOUS rate based on the window duration.
 	previousRate := rateFromResponse(peekResponses[1], drc.WindowDuration)
 	return currentRate, previousRate, nil