From 412cd66c42946947521ce5a230112f567abb3830 Mon Sep 17 00:00:00 2001 From: Jacob Sussmilch Date: Wed, 1 Jul 2026 15:06:10 +1000 Subject: [PATCH 01/12] feat: KEEP-905 Add rate-limit-aware load balancing with adaptive capacity estimation --- README.md | 110 +++++++- configs/endpoints-example.json | 18 +- internal/config/config.go | 54 ++++ internal/config/config_test.go | 151 ++++++++++ internal/health/checker.go | 31 +- internal/health/health_test.go | 122 ++++++++ internal/health/ratelimit.go | 61 ++++ internal/health/ratelimit_test.go | 166 +++++++++++ internal/helpers/helpers.go | 8 + internal/metrics/definitions.go | 56 ++++ internal/server/capacity_learning_test.go | 318 +++++++++++++++++++++ internal/server/capacity_test.go | 256 +++++++++++++++++ internal/server/failover_soak_test.go | 276 ++++++++++++++++++ internal/server/rate_limit_scheduler.go | 8 +- internal/server/server.go | 327 ++++++++++++++++++++-- internal/server/server_test.go | 284 ++++++++++++++++++- internal/store/capacity_estimate.go | 135 +++++++++ internal/store/capacity_estimate_test.go | 208 ++++++++++++++ internal/store/testutils.go | 77 ++++- internal/store/valkey.go | 76 ++++- internal/store/valkey_test.go | 73 +++++ services/health-checker/main.go | 96 ++++++- services/health-checker/main_test.go | 146 ++++++++++ 23 files changed, 2991 insertions(+), 66 deletions(-) create mode 100644 internal/health/ratelimit.go create mode 100644 internal/health/ratelimit_test.go create mode 100644 internal/server/capacity_learning_test.go create mode 100644 internal/server/capacity_test.go create mode 100644 internal/server/failover_soak_test.go create mode 100644 internal/store/capacity_estimate.go create mode 100644 internal/store/capacity_estimate_test.go diff --git a/README.md b/README.md index 07bf7de..0871f27 100644 --- a/README.md +++ b/README.md @@ -134,6 +134,8 @@ The load balancer implements intelligent retry logic with configurable timeouts: | Flag | Default | Description | | --- | --- | --- | +| `--capacity-learning-enabled` | `true` | Enable adaptive per-endpoint capacity estimation, learned from observed rate-limit hits, for endpoints with no static `capacity` configured (see [Adaptive Capacity Estimation](#adaptive-capacity-estimation)) | +| `--capacity-throttling-enabled` | `true` | Enable proactive per-endpoint capacity throttling for endpoints with a configured `capacity` (see [Proactive Capacity Throttling](#proactive-capacity-throttling)) | | `--config-file` | `configs/endpoints.json` | Path to endpoints configuration file | | `--cors-headers` | `Accept, Authorization, Content-Type, Origin, X-Requested-With` | Allowed headers for CORS requests | | `--cors-methods` | `GET, POST, OPTIONS` | Allowed HTTP methods for CORS requests | @@ -174,6 +176,8 @@ The load balancer implements intelligent retry logic with configurable timeouts: | --- | --- | --- | | `ALCHEMY_API_KEY` | - | Example API key for Alchemy RPC endpoints. **Only needed for the example config.** The name must match the variable referenced in your `configs/endpoints.json`, if you need any. | | `INFURA_API_KEY` | - | Example API key for Infura RPC endpoints. **Only needed for the example config.** The name must match the variable referenced in your `configs/endpoints.json`, if you need any. | +| `CAPACITY_LEARNING_ENABLED` | `true` | Enable adaptive per-endpoint capacity estimation, learned from observed rate-limit hits, for endpoints with no static `capacity` configured (see [Adaptive Capacity Estimation](#adaptive-capacity-estimation)) | +| `CAPACITY_THROTTLING_ENABLED` | `true` | Enable proactive per-endpoint capacity throttling for endpoints with a configured `capacity` (see [Proactive Capacity Throttling](#proactive-capacity-throttling)) | | `CONFIG_FILE` | `configs/endpoints.json` | Path to the endpoints configuration file | | `CORS_HEADERS` | `Accept, Authorization, Content-Type, Origin, X-Requested-With` | Allowed headers for CORS requests | | `CORS_METHODS` | `GET, POST, OPTIONS` | Allowed HTTP methods for CORS requests | @@ -340,8 +344,8 @@ Prevents routing traffic to pods before initial health checks complete: ### How Rate Limit Recovery Works -1. **Detection**: When a request returns a rate limit error (HTTP 429), the endpoint is automatically marked as rate-limited. -2. **Retries with "backoff"**: The system tries to reach the endpoint only after waiting for a specific amount of time, defined as a backoff, which is configurable by the user. This wait period increases each time, relative to another user-defined parameter (the backoff multiplier). +1. **Detection**: When a request returns a rate limit error, the endpoint is automatically marked as rate-limited. This covers HTTP 429, Infura's HTTP 402 daily-credit-cap convention, and rate-limit errors embedded in a JSON-RPC response body even when the HTTP status is 200 (some providers, e.g. Alchemy on batch requests, only signal rate limiting this way). When a provider sends a `Retry-After` header, it's used to size the first recovery wait precisely instead of guessing. +2. **Retries with "backoff"**: The system tries to reach the endpoint only after waiting for a specific amount of time, defined as a backoff, which is configurable by the user. This wait period increases each time, relative to another user-defined parameter (the backoff multiplier). For Infura's daily credit cap specifically, the first wait is seeded from the endpoint's own `max_backoff` rather than `initial_backoff`, since a daily quota can't be recovered by probing sooner. 3. **Automatic recovery**: The system will reintroduce the endpoint back into the load balancing pool after a certain amount of successful consecutive requests. Users can specify how many consecutive requests are required for endpoints to be marked again as healthy. 4. **Per-endpoint configuration**: Each endpoint can have its own rate limit recovery strategy tailored to the provider's limits. You can also simply rely on the system's defaults, which have been carefully set. @@ -407,6 +411,108 @@ Rate limit recovery is configured per endpoint in your `endpoints.json` file: } ``` +## Proactive Capacity Throttling + +Rate limit recovery (above) is reactive: it only kicks in once a provider has already rejected a request. None of the major RPC providers (Infura, Alchemy, dRPC) expose a live "remaining quota" header, so there's no way to know a provider's real-time budget from a response alone. Proactive capacity throttling closes that gap by letting you declare each endpoint's known ceiling and having Ætherlay self-throttle below it, so it avoids triggering the provider's limit in the first place. + +### How It Works + +1. **Configure a ceiling**: Set `capacity` on an endpoint with a request count and a window width, matching however your provider actually windows its limit. +2. **Self-imposed counting**: Ætherlay counts every dispatch attempt to that endpoint (success or failure - a rejected attempt still spent real quota) against the ceiling, in a shared counter so it works correctly across multiple load-balancer replicas. +3. **Gating**: An endpoint at its ceiling for the current window is skipped in favor of another available endpoint - the same way an already-rate-limited endpoint is skipped today. If every endpoint in a role is at its ceiling, the request falls through to the next role tier (or a 503, if none are available), exactly like today's "all endpoints unhealthy" case. +4. **Weighting**: When every candidate endpoint in a role has `capacity` configured, endpoint selection is weighted by utilization relative to each endpoint's own ceiling instead of raw request count, so a higher-capacity endpoint isn't penalized for carrying more traffic than a lower-capacity one. If any endpoint in the comparison lacks `capacity`, selection falls back to the original raw-count behavior. + +This is a self-imposed budget, separate from rate limit recovery - an endpoint can be at its configured capacity ceiling without being marked rate-limited by a provider, and vice versa. + +### Configuration + +There is no default: a numeric ceiling is only meaningful relative to your specific paid plan with that specific provider, so it must be explicitly configured. Leaving `capacity` unset (the default) disables proactive throttling for that endpoint - existing configs need no changes. + +```json +{ + "mainnet": { + "provider-1": { + "provider": "alchemy", + "role": "primary", + "type": "archive", + "http_url": "https://api.example.com", + "capacity": { + "max_requests": 190, + "window_seconds": 10 + } + } + } +} +``` + +- **`max_requests`** (`int`): Requests allowed per window. +- **`window_seconds`** (`int`): Width of the window, in seconds. Match this to how your provider actually windows its limit rather than converting everything to a per-second rate - see the worked examples below. + +### Worked Per-Provider Examples + +These are starting points, not exact conversions - size conservatively for your own call mix: + +- **Infura**: no live signal at all; use the RPS ceiling published for your specific plan tier. Example: `{"max_requests": 8, "window_seconds": 1}` for a conservative sub-10-rps free-tier budget. +- **Alchemy**: published as Compute Units over a real **10-second rolling window** (Free tier: 500 CU/s, burst up to 5000 CU/10s). Per-method CU cost isn't modeled here, so convert using your own typical call mix: `max_requests ≈ floor(window_CU_budget / avg_CU_per_call)`. Example for a mostly-simple-read workload: `{"max_requests": 190, "window_seconds": 10}` (5000 CU / ~26 CU per simple call). Lower this if your workload is CU-heavy. +- **dRPC**: free tier is roughly 120,000 CU/minute per IP (~100 eth_call/s), with a 10 CU floor per call. Example: `{"max_requests": 100, "window_seconds": 1}`, or the minute-equivalent `{"max_requests": 6000, "window_seconds": 60}`. + +### Important Limitations + +- **The counter is shared across load-balancer replicas** (via Valkey), but the check-then-dispatch isn't a single atomic reservation - under concurrent load from multiple replicas the ceiling can be slightly overshot for a given window. This is intentional: configure `max_requests` with headroom below the provider's real ceiling rather than relying on exact enforcement. +- **WebSocket capacity gates new connection attempts only, not messages on an already-open connection.** Ætherlay relays WebSocket frames without inspecting their content, so there is no way to proactively throttle traffic on a connection that's already established. + +### Disabling + +Proactive capacity throttling is on by default whenever an endpoint has `capacity` configured. To disable it globally (e.g. during an incident, without editing every endpoint's config), set `CAPACITY_THROTTLING_ENABLED=false` (or `--capacity-throttling-enabled=false`). + +## Adaptive Capacity Estimation + +Declaring a static `capacity` (above) requires knowing your provider's real ceiling ahead of time - a number you'd have to look up, guess, or re-derive whenever your plan changes. Adaptive capacity estimation removes that requirement entirely: **no configuration is needed**. Ætherlay checks for rate-limit signals first, then approximates each endpoint's safe throughput ceiling from what it actually observes, adjusting it up or down as it continues routing traffic - the same shape of control loop TCP congestion control uses (AIMD: additive increase, multiplicative decrease). + +This is the recommended default going forward. A static `capacity` still works exactly as documented above and takes priority when set; adaptive learning only ever engages for endpoints with no static `capacity` configured. + +### How It Works + +1. **No ceiling until there's evidence.** A fresh endpoint is never proactively throttled - it behaves exactly like an unconfigured endpoint today. It's still fully covered by Ætherlay's existing *reactive* rate-limit detection (429, Infura's 402 daily cap, or a JSON-RPC error embedded in a 200 response) the instant it's actually rejected. +2. **On a rate-limit hit, halve the ceiling.** The first time an endpoint is rate limited, Ætherlay looks at how many requests it had actually dispatched to that endpoint in the current window and halves that number - a conservative first estimate grounded in real observed behavior, not a guess. Every subsequent hit halves the ceiling again, always starting from whichever is more conservative: the current estimate, or what was actually observed that window. +3. **A cooldown prevents overreacting to one episode.** Several near-simultaneous rejections (e.g. multiple in-flight retries against the same endpoint) collapse into a single decrease - the cooldown is the endpoint's own learning window, so the ceiling can't be revised "worse" faster than once per window. +4. **While clean, the ceiling grows back.** For every interval of sustained clean traffic since the last decrease, the ceiling grows by a small additive step (computed once per decrease, not a flat constant, so it scales with the endpoint's own magnitude). Growth is computed on the fly from elapsed time - there's no background process running per endpoint. +5. **Gating and weighting just work.** The learned ceiling feeds the exact same proactive-gating and utilization-weighted-selection logic as a static `capacity` - an endpoint with a learned estimate is skipped at its ceiling and weighted against other endpoints (static or learned) exactly like above. + +If an estimate overshoots the provider's real limit, the very next rejection halves it back down - the AIMD process is its own safety net, converging around the real ceiling the same way a TCP congestion window sawtooths around available bandwidth. There's no separate decay or reset mechanism, and no artificial upper cap beyond overflow protection. + +### Configuration + +Nothing is required by default - adaptive learning is on out of the box (`CAPACITY_LEARNING_ENABLED=true`) for every endpoint with no static `capacity`. The AIMD tuning knobs can be overridden per endpoint if needed, following the same optional-override pattern as `rate_limit_recovery`: + +```json +{ + "mainnet": { + "provider-1": { + "provider": "alchemy", + "role": "primary", + "type": "archive", + "http_url": "https://api.example.com", + "capacity_learning": { + "decrease_factor": 0.5, + "increase_interval": 60, + "min_estimate": 1, + "window_seconds": 60 + } + } + } +} +``` + +- **`decrease_factor`** (`float`): Multiplier applied to the ceiling on a confirmed rate-limit hit (default `0.5`, halving). +- **`increase_interval`** (`int`): Seconds of sustained clean time per additive-increase step (default `60`). +- **`min_estimate`** (`int`): Floor the learned ceiling can never decrease below (default `1`). +- **`window_seconds`** (`int`): Learning window width in seconds, used to size the usage counter and the decrease cooldown (default `60`) - only takes effect for an endpoint's first-ever estimate; once seeded, the window is frozen for that endpoint even if this value later changes. + +### Disabling + +Set `CAPACITY_LEARNING_ENABLED=false` (or `--capacity-learning-enabled=false`) to turn off adaptive learning globally and return to purely reactive rate-limit recovery for any endpoint without a static `capacity`. This is independent of `CAPACITY_THROTTLING_ENABLED`, which is the overall kill switch for *all* capacity-based gating and weighting, static or learned. + ## Prometheus Metrics This service uses a **pull-based** model for metrics collection, which is standard for Prometheus. This means the application exposes a `/metrics` endpoint, and a separate Prometheus server is responsible for periodically "scraping" (or pulling) data from it. diff --git a/configs/endpoints-example.json b/configs/endpoints-example.json index acad834..1ab4831 100644 --- a/configs/endpoints-example.json +++ b/configs/endpoints-example.json @@ -19,7 +19,11 @@ "role": "primary", "type": "full", "http_url": "https://eth.drpc.org", - "ws_url": "wss://eth.drpc.org" + "ws_url": "wss://eth.drpc.org", + "capacity": { + "max_requests": 100, + "window_seconds": 1 + } }, "publicnode-1": { "provider": "publicnode", @@ -34,7 +38,13 @@ "provider": "drpc", "role": "public", "type": "full", - "http_url": "https://arbitrum.drpc.org" + "http_url": "https://arbitrum.drpc.org", + "capacity_learning": { + "decrease_factor": 0.5, + "increase_interval": 60, + "min_estimate": 1, + "window_seconds": 60 + } }, "publicnode-1": { "provider": "public_node", @@ -58,6 +68,10 @@ "max_retries": 10, "required_successes": 3, "reset_after": 3600 + }, + "capacity": { + "max_requests": 190, + "window_seconds": 10 } }, "infura-staging": { diff --git a/internal/config/config.go b/internal/config/config.go index 5fa4428..dd08ec2 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -15,11 +15,32 @@ type RateLimitRecovery struct { ResetAfter int `json:"reset_after"` // Time in seconds after which to reset backoff and start from scratch } +// CapacityLimit represents a self-imposed throughput ceiling for an endpoint, used to +// proactively throttle requests before the provider's own rate limiter would trigger. +// There is no default: a numeric ceiling is only meaningful relative to an operator's +// specific paid plan with that specific provider, so it must be explicitly configured. +type CapacityLimit struct { + MaxRequests int `json:"max_requests"` // Requests allowed per window + WindowSeconds int `json:"window_seconds"` // Width of the window, in seconds +} + +// CapacityLearning tunes the AIMD control loop used to adaptively estimate an +// endpoint's safe throughput ceiling from observed rate-limit hits, when no static +// CapacityLimit is configured. See DefaultCapacityLearning for the default values. +type CapacityLearning struct { + DecreaseFactor float64 `json:"decrease_factor"` // Multiplier applied to the ceiling on a confirmed rate-limit hit (e.g. 0.5 halves it) + IncreaseInterval int `json:"increase_interval"` // Seconds of sustained clean time per additive-increase step + MinEstimate int `json:"min_estimate"` // Floor the learned ceiling can never decrease below + WindowSeconds int `json:"window_seconds"` // Default learning window width, used when there's no static CapacityLimit to inherit one from +} + // Endpoint represents a single RPC endpoint configuration. // It contains all the necessary information to connect to and use an RPC provider. type Endpoint struct { Provider string `json:"provider"` // Name of the RPC provider (e.g., "alchemy", "infura") RateLimitRecovery *RateLimitRecovery `json:"rate_limit_recovery"` // Rate limit recovery configuration (optional) + Capacity *CapacityLimit `json:"capacity"` // Self-imposed throughput ceiling (optional; nil disables proactive throttling) + CapacityLearning *CapacityLearning `json:"capacity_learning"` // Adaptive capacity learning tuning override (optional; only used when Capacity is unset) Role string `json:"role"` // Role of the endpoint: "primary" or "fallback" SkipSyncCheck bool `json:"skip_sync_check"` // Skip eth_syncing check for this endpoint (default: false) Type string `json:"type"` // Type of node: "full" or "archive" @@ -149,3 +170,36 @@ func DefaultRateLimitRecovery() RateLimitRecovery { ResetAfter: 86400, // Reset backoff after 1 day } } + +// DefaultCapacityLearning returns the default adaptive capacity learning configuration. +func DefaultCapacityLearning() CapacityLearning { + return CapacityLearning{ + DecreaseFactor: 0.5, // Halve the estimate on a confirmed rate-limit hit + IncreaseInterval: 60, // Grow once per minute of sustained clean traffic + MinEstimate: 1, // Never learn a ceiling below 1 request/window + WindowSeconds: 60, // Default learning window when no static CapacityLimit exists + } +} + +// ResolveCapacityLearning merges an optional per-endpoint override onto the package +// defaults, replacing only the fields the operator explicitly set (non-zero) - mirrors +// the RateLimitRecovery merge pattern used in rate_limit_scheduler.go. +func ResolveCapacityLearning(override *CapacityLearning) CapacityLearning { + resolved := DefaultCapacityLearning() + if override == nil { + return resolved + } + if override.DecreaseFactor != 0 { + resolved.DecreaseFactor = override.DecreaseFactor + } + if override.IncreaseInterval != 0 { + resolved.IncreaseInterval = override.IncreaseInterval + } + if override.MinEstimate != 0 { + resolved.MinEstimate = override.MinEstimate + } + if override.WindowSeconds != 0 { + resolved.WindowSeconds = override.WindowSeconds + } + return resolved +} diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 5feec5a..0ab19f8 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -1,6 +1,7 @@ package config import ( + "encoding/json" "os" "testing" ) @@ -261,3 +262,153 @@ func TestEndpointWithRateLimitRecovery(t *testing.T) { t.Errorf("Expected ResetAfter to be 7200, got %d", endpoint.RateLimitRecovery.ResetAfter) } } + +func TestEndpointWithCapacity(t *testing.T) { + endpoint := Endpoint{ + Provider: "alchemy", + Capacity: &CapacityLimit{ + MaxRequests: 190, + WindowSeconds: 10, + }, + Role: "primary", + Type: "full", + HTTPURL: "http://test.com", + } + + if endpoint.Capacity == nil { + t.Fatal("Expected capacity configuration to be set") + } + + if endpoint.Capacity.MaxRequests != 190 { + t.Errorf("Expected MaxRequests to be 190, got %d", endpoint.Capacity.MaxRequests) + } + + if endpoint.Capacity.WindowSeconds != 10 { + t.Errorf("Expected WindowSeconds to be 10, got %d", endpoint.Capacity.WindowSeconds) + } +} + +func TestEndpointCapacityJSONRoundTrip(t *testing.T) { + original := Endpoint{ + Provider: "drpc", + Capacity: &CapacityLimit{ + MaxRequests: 100, + WindowSeconds: 1, + }, + Role: "fallback", + Type: "full", + HTTPURL: "http://test.com", + } + + data, err := json.Marshal(original) + if err != nil { + t.Fatalf("Failed to marshal endpoint: %v", err) + } + + var decoded Endpoint + if err := json.Unmarshal(data, &decoded); err != nil { + t.Fatalf("Failed to unmarshal endpoint: %v", err) + } + + if decoded.Capacity == nil { + t.Fatal("Expected capacity to survive round-trip") + } + if decoded.Capacity.MaxRequests != 100 { + t.Errorf("Expected MaxRequests to be 100, got %d", decoded.Capacity.MaxRequests) + } + if decoded.Capacity.WindowSeconds != 1 { + t.Errorf("Expected WindowSeconds to be 1, got %d", decoded.Capacity.WindowSeconds) + } + + // An endpoint with no capacity configured should round-trip to a nil pointer, + // not a zero-value struct - this is what makes the feature opt-in. + unconfigured := Endpoint{Provider: "infura", Role: "primary", Type: "full", HTTPURL: "http://test.com"} + data, err = json.Marshal(unconfigured) + if err != nil { + t.Fatalf("Failed to marshal endpoint: %v", err) + } + var decodedUnconfigured Endpoint + if err := json.Unmarshal(data, &decodedUnconfigured); err != nil { + t.Fatalf("Failed to unmarshal endpoint: %v", err) + } + if decodedUnconfigured.Capacity != nil { + t.Fatal("Expected capacity to remain nil when not configured") + } +} + +func TestDefaultCapacityLearning(t *testing.T) { + learning := DefaultCapacityLearning() + + if learning.DecreaseFactor != 0.5 { + t.Errorf("Expected DecreaseFactor to be 0.5, got %f", learning.DecreaseFactor) + } + if learning.IncreaseInterval != 60 { + t.Errorf("Expected IncreaseInterval to be 60, got %d", learning.IncreaseInterval) + } + if learning.MinEstimate != 1 { + t.Errorf("Expected MinEstimate to be 1, got %d", learning.MinEstimate) + } + if learning.WindowSeconds != 60 { + t.Errorf("Expected WindowSeconds to be 60, got %d", learning.WindowSeconds) + } +} + +func TestResolveCapacityLearning(t *testing.T) { + t.Run("nil override returns defaults", func(t *testing.T) { + resolved := ResolveCapacityLearning(nil) + if resolved != DefaultCapacityLearning() { + t.Errorf("Expected defaults, got %+v", resolved) + } + }) + + t.Run("only non-zero override fields replace defaults", func(t *testing.T) { + override := &CapacityLearning{ + DecreaseFactor: 0.25, + MinEstimate: 5, + // IncreaseInterval and WindowSeconds left zero - should fall back to defaults. + } + resolved := ResolveCapacityLearning(override) + + if resolved.DecreaseFactor != 0.25 { + t.Errorf("Expected DecreaseFactor to be 0.25, got %f", resolved.DecreaseFactor) + } + if resolved.MinEstimate != 5 { + t.Errorf("Expected MinEstimate to be 5, got %d", resolved.MinEstimate) + } + if resolved.IncreaseInterval != 60 { + t.Errorf("Expected IncreaseInterval to fall back to default 60, got %d", resolved.IncreaseInterval) + } + if resolved.WindowSeconds != 60 { + t.Errorf("Expected WindowSeconds to fall back to default 60, got %d", resolved.WindowSeconds) + } + }) +} + +func TestEndpointWithCapacityLearningOverride(t *testing.T) { + endpoint := Endpoint{ + Provider: "alchemy", + CapacityLearning: &CapacityLearning{ + DecreaseFactor: 0.75, + }, + Role: "primary", + Type: "full", + HTTPURL: "http://test.com", + } + + data, err := json.Marshal(endpoint) + if err != nil { + t.Fatalf("Failed to marshal endpoint: %v", err) + } + + var decoded Endpoint + if err := json.Unmarshal(data, &decoded); err != nil { + t.Fatalf("Failed to unmarshal endpoint: %v", err) + } + + if decoded.CapacityLearning == nil { + t.Fatal("Expected capacity_learning to survive round-trip") + } + if decoded.CapacityLearning.DecreaseFactor != 0.75 { + t.Errorf("Expected DecreaseFactor to be 0.75, got %f", decoded.CapacityLearning.DecreaseFactor) + } +} diff --git a/internal/health/checker.go b/internal/health/checker.go index d5d186c..517bd7d 100644 --- a/internal/health/checker.go +++ b/internal/health/checker.go @@ -91,7 +91,7 @@ type Checker struct { ephemeralChecksThreshold int // Rate limit handler function provided by server - HandleRateLimitFunc func(chain, endpointID, protocol string) + HandleRateLimitFunc func(chain, endpointID, protocol string, signal RateLimitSignal) // For testability: allow patching health check methods CheckHTTPHealthFunc func(ctx context.Context, chain, endpointID string, endpoint config.Endpoint) bool @@ -396,7 +396,7 @@ func (c *Checker) checkEndpoint(ctx context.Context, chain, endpointID string, e } // makeRPCCall makes a single JSON-RPC call and returns the result -func (c *Checker) makeRPCCall(ctx context.Context, url, method, chain, endpointID string) (any, error) { +func (c *Checker) makeRPCCall(ctx context.Context, url, method, chain, endpointID, provider string) (any, error) { payload := []byte(`{"jsonrpc":"2.0","method":"` + method + `","params":[],"id":1}`) req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(payload)) if err != nil { @@ -413,17 +413,18 @@ func (c *Checker) makeRPCCall(ctx context.Context, url, method, chain, endpointI // Check for "bad" HTTP status codes if resp.StatusCode < 200 || resp.StatusCode >= 300 { - // Handle 429 (Too Many Requests) specially - if resp.StatusCode == 429 && c.HandleRateLimitFunc != nil { + // Handle a rate-limit signal (429, or Infura's 402 daily cap) specially + if sig := DetectRateLimit(provider, resp.StatusCode, resp.Header, nil); sig.IsRateLimited && c.HandleRateLimitFunc != nil { log.Debug(). Str("endpoint", helpers.RedactAPIKey(url)). Str("chain", chain). Str("endpoint_id", endpointID). Int("status_code", resp.StatusCode). Str("method", method). - Msg("RPC call detected 429, handing over to rate limit handler") + Bool("daily_quota", sig.IsDailyQuota). + Msg("RPC call detected a rate-limit signal, handing over to rate limit handler") - c.HandleRateLimitFunc(chain, endpointID, "http") + c.HandleRateLimitFunc(chain, endpointID, "http", sig) } // Read and log up to 512 bytes of the failed response's body @@ -469,6 +470,20 @@ func (c *Checker) makeRPCCall(ctx context.Context, url, method, chain, endpointI return nil, err } + // A rate-limited endpoint can return HTTP 200 with the error embedded in the JSON-RPC + // body (e.g. Alchemy's -32005). Without this, a rate-limited endpoint would just be + // marked plain-unhealthy and kept polling at the normal interval - working against + // the provider's own limit instead of backing off. + if sig := DetectRateLimit(provider, http.StatusOK, resp.Header, &response); sig.IsRateLimited && c.HandleRateLimitFunc != nil { + log.Debug(). + Str("endpoint", helpers.RedactAPIKey(url)). + Str("chain", chain). + Str("endpoint_id", endpointID). + Str("method", method). + Msg("RPC call received a 200 response with an embedded rate-limit error, handing over to rate limit handler") + c.HandleRateLimitFunc(chain, endpointID, "http", sig) + } + // Check for errors inside the response if err := checkRPCError(&response, method, "HTTP", chain, endpointID, url); err != nil { return nil, err @@ -696,14 +711,14 @@ func (c *Checker) checkHTTPHealth(ctx context.Context, chain, endpointID string, log.Info().Str("chain", chain).Str("endpoint_id", endpointID).Str("url", helpers.RedactAPIKey(endpoint.HTTPURL)).Msg("Running HTTP health check") // Always make the eth_blockNumber call - blockResult, blockErr := c.makeRPCCall(ctx, endpoint.HTTPURL, "eth_blockNumber", chain, endpointID) + blockResult, blockErr := c.makeRPCCall(ctx, endpoint.HTTPURL, "eth_blockNumber", chain, endpointID, endpoint.Provider) c.incrementHealthRequestCount(ctx, chain, endpointID) // Only make the eth_syncing call if sync status checking is enabled and not skipped for this endpoint var syncResult any var syncErr error if c.healthCheckSyncStatus && !endpoint.SkipSyncCheck { - syncResult, syncErr = c.makeRPCCall(ctx, endpoint.HTTPURL, "eth_syncing", chain, endpointID) + syncResult, syncErr = c.makeRPCCall(ctx, endpoint.HTTPURL, "eth_syncing", chain, endpointID, endpoint.Provider) c.incrementHealthRequestCount(ctx, chain, endpointID) } diff --git a/internal/health/health_test.go b/internal/health/health_test.go index a664214..359b1e5 100644 --- a/internal/health/health_test.go +++ b/internal/health/health_test.go @@ -164,6 +164,128 @@ func TestCheckHealthWithTimeout(t *testing.T) { } } +// TestMakeRPCCallDetects429WithRetryAfter tests that makeRPCCall's 429 handling parses a +// Retry-After header and forwards it (and chain/endpointID/protocol) to HandleRateLimitFunc. +func TestMakeRPCCallDetects429WithRetryAfter(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Retry-After", "9") + w.WriteHeader(http.StatusTooManyRequests) + })) + defer server.Close() + + var gotChain, gotEndpointID, gotProtocol string + var gotSignal RateLimitSignal + checker := &Checker{ + valkeyClient: store.NewMockValkeyClient(), + HandleRateLimitFunc: func(chain, endpointID, protocol string, signal RateLimitSignal) { + gotChain, gotEndpointID, gotProtocol, gotSignal = chain, endpointID, protocol, signal + }, + } + + _, err := checker.makeRPCCall(context.Background(), server.URL, "eth_blockNumber", "ethereum", "ep1", "alchemy") + if err == nil { + t.Fatal("Expected error from 429 response") + } + + if gotChain != "ethereum" || gotEndpointID != "ep1" || gotProtocol != "http" { + t.Errorf("HandleRateLimitFunc called with unexpected args: chain=%s endpointID=%s protocol=%s", gotChain, gotEndpointID, gotProtocol) + } + if !gotSignal.IsRateLimited { + t.Error("Expected IsRateLimited to be true") + } + if gotSignal.RetryAfter != 9*time.Second { + t.Errorf("Expected RetryAfter to be 9s, got %v", gotSignal.RetryAfter) + } +} + +// TestMakeRPCCallDetects402DailyCapForInfura tests that a 402 from an Infura endpoint is +// classified as a daily-quota rate-limit signal. +func TestMakeRPCCallDetects402DailyCapForInfura(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusPaymentRequired) + })) + defer server.Close() + + var gotSignal RateLimitSignal + called := false + checker := &Checker{ + valkeyClient: store.NewMockValkeyClient(), + HandleRateLimitFunc: func(chain, endpointID, protocol string, signal RateLimitSignal) { + called = true + gotSignal = signal + }, + } + + _, err := checker.makeRPCCall(context.Background(), server.URL, "eth_blockNumber", "ethereum", "ep1", "infura") + if err == nil { + t.Fatal("Expected error from 402 response") + } + if !called { + t.Fatal("Expected HandleRateLimitFunc to be called") + } + if !gotSignal.IsDailyQuota { + t.Error("Expected IsDailyQuota to be true for an Infura 402") + } +} + +// TestMakeRPCCallDetectsEmbeddedRateLimitErrorIn200Response tests the periodic +// health-check path's version of the same blind spot closed on the live proxy path: a +// 200 response whose JSON-RPC body carries a rate-limit error code. +func TestMakeRPCCallDetectsEmbeddedRateLimitErrorIn200Response(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"jsonrpc":"2.0","id":1,"error":{"code":-32005,"message":"Request limit exceeded"}}`)) + })) + defer server.Close() + + called := false + checker := &Checker{ + valkeyClient: store.NewMockValkeyClient(), + HandleRateLimitFunc: func(chain, endpointID, protocol string, signal RateLimitSignal) { + called = true + }, + } + + _, err := checker.makeRPCCall(context.Background(), server.URL, "eth_blockNumber", "ethereum", "ep1", "alchemy") + if err == nil { + t.Fatal("Expected an error for the embedded JSON-RPC error") + } + if !called { + t.Error("Expected HandleRateLimitFunc to be called for a 200 response with an embedded rate-limit error") + } +} + +// TestMakeRPCCallSuccessPathUnaffected is a regression guard: a clean success response +// must not trigger HandleRateLimitFunc and must still return the parsed result. +func TestMakeRPCCallSuccessPathUnaffected(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"jsonrpc":"2.0","id":1,"result":"0x123"}`)) + })) + defer server.Close() + + called := false + checker := &Checker{ + valkeyClient: store.NewMockValkeyClient(), + HandleRateLimitFunc: func(chain, endpointID, protocol string, signal RateLimitSignal) { + called = true + }, + } + + result, err := checker.makeRPCCall(context.Background(), server.URL, "eth_blockNumber", "ethereum", "ep1", "alchemy") + if err != nil { + t.Fatalf("Expected no error for a clean success response, got %v", err) + } + if called { + t.Error("Expected HandleRateLimitFunc NOT to be called for a clean success response") + } + if result != "0x123" { + t.Errorf("Expected result '0x123', got %v", result) + } +} + func TestStartEphemeralChecksDisabled(t *testing.T) { checker := &Checker{ ephemeralChecksEnabled: false, diff --git a/internal/health/ratelimit.go b/internal/health/ratelimit.go new file mode 100644 index 0000000..7e87bcc --- /dev/null +++ b/internal/health/ratelimit.go @@ -0,0 +1,61 @@ +package health + +import ( + "net/http" + "strconv" + "strings" + "time" +) + +// RateLimitSignal describes a detected rate-limit condition and any recovery timing +// hint the provider gave, so callers can seed a precise backoff instead of guessing. +type RateLimitSignal struct { + IsRateLimited bool + IsDailyQuota bool // true only for Infura's HTTP 402 daily credit cap - can't be sped up by probing + RetryAfter time.Duration // 0 if absent/unparseable +} + +// IsJSONRPCRateLimitCode reports whether a JSON-RPC error code indicates rate limiting. +// -32005 is the standard "Request limit exceeded" code used by Infura, Alchemy, and others. +func IsJSONRPCRateLimitCode(code int) bool { + return code == -32005 +} + +// DetectRateLimit inspects an HTTP response (status code + headers) and, when available, +// a parsed JSON-RPC response body, to determine whether a request was rate limited and +// what recovery timing hint (if any) the provider gave. +// +// This has exactly one provider-specific branch (Infura's HTTP 402 daily-credit-cap +// convention). Alchemy's unreliable Retry-After and dRPC's total absence of rate-limit +// headers are both handled by the same generic path - a plugin/adapter system isn't +// warranted for a single behavioral axis across the providers this was built against. +func DetectRateLimit(provider string, statusCode int, headers http.Header, rpcResp *RpcResponse) RateLimitSignal { + var sig RateLimitSignal + + switch { + case statusCode == http.StatusTooManyRequests: + sig.IsRateLimited = true + case statusCode == http.StatusPaymentRequired && strings.EqualFold(provider, "infura"): + // Infura-documented behavior: 402 means the daily credit cap is exhausted for + // the rest of the day, not a short burst limit - kept distinct so recovery + // doesn't re-probe on the same short interval as a 429 burst limit. + sig.IsRateLimited = true + sig.IsDailyQuota = true + case rpcResp != nil && rpcResp.Error != nil && IsJSONRPCRateLimitCode(rpcResp.Error.Code): + sig.IsRateLimited = true + } + + if sig.IsRateLimited && headers != nil { + if ra := headers.Get("Retry-After"); ra != "" { + if secs, err := strconv.Atoi(strings.TrimSpace(ra)); err == nil && secs > 0 { + sig.RetryAfter = time.Duration(secs) * time.Second + } else if t, err := http.ParseTime(ra); err == nil { + if d := time.Until(t); d > 0 { + sig.RetryAfter = d + } + } + } + } + + return sig +} diff --git a/internal/health/ratelimit_test.go b/internal/health/ratelimit_test.go new file mode 100644 index 0000000..fcb1561 --- /dev/null +++ b/internal/health/ratelimit_test.go @@ -0,0 +1,166 @@ +package health + +import ( + "net/http" + "testing" + "time" +) + +func TestIsJSONRPCRateLimitCode(t *testing.T) { + tests := []struct { + name string + code int + expected bool + }{ + {"rate limit code -32005", -32005, true}, + {"method not found code", -32601, false}, + {"generic error code", -32000, false}, + {"zero code", 0, false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := IsJSONRPCRateLimitCode(tt.code); got != tt.expected { + t.Errorf("IsJSONRPCRateLimitCode(%d) = %v, want %v", tt.code, got, tt.expected) + } + }) + } +} + +func TestDetectRateLimit(t *testing.T) { + rateLimitedRPCResp := &RpcResponse{Error: &struct { + Code int `json:"code"` + Message string `json:"message"` + }{Code: -32005, Message: "Request limit exceeded"}} + + methodNotFoundRPCResp := &RpcResponse{Error: &struct { + Code int `json:"code"` + Message string `json:"message"` + }{Code: -32601, Message: "Method not found"}} + + tests := []struct { + name string + provider string + statusCode int + headers http.Header + rpcResp *RpcResponse + wantRateLimited bool + wantDailyQuota bool + wantRetryAfterEq time.Duration + }{ + { + name: "429 for any provider is rate limited", + provider: "alchemy", + statusCode: 429, + wantRateLimited: true, + }, + { + name: "402 for infura is rate limited and daily quota", + provider: "infura", + statusCode: 402, + wantRateLimited: true, + wantDailyQuota: true, + }, + { + name: "402 for other provider is not rate limited", + provider: "alchemy", + statusCode: 402, + wantRateLimited: false, + }, + { + name: "402 for infura with different case still matches", + provider: "Infura", + statusCode: 402, + wantRateLimited: true, + wantDailyQuota: true, + }, + { + name: "200 with -32005 JSON-RPC error is rate limited", + provider: "alchemy", + statusCode: 200, + rpcResp: rateLimitedRPCResp, + wantRateLimited: true, + }, + { + name: "200 with unrelated JSON-RPC error is not rate limited", + provider: "alchemy", + statusCode: 200, + rpcResp: methodNotFoundRPCResp, + wantRateLimited: false, + }, + { + name: "200 with no error is not rate limited", + provider: "alchemy", + statusCode: 200, + wantRateLimited: false, + }, + { + name: "429 with Retry-After as integer seconds", + provider: "alchemy", + statusCode: 429, + headers: http.Header{"Retry-After": []string{"5"}}, + wantRateLimited: true, + wantRetryAfterEq: 5 * time.Second, + }, + { + name: "429 with missing Retry-After", + provider: "alchemy", + statusCode: 429, + headers: http.Header{}, + wantRateLimited: true, + wantRetryAfterEq: 0, + }, + { + name: "429 with garbage Retry-After", + provider: "alchemy", + statusCode: 429, + headers: http.Header{"Retry-After": []string{"not-a-number"}}, + wantRateLimited: true, + wantRetryAfterEq: 0, + }, + { + name: "429 with negative Retry-After is ignored", + provider: "alchemy", + statusCode: 429, + headers: http.Header{"Retry-After": []string{"-5"}}, + wantRateLimited: true, + wantRetryAfterEq: 0, + }, + { + name: "200 success does not surface a stray Retry-After header", + provider: "alchemy", + statusCode: 200, + headers: http.Header{"Retry-After": []string{"30"}}, + wantRateLimited: false, + wantRetryAfterEq: 0, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + sig := DetectRateLimit(tt.provider, tt.statusCode, tt.headers, tt.rpcResp) + if sig.IsRateLimited != tt.wantRateLimited { + t.Errorf("IsRateLimited = %v, want %v", sig.IsRateLimited, tt.wantRateLimited) + } + if sig.IsDailyQuota != tt.wantDailyQuota { + t.Errorf("IsDailyQuota = %v, want %v", sig.IsDailyQuota, tt.wantDailyQuota) + } + if sig.RetryAfter != tt.wantRetryAfterEq { + t.Errorf("RetryAfter = %v, want %v", sig.RetryAfter, tt.wantRetryAfterEq) + } + }) + } + + t.Run("429 with Retry-After as HTTP-date", func(t *testing.T) { + future := time.Now().Add(10 * time.Second).UTC().Truncate(time.Second) + headers := http.Header{"Retry-After": []string{future.Format(http.TimeFormat)}} + sig := DetectRateLimit("alchemy", 429, headers, nil) + if !sig.IsRateLimited { + t.Fatal("expected IsRateLimited to be true") + } + // Allow a small delta since time.Until(future) is computed after truncation to the second. + if sig.RetryAfter <= 0 || sig.RetryAfter > 11*time.Second { + t.Errorf("RetryAfter = %v, want roughly 10s", sig.RetryAfter) + } + }) +} diff --git a/internal/helpers/helpers.go b/internal/helpers/helpers.go index 2f9ad7d..5f9f14c 100644 --- a/internal/helpers/helpers.go +++ b/internal/helpers/helpers.go @@ -12,6 +12,8 @@ import ( // Config holds all CLI flags and their values type Config struct { + CapacityLearningEnabled bool + CapacityThrottlingEnabled bool ConfigFile string CorsHeaders string CorsOrigin string @@ -50,6 +52,8 @@ func ParseFlags() *Config { config := &Config{} // Define all flags with proper defaults + flag.BoolVar(&config.CapacityLearningEnabled, "capacity-learning-enabled", true, "Enable adaptive per-endpoint capacity estimation, learned from observed rate-limit hits, for endpoints with no static capacity configured") + flag.BoolVar(&config.CapacityThrottlingEnabled, "capacity-throttling-enabled", true, "Enable proactive per-endpoint capacity throttling based on configured request ceilings") flag.StringVar(&config.ConfigFile, "config-file", "configs/endpoints.json", "Configuration file path") flag.StringVar(&config.CorsHeaders, "cors-headers", "Accept, Authorization, Content-Type, Origin, X-Requested-With", "CORS allowed headers") flag.StringVar(&config.CorsMethods, "cors-methods", "GET, POST, OPTIONS", "CORS allowed methods") @@ -126,6 +130,8 @@ func (c *Config) GetBoolValue(flagName string, flagValue bool, envKey string, de // LoadConfiguration loads all configuration values with proper precedence func (c *Config) LoadConfiguration() *LoadedConfig { cfg := &LoadedConfig{ + CapacityLearningEnabled: c.GetBoolValue("capacity-learning-enabled", c.CapacityLearningEnabled, "CAPACITY_LEARNING_ENABLED", true), + CapacityThrottlingEnabled: c.GetBoolValue("capacity-throttling-enabled", c.CapacityThrottlingEnabled, "CAPACITY_THROTTLING_ENABLED", true), ConfigFile: c.GetStringValue("config-file", c.ConfigFile, "CONFIG_FILE", "configs/endpoints.json"), CorsHeaders: c.GetStringValue("cors-headers", c.CorsHeaders, "CORS_HEADERS", "Accept, Authorization, Content-Type, Origin, X-Requested-With"), CorsMethods: c.GetStringValue("cors-methods", c.CorsMethods, "CORS_METHODS", "GET, POST, OPTIONS"), @@ -171,6 +177,8 @@ func (c *Config) LoadConfiguration() *LoadedConfig { // LoadedConfig contains the final resolved configuration values type LoadedConfig struct { + CapacityLearningEnabled bool + CapacityThrottlingEnabled bool ConfigFile string CorsHeaders string CorsMethods string diff --git a/internal/metrics/definitions.go b/internal/metrics/definitions.go index e06b76a..cdaaf91 100644 --- a/internal/metrics/definitions.go +++ b/internal/metrics/definitions.go @@ -30,6 +30,18 @@ var ( var ( // EndpointProxyRequestsTotal counts the total number of requests successfully forwarded to each endpoint. EndpointProxyRequestsTotal *prometheus.CounterVec + // EndpointCapacityUtilization tracks an endpoint's self-imposed capacity usage as a + // fraction (count/max_requests) of its configured ceiling for the current window. + EndpointCapacityUtilization *prometheus.GaugeVec + // EndpointCapacitySkippedTotal counts how often an endpoint was excluded from + // selection for being at its configured capacity ceiling. + EndpointCapacitySkippedTotal *prometheus.CounterVec + // EndpointCapacityEstimatedCeiling tracks the currently learned safe throughput + // ceiling (requests/window) for endpoints with no static capacity configured. + EndpointCapacityEstimatedCeiling *prometheus.GaugeVec + // EndpointCapacityEstimateDecreasedTotal counts how often the learned capacity + // estimate was decreased in response to a confirmed rate-limit hit. + EndpointCapacityEstimateDecreasedTotal *prometheus.CounterVec ) // init initializes all metrics with error handling @@ -87,6 +99,50 @@ func initEndpointMetrics() { if EndpointProxyRequestsTotal == nil { log.Warn().Msg("Failed to register metric aetherlay_endpoint_proxy_requests_total") } + + EndpointCapacityUtilization = promauto.NewGaugeVec( + prometheus.GaugeOpts{ + Name: "aetherlay_endpoint_capacity_utilization", + Help: "Fraction of an endpoint's configured capacity ceiling used in the current window.", + }, + []string{"chain", "endpoint"}, + ) + if EndpointCapacityUtilization == nil { + log.Warn().Msg("Failed to register metric aetherlay_endpoint_capacity_utilization") + } + + EndpointCapacitySkippedTotal = promauto.NewCounterVec( + prometheus.CounterOpts{ + Name: "aetherlay_endpoint_capacity_skipped_total", + Help: "Total number of times an endpoint was excluded from selection for being at its configured capacity ceiling.", + }, + []string{"chain", "endpoint"}, + ) + if EndpointCapacitySkippedTotal == nil { + log.Warn().Msg("Failed to register metric aetherlay_endpoint_capacity_skipped_total") + } + + EndpointCapacityEstimatedCeiling = promauto.NewGaugeVec( + prometheus.GaugeOpts{ + Name: "aetherlay_endpoint_capacity_estimated_ceiling", + Help: "Currently learned safe throughput ceiling (requests/window) for an endpoint with no static capacity configured.", + }, + []string{"chain", "endpoint"}, + ) + if EndpointCapacityEstimatedCeiling == nil { + log.Warn().Msg("Failed to register metric aetherlay_endpoint_capacity_estimated_ceiling") + } + + EndpointCapacityEstimateDecreasedTotal = promauto.NewCounterVec( + prometheus.CounterOpts{ + Name: "aetherlay_endpoint_capacity_estimate_decreased_total", + Help: "Total number of times the learned capacity estimate was decreased in response to a confirmed rate-limit hit.", + }, + []string{"chain", "endpoint"}, + ) + if EndpointCapacityEstimateDecreasedTotal == nil { + log.Warn().Msg("Failed to register metric aetherlay_endpoint_capacity_estimate_decreased_total") + } } // initHealthMetrics initializes health check-related metrics diff --git a/internal/server/capacity_learning_test.go b/internal/server/capacity_learning_test.go new file mode 100644 index 0000000..66a0059 --- /dev/null +++ b/internal/server/capacity_learning_test.go @@ -0,0 +1,318 @@ +package server + +import ( + "context" + "testing" + "time" + + "aetherlay/internal/config" + "aetherlay/internal/helpers" + "aetherlay/internal/store" +) + +// learningTestConfig returns a LoadedConfig with capacity throttling and adaptive +// learning both explicitly enabled, since createTestConfig()'s struct literal leaves +// both at their Go zero value (false) - matching how existing, pre-adaptive-learning +// tests continue to see no behavior change unless a test opts in explicitly. +func learningTestConfig() *helpers.LoadedConfig { + cfg := createTestConfig() + cfg.CapacityThrottlingEnabled = true + cfg.CapacityLearningEnabled = true + return cfg +} + +func TestApplyLearnedCapacityDecreaseSeedsFromObservedCount(t *testing.T) { + cfg := &config.Config{ + Endpoints: map[string]config.ChainEndpoints{ + "ethereum": { + "ep1": config.Endpoint{Provider: "alchemy", Role: "primary", Type: "full", HTTPURL: "http://ep1"}, + }, + }, + } + valkeyClient := store.NewMockValkeyClient() + ctx := context.Background() + + // Simulate 10 requests already dispatched in the current (default 60s) learning window. + for i := 0; i < 10; i++ { + if _, err := valkeyClient.IncrementCapacityCount(ctx, "ethereum", "ep1", 60); err != nil { + t.Fatalf("IncrementCapacityCount failed: %v", err) + } + } + + server := NewServer(cfg, valkeyClient, learningTestConfig()) + server.applyLearnedCapacityDecrease("ethereum", "ep1", cfg.Endpoints["ethereum"]["ep1"]) + + estimate, err := valkeyClient.GetCapacityEstimate(ctx, "ethereum", "ep1") + if err != nil { + t.Fatalf("GetCapacityEstimate failed: %v", err) + } + if !estimate.HasEstimate { + t.Fatal("Expected an estimate to be seeded") + } + if estimate.MaxRequests != 5 { + t.Errorf("Expected MaxRequests to be 5 (10 observed * 0.5), got %d", estimate.MaxRequests) + } + if estimate.WindowSeconds != 60 { + t.Errorf("Expected WindowSeconds to be 60 (default learning window), got %d", estimate.WindowSeconds) + } +} + +func TestApplyLearnedCapacityDecreaseSkipsWithinCooldown(t *testing.T) { + cfg := &config.Config{ + Endpoints: map[string]config.ChainEndpoints{ + "ethereum": { + "ep1": config.Endpoint{Provider: "alchemy", Role: "primary", Type: "full", HTTPURL: "http://ep1"}, + }, + }, + } + valkeyClient := store.NewMockValkeyClient() + ctx := context.Background() + + for i := 0; i < 10; i++ { + valkeyClient.IncrementCapacityCount(ctx, "ethereum", "ep1", 60) + } + + server := NewServer(cfg, valkeyClient, learningTestConfig()) + ep := cfg.Endpoints["ethereum"]["ep1"] + + server.applyLearnedCapacityDecrease("ethereum", "ep1", ep) + first, _ := valkeyClient.GetCapacityEstimate(ctx, "ethereum", "ep1") + if first.MaxRequests != 5 { + t.Fatalf("Expected first decrease to seed 5, got %d", first.MaxRequests) + } + + // Second hit immediately after - within the cooldown (the learning window itself) - + // must not decrease again even though the observed count hasn't changed. + server.applyLearnedCapacityDecrease("ethereum", "ep1", ep) + second, _ := valkeyClient.GetCapacityEstimate(ctx, "ethereum", "ep1") + if second.MaxRequests != 5 { + t.Errorf("Expected MaxRequests to remain 5 (no double decrease within cooldown), got %d", second.MaxRequests) + } +} + +func TestApplyLearnedCapacityDecreaseAppliesAgainAfterCooldownElapsed(t *testing.T) { + cfg := &config.Config{ + Endpoints: map[string]config.ChainEndpoints{ + "ethereum": { + "ep1": config.Endpoint{Provider: "alchemy", Role: "primary", Type: "full", HTTPURL: "http://ep1"}, + }, + }, + } + valkeyClient := store.NewMockValkeyClient() + ctx := context.Background() + ep := cfg.Endpoints["ethereum"]["ep1"] + + for i := 0; i < 10; i++ { + valkeyClient.IncrementCapacityCount(ctx, "ethereum", "ep1", 60) + } + + server := NewServer(cfg, valkeyClient, learningTestConfig()) + server.applyLearnedCapacityDecrease("ethereum", "ep1", ep) + + seeded, _ := valkeyClient.GetCapacityEstimate(ctx, "ethereum", "ep1") + if seeded.MaxRequests != 5 { + t.Fatalf("Expected seed to be 5, got %d", seeded.MaxRequests) + } + + // Backdate LastDecreaseAt past both the cooldown and two IncreaseIntervals (60s each, + // default), so growth applies before the next decrease grounds itself in whichever of + // the grown estimate or fresh observed count is lower. + seeded.LastDecreaseAt = time.Now().Add(-125 * time.Second) + if err := valkeyClient.SetCapacityEstimate(ctx, "ethereum", "ep1", *seeded); err != nil { + t.Fatalf("SetCapacityEstimate failed: %v", err) + } + + // More traffic since the (backdated) last decrease: 6 additional requests, 16 total + // in the window bucket. + for i := 0; i < 6; i++ { + valkeyClient.IncrementCapacityCount(ctx, "ethereum", "ep1", 60) + } + + server.applyLearnedCapacityDecrease("ethereum", "ep1", ep) + + final, err := valkeyClient.GetCapacityEstimate(ctx, "ethereum", "ep1") + if err != nil { + t.Fatalf("GetCapacityEstimate failed: %v", err) + } + // effectiveNow = 5 (seed) + 2 steps * 1 (step, max(1, 5/10)) = 7 + // observedCount = 16 + // base = min(7, 16) = 7 -> newMax = floor(7*0.5) = 3 + if final.MaxRequests != 3 { + t.Errorf("Expected second decrease to be grounded in the grown estimate (7) not raw observed count (16): expected 3, got %d", final.MaxRequests) + } +} + +func TestApplyLearnedCapacityDecreaseSkipsWhenStaticCapacityConfigured(t *testing.T) { + cfg := &config.Config{ + Endpoints: map[string]config.ChainEndpoints{ + "ethereum": { + "ep1": config.Endpoint{ + Provider: "alchemy", Role: "primary", Type: "full", HTTPURL: "http://ep1", + Capacity: &config.CapacityLimit{MaxRequests: 100, WindowSeconds: 10}, + }, + }, + }, + } + valkeyClient := store.NewMockValkeyClient() + ctx := context.Background() + + server := NewServer(cfg, valkeyClient, learningTestConfig()) + server.applyLearnedCapacityDecrease("ethereum", "ep1", cfg.Endpoints["ethereum"]["ep1"]) + + estimate, _ := valkeyClient.GetCapacityEstimate(ctx, "ethereum", "ep1") + if estimate.HasEstimate { + t.Error("Expected no learned estimate when a static Capacity is configured") + } +} + +func TestApplyLearnedCapacityDecreaseSkipsWhenLearningDisabled(t *testing.T) { + cfg := &config.Config{ + Endpoints: map[string]config.ChainEndpoints{ + "ethereum": { + "ep1": config.Endpoint{Provider: "alchemy", Role: "primary", Type: "full", HTTPURL: "http://ep1"}, + }, + }, + } + valkeyClient := store.NewMockValkeyClient() + ctx := context.Background() + + appConfig := learningTestConfig() + appConfig.CapacityLearningEnabled = false + server := NewServer(cfg, valkeyClient, appConfig) + server.applyLearnedCapacityDecrease("ethereum", "ep1", cfg.Endpoints["ethereum"]["ep1"]) + + estimate, _ := valkeyClient.GetCapacityEstimate(ctx, "ethereum", "ep1") + if estimate.HasEstimate { + t.Error("Expected no learned estimate when CapacityLearningEnabled is false") + } +} + +func TestEffectiveCapacityCeilingNoBlackHoleBeforeEvidence(t *testing.T) { + cfg := &config.Config{} + valkeyClient := store.NewMockValkeyClient() + server := NewServer(cfg, valkeyClient, learningTestConfig()) + + ep := config.Endpoint{Provider: "alchemy"} + _, _, ok := server.effectiveCapacityCeiling("ethereum", "fresh-endpoint", ep) + if ok { + t.Error("Expected no ceiling for an endpoint with no static Capacity and no learned estimate yet") + } +} + +func TestGetAvailableEndpointsSkipsEndpointOnceLearnedEstimateExhausted(t *testing.T) { + cfg := &config.Config{ + Endpoints: map[string]config.ChainEndpoints{ + "ethereum": { + "learned-tight": config.Endpoint{Provider: "alchemy", Role: "primary", Type: "full", HTTPURL: "http://tight"}, + "unbounded": config.Endpoint{Provider: "alchemy", Role: "primary", Type: "full", HTTPURL: "http://unbounded"}, + }, + }, + } + valkeyClient := store.NewMockValkeyClient() + valkeyClient.PopulateStatuses(map[string]*store.EndpointStatus{ + "ethereum:learned-tight": {HasHTTP: true, HealthyHTTP: true}, + "ethereum:unbounded": {HasHTTP: true, HealthyHTTP: true}, + }) + ctx := context.Background() + + // Seed a learned estimate for "learned-tight" and drive its window count up to it. + if err := valkeyClient.SetCapacityEstimate(ctx, "ethereum", "learned-tight", store.CapacityEstimate{ + HasEstimate: true, MaxRequests: 2, IncreaseStep: 1, WindowSeconds: 60, LastDecreaseAt: time.Now(), + }); err != nil { + t.Fatalf("SetCapacityEstimate failed: %v", err) + } + valkeyClient.IncrementCapacityCount(ctx, "ethereum", "learned-tight", 60) + valkeyClient.IncrementCapacityCount(ctx, "ethereum", "learned-tight", 60) + + server := NewServer(cfg, valkeyClient, learningTestConfig()) + endpoints := server.getAvailableEndpoints("ethereum", false, false) + + if len(endpoints) != 1 { + t.Fatalf("Expected 1 available endpoint, got %d", len(endpoints)) + } + if endpoints[0].ID != "unbounded" { + t.Errorf("Expected 'unbounded' (no black hole before evidence), got %s", endpoints[0].ID) + } +} + +func TestSelectBestEndpointByRoleWeightsLearnedAgainstStatic(t *testing.T) { + valkeyClient := store.NewMockValkeyClient() + ctx := context.Background() + + for i := 0; i < 10; i++ { + valkeyClient.IncrementRequestCount(ctx, "ethereum", "static-large", "proxy_requests") + valkeyClient.IncrementRequestCount(ctx, "ethereum", "learned-small", "proxy_requests") + } + + if err := valkeyClient.SetCapacityEstimate(ctx, "ethereum", "learned-small", store.CapacityEstimate{ + HasEstimate: true, MaxRequests: 10, IncreaseStep: 1, WindowSeconds: 60, LastDecreaseAt: time.Now(), + }); err != nil { + t.Fatalf("SetCapacityEstimate failed: %v", err) + } + + endpoints := []EndpointWithID{ + {ID: "static-large", Endpoint: config.Endpoint{ + Role: "primary", Type: "full", HTTPURL: "http://large", + Capacity: &config.CapacityLimit{MaxRequests: 1000, WindowSeconds: 1}, + }}, + {ID: "learned-small", Endpoint: config.Endpoint{ + Role: "primary", Type: "full", HTTPURL: "http://small", + }}, + } + + server := NewServer(&config.Config{}, valkeyClient, learningTestConfig()) + best := server.selectBestEndpointByRole("ethereum", endpoints, "primary") + + if best == nil { + t.Fatal("Expected a best endpoint") + } + // Equal raw counts, but static-large's huge daily-equivalent budget (1000 req/s) + // gives it a far lower utilization ratio than learned-small's modest learned ceiling. + if best.ID != "static-large" { + t.Errorf("Expected static-large (lower utilization ratio), got %s", best.ID) + } +} + +func TestEffectiveCapacityCeilingReflectsLazyGrowth(t *testing.T) { + cfg := &config.Config{} + valkeyClient := store.NewMockValkeyClient() + ctx := context.Background() + + if err := valkeyClient.SetCapacityEstimate(ctx, "ethereum", "ep1", store.CapacityEstimate{ + HasEstimate: true, MaxRequests: 5, IncreaseStep: 1, WindowSeconds: 60, + LastDecreaseAt: time.Now().Add(-185 * time.Second), // 3 intervals of 60s elapsed + }); err != nil { + t.Fatalf("SetCapacityEstimate failed: %v", err) + } + + server := NewServer(cfg, valkeyClient, learningTestConfig()) + maxRequests, windowSeconds, ok := server.effectiveCapacityCeiling("ethereum", "ep1", config.Endpoint{Provider: "alchemy"}) + + if !ok { + t.Fatal("Expected a resolvable ceiling") + } + if maxRequests != 8 { // 5 + floor(185/60)=3 steps * 1 + t.Errorf("Expected grown ceiling of 8, got %d", maxRequests) + } + if windowSeconds != 60 { + t.Errorf("Expected WindowSeconds to be 60, got %d", windowSeconds) + } +} + +// Regression guard: existing static-capacity-only tests in capacity_test.go must see no +// behavior change now that CapacityLearningEnabled exists and defaults to true at +// runtime - createTestConfig() leaves it at Go's zero value (false) unless a test opts +// in, so effectiveCapacityCeiling never consults a learned estimate for them. +func TestEffectiveCapacityCeilingUnaffectedWhenLearningNotOptedIn(t *testing.T) { + cfg := &config.Config{} + valkeyClient := store.NewMockValkeyClient() + appConfig := createTestConfig() + appConfig.CapacityThrottlingEnabled = true + // appConfig.CapacityLearningEnabled left at zero-value false, as in every pre-existing test. + server := NewServer(cfg, valkeyClient, appConfig) + + _, _, ok := server.effectiveCapacityCeiling("ethereum", "ep1", config.Endpoint{Provider: "alchemy"}) + if ok { + t.Error("Expected no ceiling for an unconfigured endpoint when learning is not opted in") + } +} diff --git a/internal/server/capacity_test.go b/internal/server/capacity_test.go new file mode 100644 index 0000000..6e4689c --- /dev/null +++ b/internal/server/capacity_test.go @@ -0,0 +1,256 @@ +package server + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + + "aetherlay/internal/config" + "aetherlay/internal/store" +) + +// TestGetAvailableEndpointsSkipsAtCapacityEndpoint tests that an endpoint at its +// configured capacity ceiling for the current window is excluded from selection, +// independent of the (unrelated) provider-triggered RateLimited state. +func TestGetAvailableEndpointsSkipsAtCapacityEndpoint(t *testing.T) { + cfg := &config.Config{ + Endpoints: map[string]config.ChainEndpoints{ + "ethereum": { + "under-capacity": config.Endpoint{ + Provider: "alchemy", + Role: "primary", + Type: "full", + HTTPURL: "http://under-capacity.com", + Capacity: &config.CapacityLimit{MaxRequests: 5, WindowSeconds: 10}, + }, + "at-capacity": config.Endpoint{ + Provider: "alchemy", + Role: "primary", + Type: "full", + HTTPURL: "http://at-capacity.com", + Capacity: &config.CapacityLimit{MaxRequests: 2, WindowSeconds: 10}, + }, + }, + }, + } + valkeyClient := store.NewMockValkeyClient() + valkeyClient.PopulateStatuses(map[string]*store.EndpointStatus{ + "ethereum:under-capacity": {HasHTTP: true, HealthyHTTP: true}, + "ethereum:at-capacity": {HasHTTP: true, HealthyHTTP: true}, + }) + ctx := context.Background() + valkeyClient.IncrementCapacityCount(ctx, "ethereum", "at-capacity", 10) + valkeyClient.IncrementCapacityCount(ctx, "ethereum", "at-capacity", 10) + + appConfig := createTestConfig() + appConfig.CapacityThrottlingEnabled = true + server := NewServer(cfg, valkeyClient, appConfig) + + endpoints := server.getAvailableEndpoints("ethereum", false, false) + + if len(endpoints) != 1 { + t.Fatalf("Expected 1 available endpoint, got %d", len(endpoints)) + } + if endpoints[0].ID != "under-capacity" { + t.Errorf("Expected under-capacity endpoint, got %s", endpoints[0].ID) + } +} + +// TestGetAvailableEndpointsCapacityKillSwitch tests that CapacityThrottlingEnabled=false +// fully bypasses the capacity gate, even for an endpoint at its configured ceiling. +func TestGetAvailableEndpointsCapacityKillSwitch(t *testing.T) { + cfg := &config.Config{ + Endpoints: map[string]config.ChainEndpoints{ + "ethereum": { + "at-capacity": config.Endpoint{ + Provider: "alchemy", + Role: "primary", + Type: "full", + HTTPURL: "http://at-capacity.com", + Capacity: &config.CapacityLimit{MaxRequests: 2, WindowSeconds: 10}, + }, + }, + }, + } + valkeyClient := store.NewMockValkeyClient() + valkeyClient.PopulateStatuses(map[string]*store.EndpointStatus{ + "ethereum:at-capacity": {HasHTTP: true, HealthyHTTP: true}, + }) + ctx := context.Background() + valkeyClient.IncrementCapacityCount(ctx, "ethereum", "at-capacity", 10) + valkeyClient.IncrementCapacityCount(ctx, "ethereum", "at-capacity", 10) + + appConfig := createTestConfig() + appConfig.CapacityThrottlingEnabled = false + server := NewServer(cfg, valkeyClient, appConfig) + + endpoints := server.getAvailableEndpoints("ethereum", false, false) + + if len(endpoints) != 1 { + t.Fatalf("Expected the at-capacity endpoint to still be available with the kill switch off, got %d endpoints", len(endpoints)) + } +} + +// TestSelectBestEndpointByRoleFallsBackWhenNotAllHaveCapacity tests that selection falls +// back to raw 24h-count comparison (today's behavior) when not every candidate endpoint +// in the role has a configured Capacity - avoiding comparing endpoints on incompatible units. +func TestSelectBestEndpointByRoleFallsBackWhenNotAllHaveCapacity(t *testing.T) { + valkeyClient := store.NewMockValkeyClient() + ctx := context.Background() + + valkeyClient.IncrementRequestCount(ctx, "ethereum", "ep-a", "proxy_requests") + for i := 0; i < 5; i++ { + valkeyClient.IncrementRequestCount(ctx, "ethereum", "ep-b", "proxy_requests") + } + + endpoints := []EndpointWithID{ + {ID: "ep-a", Endpoint: config.Endpoint{Role: "primary", Type: "full", HTTPURL: "http://a"}}, + {ID: "ep-b", Endpoint: config.Endpoint{ + Role: "primary", Type: "full", HTTPURL: "http://b", + Capacity: &config.CapacityLimit{MaxRequests: 1000000, WindowSeconds: 1}, + }}, + } + + server := NewServer(&config.Config{}, valkeyClient, createTestConfig()) + best := server.selectBestEndpointByRole("ethereum", endpoints, "primary") + + if best == nil { + t.Fatal("Expected a best endpoint") + } + if best.ID != "ep-a" { + t.Errorf("Expected ep-a (lowest raw count, fallback behavior since not all candidates have Capacity), got %s", best.ID) + } +} + +// TestSelectBestEndpointByRoleWeightsByCapacityWhenAllConfigured tests that, once every +// candidate endpoint in the role has a configured Capacity, selection is weighted by +// utilization relative to each endpoint's own ceiling rather than raw request count. +func TestSelectBestEndpointByRoleWeightsByCapacityWhenAllConfigured(t *testing.T) { + valkeyClient := store.NewMockValkeyClient() + ctx := context.Background() + + for i := 0; i < 10; i++ { + valkeyClient.IncrementRequestCount(ctx, "ethereum", "low-capacity", "proxy_requests") + valkeyClient.IncrementRequestCount(ctx, "ethereum", "high-capacity", "proxy_requests") + } + + endpoints := []EndpointWithID{ + {ID: "low-capacity", Endpoint: config.Endpoint{ + Role: "primary", Type: "full", HTTPURL: "http://low", + Capacity: &config.CapacityLimit{MaxRequests: 10, WindowSeconds: 1}, + }}, + {ID: "high-capacity", Endpoint: config.Endpoint{ + Role: "primary", Type: "full", HTTPURL: "http://high", + Capacity: &config.CapacityLimit{MaxRequests: 1000, WindowSeconds: 1}, + }}, + } + + server := NewServer(&config.Config{}, valkeyClient, createTestConfig()) + best := server.selectBestEndpointByRole("ethereum", endpoints, "primary") + + if best == nil { + t.Fatal("Expected a best endpoint") + } + // Equal raw counts, but low-capacity's small ceiling means a much higher utilization + // ratio than high-capacity's - so high-capacity should be preferred. + if best.ID != "high-capacity" { + t.Errorf("Expected high-capacity (lower utilization ratio despite equal raw count), got %s", best.ID) + } +} + +// TestRecordCapacityUsageIncrementsRegardlessOfOutcome tests that recordCapacityUsage +// increments the counter on every call - the caller is responsible for calling it on +// every dispatch attempt whether it succeeds or fails, since a rejected attempt still +// spends a real unit of the provider's quota. +func TestRecordCapacityUsageIncrementsRegardlessOfOutcome(t *testing.T) { + valkeyClient := store.NewMockValkeyClient() + appConfig := createTestConfig() + appConfig.CapacityThrottlingEnabled = true + server := NewServer(&config.Config{}, valkeyClient, appConfig) + + ep := config.Endpoint{Capacity: &config.CapacityLimit{MaxRequests: 100, WindowSeconds: 10}} + server.recordCapacityUsage("ethereum", ep, "ep1") + server.recordCapacityUsage("ethereum", ep, "ep1") + + count, err := valkeyClient.GetCapacityCount(context.Background(), "ethereum", "ep1", 10) + if err != nil { + t.Fatalf("GetCapacityCount failed: %v", err) + } + if count != 2 { + t.Errorf("Expected capacity count 2, got %d", count) + } +} + +// TestRecordCapacityUsageNoop tests the two conditions under which recordCapacityUsage +// must be a no-op: no Capacity configured, and the CapacityThrottlingEnabled kill switch. +func TestRecordCapacityUsageNoop(t *testing.T) { + t.Run("no Capacity configured", func(t *testing.T) { + valkeyClient := store.NewMockValkeyClient() + appConfig := createTestConfig() + appConfig.CapacityThrottlingEnabled = true + server := NewServer(&config.Config{}, valkeyClient, appConfig) + + server.recordCapacityUsage("ethereum", config.Endpoint{}, "no-cap") + + count, _ := valkeyClient.GetCapacityCount(context.Background(), "ethereum", "no-cap", 10) + if count != 0 { + t.Errorf("Expected no increment without Capacity configured, got %d", count) + } + }) + + t.Run("kill switch disabled", func(t *testing.T) { + valkeyClient := store.NewMockValkeyClient() + appConfig := createTestConfig() + appConfig.CapacityThrottlingEnabled = false + server := NewServer(&config.Config{}, valkeyClient, appConfig) + + ep := config.Endpoint{Capacity: &config.CapacityLimit{MaxRequests: 100, WindowSeconds: 10}} + server.recordCapacityUsage("ethereum", ep, "disabled-ep") + + count, _ := valkeyClient.GetCapacityCount(context.Background(), "ethereum", "disabled-ep", 10) + if count != 0 { + t.Errorf("Expected no increment when CapacityThrottlingEnabled is false, got %d", count) + } + }) +} + +// TestHandleRequestHTTPRecordsCapacityUsagePerAttempt tests that a full HTTP request +// dispatched through the router increments the endpoint's capacity counter, confirming +// the dispatch-point wiring (not just the helper function in isolation). +func TestHandleRequestHTTPRecordsCapacityUsagePerAttempt(t *testing.T) { + cfg := &config.Config{ + Endpoints: map[string]config.ChainEndpoints{ + "ethereum": { + "ep1": config.Endpoint{ + Provider: "alchemy", Role: "primary", Type: "full", HTTPURL: "http://ep1", + Capacity: &config.CapacityLimit{MaxRequests: 100, WindowSeconds: 10}, + }, + }, + }, + } + valkeyClient := store.NewMockValkeyClient() + valkeyClient.PopulateStatuses(map[string]*store.EndpointStatus{ + "ethereum:ep1": {HasHTTP: true, HealthyHTTP: true}, + }) + appConfig := createTestConfig() + appConfig.CapacityThrottlingEnabled = true + server := NewServer(cfg, valkeyClient, appConfig) + server.forwardRequestWithBody = stubForwardRequestWithBody + + req := httptest.NewRequest("POST", "/ethereum", nil) + w := httptest.NewRecorder() + server.router.ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("Expected 200 from stubbed forward, got %d", w.Code) + } + + count, err := valkeyClient.GetCapacityCount(context.Background(), "ethereum", "ep1", 10) + if err != nil { + t.Fatalf("GetCapacityCount failed: %v", err) + } + if count != 1 { + t.Errorf("Expected capacity count 1 after one successful dispatch, got %d", count) + } +} diff --git a/internal/server/failover_soak_test.go b/internal/server/failover_soak_test.go new file mode 100644 index 0000000..219ee0d --- /dev/null +++ b/internal/server/failover_soak_test.go @@ -0,0 +1,276 @@ +package server + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + + "aetherlay/internal/config" + "aetherlay/internal/store" +) + +// jsonRPCSuccessBody is a minimal, valid JSON-RPC success response body used by the +// "healthy" backends in the failover soak tests below. +const jsonRPCSuccessBody = `{"jsonrpc":"2.0","id":1,"result":"0x1"}` + +// alwaysRateLimitedBackend returns a real httptest.Server that always responds with a +// 429, and a pointer to a counter tracking how many times it was actually dispatched to. +func alwaysRateLimitedBackend(t *testing.T) (*httptest.Server, *int64) { + t.Helper() + var calls int64 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + atomic.AddInt64(&calls, 1) + w.WriteHeader(http.StatusTooManyRequests) + })) + t.Cleanup(server.Close) + return server, &calls +} + +// alwaysErroringBackend returns a real httptest.Server that always responds with a +// generic 500 (the "or something else" non-rate-limit failure mode), and a call counter. +func alwaysErroringBackend(t *testing.T) (*httptest.Server, *int64) { + t.Helper() + var calls int64 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + atomic.AddInt64(&calls, 1) + w.WriteHeader(http.StatusInternalServerError) + })) + t.Cleanup(server.Close) + return server, &calls +} + +// alwaysHealthyBackend returns a real httptest.Server that always responds with a valid +// 200 JSON-RPC success body, and a call counter. +func alwaysHealthyBackend(t *testing.T) (*httptest.Server, *int64) { + t.Helper() + var calls int64 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + atomic.AddInt64(&calls, 1) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + w.Write([]byte(jsonRPCSuccessBody)) + })) + t.Cleanup(server.Close) + return server, &calls +} + +// postRequest fires one real POST request through the server's actual router (the full +// retry loop, endpoint selection, and dispatch path - no stubbed forwarder). +func postRequest(server *Server, chain string) *httptest.ResponseRecorder { + req := httptest.NewRequest("POST", "/"+chain, strings.NewReader(`{"jsonrpc":"2.0","method":"eth_blockNumber","params":[],"id":1}`)) + w := httptest.NewRecorder() + server.router.ServeHTTP(w, req) + return w +} + +// TestManyRequestsAllServedWhenOneEndpointAlwaysRateLimited fires many requests at a +// two-endpoint chain where one endpoint always returns 429. Every single client-facing +// request must still succeed - the load balancer must fail over to the healthy endpoint, +// both via the internal per-request retry (for whichever request first tries the bad +// endpoint) and by excluding it globally afterward. Neither endpoint has a static +// `capacity` configured, so this also exercises adaptive capacity learning bootstrapping +// for the endpoint that gets rate limited. +func TestManyRequestsAllServedWhenOneEndpointAlwaysRateLimited(t *testing.T) { + const rounds = 100 + + flakyServer, flakyCalls := alwaysRateLimitedBackend(t) + reliableServer, reliableCalls := alwaysHealthyBackend(t) + + cfg := &config.Config{ + Endpoints: map[string]config.ChainEndpoints{ + "ethereum": { + "flaky": config.Endpoint{Provider: "alchemy", Role: "primary", Type: "full", HTTPURL: flakyServer.URL}, + "reliable": config.Endpoint{Provider: "drpc", Role: "primary", Type: "full", HTTPURL: reliableServer.URL}, + }, + }, + } + valkeyClient := store.NewMockValkeyClient() + valkeyClient.PopulateStatuses(map[string]*store.EndpointStatus{ + "ethereum:flaky": {HasHTTP: true, HealthyHTTP: true}, + "ethereum:reliable": {HasHTTP: true, HealthyHTTP: true}, + }) + + appConfig := createTestConfig() + appConfig.CapacityThrottlingEnabled = true + appConfig.CapacityLearningEnabled = true + server := NewServer(cfg, valkeyClient, appConfig) + + for i := 0; i < rounds; i++ { + w := postRequest(server, "ethereum") + if w.Code != http.StatusOK { + t.Fatalf("request %d: expected 200, got %d (body: %s)", i, w.Code, w.Body.String()) + } + } + + if atomic.LoadInt64(flakyCalls) < 1 { + t.Error("Expected the rate-limited endpoint to have been dispatched to at least once, to prove real failover (not just never selected)") + } + if got := atomic.LoadInt64(reliableCalls); got != rounds { + t.Errorf("Expected the reliable endpoint to have served all %d successful responses, got %d", rounds, got) + } + + state, err := valkeyClient.GetRateLimitState(context.Background(), "ethereum", "flaky") + if err != nil { + t.Fatalf("GetRateLimitState failed: %v", err) + } + if !state.RateLimited { + t.Error("Expected the flaky endpoint to end up marked as rate limited") + } + + estimate, err := valkeyClient.GetCapacityEstimate(context.Background(), "ethereum", "flaky") + if err != nil { + t.Fatalf("GetCapacityEstimate failed: %v", err) + } + if !estimate.HasEstimate { + t.Error("Expected an adaptive capacity estimate to have been learned for the rate-limited endpoint") + } +} + +// TestManyRequestsAllServedWhenOneEndpointHasGenericErrors covers the "or something +// else" failure mode: a plain 500, which excludes an endpoint via the generic +// failure-threshold path rather than rate-limit detection. Every request must still succeed. +func TestManyRequestsAllServedWhenOneEndpointHasGenericErrors(t *testing.T) { + const rounds = 100 + + brokenServer, brokenCalls := alwaysErroringBackend(t) + reliableServer, reliableCalls := alwaysHealthyBackend(t) + + cfg := &config.Config{ + Endpoints: map[string]config.ChainEndpoints{ + "ethereum": { + "broken": config.Endpoint{Provider: "alchemy", Role: "primary", Type: "full", HTTPURL: brokenServer.URL}, + "reliable": config.Endpoint{Provider: "drpc", Role: "primary", Type: "full", HTTPURL: reliableServer.URL}, + }, + }, + } + valkeyClient := store.NewMockValkeyClient() + valkeyClient.PopulateStatuses(map[string]*store.EndpointStatus{ + "ethereum:broken": {HasHTTP: true, HealthyHTTP: true}, + "ethereum:reliable": {HasHTTP: true, HealthyHTTP: true}, + }) + + server := NewServer(cfg, valkeyClient, createTestConfig()) + + for i := 0; i < rounds; i++ { + w := postRequest(server, "ethereum") + if w.Code != http.StatusOK { + t.Fatalf("request %d: expected 200, got %d (body: %s)", i, w.Code, w.Body.String()) + } + } + + if atomic.LoadInt64(brokenCalls) < 1 { + t.Error("Expected the broken endpoint to have been dispatched to at least once, to prove real failover") + } + if got := atomic.LoadInt64(reliableCalls); got != rounds { + t.Errorf("Expected the reliable endpoint to have served all %d successful responses, got %d", rounds, got) + } + + status, err := valkeyClient.GetEndpointStatus(context.Background(), "ethereum", "broken") + if err != nil { + t.Fatalf("GetEndpointStatus failed: %v", err) + } + if status.HealthyHTTP { + t.Error("Expected the broken endpoint to end up marked unhealthy") + } +} + +// TestManyRequestsAllServedWithMixedSimultaneousFailureModes combines a rate-limited +// endpoint, a generically-erroring endpoint, and a healthy one in the same chain. Every +// request must still succeed despite two of the three providers misbehaving at once. +func TestManyRequestsAllServedWithMixedSimultaneousFailureModes(t *testing.T) { + const rounds = 150 + + rateLimitedServer, _ := alwaysRateLimitedBackend(t) + brokenServer, _ := alwaysErroringBackend(t) + reliableServer, reliableCalls := alwaysHealthyBackend(t) + + cfg := &config.Config{ + Endpoints: map[string]config.ChainEndpoints{ + "ethereum": { + "rate-limited": config.Endpoint{Provider: "alchemy", Role: "primary", Type: "full", HTTPURL: rateLimitedServer.URL}, + "broken": config.Endpoint{Provider: "infura", Role: "primary", Type: "full", HTTPURL: brokenServer.URL}, + "reliable": config.Endpoint{Provider: "drpc", Role: "primary", Type: "full", HTTPURL: reliableServer.URL}, + }, + }, + } + valkeyClient := store.NewMockValkeyClient() + valkeyClient.PopulateStatuses(map[string]*store.EndpointStatus{ + "ethereum:rate-limited": {HasHTTP: true, HealthyHTTP: true}, + "ethereum:broken": {HasHTTP: true, HealthyHTTP: true}, + "ethereum:reliable": {HasHTTP: true, HealthyHTTP: true}, + }) + + appConfig := createTestConfig() + appConfig.ProxyMaxRetries = 3 // must cover up to 2 failing endpoints before reaching the healthy one + server := NewServer(cfg, valkeyClient, appConfig) + + for i := 0; i < rounds; i++ { + w := postRequest(server, "ethereum") + if w.Code != http.StatusOK { + t.Fatalf("request %d: expected 200, got %d (body: %s)", i, w.Code, w.Body.String()) + } + } + + if got := atomic.LoadInt64(reliableCalls); got != rounds { + t.Errorf("Expected the reliable endpoint to have served all %d successful responses, got %d", rounds, got) + } + + rlState, _ := valkeyClient.GetRateLimitState(context.Background(), "ethereum", "rate-limited") + if !rlState.RateLimited { + t.Error("Expected the rate-limited endpoint to end up marked as rate limited") + } + brokenStatus, _ := valkeyClient.GetEndpointStatus(context.Background(), "ethereum", "broken") + if brokenStatus.HealthyHTTP { + t.Error("Expected the broken endpoint to end up marked unhealthy") + } +} + +// TestManyRequestsFailoverWhenStaticCapacityExhausted confirms the proactive +// static-capacity gate (no provider error involved at all) also fails over correctly: +// once a configured endpoint hits its self-imposed ceiling, later requests route to the +// other endpoint, and every request still succeeds. +func TestManyRequestsFailoverWhenStaticCapacityExhausted(t *testing.T) { + const rounds = 50 + const ceilingLimit = 5 + + cappedServer, cappedCalls := alwaysHealthyBackend(t) + reliableServer, reliableCalls := alwaysHealthyBackend(t) + + cfg := &config.Config{ + Endpoints: map[string]config.ChainEndpoints{ + "ethereum": { + "capped": config.Endpoint{ + Provider: "alchemy", Role: "primary", Type: "full", HTTPURL: cappedServer.URL, + Capacity: &config.CapacityLimit{MaxRequests: ceilingLimit, WindowSeconds: 60}, + }, + "reliable": config.Endpoint{Provider: "drpc", Role: "primary", Type: "full", HTTPURL: reliableServer.URL}, + }, + }, + } + valkeyClient := store.NewMockValkeyClient() + valkeyClient.PopulateStatuses(map[string]*store.EndpointStatus{ + "ethereum:capped": {HasHTTP: true, HealthyHTTP: true}, + "ethereum:reliable": {HasHTTP: true, HealthyHTTP: true}, + }) + + appConfig := createTestConfig() + appConfig.CapacityThrottlingEnabled = true + server := NewServer(cfg, valkeyClient, appConfig) + + for i := 0; i < rounds; i++ { + w := postRequest(server, "ethereum") + if w.Code != http.StatusOK { + t.Fatalf("request %d: expected 200, got %d (body: %s)", i, w.Code, w.Body.String()) + } + } + + if got := atomic.LoadInt64(cappedCalls); got != ceilingLimit { + t.Errorf("Expected the capped endpoint to be dispatched to exactly %d times (its configured ceiling), got %d", ceilingLimit, got) + } + if got := atomic.LoadInt64(reliableCalls); got != rounds-ceilingLimit { + t.Errorf("Expected the reliable endpoint to serve the remaining %d requests, got %d", rounds-ceilingLimit, got) + } +} diff --git a/internal/server/rate_limit_scheduler.go b/internal/server/rate_limit_scheduler.go index 2a51532..63e0d62 100644 --- a/internal/server/rate_limit_scheduler.go +++ b/internal/server/rate_limit_scheduler.go @@ -375,7 +375,7 @@ func (rls *RateLimitScheduler) checkEndpointHealth(ctx context.Context, endpoint Str("url", helpers.RedactAPIKey(endpoint.HTTPURL)). Int("code", rpcResp.Error.Code). Str("message", rpcResp.Error.Message) - if isJSONRPCRateLimitCode(rpcResp.Error.Code) { + if health.IsJSONRPCRateLimitCode(rpcResp.Error.Code) { evt.Msg("Recovery check received JSON-RPC rate-limit error") } else { evt.Msg("Recovery check received JSON-RPC error") @@ -394,12 +394,6 @@ func (rls *RateLimitScheduler) checkEndpointHealth(ctx context.Context, endpoint return true } -// isJSONRPCRateLimitCode reports whether a JSON-RPC error code indicates rate limiting. -// -32005 is the standard "Request limit exceeded" code used by Infura, Alchemy, and others. -func isJSONRPCRateLimitCode(code int) bool { - return code == -32005 -} - // shouldResetBackoff determines if the backoff cycle should be reset func (rls *RateLimitScheduler) shouldResetBackoff(state *store.RateLimitState, config config.RateLimitRecovery) bool { if state.FirstRateLimited.IsZero() { diff --git a/internal/server/server.go b/internal/server/server.go index 7353295..c024339 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -18,6 +18,7 @@ import ( "aetherlay/internal/cache" "aetherlay/internal/config" + "aetherlay/internal/health" "aetherlay/internal/helpers" "aetherlay/internal/metrics" "aetherlay/internal/store" @@ -31,6 +32,7 @@ import ( type RateLimitError struct { StatusCode int Message string + Signal health.RateLimitSignal } // Error returns the rate limit error message. @@ -472,6 +474,8 @@ func (s *Server) handleRequestHTTP(chain string) http.HandlerFunc { // Create per-try timeout context that respects the overall timeout tryCtx, tryCancel := context.WithTimeout(ctx, s.requestTimeoutPerTry) + s.recordCapacityUsage(chain, endpoint.Endpoint, endpoint.ID) + // Create a fresh request with a new body reader for each retry attempt err := s.forwardRequestWithBody(w, tryCtx, r.Method, endpoint.Endpoint.HTTPURL, bodyBytes, r.Header) tryCancel() // Always cancel the per-try context @@ -632,6 +636,8 @@ func (s *Server) handleRequestWS(chain string) http.HandlerFunc { tryCtx, tryCancel := context.WithTimeout(ctx, s.requestTimeoutPerTry) reqWithCtx := r.WithContext(tryCtx) + s.recordCapacityUsage(chain, endpoint.Endpoint, endpoint.ID) + err := s.proxyWebSocket(w, reqWithCtx, endpoint.Endpoint.WSURL) tryCancel() // Always cancel the per-try context @@ -692,10 +698,10 @@ func (s *Server) handleRequestWS(chain string) http.HandlerFunc { break } - // Check if this is a 429 rate limiting error during handshake - if _, ok := err.(*RateLimitError); ok { + // Check if this is a rate limiting error during handshake (429, or Infura's 402 daily cap) + if rlErr, ok := err.(*RateLimitError); ok { log.Debug().Str("chain", chain).Str("endpoint", endpoint.ID).Int("retry", retryCount).Msg("WebSocket handshake rate limited") - s.handleRateLimit(chain, endpoint.ID, "ws") + s.handleRateLimit(chain, endpoint.ID, "ws", rlErr.Signal) // Remove the rate-limited endpoint from the list var remainingEndpoints []EndpointWithID for _, ep := range allEndpoints { @@ -872,6 +878,29 @@ func (s *Server) getEndpointsByRole(chainEndpoints config.ChainEndpoints, role s continue } + // Proactively skip an endpoint that has hit its capacity ceiling (static + // or learned) for the current window, so selection routes to another + // endpoint before the provider's own rate limiter would trigger. + // Independent of RateLimited above: this is a self-imposed budget, not + // a provider signal. + if s.appConfig.CapacityThrottlingEnabled { + if maxRequests, windowSeconds, hasCeiling := s.effectiveCapacityCeiling(chain, endpointID, endpoint); hasCeiling { + count, err := s.valkeyClient.GetCapacityCount(context.Background(), chain, endpointID, windowSeconds) + if err == nil { + if metrics.EndpointCapacityUtilization != nil { + metrics.EndpointCapacityUtilization.WithLabelValues(chain, endpointID).Set(float64(count) / float64(maxRequests)) + } + if count >= maxRequests { + log.Debug().Str("chain", chain).Str("endpoint", endpointID).Str("role", role).Msg("Skipping endpoint at its capacity ceiling") + if metrics.EndpointCapacitySkippedTotal != nil { + metrics.EndpointCapacitySkippedTotal.WithLabelValues(chain, endpointID).Inc() + } + continue + } + } + } + } + if ws { if status.HasWS && status.HealthyWS { endpoints = append(endpoints, EndpointWithID{ID: endpointID, Endpoint: endpoint}) @@ -913,26 +942,62 @@ func (s *Server) selectBestEndpoint(chain string, endpoints []EndpointWithID) *E return nil } -// selectBestEndpointByRole selects the best endpoint of a specific role based on request counts -func (s *Server) selectBestEndpointByRole(chain string, endpoints []EndpointWithID, role string) *EndpointWithID { - var bestEndpoint *EndpointWithID - var minRequests int64 = -1 +// endpointCeiling caches a resolved capacity ceiling (static or learned) for one +// candidate, so selectBestEndpointByRole doesn't re-resolve it (and re-read Valkey for +// the learned case) once to decide allHaveCeiling and again to compute its score. +type endpointCeiling struct { + maxRequests int64 + windowSeconds int + ok bool +} +// selectBestEndpointByRole selects the best endpoint of a specific role based on request counts. +// When every candidate endpoint in this role has a resolvable capacity ceiling - static +// Capacity, or a learned estimate once adaptive learning has evidence for it - selection +// is instead weighted by utilization relative to each endpoint's own ceiling - projected +// out to a 24h-equivalent budget so it's comparable to the existing r24h counter - so a +// higher-capacity endpoint isn't penalized for having a higher raw request count than a +// lower-capacity one. If any candidate has no ceiling at all yet, this falls back to the +// original behavior (lowest raw 24h count wins) to avoid comparing endpoints on +// incompatible units. +func (s *Server) selectBestEndpointByRole(chain string, endpoints []EndpointWithID, role string) *EndpointWithID { + var candidateIndices []int + ceilings := make(map[int]endpointCeiling) + allHaveCeiling := true for i := range endpoints { - // Skip endpoints that don't match the requested role if endpoints[i].Endpoint.Role != role { continue } + candidateIndices = append(candidateIndices, i) + maxRequests, windowSeconds, ok := s.effectiveCapacityCeiling(chain, endpoints[i].ID, endpoints[i].Endpoint) + ceilings[i] = endpointCeiling{maxRequests: maxRequests, windowSeconds: windowSeconds, ok: ok} + if !ok { + allHaveCeiling = false + } + } + + var bestEndpoint *EndpointWithID + var minScore float64 = -1 + for _, i := range candidateIndices { r24h, _, _, err := s.valkeyClient.GetCombinedRequestCounts(context.Background(), chain, endpoints[i].ID) // Skip endpoints where we can't get request count data if err != nil { continue } - // Select endpoint with lowest 24h request count (or first one if minRequests is uninitialized) - if minRequests == -1 || r24h < minRequests { - minRequests = r24h + score := float64(r24h) + if allHaveCeiling { + c := ceilings[i] + dailyBudget := float64(c.maxRequests) * (86400.0 / float64(c.windowSeconds)) + if dailyBudget > 0 { + score = float64(r24h) / dailyBudget + } + } + + // Select endpoint with the lowest score (or first one if minScore is uninitialized) + if minScore == -1 || score < minScore { + minScore = score bestEndpoint = &endpoints[i] } } @@ -940,6 +1005,30 @@ func (s *Server) selectBestEndpointByRole(chain string, endpoints []EndpointWith return bestEndpoint } +// effectiveCapacityCeiling resolves the ceiling (requests/window) to gate and weight an +// endpoint against, whichever source is authoritative: +// 1. Static Capacity configured -> today's exact values, unchanged - adaptive learning +// never engages for this endpoint. +// 2. No static Capacity, adaptive learning enabled, and a learned estimate exists -> +// the estimate's ceiling, grown lazily via store.EffectiveMaxRequests. +// 3. Otherwise -> ok=false: no ceiling at all, never proactively skipped. Absence of +// evidence isn't evidence of a limit - the endpoint remains fully covered by the +// independent, reactive RateLimitState check the instant it's actually rate limited. +func (s *Server) effectiveCapacityCeiling(chain, endpointID string, ep config.Endpoint) (maxRequests int64, windowSeconds int, ok bool) { + if ep.Capacity != nil { + return int64(ep.Capacity.MaxRequests), ep.Capacity.WindowSeconds, true + } + if !s.appConfig.CapacityLearningEnabled { + return 0, 0, false + } + estimate, err := s.valkeyClient.GetCapacityEstimate(context.Background(), chain, endpointID) + if err != nil || !estimate.HasEstimate { + return 0, 0, false + } + params := config.ResolveCapacityLearning(ep.CapacityLearning) + return store.EffectiveMaxRequests(*estimate, params, time.Now()), estimate.WindowSeconds, true +} + // removeEndpointByID removes an endpoint from a slice by its ID func removeEndpointByID(endpoints []EndpointWithID, id string) []EndpointWithID { var remaining []EndpointWithID @@ -1151,11 +1240,11 @@ func (s *Server) defaultForwardRequestWithBodyFunc(w http.ResponseWriter, ctx co // For all other non-2xx responses (400 already handled above), mark endpoint as unhealthy if found { - if resp.StatusCode == 429 { - // For 429 (Too Many Requests), use the rate limit handler + sig := health.DetectRateLimit(s.providerForEndpoint(chain, endpointID), resp.StatusCode, resp.Header, nil) + if sig.IsRateLimited { s.markEndpointUnhealthyProtocol(chain, endpointID, "http") - s.handleRateLimit(chain, endpointID, "http") - log.Debug().Str("url", helpers.RedactAPIKey(targetURL)).Int("status_code", resp.StatusCode).Msg("Endpoint returned 429 (Too Many Requests), handling rate limit") + s.handleRateLimit(chain, endpointID, "http", sig) + log.Debug().Str("url", helpers.RedactAPIKey(targetURL)).Int("status_code", resp.StatusCode).Bool("daily_quota", sig.IsDailyQuota).Msg("Endpoint returned a rate-limit signal, handling rate limit") } else { s.markEndpointUnhealthyProtocol(chain, endpointID, "http") log.Debug().Str("url", helpers.RedactAPIKey(targetURL)).Int("status_code", resp.StatusCode).Msg("Endpoint returned non-2xx status, marked unhealthy") @@ -1177,6 +1266,35 @@ func (s *Server) defaultForwardRequestWithBodyFunc(w http.ResponseWriter, ctx co } } + // For JSON responses within a bounded size, buffer and scan the body for an embedded + // rate-limit error before forwarding. Some providers (e.g. Alchemy on batch requests) + // return HTTP 200 with the rate-limit signal only inside the JSON-RPC body - invisible + // to a status-code-only check. Larger bodies are streamed through unmodified and + // uninspected, to bound memory use on large responses (e.g. eth_getLogs). + if isJSONContentType(resp.Header.Get("Content-Type")) { + buffered, readErr := io.ReadAll(io.LimitReader(resp.Body, maxRateLimitScanBodyBytes+1)) + if readErr == nil && int64(len(buffered)) <= maxRateLimitScanBodyBytes { + if chain, endpointID, found := s.findChainAndEndpointByURL(targetURL); found { + if sig := bodyCarriesRateLimitSignal(s.providerForEndpoint(chain, endpointID), buffered, resp.Header); sig.IsRateLimited { + // The HTTP transaction itself succeeded (2xx) - don't mark the endpoint + // unhealthy, just flag it as rate limited so selection avoids it, and + // retry the same buffered request body against a different endpoint. + s.handleRateLimit(chain, endpointID, "http", sig) + log.Debug().Str("url", helpers.RedactAPIKey(targetURL)).Bool("daily_quota", sig.IsDailyQuota).Msg("2xx response carried an embedded rate-limit signal, retrying with a different endpoint") + return fmt.Errorf("rate limited: embedded JSON-RPC rate-limit error in 2xx response") + } + } + w.WriteHeader(resp.StatusCode) + _, err = w.Write(buffered) + return err + } + // Exceeded the scan cap (or a read error occurred): stream the buffered prefix + // plus whatever remains, unmodified and uninspected. + w.WriteHeader(resp.StatusCode) + _, err = io.Copy(w, io.MultiReader(bytes.NewReader(buffered), resp.Body)) + return err + } + // Set response status w.WriteHeader(resp.StatusCode) @@ -1185,6 +1303,132 @@ func (s *Server) defaultForwardRequestWithBodyFunc(w http.ResponseWriter, ctx co return err } +// maxRateLimitScanBodyBytes bounds how much of a 2xx response body is buffered to scan +// for an embedded JSON-RPC rate-limit error. Single-request rate-limit bodies are tiny; +// this cap exists so large legitimate responses aren't fully buffered in memory. +const maxRateLimitScanBodyBytes = 10 * 1024 * 1024 + +// isJSONContentType reports whether a Content-Type header value looks like JSON. +func isJSONContentType(contentType string) bool { + return strings.Contains(contentType, "application/json") || strings.Contains(contentType, "text/json") +} + +// providerForEndpoint looks up the configured provider name for a chain/endpoint, used to +// select provider-specific rate-limit detection behavior (see health.DetectRateLimit). +func (s *Server) providerForEndpoint(chain, endpointID string) string { + chainEndpoints, ok := s.config.GetEndpointsForChain(chain) + if !ok { + return "" + } + return chainEndpoints[endpointID].Provider +} + +// capacityWindowSeconds resolves the window width to track usage against: the static +// CapacityLimit's window if configured, else the (possibly overridden) adaptive +// learning window - so an endpoint with no static Capacity still accumulates real +// GetCapacityCount evidence before it's ever rate limited, for applyLearnedCapacityDecrease +// to seed an estimate from on the first hit. +func capacityWindowSeconds(endpoint config.Endpoint) int { + if endpoint.Capacity != nil { + return endpoint.Capacity.WindowSeconds + } + return config.ResolveCapacityLearning(endpoint.CapacityLearning).WindowSeconds +} + +// recordCapacityUsage increments an endpoint's self-imposed capacity counter for the +// current window, on every dispatch attempt regardless of outcome - a failed/429'd +// attempt still spent a real unit of the provider's quota. No-op when capacity +// throttling is disabled entirely, or when the endpoint has no static Capacity and +// adaptive learning is also disabled. +func (s *Server) recordCapacityUsage(chain string, endpoint config.Endpoint, endpointID string) { + if !s.appConfig.CapacityThrottlingEnabled { + return + } + if endpoint.Capacity == nil && !s.appConfig.CapacityLearningEnabled { + return + } + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + if _, err := s.valkeyClient.IncrementCapacityCount(ctx, chain, endpointID, capacityWindowSeconds(endpoint)); err != nil { + log.Debug().Err(err).Str("chain", chain).Str("endpoint", endpointID).Msg("Failed to record capacity usage") + } +} + +// applyLearnedCapacityDecrease is called from handleRateLimit on every confirmed +// rate-limit signal. It only ever engages for endpoints with no static Capacity +// configured (static config, if present, is left completely untouched - see +// effectiveCapacityCeiling). A cooldown (the learning window itself) collapses several +// near-simultaneous hits from one episode into a single decrease. +func (s *Server) applyLearnedCapacityDecrease(chain, endpointID string, endpoint config.Endpoint) { + if !s.appConfig.CapacityThrottlingEnabled || !s.appConfig.CapacityLearningEnabled || endpoint.Capacity != nil { + return + } + + params := config.ResolveCapacityLearning(endpoint.CapacityLearning) + now := time.Now() + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + prior, err := s.valkeyClient.GetCapacityEstimate(ctx, chain, endpointID) + if err != nil { + log.Debug().Err(err).Str("chain", chain).Str("endpoint", endpointID).Msg("Failed to get capacity estimate") + return + } + + if !store.ShouldApplyCapacityDecrease(*prior, params.WindowSeconds, now) { + log.Debug().Str("chain", chain).Str("endpoint", endpointID).Msg("Skipping capacity estimate decrease, within cooldown of the last decrease") + return + } + + observedCount, err := s.valkeyClient.GetCapacityCount(ctx, chain, endpointID, params.WindowSeconds) + if err != nil { + log.Debug().Err(err).Str("chain", chain).Str("endpoint", endpointID).Msg("Failed to get capacity count for estimate decrease") + return + } + + effectiveNow := store.EffectiveMaxRequests(*prior, params, now) + newEstimate := store.ApplyCapacityDecrease(*prior, effectiveNow, observedCount, params.WindowSeconds, params, now) + + if err := s.valkeyClient.SetCapacityEstimate(ctx, chain, endpointID, newEstimate); err != nil { + log.Debug().Err(err).Str("chain", chain).Str("endpoint", endpointID).Msg("Failed to set capacity estimate") + return + } + + log.Info().Str("chain", chain).Str("endpoint", endpointID).Int64("new_estimate", newEstimate.MaxRequests).Int("window_seconds", newEstimate.WindowSeconds).Msg("Decreased learned capacity estimate after a rate-limit hit") + if metrics.EndpointCapacityEstimatedCeiling != nil { + metrics.EndpointCapacityEstimatedCeiling.WithLabelValues(chain, endpointID).Set(float64(newEstimate.MaxRequests)) + } + if metrics.EndpointCapacityEstimateDecreasedTotal != nil { + metrics.EndpointCapacityEstimateDecreasedTotal.WithLabelValues(chain, endpointID).Inc() + } +} + +// bodyCarriesRateLimitSignal reports whether a 2xx JSON-RPC response body (a single +// object or a batch array) contains an embedded rate-limit error - the case that makes +// Alchemy's HTTP-200-with-one-batch-item-429 blind spot visible, since HTTP status alone +// can't see it. Returns a zero-value signal if the body isn't a recognizable JSON-RPC +// response or carries no rate-limit error. +func bodyCarriesRateLimitSignal(provider string, body []byte, headers http.Header) health.RateLimitSignal { + var single health.RpcResponse + if err := json.Unmarshal(body, &single); err == nil && single.Error != nil { + return health.DetectRateLimit(provider, http.StatusOK, headers, &single) + } + + var batch []health.RpcResponse + if err := json.Unmarshal(body, &batch); err == nil { + for i := range batch { + if batch[i].Error == nil { + continue + } + if sig := health.DetectRateLimit(provider, http.StatusOK, headers, &batch[i]); sig.IsRateLimited { + return sig + } + } + } + + return health.RateLimitSignal{} +} + // proxyWebSocketCopy copies messages from src to dst, forwarding close frames // to the destination so both peers receive a proper WebSocket close handshake. // It returns the first error and a bool indicating whether the error came from @@ -1259,15 +1503,16 @@ func (s *Server) defaultProxyWebSocket(w http.ResponseWriter, r *http.Request, b // Check for non-2xx status codes during handshake if resp != nil && (resp.StatusCode < 200 || resp.StatusCode >= 300) { - if resp.StatusCode == 429 { - // For 429 (Too Many Requests), mark unhealthy and return RateLimitError as signal + if sig := health.DetectRateLimit(s.providerForEndpoint(chain, endpointID), resp.StatusCode, resp.Header, nil); sig.IsRateLimited { + // For a rate-limit signal (429, or Infura's 402 daily cap), mark unhealthy and return RateLimitError as signal if found { s.markEndpointUnhealthyProtocol(chain, endpointID, "ws") } - log.Debug().Str("url", helpers.RedactAPIKey(backendURL)).Int("status_code", resp.StatusCode).Msg("WebSocket handshake rate limited") + log.Debug().Str("url", helpers.RedactAPIKey(backendURL)).Int("status_code", resp.StatusCode).Bool("daily_quota", sig.IsDailyQuota).Msg("WebSocket handshake rate limited") return &RateLimitError{ StatusCode: resp.StatusCode, Message: fmt.Sprintf("WebSocket handshake was rate-limited: HTTP %d", resp.StatusCode), + Signal: sig, } } @@ -1411,12 +1656,14 @@ func (s *Server) defaultProxyWebSocket(w http.ResponseWriter, r *http.Request, b } // GetRateLimitHandler returns the rate limit handler function for the health checker -func (s *Server) GetRateLimitHandler() func(chain, endpointID, protocol string) { +func (s *Server) GetRateLimitHandler() func(chain, endpointID, protocol string, signal health.RateLimitSignal) { return s.handleRateLimit } -// handleRateLimit handles rate limiting for an endpoint -func (s *Server) handleRateLimit(chain, endpointID, protocol string) { +// handleRateLimit handles rate limiting for an endpoint. The signal carries whatever +// recovery timing hint the provider gave (a Retry-After header, or Infura's daily-quota +// distinction), so the backoff can be seeded precisely instead of guessed. +func (s *Server) handleRateLimit(chain, endpointID, protocol string, signal health.RateLimitSignal) { log.Debug().Str("chain", chain).Str("endpoint", endpointID).Str("protocol", protocol).Msg("Handling rate limit") // Set the endpoint as rate limited in Valkey @@ -1432,7 +1679,7 @@ func (s *Server) handleRateLimit(chain, endpointID, protocol string) { state.RecoveryAttempts = 0 state.LastRecoveryCheck = now state.ConsecutiveSuccess = 0 - state.CurrentBackoff = 0 // Will be set to initial backoff on first attempt + state.CurrentBackoff = s.initialBackoffForSignal(chain, endpointID, signal) // Set first rate limited time if this is the first time if state.FirstRateLimited.IsZero() { @@ -1444,8 +1691,42 @@ func (s *Server) handleRateLimit(chain, endpointID, protocol string) { return } - log.Info().Str("chain", chain).Str("endpoint", endpointID).Str("protocol", protocol).Msg("Endpoint marked as rate limited") + log.Info().Str("chain", chain).Str("endpoint", endpointID).Str("protocol", protocol).Int("current_backoff", state.CurrentBackoff).Msg("Endpoint marked as rate limited") + + // Check for rate limits first (this signal), then approximate the endpoint's safe + // throughput ceiling from it - only engages for endpoints with no static Capacity. + if chainEndpoints, ok := s.config.GetEndpointsForChain(chain); ok { + if ep, ok := chainEndpoints[endpointID]; ok { + s.applyLearnedCapacityDecrease(chain, endpointID, ep) + } + } // Start rate limit recovery monitoring s.rateLimitScheduler.StartMonitoring(chain, endpointID) } + +// initialBackoffForSignal computes the first recovery-check wait time from whatever the +// provider actually told us, instead of always guessing via the scheduler's InitialBackoff: +// - A parsed Retry-After is the most precise signal available and is used directly. +// - Infura's daily credit cap (402) can't be sped up by probing sooner, since it only +// resets once the day rolls over - Infura's docs don't guarantee an exact reset +// boundary, so rather than assume one, this seeds at the endpoint's own configured +// (or default) MaxBackoff, so a known-exhausted daily quota isn't re-probed on a +// short cycle. +// - Otherwise 0, which leaves rate_limit_scheduler.go's calculateNextBackoff to fall +// back to InitialBackoff, exactly as it did before this signal existed. +func (s *Server) initialBackoffForSignal(chain, endpointID string, signal health.RateLimitSignal) int { + if signal.RetryAfter > 0 { + return int(signal.RetryAfter.Seconds()) + } + if signal.IsDailyQuota { + rlc := config.DefaultRateLimitRecovery() + if chainEndpoints, ok := s.config.GetEndpointsForChain(chain); ok { + if ep, ok := chainEndpoints[endpointID]; ok && ep.RateLimitRecovery != nil && ep.RateLimitRecovery.MaxBackoff != 0 { + rlc.MaxBackoff = ep.RateLimitRecovery.MaxBackoff + } + } + return rlc.MaxBackoff + } + return 0 +} diff --git a/internal/server/server_test.go b/internal/server/server_test.go index 0568233..54f97a5 100644 --- a/internal/server/server_test.go +++ b/internal/server/server_test.go @@ -6,10 +6,12 @@ import ( "io" "net/http" "net/http/httptest" + "strings" "testing" "time" "aetherlay/internal/config" + "aetherlay/internal/health" "aetherlay/internal/helpers" "aetherlay/internal/store" @@ -589,7 +591,7 @@ func TestHandleRateLimit(t *testing.T) { server := NewServer(cfg, mockValkey, appConfig) // Test handling rate limit - server.handleRateLimit("ethereum", "test-endpoint", "http") + server.handleRateLimit("ethereum", "test-endpoint", "http", health.RateLimitSignal{IsRateLimited: true}) // Verify rate limit state was set state, err := mockValkey.GetRateLimitState(context.Background(), "ethereum", "test-endpoint") @@ -732,7 +734,7 @@ func TestServerGetRateLimitHandler(t *testing.T) { } // Test that handler works - handler("ethereum", "test-endpoint", "http") + handler("ethereum", "test-endpoint", "http", health.RateLimitSignal{IsRateLimited: true}) // Verify rate limit state was set state, err := mockValkey.GetRateLimitState(context.Background(), "ethereum", "test-endpoint") @@ -811,3 +813,281 @@ func TestEphemeralChecksEnabled(t *testing.T) { } }) } + +// TestInitialBackoffForSignal tests that handleRateLimit's backoff seeding prefers a +// parsed Retry-After, falls back to the endpoint's own (or default) MaxBackoff for a +// daily-quota signal, and leaves 0 (scheduler default) for a plain signal. +func TestInitialBackoffForSignal(t *testing.T) { + cfg := &config.Config{ + Endpoints: map[string]config.ChainEndpoints{ + "chainA": { + "with-override": config.Endpoint{ + Provider: "infura", + Role: "primary", + Type: "full", + HTTPURL: "http://with-override", + RateLimitRecovery: &config.RateLimitRecovery{MaxBackoff: 999}, + }, + "no-override": config.Endpoint{ + Provider: "infura", + Role: "primary", + Type: "full", + HTTPURL: "http://no-override", + }, + }, + }, + } + valkeyClient := store.NewMockValkeyClient() + server := NewServer(cfg, valkeyClient, createTestConfig()) + + tests := []struct { + name string + endpointID string + signal health.RateLimitSignal + expected int + }{ + {"retry-after takes priority over daily quota", "no-override", health.RateLimitSignal{RetryAfter: 42 * time.Second, IsDailyQuota: true}, 42}, + {"daily quota uses endpoint's own MaxBackoff override", "with-override", health.RateLimitSignal{IsDailyQuota: true}, 999}, + {"daily quota without override uses default MaxBackoff", "no-override", health.RateLimitSignal{IsDailyQuota: true}, config.DefaultRateLimitRecovery().MaxBackoff}, + {"plain rate limit signal leaves 0 for scheduler to guess InitialBackoff", "no-override", health.RateLimitSignal{IsRateLimited: true}, 0}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := server.initialBackoffForSignal("chainA", tt.endpointID, tt.signal); got != tt.expected { + t.Errorf("initialBackoffForSignal() = %d, want %d", got, tt.expected) + } + }) + } +} + +// TestBodyCarriesRateLimitSignal tests the pure JSON-RPC body scanner used to close the +// blind spot where a provider (e.g. Alchemy on batch requests) signals rate limiting +// only inside a 200 response body, not via HTTP status. +func TestBodyCarriesRateLimitSignal(t *testing.T) { + tests := []struct { + name string + body string + expected bool + }{ + {"single object with rate limit error", `{"jsonrpc":"2.0","id":1,"error":{"code":-32005,"message":"Request limit exceeded"}}`, true}, + {"single object success", `{"jsonrpc":"2.0","id":1,"result":"0x1"}`, false}, + {"batch with one rate limited element", `[{"jsonrpc":"2.0","id":1,"result":"0x1"},{"jsonrpc":"2.0","id":2,"error":{"code":-32005,"message":"limit"}}]`, true}, + {"batch all success", `[{"jsonrpc":"2.0","id":1,"result":"0x1"},{"jsonrpc":"2.0","id":2,"result":"0x2"}]`, false}, + {"batch with unrelated error", `[{"jsonrpc":"2.0","id":1,"error":{"code":-32601,"message":"method not found"}}]`, false}, + {"non-JSON body", `not json at all`, false}, + {"unrelated single JSON-RPC error", `{"jsonrpc":"2.0","id":1,"error":{"code":-32601,"message":"method not found"}}`, false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + sig := bodyCarriesRateLimitSignal("alchemy", []byte(tt.body), nil) + if sig.IsRateLimited != tt.expected { + t.Errorf("bodyCarriesRateLimitSignal() = %v, want %v", sig.IsRateLimited, tt.expected) + } + }) + } +} + +// TestDefaultForwardRequestWithBodyFunc429ParsesRetryAfter tests that a 429 response +// carrying a Retry-After header seeds RateLimitState.CurrentBackoff precisely instead of +// leaving it to the scheduler's guessed InitialBackoff. +func TestDefaultForwardRequestWithBodyFunc429ParsesRetryAfter(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Retry-After", "5") + w.WriteHeader(http.StatusTooManyRequests) + })) + defer ts.Close() + + cfg := &config.Config{ + Endpoints: map[string]config.ChainEndpoints{ + "chainA": { + "ep1": config.Endpoint{Provider: "alchemy", HTTPURL: ts.URL, Role: "primary", Type: "full"}, + }, + }, + } + valkeyClient := store.NewMockValkeyClient() + server := NewServer(cfg, valkeyClient, createTestConfig()) + + err := server.defaultForwardRequestWithBodyFunc(httptest.NewRecorder(), context.Background(), "POST", ts.URL, nil, http.Header{}) + if err == nil { + t.Fatal("Expected error from 429 response") + } + + state, stateErr := valkeyClient.GetRateLimitState(context.Background(), "chainA", "ep1") + if stateErr != nil { + t.Fatalf("Failed to get rate limit state: %v", stateErr) + } + if !state.RateLimited { + t.Error("Expected endpoint to be marked as rate limited") + } + if state.CurrentBackoff != 5 { + t.Errorf("Expected CurrentBackoff to be seeded to 5 from Retry-After, got %d", state.CurrentBackoff) + } +} + +// TestDefaultForwardRequestWithBodyFunc402InfuraSeedsFromMaxBackoff tests that Infura's +// HTTP 402 daily-credit-cap signal seeds CurrentBackoff from the endpoint's configured +// MaxBackoff rather than the scheduler's normal short InitialBackoff. +func TestDefaultForwardRequestWithBodyFunc402InfuraSeedsFromMaxBackoff(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusPaymentRequired) + })) + defer ts.Close() + + cfg := &config.Config{ + Endpoints: map[string]config.ChainEndpoints{ + "chainA": { + "ep1": config.Endpoint{ + Provider: "infura", + HTTPURL: ts.URL, + Role: "primary", + Type: "full", + RateLimitRecovery: &config.RateLimitRecovery{MaxBackoff: 12345}, + }, + }, + }, + } + valkeyClient := store.NewMockValkeyClient() + server := NewServer(cfg, valkeyClient, createTestConfig()) + + err := server.defaultForwardRequestWithBodyFunc(httptest.NewRecorder(), context.Background(), "POST", ts.URL, nil, http.Header{}) + if err == nil { + t.Fatal("Expected error from 402 response") + } + + state, stateErr := valkeyClient.GetRateLimitState(context.Background(), "chainA", "ep1") + if stateErr != nil { + t.Fatalf("Failed to get rate limit state: %v", stateErr) + } + if !state.RateLimited { + t.Error("Expected endpoint to be marked as rate limited") + } + if state.CurrentBackoff != 12345 { + t.Errorf("Expected CurrentBackoff to be seeded from the endpoint's MaxBackoff override (12345), got %d", state.CurrentBackoff) + } +} + +// TestDefaultForwardRequestWithBodyFunc402NonInfuraNotRateLimited tests that a 402 from a +// non-Infura provider is NOT treated as a rate-limit signal, since 402 is only a documented +// Infura convention. +func TestDefaultForwardRequestWithBodyFunc402NonInfuraNotRateLimited(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusPaymentRequired) + })) + defer ts.Close() + + cfg := &config.Config{ + Endpoints: map[string]config.ChainEndpoints{ + "chainA": { + "ep1": config.Endpoint{Provider: "alchemy", HTTPURL: ts.URL, Role: "primary", Type: "full"}, + }, + }, + } + valkeyClient := store.NewMockValkeyClient() + server := NewServer(cfg, valkeyClient, createTestConfig()) + + err := server.defaultForwardRequestWithBodyFunc(httptest.NewRecorder(), context.Background(), "POST", ts.URL, nil, http.Header{}) + if err == nil { + t.Fatal("Expected error from 402 response") + } + + state, stateErr := valkeyClient.GetRateLimitState(context.Background(), "chainA", "ep1") + if stateErr != nil { + t.Fatalf("Failed to get rate limit state: %v", stateErr) + } + if state.RateLimited { + t.Error("Expected a non-Infura 402 to NOT be treated as a rate-limit signal") + } +} + +// TestDefaultForwardRequestWithBodyFuncDetectsEmbeddedBatchRateLimit tests that a 200 +// response carrying a rate-limit error embedded in one element of a JSON-RPC batch array +// is treated as a failed attempt (retried against a different endpoint), rather than +// forwarded to the client with the embedded error silently mixed in. +func TestDefaultForwardRequestWithBodyFuncDetectsEmbeddedBatchRateLimit(t *testing.T) { + batchBody := `[{"jsonrpc":"2.0","id":1,"result":"0x1"},{"jsonrpc":"2.0","id":2,"error":{"code":-32005,"message":"Request limit exceeded"}}]` + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + w.Write([]byte(batchBody)) + })) + defer ts.Close() + + cfg := &config.Config{ + Endpoints: map[string]config.ChainEndpoints{ + "chainA": { + "ep1": config.Endpoint{Provider: "alchemy", HTTPURL: ts.URL, Role: "primary", Type: "full"}, + }, + }, + } + valkeyClient := store.NewMockValkeyClient() + // Pre-populate as healthy so we can confirm the 2xx transaction is NOT marked + // unhealthy (only flagged as rate limited) - it genuinely succeeded at the HTTP level. + valkeyClient.PopulateStatuses(map[string]*store.EndpointStatus{ + "chainA:ep1": {HasHTTP: true, HealthyHTTP: true}, + }) + appConfig := createTestConfig() + appConfig.EndpointFailureThreshold = 1 + server := NewServer(cfg, valkeyClient, appConfig) + + rec := httptest.NewRecorder() + err := server.defaultForwardRequestWithBodyFunc(rec, context.Background(), "POST", ts.URL, nil, http.Header{}) + if err == nil { + t.Fatal("Expected an error so the caller retries with a different endpoint") + } + if rec.Body.Len() != 0 { + t.Errorf("Expected nothing written to the client, got %q", rec.Body.String()) + } + + state, stateErr := valkeyClient.GetRateLimitState(context.Background(), "chainA", "ep1") + if stateErr != nil { + t.Fatalf("Failed to get rate limit state: %v", stateErr) + } + if !state.RateLimited { + t.Error("Expected endpoint to be marked rate limited from the embedded batch error") + } + + status, statusErr := valkeyClient.GetEndpointStatus(context.Background(), "chainA", "ep1") + if statusErr != nil { + t.Fatalf("Failed to get endpoint status: %v", statusErr) + } + if !status.HealthyHTTP { + t.Error("Expected HealthyHTTP to remain true - a 2xx response should not be marked unhealthy, only rate limited") + } +} + +// TestDefaultProxyWebSocketHandshake429ParsesRetryAfter tests that a 429 during the WS +// handshake dial (before any client upgrade) carries a Retry-After header through to the +// returned RateLimitError's Signal field. +func TestDefaultProxyWebSocketHandshake429ParsesRetryAfter(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Retry-After", "7") + w.WriteHeader(http.StatusTooManyRequests) + })) + defer ts.Close() + + backendURL := "ws" + strings.TrimPrefix(ts.URL, "http") + + cfg := &config.Config{ + Endpoints: map[string]config.ChainEndpoints{ + "chainA": { + "ep1": config.Endpoint{Provider: "alchemy", WSURL: backendURL, Role: "primary", Type: "full"}, + }, + }, + } + valkeyClient := store.NewMockValkeyClient() + server := NewServer(cfg, valkeyClient, createTestConfig()) + + err := server.defaultProxyWebSocket(httptest.NewRecorder(), httptest.NewRequest("GET", "/", nil), backendURL) + if err == nil { + t.Fatal("Expected an error from the failed WS handshake") + } + rlErr, ok := err.(*RateLimitError) + if !ok { + t.Fatalf("Expected *RateLimitError, got %T: %v", err, err) + } + if rlErr.Signal.RetryAfter != 7*time.Second { + t.Errorf("Expected RetryAfter to be 7s, got %v", rlErr.Signal.RetryAfter) + } +} diff --git a/internal/store/capacity_estimate.go b/internal/store/capacity_estimate.go new file mode 100644 index 0000000..ddc3f15 --- /dev/null +++ b/internal/store/capacity_estimate.go @@ -0,0 +1,135 @@ +package store + +import ( + "context" + "encoding/json" + "time" + + "aetherlay/internal/config" + + "github.com/valkey-io/valkey-go" +) + +// CapacityEstimate represents the AIMD-learned safe throughput ceiling for an endpoint, +// derived from observed rate-limit hits rather than an operator-declared CapacityLimit. +// Unlike RateLimitState (a binary "currently blocked" flag with a 24h TTL that must +// self-heal), this is valuable precisely because it persists indefinitely - the AIMD +// decrease step is itself the safety bound against a stale-high estimate. +type CapacityEstimate struct { + HasEstimate bool `json:"has_estimate"` // Whether any evidence has been observed yet + MaxRequests int64 `json:"max_requests"` // Learned ceiling as of the last decrease/seed + IncreaseStep int64 `json:"increase_step"` // Additive step size, recomputed at each decrease + WindowSeconds int `json:"window_seconds"` // Frozen at seed time, stable even if config later changes + LastDecreaseAt time.Time `json:"last_decrease_at"` // Anchors both the decrease cooldown and the increase clock +} + +// ApplyCapacityDecrease computes the next CapacityEstimate after a confirmed rate-limit +// hit. effectiveNow is the currently effective ceiling (after any lazy growth since the +// last decrease, see EffectiveMaxRequests) - passing 0 when prior.HasEstimate is false is +// fine, since it's only consulted when there is a prior estimate to compare against. +// observedCount is the real number of requests dispatched to this endpoint in the +// current window at the moment of the hit (from GetCapacityCount), which already +// includes the rejected attempt itself. +// +// The next ceiling is grounded in whichever of the two is more conservative: if growth +// had crept the effective ceiling above what was actually observed this window, trust +// the observed count instead of compounding off a purely speculative grown value. When +// there is no prior estimate, this collapses to using the observed count directly - the +// seed and every subsequent decrease are the same operation. +func ApplyCapacityDecrease(prior CapacityEstimate, effectiveNow int64, observedCount int64, windowSeconds int, params config.CapacityLearning, now time.Time) CapacityEstimate { + base := observedCount + if prior.HasEstimate && effectiveNow < base { + base = effectiveNow + } + + newMax := int64(float64(base) * params.DecreaseFactor) + if newMax < int64(params.MinEstimate) { + newMax = int64(params.MinEstimate) + } + + step := newMax / 10 + if step < 1 { + step = 1 + } + + return CapacityEstimate{ + HasEstimate: true, + MaxRequests: newMax, + IncreaseStep: step, + WindowSeconds: windowSeconds, + LastDecreaseAt: now, + } +} + +// ShouldApplyCapacityDecrease reports whether enough time has passed since the last +// decrease to apply another one. It uses the window width itself as the cooldown, since +// a ceiling can't meaningfully be learned to be "worse" faster than one window boundary, +// and this collapses multiple near-simultaneous rate-limit hits from a single episode +// (e.g. several in-flight retries all rejected within the same window) into one decrease. +func ShouldApplyCapacityDecrease(estimate CapacityEstimate, windowSeconds int, now time.Time) bool { + if !estimate.HasEstimate { + return true + } + return now.Sub(estimate.LastDecreaseAt) >= time.Duration(windowSeconds)*time.Second +} + +// EffectiveMaxRequests computes the currently believed-safe ceiling, growing it +// additively based purely on elapsed wall-clock time since the last decrease. This is +// deliberately stateless between reads - identical on every replica from the same +// persisted estimate and current time, with no background goroutine and no counter to +// race on. +func EffectiveMaxRequests(estimate CapacityEstimate, params config.CapacityLearning, now time.Time) int64 { + if !estimate.HasEstimate { + return 0 + } + if params.IncreaseInterval <= 0 { + return estimate.MaxRequests + } + elapsed := now.Sub(estimate.LastDecreaseAt) + if elapsed <= 0 { + return estimate.MaxRequests + } + steps := int64(elapsed / (time.Duration(params.IncreaseInterval) * time.Second)) + return estimate.MaxRequests + steps*estimate.IncreaseStep +} + +// GetCapacityEstimate retrieves the learned capacity estimate for an endpoint from +// Valkey. Returns a zero-value estimate (HasEstimate: false) if none has been recorded. +func (r *ValkeyClient) GetCapacityEstimate(ctx context.Context, chain, endpoint string) (*CapacityEstimate, error) { + key := capacityEstimatePrefix + chain + ":" + endpoint + cmd := r.client.B().Get().Key(key).Build() + result := r.client.Do(ctx, cmd) + + if valkey.IsValkeyNil(result.Error()) { + return &CapacityEstimate{}, nil + } + + data, err := result.AsBytes() + if err != nil { + return nil, err + } + + var estimate CapacityEstimate + if err := json.Unmarshal(data, &estimate); err != nil { + return nil, err + } + return &estimate, nil +} + +// SetCapacityEstimate stores the learned capacity estimate for an endpoint in Valkey, +// with no expiration - unlike RateLimitState, a learned throughput number is valuable +// because it persists through long clean periods; the AIMD decrease step at the next +// real hit is the safety bound against a stale-high value, not a TTL. +// Uses last-write-wins semantics: a concurrent double-decrease across replicas only +// errs toward a more conservative estimate, which the additive-increase step corrects. +func (r *ValkeyClient) SetCapacityEstimate(ctx context.Context, chain, endpoint string, estimate CapacityEstimate) error { + key := capacityEstimatePrefix + chain + ":" + endpoint + + jsonBytes, err := json.Marshal(estimate) + if err != nil { + return err + } + + cmd := r.client.B().Set().Key(key).Value(string(jsonBytes)).Build() + return r.client.Do(ctx, cmd).Error() +} diff --git a/internal/store/capacity_estimate_test.go b/internal/store/capacity_estimate_test.go new file mode 100644 index 0000000..5ff79be --- /dev/null +++ b/internal/store/capacity_estimate_test.go @@ -0,0 +1,208 @@ +package store + +import ( + "context" + "testing" + "time" + + "aetherlay/internal/config" +) + +func TestApplyCapacityDecrease(t *testing.T) { + params := config.CapacityLearning{DecreaseFactor: 0.5, MinEstimate: 1, IncreaseInterval: 60, WindowSeconds: 60} + now := time.Now() + + tests := []struct { + name string + prior CapacityEstimate + effectiveNow int64 + observedCount int64 + windowSeconds int + expectedMax int64 + expectedStep int64 + expectedWindow int + }{ + { + name: "no prior estimate seeds from observed count", + prior: CapacityEstimate{}, + effectiveNow: 0, + observedCount: 100, + windowSeconds: 10, + expectedMax: 50, + expectedStep: 5, + expectedWindow: 10, + }, + { + name: "prior lower than observed - trust the prior (more conservative)", + prior: CapacityEstimate{HasEstimate: true, MaxRequests: 40}, + effectiveNow: 40, + observedCount: 100, + windowSeconds: 10, + expectedMax: 20, + expectedStep: 2, + expectedWindow: 10, + }, + { + name: "observed lower than prior effective ceiling - trust the observed count", + prior: CapacityEstimate{HasEstimate: true, MaxRequests: 150}, + effectiveNow: 150, + observedCount: 80, + windowSeconds: 10, + expectedMax: 40, + expectedStep: 4, + expectedWindow: 10, + }, + { + name: "floor enforcement - result clamped to MinEstimate", + prior: CapacityEstimate{}, + effectiveNow: 0, + observedCount: 1, + windowSeconds: 10, + expectedMax: 1, + expectedStep: 1, + expectedWindow: 10, + }, + { + name: "zero observed count clamps to floor", + prior: CapacityEstimate{}, + effectiveNow: 0, + observedCount: 0, + windowSeconds: 10, + expectedMax: 1, + expectedStep: 1, + expectedWindow: 10, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := ApplyCapacityDecrease(tt.prior, tt.effectiveNow, tt.observedCount, tt.windowSeconds, params, now) + + if !result.HasEstimate { + t.Error("Expected HasEstimate to be true after a decrease") + } + if result.MaxRequests != tt.expectedMax { + t.Errorf("MaxRequests = %d, want %d", result.MaxRequests, tt.expectedMax) + } + if result.IncreaseStep != tt.expectedStep { + t.Errorf("IncreaseStep = %d, want %d", result.IncreaseStep, tt.expectedStep) + } + if result.WindowSeconds != tt.expectedWindow { + t.Errorf("WindowSeconds = %d, want %d", result.WindowSeconds, tt.expectedWindow) + } + if !result.LastDecreaseAt.Equal(now) { + t.Errorf("LastDecreaseAt = %v, want %v", result.LastDecreaseAt, now) + } + }) + } +} + +func TestShouldApplyCapacityDecrease(t *testing.T) { + now := time.Now() + + tests := []struct { + name string + estimate CapacityEstimate + window int + expected bool + }{ + {"never decreased", CapacityEstimate{HasEstimate: false}, 10, true}, + {"within cooldown", CapacityEstimate{HasEstimate: true, LastDecreaseAt: now.Add(-5 * time.Second)}, 10, false}, + {"exactly at cooldown boundary", CapacityEstimate{HasEstimate: true, LastDecreaseAt: now.Add(-10 * time.Second)}, 10, true}, + {"after cooldown elapsed", CapacityEstimate{HasEstimate: true, LastDecreaseAt: now.Add(-30 * time.Second)}, 10, true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := ShouldApplyCapacityDecrease(tt.estimate, tt.window, now); got != tt.expected { + t.Errorf("ShouldApplyCapacityDecrease() = %v, want %v", got, tt.expected) + } + }) + } +} + +func TestEffectiveMaxRequests(t *testing.T) { + now := time.Now() + params := config.CapacityLearning{IncreaseInterval: 60} + + tests := []struct { + name string + estimate CapacityEstimate + params config.CapacityLearning + expected int64 + }{ + { + name: "no estimate returns 0", + estimate: CapacityEstimate{HasEstimate: false}, + params: params, + expected: 0, + }, + { + name: "zero elapsed returns base ceiling unchanged", + estimate: CapacityEstimate{HasEstimate: true, MaxRequests: 50, IncreaseStep: 5, LastDecreaseAt: now}, + params: params, + expected: 50, + }, + { + name: "one interval elapsed adds one step", + estimate: CapacityEstimate{HasEstimate: true, MaxRequests: 50, IncreaseStep: 5, LastDecreaseAt: now.Add(-60 * time.Second)}, + params: params, + expected: 55, + }, + { + name: "multiple intervals elapsed adds proportional steps", + estimate: CapacityEstimate{HasEstimate: true, MaxRequests: 50, IncreaseStep: 5, LastDecreaseAt: now.Add(-185 * time.Second)}, + params: params, + expected: 65, // floor(185/60) = 3 steps * 5 = 15 + }, + { + name: "IncreaseInterval <= 0 guard returns base ceiling unchanged", + estimate: CapacityEstimate{HasEstimate: true, MaxRequests: 50, IncreaseStep: 5, LastDecreaseAt: now.Add(-1 * time.Hour)}, + params: config.CapacityLearning{IncreaseInterval: 0}, + expected: 50, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := EffectiveMaxRequests(tt.estimate, tt.params, now); got != tt.expected { + t.Errorf("EffectiveMaxRequests() = %d, want %d", got, tt.expected) + } + }) + } +} + +func TestGetAndSetCapacityEstimateMock(t *testing.T) { + client := NewMockValkeyClient() + ctx := context.Background() + chain := "ethereum" + endpoint := "ep1" + + // Absent estimate returns a zero-value, non-nil result. + estimate, err := client.GetCapacityEstimate(ctx, chain, endpoint) + if err != nil { + t.Fatalf("GetCapacityEstimate failed: %v", err) + } + if estimate.HasEstimate { + t.Error("Expected HasEstimate to be false for an unseen endpoint") + } + + stored := CapacityEstimate{ + HasEstimate: true, + MaxRequests: 42, + IncreaseStep: 4, + WindowSeconds: 60, + LastDecreaseAt: time.Now(), + } + if err := client.SetCapacityEstimate(ctx, chain, endpoint, stored); err != nil { + t.Fatalf("SetCapacityEstimate failed: %v", err) + } + + retrieved, err := client.GetCapacityEstimate(ctx, chain, endpoint) + if err != nil { + t.Fatalf("GetCapacityEstimate failed: %v", err) + } + if !retrieved.HasEstimate || retrieved.MaxRequests != 42 || retrieved.IncreaseStep != 4 { + t.Errorf("Expected stored estimate to round-trip, got %+v", retrieved) + } +} diff --git a/internal/store/testutils.go b/internal/store/testutils.go index 4614c1b..bb32640 100644 --- a/internal/store/testutils.go +++ b/internal/store/testutils.go @@ -9,20 +9,29 @@ import ( // MockValkeyClient is a mock implementation of ValkeyClientIface for testing. // It supports in-memory endpoint status storage and is safe for concurrent use. type MockValkeyClient struct { - rateLimitStates map[string]*RateLimitState - requestCounts map[string]map[string]map[string][3]int64 // [0]=24h, [1]=1m, [2]=all - statuses map[string]*EndpointStatus - values map[string]string // Generic key-value storage for Set/Get - mu sync.RWMutex + rateLimitStates map[string]*RateLimitState + requestCounts map[string]map[string]map[string][3]int64 // [0]=24h, [1]=1m, [2]=all + capacityCounts map[string]map[int64]int64 // "chain:endpoint" -> bucket -> count + capacityEstimates map[string]*CapacityEstimate // "chain:endpoint" -> learned estimate + statuses map[string]*EndpointStatus + values map[string]string // Generic key-value storage for Set/Get + mu sync.RWMutex + + // NowFunc lets tests deterministically simulate capacity-window rollover + // without a real sleep. Defaults to time.Now. + NowFunc func() time.Time } // NewMockValkeyClient creates a new MockValkeyClient with empty state. func NewMockValkeyClient() *MockValkeyClient { return &MockValkeyClient{ - rateLimitStates: make(map[string]*RateLimitState), - requestCounts: make(map[string]map[string]map[string][3]int64), - statuses: make(map[string]*EndpointStatus), - values: make(map[string]string), + rateLimitStates: make(map[string]*RateLimitState), + requestCounts: make(map[string]map[string]map[string][3]int64), + capacityCounts: make(map[string]map[int64]int64), + capacityEstimates: make(map[string]*CapacityEstimate), + statuses: make(map[string]*EndpointStatus), + values: make(map[string]string), + NowFunc: time.Now, } } @@ -133,6 +142,56 @@ func (m *MockValkeyClient) SetRateLimitState(_ context.Context, chain, endpoint return nil } +// IncrementCapacityCount increments the mock's in-memory capacity counter for the +// endpoint's current fixed window (bucketed by NowFunc().Unix()/windowSeconds, mirroring +// the real ValkeyClient's bucketing) and returns the new count. +func (m *MockValkeyClient) IncrementCapacityCount(_ context.Context, chain, endpoint string, windowSeconds int) (int64, error) { + m.mu.Lock() + defer m.mu.Unlock() + key := chain + ":" + endpoint + bucket := m.NowFunc().Unix() / int64(windowSeconds) + if _, ok := m.capacityCounts[key]; !ok { + m.capacityCounts[key] = make(map[int64]int64) + } + m.capacityCounts[key][bucket]++ + return m.capacityCounts[key][bucket], nil +} + +// GetCapacityCount returns the mock's in-memory capacity counter for the endpoint's +// current fixed window, or 0 if nothing has been recorded in that window yet. +func (m *MockValkeyClient) GetCapacityCount(_ context.Context, chain, endpoint string, windowSeconds int) (int64, error) { + m.mu.RLock() + defer m.mu.RUnlock() + key := chain + ":" + endpoint + bucket := m.NowFunc().Unix() / int64(windowSeconds) + if buckets, ok := m.capacityCounts[key]; ok { + return buckets[bucket], nil + } + return 0, nil +} + +// GetCapacityEstimate returns the mock's in-memory learned capacity estimate for an +// endpoint, or a zero-value estimate (HasEstimate: false) if none has been recorded. +func (m *MockValkeyClient) GetCapacityEstimate(_ context.Context, chain, endpoint string) (*CapacityEstimate, error) { + m.mu.RLock() + defer m.mu.RUnlock() + key := chain + ":" + endpoint + estimate, ok := m.capacityEstimates[key] + if !ok { + return &CapacityEstimate{}, nil + } + return estimate, nil +} + +// SetCapacityEstimate stores the mock's in-memory learned capacity estimate for an endpoint. +func (m *MockValkeyClient) SetCapacityEstimate(_ context.Context, chain, endpoint string, estimate CapacityEstimate) error { + m.mu.Lock() + defer m.mu.Unlock() + key := chain + ":" + endpoint + m.capacityEstimates[key] = &estimate + return nil +} + // PopulateStatuses allows tests to pre-populate endpoint statuses in the mock. func (m *MockValkeyClient) PopulateStatuses(statuses map[string]*EndpointStatus) { m.mu.Lock() diff --git a/internal/store/valkey.go b/internal/store/valkey.go index 0a3fce6..bf4c32b 100644 --- a/internal/store/valkey.go +++ b/internal/store/valkey.go @@ -4,6 +4,7 @@ import ( "context" "crypto/tls" "encoding/json" + "strconv" "strings" "time" @@ -12,14 +13,16 @@ import ( const ( // Key prefixes for Valkey storage - healthPrefix = "health:" - metricsPrefix = "metrics:" - rateLimitPrefix = "rate_limit:" - proxyRequests = "proxy_requests" - healthRequests = "health_requests" - requests24hKey = "requests_24h" - requests1mKey = "requests_1m" - requestsAllKey = "requests_all" + healthPrefix = "health:" + metricsPrefix = "metrics:" + rateLimitPrefix = "rate_limit:" + capacityPrefix = "capacity:" + capacityEstimatePrefix = "capacity_estimate:" + proxyRequests = "proxy_requests" + healthRequests = "health_requests" + requests24hKey = "requests_24h" + requests1mKey = "requests_1m" + requestsAllKey = "requests_all" ) // EndpointStatus represents the health status and metrics of an endpoint. @@ -66,6 +69,10 @@ type ValkeyClientIface interface { GetCombinedRequestCounts(ctx context.Context, chain, endpoint string) (int64, int64, int64, error) GetRateLimitState(ctx context.Context, chain, endpoint string) (*RateLimitState, error) SetRateLimitState(ctx context.Context, chain, endpoint string, state RateLimitState) error + IncrementCapacityCount(ctx context.Context, chain, endpoint string, windowSeconds int) (int64, error) + GetCapacityCount(ctx context.Context, chain, endpoint string, windowSeconds int) (int64, error) + GetCapacityEstimate(ctx context.Context, chain, endpoint string) (*CapacityEstimate, error) + SetCapacityEstimate(ctx context.Context, chain, endpoint string, estimate CapacityEstimate) error CleanupStaleEndpoints(ctx context.Context, activeEndpoints map[string][]string) (int, error) Ping(ctx context.Context) error Close() error @@ -262,7 +269,7 @@ func (r *ValkeyClient) CleanupStaleEndpoints(ctx context.Context, activeEndpoint } } - prefixes := []string{healthPrefix, metricsPrefix, rateLimitPrefix} + prefixes := []string{healthPrefix, metricsPrefix, rateLimitPrefix, capacityEstimatePrefix} var staleKeys []string for _, prefix := range prefixes { @@ -381,3 +388,54 @@ func (r *ValkeyClient) SetRateLimitState(ctx context.Context, chain, endpoint st cmd := r.client.B().Set().Key(key).Value(string(jsonBytes)).Ex(24 * time.Hour).Build() return r.client.Do(ctx, cmd).Error() } + +// capacityBucketKey returns the Valkey key for the current fixed window of width +// windowSeconds, e.g. window 10 buckets time into 10-second slices. The window +// resets every windowSeconds because each slice gets its own key - unlike +// IncrementRequestCount's rolling TTL, this key naturally stops being written to +// once the window elapses, so a fresh window always starts at zero. +func capacityBucketKey(chain, endpoint string, windowSeconds int) string { + bucket := time.Now().Unix() / int64(windowSeconds) + return capacityPrefix + chain + ":" + endpoint + ":" + strconv.FormatInt(bucket, 10) +} + +// IncrementCapacityCount increments the self-imposed capacity counter for an endpoint +// within the current fixed window of width windowSeconds, and returns the new count. +// Used to proactively throttle requests below a configured ceiling, independent of +// any provider-reported rate limit state. +func (r *ValkeyClient) IncrementCapacityCount(ctx context.Context, chain, endpoint string, windowSeconds int) (int64, error) { + key := capacityBucketKey(chain, endpoint, windowSeconds) + + cmds := []valkey.Completed{ + r.client.B().Incr().Key(key).Build(), + r.client.B().Expire().Key(key).Seconds(int64(2 * windowSeconds)).Build(), + } + + results := r.client.DoMulti(ctx, cmds...) + if err := results[0].Error(); err != nil { + return 0, err + } + count, err := results[0].AsInt64() + if err != nil { + return 0, err + } + if err := results[1].Error(); err != nil { + return 0, err + } + return count, nil +} + +// GetCapacityCount returns the current count for an endpoint's capacity window, +// or 0 if nothing has been recorded in the current window yet. +func (r *ValkeyClient) GetCapacityCount(ctx context.Context, chain, endpoint string, windowSeconds int) (int64, error) { + key := capacityBucketKey(chain, endpoint, windowSeconds) + result := r.client.Do(ctx, r.client.B().Get().Key(key).Build()) + + if valkey.IsValkeyNil(result.Error()) { + return 0, nil + } + if err := result.Error(); err != nil { + return 0, err + } + return result.AsInt64() +} diff --git a/internal/store/valkey_test.go b/internal/store/valkey_test.go index 5ebaab9..02f2e32 100644 --- a/internal/store/valkey_test.go +++ b/internal/store/valkey_test.go @@ -286,6 +286,79 @@ func TestCombinedRequestCounts(t *testing.T) { } } +func TestIncrementAndGetCapacityCount(t *testing.T) { + client := NewMockValkeyClient() + + ctx := context.Background() + chain := "test-chain" + endpoint := uniqueTestKey("https://test.example.com") + + for i := 0; i < 3; i++ { + count, err := client.IncrementCapacityCount(ctx, chain, endpoint, 10) + if err != nil { + t.Fatalf("Increment failed: %v", err) + } + if count != int64(i+1) { + t.Errorf("Expected count %d, got %d", i+1, count) + } + } + + count, err := client.GetCapacityCount(ctx, chain, endpoint, 10) + if err != nil { + t.Fatalf("Get capacity count failed: %v", err) + } + if count != 3 { + t.Errorf("Expected capacity count to be 3, got %d", count) + } +} + +func TestGetCapacityCountForUnusedEndpointReturnsZero(t *testing.T) { + client := NewMockValkeyClient() + + ctx := context.Background() + count, err := client.GetCapacityCount(ctx, "test-chain", "https://unused.example.com", 10) + if err != nil { + t.Fatalf("Get capacity count failed: %v", err) + } + if count != 0 { + t.Errorf("Expected capacity count to be 0, got %d", count) + } +} + +func TestCapacityCountWindowRollover(t *testing.T) { + client := NewMockValkeyClient() + ctx := context.Background() + chain := "test-chain" + endpoint := uniqueTestKey("https://test.example.com") + + windowStart := time.Unix(1_700_000_000, 0) + client.NowFunc = func() time.Time { return windowStart } + + for i := 0; i < 2; i++ { + if _, err := client.IncrementCapacityCount(ctx, chain, endpoint, 10); err != nil { + t.Fatalf("Increment failed: %v", err) + } + } + count, err := client.GetCapacityCount(ctx, chain, endpoint, 10) + if err != nil { + t.Fatalf("Get capacity count failed: %v", err) + } + if count != 2 { + t.Errorf("Expected count 2 within the window, got %d", count) + } + + // Jump forward past the window boundary (window width 10s) without a real sleep. + client.NowFunc = func() time.Time { return windowStart.Add(11 * time.Second) } + + count, err = client.GetCapacityCount(ctx, chain, endpoint, 10) + if err != nil { + t.Fatalf("Get capacity count failed: %v", err) + } + if count != 0 { + t.Errorf("Expected count to reset to 0 in the new window, got %d", count) + } +} + // TestNewValkeyClientTLSConfig is an integration test that checks the TLS configuration. // It requires a running Valkey server with TLS enabled on port 6380 and non-TLS on 6379. func TestNewValkeyClientTLSConfig(t *testing.T) { diff --git a/services/health-checker/main.go b/services/health-checker/main.go index c9be4ac..602d167 100644 --- a/services/health-checker/main.go +++ b/services/health-checker/main.go @@ -36,9 +36,87 @@ var testExitAfterSetup bool // exitCode is used to track the exit code for the process var exitCode int +// standaloneInitialBackoff mirrors server.Server.initialBackoffForSignal: seed the first +// recovery-check wait from whatever the provider actually told us, instead of always +// guessing. A parsed Retry-After is used directly; Infura's daily credit cap (402) seeds +// from the endpoint's own (or default) MaxBackoff, since the docs don't guarantee an +// exact reset boundary. Otherwise 0, leaving the load balancer's scheduler to fall back +// to InitialBackoff once it picks up recovery monitoring. +func standaloneInitialBackoff(cfg *config.Config, chain, endpointID string, signal health.RateLimitSignal) int { + if signal.RetryAfter > 0 { + return int(signal.RetryAfter.Seconds()) + } + if signal.IsDailyQuota { + rlc := config.DefaultRateLimitRecovery() + if chainEndpoints, ok := cfg.GetEndpointsForChain(chain); ok { + if ep, ok := chainEndpoints[endpointID]; ok && ep.RateLimitRecovery != nil && ep.RateLimitRecovery.MaxBackoff != 0 { + rlc.MaxBackoff = ep.RateLimitRecovery.MaxBackoff + } + } + return rlc.MaxBackoff + } + return 0 +} + +// applyStandaloneLearnedCapacityDecrease mirrors server.Server.applyLearnedCapacityDecrease +// exactly, using the same shared store.ApplyCapacityDecrease/EffectiveMaxRequests math, so +// the load balancer and the standalone health checker - two separate processes mutating +// the same Valkey-persisted estimate - never diverge. Only engages for endpoints with no +// static Capacity configured, mirroring effectiveCapacityCeiling's rule in server.go. +func applyStandaloneLearnedCapacityDecrease(cfg *config.Config, valkeyClient store.ValkeyClientIface, capacityThrottlingEnabled, capacityLearningEnabled bool, chain, endpointID string) { + chainEndpoints, ok := cfg.GetEndpointsForChain(chain) + if !ok { + return + } + ep, ok := chainEndpoints[endpointID] + if !ok { + return + } + if !capacityThrottlingEnabled || !capacityLearningEnabled || ep.Capacity != nil { + return + } + + params := config.ResolveCapacityLearning(ep.CapacityLearning) + now := time.Now() + ctx := context.Background() + + prior, err := valkeyClient.GetCapacityEstimate(ctx, chain, endpointID) + if err != nil { + log.Debug().Err(err).Str("chain", chain).Str("endpoint", endpointID).Msg("Failed to get capacity estimate in standalone health checker") + return + } + + if !store.ShouldApplyCapacityDecrease(*prior, params.WindowSeconds, now) { + log.Debug().Str("chain", chain).Str("endpoint", endpointID).Msg("Skipping capacity estimate decrease in standalone health checker, within cooldown") + return + } + + observedCount, err := valkeyClient.GetCapacityCount(ctx, chain, endpointID, params.WindowSeconds) + if err != nil { + log.Debug().Err(err).Str("chain", chain).Str("endpoint", endpointID).Msg("Failed to get capacity count in standalone health checker") + return + } + + effectiveNow := store.EffectiveMaxRequests(*prior, params, now) + newEstimate := store.ApplyCapacityDecrease(*prior, effectiveNow, observedCount, params.WindowSeconds, params, now) + + if err := valkeyClient.SetCapacityEstimate(ctx, chain, endpointID, newEstimate); err != nil { + log.Debug().Err(err).Str("chain", chain).Str("endpoint", endpointID).Msg("Failed to set capacity estimate in standalone health checker") + return + } + + log.Info().Str("chain", chain).Str("endpoint", endpointID).Int64("new_estimate", newEstimate.MaxRequests).Msg("Standalone health checker decreased learned capacity estimate after a rate-limit hit") + if metrics.EndpointCapacityEstimatedCeiling != nil { + metrics.EndpointCapacityEstimatedCeiling.WithLabelValues(chain, endpointID).Set(float64(newEstimate.MaxRequests)) + } + if metrics.EndpointCapacityEstimateDecreasedTotal != nil { + metrics.EndpointCapacityEstimateDecreasedTotal.WithLabelValues(chain, endpointID).Inc() + } +} + // createStandaloneRateLimitHandler creates a simple rate limit handler for the standalone health checker -func createStandaloneRateLimitHandler(valkeyClient store.ValkeyClientIface) func(chain, endpointID, protocol string) { - return func(chain, endpointID, protocol string) { +func createStandaloneRateLimitHandler(cfg *config.Config, valkeyClient store.ValkeyClientIface, capacityThrottlingEnabled, capacityLearningEnabled bool) func(chain, endpointID, protocol string, signal health.RateLimitSignal) { + return func(chain, endpointID, protocol string, signal health.RateLimitSignal) { log.Debug().Str("chain", chain).Str("endpoint", endpointID).Str("protocol", protocol).Msg("Standalone health checker detected rate limit") // Get current rate limit state @@ -52,6 +130,7 @@ func createStandaloneRateLimitHandler(valkeyClient store.ValkeyClientIface) func now := time.Now() state.RateLimited = true state.LastRecoveryCheck = now + state.CurrentBackoff = standaloneInitialBackoff(cfg, chain, endpointID, signal) // Set first rate limited time if this is the first time if state.FirstRateLimited.IsZero() { @@ -63,7 +142,12 @@ func createStandaloneRateLimitHandler(valkeyClient store.ValkeyClientIface) func return } - log.Info().Str("chain", chain).Str("endpoint", endpointID).Str("protocol", protocol).Msg("Standalone health checker marked endpoint as rate limited") + log.Info().Str("chain", chain).Str("endpoint", endpointID).Str("protocol", protocol).Int("current_backoff", state.CurrentBackoff).Msg("Standalone health checker marked endpoint as rate limited") + + // Check for rate limits first (this signal), then approximate the endpoint's + // safe throughput ceiling from it - only engages for endpoints with no static + // Capacity, and shares the exact same math as the load balancer's own path. + applyStandaloneLearnedCapacityDecrease(cfg, valkeyClient, capacityThrottlingEnabled, capacityLearningEnabled, chain, endpointID) } } @@ -73,6 +157,8 @@ func RunHealthChecker( corsHeaders string, corsMethods string, corsOrigin string, + capacityLearningEnabled bool, + capacityThrottlingEnabled bool, ephemeralChecksEnabled bool, ephemeralChecksHealthyThreshold int, ephemeralChecksInterval int, @@ -170,7 +256,7 @@ func RunHealthChecker( checker := health.NewChecker(cfg, valkeyClient, time.Duration(healthCheckInterval)*time.Second, time.Duration(ephemeralChecksInterval)*time.Second, ephemeralChecksHealthyThreshold, healthCheckSyncStatus, healthCheckConcurrency, ephemeralChecksEnabled) // Set up simple rate limit handler for standalone health checker - checker.HandleRateLimitFunc = createStandaloneRateLimitHandler(valkeyClient) + checker.HandleRateLimitFunc = createStandaloneRateLimitHandler(cfg, valkeyClient, capacityThrottlingEnabled, capacityLearningEnabled) if testCheckerPatch != nil { testCheckerPatch(checker) @@ -237,6 +323,8 @@ func main() { config.CorsHeaders, config.CorsMethods, config.CorsOrigin, + config.CapacityLearningEnabled, + config.CapacityThrottlingEnabled, config.EphemeralChecksEnabled, config.EphemeralChecksHealthyThreshold, config.EphemeralChecksInterval, diff --git a/services/health-checker/main_test.go b/services/health-checker/main_test.go index e2320a6..18270e2 100644 --- a/services/health-checker/main_test.go +++ b/services/health-checker/main_test.go @@ -6,6 +6,7 @@ import ( "aetherlay/internal/store" "context" "testing" + "time" ) // mockConfig returns a minimal valid *config.Config for testing @@ -64,6 +65,8 @@ func TestRunHealthCheckerFromEnv_Standalone(t *testing.T) { "Accept, Authorization, Content-Type, Origin, X-Requested-With", // corsHeaders "GET, POST, OPTIONS", // corsMethods "*", // corsOrigin + true, // capacityLearningEnabled + true, // capacityThrottlingEnabled true, // ephemeralChecksEnabled 3, // ephemeralChecksHealthyThreshold 30, // ephemeralChecksInterval @@ -125,6 +128,8 @@ func TestRunHealthCheckerFromEnv_Ephemeral(t *testing.T) { "Accept, Authorization, Content-Type, Origin, X-Requested-With", // corsHeaders "GET, POST, OPTIONS", // corsMethods "*", // corsOrigin + true, // capacityLearningEnabled + true, // capacityThrottlingEnabled true, // ephemeralChecksEnabled 3, // ephemeralChecksHealthyThreshold 30, // ephemeralChecksInterval @@ -186,6 +191,8 @@ func TestRunHealthCheckerFromEnv_Disabled(t *testing.T) { "Accept, Authorization, Content-Type, Origin, X-Requested-With", // corsHeaders "GET, POST, OPTIONS", // corsMethods "*", // corsOrigin + true, // capacityLearningEnabled + true, // capacityThrottlingEnabled true, // ephemeralChecksEnabled 3, // ephemeralChecksHealthyThreshold 30, // ephemeralChecksInterval @@ -207,3 +214,142 @@ func TestRunHealthCheckerFromEnv_Disabled(t *testing.T) { t.Errorf("Expected mode 'disabled', got '%s'", detectedMode) } } + +// TestStandaloneInitialBackoff tests that standaloneInitialBackoff mirrors +// server.Server.initialBackoffForSignal's precedence: Retry-After first, then the +// endpoint's own (or default) MaxBackoff for a daily-quota signal, else 0. +func TestStandaloneInitialBackoff(t *testing.T) { + cfg := &config.Config{ + Endpoints: map[string]config.ChainEndpoints{ + "mainnet": { + "with-override": config.Endpoint{ + Provider: "infura", + Role: "primary", + Type: "full", + HTTPURL: "http://with-override", + RateLimitRecovery: &config.RateLimitRecovery{MaxBackoff: 555}, + }, + "no-override": config.Endpoint{ + Provider: "infura", + Role: "primary", + Type: "full", + HTTPURL: "http://no-override", + }, + }, + }, + } + + tests := []struct { + name string + endpointID string + signal health.RateLimitSignal + expected int + }{ + {"retry-after takes priority over daily quota", "no-override", health.RateLimitSignal{RetryAfter: 20 * time.Second, IsDailyQuota: true}, 20}, + {"daily quota uses endpoint's own MaxBackoff override", "with-override", health.RateLimitSignal{IsDailyQuota: true}, 555}, + {"daily quota without override uses default MaxBackoff", "no-override", health.RateLimitSignal{IsDailyQuota: true}, config.DefaultRateLimitRecovery().MaxBackoff}, + {"plain rate limit signal leaves 0", "no-override", health.RateLimitSignal{IsRateLimited: true}, 0}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := standaloneInitialBackoff(cfg, "mainnet", tt.endpointID, tt.signal); got != tt.expected { + t.Errorf("standaloneInitialBackoff() = %d, want %d", got, tt.expected) + } + }) + } +} + +// TestCreateStandaloneRateLimitHandlerSeedsBackoffFromSignal tests that the standalone +// health checker's rate limit handler seeds CurrentBackoff from the signal, not just +// marking the endpoint rate limited with a zero backoff. +func TestCreateStandaloneRateLimitHandlerSeedsBackoffFromSignal(t *testing.T) { + cfg := mockConfig() + valkeyClient := store.NewMockValkeyClient() + handler := createStandaloneRateLimitHandler(cfg, valkeyClient, true, true) + + handler("mainnet", "mock", "http", health.RateLimitSignal{RetryAfter: 15 * time.Second}) + + state, err := valkeyClient.GetRateLimitState(context.Background(), "mainnet", "mock") + if err != nil { + t.Fatalf("Failed to get rate limit state: %v", err) + } + if !state.RateLimited { + t.Error("Expected endpoint to be marked rate limited") + } + if state.CurrentBackoff != 15 { + t.Errorf("Expected CurrentBackoff to be seeded to 15 from Retry-After, got %d", state.CurrentBackoff) + } +} + +// TestApplyStandaloneLearnedCapacityDecreaseSeedsFromObservedCount confirms the +// standalone health checker's decrease path shares the exact same store.ApplyCapacityDecrease +// math as the load balancer's own applyLearnedCapacityDecrease - load-bearing, since +// these run as separate processes mutating the same Valkey-persisted estimate. +func TestApplyStandaloneLearnedCapacityDecreaseSeedsFromObservedCount(t *testing.T) { + cfg := mockConfig() + valkeyClient := store.NewMockValkeyClient() + ctx := context.Background() + + for i := 0; i < 10; i++ { + valkeyClient.IncrementCapacityCount(ctx, "mainnet", "mock", 60) + } + + applyStandaloneLearnedCapacityDecrease(cfg, valkeyClient, true, true, "mainnet", "mock") + + estimate, err := valkeyClient.GetCapacityEstimate(ctx, "mainnet", "mock") + if err != nil { + t.Fatalf("GetCapacityEstimate failed: %v", err) + } + if !estimate.HasEstimate || estimate.MaxRequests != 5 { + t.Errorf("Expected a learned estimate of 5 (10 observed * 0.5), got %+v", estimate) + } +} + +// TestApplyStandaloneLearnedCapacityDecreaseSkipsWhenStaticCapacityConfigured mirrors +// server.effectiveCapacityCeiling's rule: adaptive learning never engages for an +// endpoint that already has a static Capacity configured. +func TestApplyStandaloneLearnedCapacityDecreaseSkipsWhenStaticCapacityConfigured(t *testing.T) { + cfg := &config.Config{ + Endpoints: map[string]config.ChainEndpoints{ + "mainnet": { + "mock": config.Endpoint{ + Provider: "mock", Role: "primary", Type: "full", HTTPURL: "http://mock", + Capacity: &config.CapacityLimit{MaxRequests: 100, WindowSeconds: 10}, + }, + }, + }, + } + valkeyClient := store.NewMockValkeyClient() + ctx := context.Background() + + applyStandaloneLearnedCapacityDecrease(cfg, valkeyClient, true, true, "mainnet", "mock") + + estimate, _ := valkeyClient.GetCapacityEstimate(ctx, "mainnet", "mock") + if estimate.HasEstimate { + t.Error("Expected no learned estimate when a static Capacity is configured") + } +} + +// TestCreateStandaloneRateLimitHandlerAlsoSeedsCapacityEstimate confirms the handler +// wires both the reactive backoff and the proactive learned-estimate decrease together. +func TestCreateStandaloneRateLimitHandlerAlsoSeedsCapacityEstimate(t *testing.T) { + cfg := mockConfig() + valkeyClient := store.NewMockValkeyClient() + ctx := context.Background() + + for i := 0; i < 8; i++ { + valkeyClient.IncrementCapacityCount(ctx, "mainnet", "mock", 60) + } + + handler := createStandaloneRateLimitHandler(cfg, valkeyClient, true, true) + handler("mainnet", "mock", "http", health.RateLimitSignal{IsRateLimited: true}) + + estimate, err := valkeyClient.GetCapacityEstimate(ctx, "mainnet", "mock") + if err != nil { + t.Fatalf("GetCapacityEstimate failed: %v", err) + } + if !estimate.HasEstimate || estimate.MaxRequests != 4 { + t.Errorf("Expected a learned estimate of 4 (8 observed * 0.5), got %+v", estimate) + } +} From 80f7f5b7a299d6f011cae6a031f99bd9a93a2c6f Mon Sep 17 00:00:00 2001 From: Jacob Sussmilch Date: Wed, 1 Jul 2026 15:34:37 +1000 Subject: [PATCH 02/12] fix: KEEP-905 Guard against divide-by-zero on non-positive capacity window_seconds An endpoint's static capacity.window_seconds is used as a divisor when computing the Valkey bucket key for capacity counting. Omitting the field (JSON zero value) or setting it negative caused a panic on the endpoint's first request, and since capacity gating evaluates every endpoint in a role before selection, this took down the entire chain, not just the misconfigured endpoint. LoadConfig now validates window_seconds at load time, warning and disabling capacity throttling for that endpoint rather than crashing. capacityBucketKey and MockValkeyClient's capacity methods also clamp defensively, since they're reachable through the public ValkeyClientIface beyond just the config path. --- internal/config/config.go | 25 +++++++++++++++ internal/config/config_test.go | 58 ++++++++++++++++++++++++++++++++++ internal/store/testutils.go | 6 ++++ internal/store/valkey.go | 6 ++++ internal/store/valkey_test.go | 18 +++++++++++ 5 files changed, 113 insertions(+) diff --git a/internal/config/config.go b/internal/config/config.go index dd08ec2..9705b7a 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -3,6 +3,8 @@ package config import ( "encoding/json" "os" + + "github.com/rs/zerolog/log" ) // RateLimitRecovery represents the configuration for rate limit recovery @@ -91,6 +93,7 @@ func LoadConfig(path string) (*Config, error) { for chainName, chainEndpoints := range config.Endpoints { for endpointID, endpoint := range chainEndpoints { substituteEnvVarsInEndpoint(&endpoint) + validateEndpointCapacity(chainName, endpointID, &endpoint) config.Endpoints[chainName][endpointID] = endpoint } } @@ -98,6 +101,28 @@ func LoadConfig(path string) (*Config, error) { return &config, nil } +// validateEndpointCapacity catches an endpoint's static Capacity being configured with a +// non-positive WindowSeconds (e.g. omitted entirely, which JSON-decodes to the int zero +// value). WindowSeconds ends up as a divisor when computing the Valkey bucket key for +// capacity counting (see capacityBucketKey in internal/store/valkey.go), so a zero or +// negative value would panic on the endpoint's very first request. Rather than crash, +// disable proactive capacity throttling for that endpoint (matching today's behavior for +// an endpoint with no capacity configured at all) and warn loudly so the operator notices +// the misconfiguration. +func validateEndpointCapacity(chain, endpointID string, endpoint *Endpoint) { + if endpoint.Capacity == nil { + return + } + if endpoint.Capacity.WindowSeconds <= 0 { + log.Warn(). + Str("chain", chain). + Str("endpoint", endpointID). + Int("window_seconds", endpoint.Capacity.WindowSeconds). + Msg("Endpoint's capacity.window_seconds must be positive - disabling proactive capacity throttling for this endpoint") + endpoint.Capacity = nil + } +} + // GetEndpointsForChain returns all endpoints for a specific chain. // Returns the endpoints and a boolean indicating if the chain exists. func (c *Config) GetEndpointsForChain(chain string) (ChainEndpoints, bool) { diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 0ab19f8..f50a8cb 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -152,6 +152,64 @@ func TestLoadConfigInvalidJSON(t *testing.T) { } } +// TestLoadConfigDisablesCapacityWithNonPositiveWindowSeconds guards against a regression +// of a divide-by-zero panic: WindowSeconds is used as a divisor when computing the Valkey +// capacity bucket key, so a zero (e.g. omitted from the JSON) or negative value must never +// reach the rest of the system - LoadConfig should catch it and disable capacity for that +// endpoint instead of letting it through. +func TestLoadConfigDisablesCapacityWithNonPositiveWindowSeconds(t *testing.T) { + tmpFile := "test_zero_window.json" + content := `{ + "ethereum": { + "zero-window": { + "provider": "alchemy", + "role": "primary", + "type": "full", + "http_url": "http://test.com", + "capacity": {"max_requests": 100} + }, + "negative-window": { + "provider": "alchemy", + "role": "primary", + "type": "full", + "http_url": "http://test2.com", + "capacity": {"max_requests": 100, "window_seconds": -5} + }, + "valid-window": { + "provider": "alchemy", + "role": "primary", + "type": "full", + "http_url": "http://test3.com", + "capacity": {"max_requests": 100, "window_seconds": 10} + } + } + }` + if err := os.WriteFile(tmpFile, []byte(content), 0644); err != nil { + t.Fatalf("Failed to create test file: %v", err) + } + defer os.Remove(tmpFile) + + cfg, err := LoadConfig(tmpFile) + if err != nil { + t.Fatalf("LoadConfig failed: %v", err) + } + + endpoints := cfg.Endpoints["ethereum"] + + if endpoints["zero-window"].Capacity != nil { + t.Error("Expected Capacity to be disabled (nil) for an endpoint with omitted (zero) window_seconds") + } + if endpoints["negative-window"].Capacity != nil { + t.Error("Expected Capacity to be disabled (nil) for an endpoint with negative window_seconds") + } + if endpoints["valid-window"].Capacity == nil { + t.Fatal("Expected Capacity to remain set for an endpoint with a valid window_seconds") + } + if endpoints["valid-window"].Capacity.WindowSeconds != 10 { + t.Errorf("Expected valid endpoint's WindowSeconds to be untouched (10), got %d", endpoints["valid-window"].Capacity.WindowSeconds) + } +} + func TestDefaultRateLimitRecovery(t *testing.T) { config := DefaultRateLimitRecovery() diff --git a/internal/store/testutils.go b/internal/store/testutils.go index bb32640..6d05449 100644 --- a/internal/store/testutils.go +++ b/internal/store/testutils.go @@ -148,6 +148,9 @@ func (m *MockValkeyClient) SetRateLimitState(_ context.Context, chain, endpoint func (m *MockValkeyClient) IncrementCapacityCount(_ context.Context, chain, endpoint string, windowSeconds int) (int64, error) { m.mu.Lock() defer m.mu.Unlock() + if windowSeconds <= 0 { + windowSeconds = 1 + } key := chain + ":" + endpoint bucket := m.NowFunc().Unix() / int64(windowSeconds) if _, ok := m.capacityCounts[key]; !ok { @@ -162,6 +165,9 @@ func (m *MockValkeyClient) IncrementCapacityCount(_ context.Context, chain, endp func (m *MockValkeyClient) GetCapacityCount(_ context.Context, chain, endpoint string, windowSeconds int) (int64, error) { m.mu.RLock() defer m.mu.RUnlock() + if windowSeconds <= 0 { + windowSeconds = 1 + } key := chain + ":" + endpoint bucket := m.NowFunc().Unix() / int64(windowSeconds) if buckets, ok := m.capacityCounts[key]; ok { diff --git a/internal/store/valkey.go b/internal/store/valkey.go index bf4c32b..ee07403 100644 --- a/internal/store/valkey.go +++ b/internal/store/valkey.go @@ -395,6 +395,12 @@ func (r *ValkeyClient) SetRateLimitState(ctx context.Context, chain, endpoint st // IncrementRequestCount's rolling TTL, this key naturally stops being written to // once the window elapses, so a fresh window always starts at zero. func capacityBucketKey(chain, endpoint string, windowSeconds int) string { + // windowSeconds is a divisor below; config.LoadConfig is the primary guard against a + // non-positive value reaching here, but this function is reachable through the public + // ValkeyClientIface, so it defends itself too rather than panicking on bad input. + if windowSeconds <= 0 { + windowSeconds = 1 + } bucket := time.Now().Unix() / int64(windowSeconds) return capacityPrefix + chain + ":" + endpoint + ":" + strconv.FormatInt(bucket, 10) } diff --git a/internal/store/valkey_test.go b/internal/store/valkey_test.go index 02f2e32..4c6843c 100644 --- a/internal/store/valkey_test.go +++ b/internal/store/valkey_test.go @@ -325,6 +325,24 @@ func TestGetCapacityCountForUnusedEndpointReturnsZero(t *testing.T) { } } +// TestCapacityCountDoesNotPanicOnNonPositiveWindowSeconds guards against a regression of +// a divide-by-zero panic: windowSeconds is used as a divisor to compute the capacity +// bucket, so a caller passing 0 or a negative value (e.g. from a bug elsewhere, since +// config.LoadConfig is the primary guard, not this function) must never crash the process. +func TestCapacityCountDoesNotPanicOnNonPositiveWindowSeconds(t *testing.T) { + client := NewMockValkeyClient() + ctx := context.Background() + + for _, windowSeconds := range []int{0, -1, -100} { + if _, err := client.IncrementCapacityCount(ctx, "test-chain", "ep1", windowSeconds); err != nil { + t.Fatalf("IncrementCapacityCount(windowSeconds=%d) failed: %v", windowSeconds, err) + } + if _, err := client.GetCapacityCount(ctx, "test-chain", "ep1", windowSeconds); err != nil { + t.Fatalf("GetCapacityCount(windowSeconds=%d) failed: %v", windowSeconds, err) + } + } +} + func TestCapacityCountWindowRollover(t *testing.T) { client := NewMockValkeyClient() ctx := context.Background() From 986772b864127eb082531eeca23e2bb898335587 Mon Sep 17 00:00:00 2001 From: Jacob Sussmilch Date: Wed, 1 Jul 2026 15:39:11 +1000 Subject: [PATCH 03/12] fix: KEEP-905 Fix capacity window desync between write and read paths recordCapacityUsage and applyLearnedCapacityDecrease resolved an endpoint's learning window from the live config on every call, while effectiveCapacityCeiling's gating/weighting read path used the window frozen into the endpoint's CapacityEstimate at last-seed time. Since the Valkey bucket key is derived from the window value itself, a config change to an endpoint's capacity_learning.window_seconds after it already had a learned estimate caused writes and reads to target different buckets - the gate would then read a permanently stale/zero count and never trip again, silently defeating the proactive throttle until another real rate-limit hit happened to re-sync the frozen value. Both the load balancer and the standalone health checker now consult the estimate's frozen window (once one exists) for the cooldown check and observed-count read, falling back to the live-resolved default only before any estimate has been seeded. --- internal/server/capacity_learning_test.go | 137 ++++++++++++++++++++++ internal/server/server.go | 33 +++++- services/health-checker/main.go | 14 ++- 3 files changed, 178 insertions(+), 6 deletions(-) diff --git a/internal/server/capacity_learning_test.go b/internal/server/capacity_learning_test.go index 66a0059..a4dfa80 100644 --- a/internal/server/capacity_learning_test.go +++ b/internal/server/capacity_learning_test.go @@ -316,3 +316,140 @@ func TestEffectiveCapacityCeilingUnaffectedWhenLearningNotOptedIn(t *testing.T) t.Error("Expected no ceiling for an unconfigured endpoint when learning is not opted in") } } + +// TestCapacityWindowSecondsUsesFrozenEstimateWindow guards against a regression where +// the write path (recordCapacityUsage) resolved window_seconds from the live config while +// the read/gating path (effectiveCapacityCeiling) used the estimate's frozen window - +// causing writes and reads to target different Valkey bucket keys and silently +// defeating the proactive gate. Once an estimate is seeded, both paths must agree. +func TestCapacityWindowSecondsUsesFrozenEstimateWindow(t *testing.T) { + cfg := &config.Config{} + valkeyClient := store.NewMockValkeyClient() + server := NewServer(cfg, valkeyClient, learningTestConfig()) + + // Endpoint has no static Capacity; its live-resolved learning window (30s) differs + // from the frozen window already recorded on its estimate (60s) - simulating an + // operator changing capacity_learning.window_seconds after the estimate was seeded. + ep := config.Endpoint{ + Provider: "alchemy", + CapacityLearning: &config.CapacityLearning{WindowSeconds: 30}, + } + if err := valkeyClient.SetCapacityEstimate(context.Background(), "ethereum", "ep1", store.CapacityEstimate{ + HasEstimate: true, MaxRequests: 10, IncreaseStep: 1, WindowSeconds: 60, LastDecreaseAt: time.Now(), + }); err != nil { + t.Fatalf("SetCapacityEstimate failed: %v", err) + } + + if got := server.capacityWindowSeconds("ethereum", "ep1", ep); got != 60 { + t.Errorf("Expected capacityWindowSeconds to return the frozen estimate window (60), got %d", got) + } + + // Bootstrap case: no estimate exists yet for this endpoint - falls back to the live + // config, since there's no frozen value yet to match. + if got := server.capacityWindowSeconds("ethereum", "no-estimate-yet", ep); got != 30 { + t.Errorf("Expected capacityWindowSeconds to fall back to the live config (30) before any estimate exists, got %d", got) + } +} + +// TestRecordCapacityUsageWritesToFrozenWindowBucket confirms recordCapacityUsage's write +// lands in the same Valkey bucket effectiveCapacityCeiling's gating check reads from, +// even when the endpoint's live-resolved learning window has since changed. +func TestRecordCapacityUsageWritesToFrozenWindowBucket(t *testing.T) { + cfg := &config.Config{} + valkeyClient := store.NewMockValkeyClient() + appConfig := learningTestConfig() + server := NewServer(cfg, valkeyClient, appConfig) + ctx := context.Background() + + ep := config.Endpoint{ + Provider: "alchemy", + CapacityLearning: &config.CapacityLearning{WindowSeconds: 30}, // live value, deliberately different from frozen + } + if err := valkeyClient.SetCapacityEstimate(ctx, "ethereum", "ep1", store.CapacityEstimate{ + HasEstimate: true, MaxRequests: 10, IncreaseStep: 1, WindowSeconds: 60, LastDecreaseAt: time.Now(), + }); err != nil { + t.Fatalf("SetCapacityEstimate failed: %v", err) + } + + server.recordCapacityUsage("ethereum", ep, "ep1") + + frozenWindowCount, err := valkeyClient.GetCapacityCount(ctx, "ethereum", "ep1", 60) + if err != nil { + t.Fatalf("GetCapacityCount failed: %v", err) + } + if frozenWindowCount != 1 { + t.Errorf("Expected the write to land in the frozen-window (60s) bucket, got count %d there", frozenWindowCount) + } + + liveWindowCount, err := valkeyClient.GetCapacityCount(ctx, "ethereum", "ep1", 30) + if err != nil { + t.Fatalf("GetCapacityCount failed: %v", err) + } + if liveWindowCount != 0 { + t.Errorf("Expected nothing written to the stale live-window (30s) bucket, got count %d there", liveWindowCount) + } + + // The gating path must observe the same count recordCapacityUsage just wrote. + maxRequests, windowSeconds, ok := server.effectiveCapacityCeiling("ethereum", "ep1", ep) + if !ok { + t.Fatal("Expected a resolvable ceiling") + } + gatedCount, err := valkeyClient.GetCapacityCount(ctx, "ethereum", "ep1", windowSeconds) + if err != nil { + t.Fatalf("GetCapacityCount failed: %v", err) + } + if gatedCount != 1 { + t.Errorf("Expected the gating path to observe the write (count=1) via its own resolved window (%d), got count %d", windowSeconds, gatedCount) + } + _ = maxRequests +} + +// TestApplyLearnedCapacityDecreaseReadsFrozenWindowNotLiveConfig confirms the decrease +// path's cooldown check and observed-count read are grounded in the estimate's frozen +// window, not a live config value that may have since diverged - otherwise the decrease +// would be seeded from an empty/wrong bucket's count. +func TestApplyLearnedCapacityDecreaseReadsFrozenWindowNotLiveConfig(t *testing.T) { + cfg := &config.Config{ + Endpoints: map[string]config.ChainEndpoints{ + "ethereum": { + "ep1": config.Endpoint{ + Provider: "alchemy", + Role: "primary", + Type: "full", + HTTPURL: "http://ep1", + CapacityLearning: &config.CapacityLearning{WindowSeconds: 30}, // live value, differs from frozen + }, + }, + }, + } + valkeyClient := store.NewMockValkeyClient() + ctx := context.Background() + ep := cfg.Endpoints["ethereum"]["ep1"] + + // Seed an estimate frozen at 60s, with its cooldown already elapsed. + if err := valkeyClient.SetCapacityEstimate(ctx, "ethereum", "ep1", store.CapacityEstimate{ + HasEstimate: true, MaxRequests: 10, IncreaseStep: 1, WindowSeconds: 60, + LastDecreaseAt: time.Now().Add(-120 * time.Second), + }); err != nil { + t.Fatalf("SetCapacityEstimate failed: %v", err) + } + + // Real traffic landed in the frozen (60s) bucket, as recordCapacityUsage would write it. + for i := 0; i < 8; i++ { + valkeyClient.IncrementCapacityCount(ctx, "ethereum", "ep1", 60) + } + + server := NewServer(cfg, valkeyClient, learningTestConfig()) + server.applyLearnedCapacityDecrease("ethereum", "ep1", ep) + + final, err := valkeyClient.GetCapacityEstimate(ctx, "ethereum", "ep1") + if err != nil { + t.Fatalf("GetCapacityEstimate failed: %v", err) + } + // If the decrease had incorrectly read the live 30s bucket (count=0 there), it would + // seed from 0 and clamp to MinEstimate (1). Reading the correct frozen 60s bucket + // (count=8) grounds the decrease at floor(8*0.5)=4 instead. + if final.MaxRequests != 4 { + t.Errorf("Expected decrease to be grounded in the frozen window's observed count (8 -> 4), got MaxRequests=%d", final.MaxRequests) + } +} diff --git a/internal/server/server.go b/internal/server/server.go index c024339..bde3777 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -1328,10 +1328,24 @@ func (s *Server) providerForEndpoint(chain, endpointID string) string { // learning window - so an endpoint with no static Capacity still accumulates real // GetCapacityCount evidence before it's ever rate limited, for applyLearnedCapacityDecrease // to seed an estimate from on the first hit. -func capacityWindowSeconds(endpoint config.Endpoint) int { +// capacityWindowSeconds resolves the window width to track usage against for the WRITE +// path (recordCapacityUsage's usage counter). It must agree with whatever window +// effectiveCapacityCeiling's READ path is watching, or gating silently stops working - +// writes and reads would land in different Valkey bucket keys (see capacityBucketKey, +// which derives the bucket from windowSeconds itself). So: static Capacity's window if +// configured (unchanged, always stable); else the frozen window already recorded on a +// learned estimate, if one has been seeded (matching effectiveCapacityCeiling exactly); +// else the live-resolved adaptive-learning window as a bootstrap default before any +// estimate exists yet to freeze against. +func (s *Server) capacityWindowSeconds(chain, endpointID string, endpoint config.Endpoint) int { if endpoint.Capacity != nil { return endpoint.Capacity.WindowSeconds } + if s.appConfig.CapacityLearningEnabled { + if estimate, err := s.valkeyClient.GetCapacityEstimate(context.Background(), chain, endpointID); err == nil && estimate.HasEstimate { + return estimate.WindowSeconds + } + } return config.ResolveCapacityLearning(endpoint.CapacityLearning).WindowSeconds } @@ -1349,7 +1363,7 @@ func (s *Server) recordCapacityUsage(chain string, endpoint config.Endpoint, end } ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - if _, err := s.valkeyClient.IncrementCapacityCount(ctx, chain, endpointID, capacityWindowSeconds(endpoint)); err != nil { + if _, err := s.valkeyClient.IncrementCapacityCount(ctx, chain, endpointID, s.capacityWindowSeconds(chain, endpointID, endpoint)); err != nil { log.Debug().Err(err).Str("chain", chain).Str("endpoint", endpointID).Msg("Failed to record capacity usage") } } @@ -1375,12 +1389,23 @@ func (s *Server) applyLearnedCapacityDecrease(chain, endpointID string, endpoint return } - if !store.ShouldApplyCapacityDecrease(*prior, params.WindowSeconds, now) { + // Read against whichever window the usage counter is actually being written to right + // now (see capacityWindowSeconds): the estimate's own frozen window once one exists, + // not the live-resolved config - otherwise this read targets a different Valkey + // bucket key than recordCapacityUsage's writes, and observedCount below would always + // be stale/zero. Only before any estimate has been seeded is there no frozen value to + // match, so the live-resolved window is the correct (and only) choice at that point. + readWindowSeconds := params.WindowSeconds + if prior.HasEstimate { + readWindowSeconds = prior.WindowSeconds + } + + if !store.ShouldApplyCapacityDecrease(*prior, readWindowSeconds, now) { log.Debug().Str("chain", chain).Str("endpoint", endpointID).Msg("Skipping capacity estimate decrease, within cooldown of the last decrease") return } - observedCount, err := s.valkeyClient.GetCapacityCount(ctx, chain, endpointID, params.WindowSeconds) + observedCount, err := s.valkeyClient.GetCapacityCount(ctx, chain, endpointID, readWindowSeconds) if err != nil { log.Debug().Err(err).Str("chain", chain).Str("endpoint", endpointID).Msg("Failed to get capacity count for estimate decrease") return diff --git a/services/health-checker/main.go b/services/health-checker/main.go index 602d167..ef0b7ac 100644 --- a/services/health-checker/main.go +++ b/services/health-checker/main.go @@ -86,12 +86,22 @@ func applyStandaloneLearnedCapacityDecrease(cfg *config.Config, valkeyClient sto return } - if !store.ShouldApplyCapacityDecrease(*prior, params.WindowSeconds, now) { + // Read against the estimate's own frozen window once one exists (matching + // server.Server.capacityWindowSeconds/effectiveCapacityCeiling exactly), not the + // live-resolved config - otherwise this read targets a different Valkey bucket key + // than the load balancer's writes. Only before any estimate has been seeded is there + // no frozen value to match. + readWindowSeconds := params.WindowSeconds + if prior.HasEstimate { + readWindowSeconds = prior.WindowSeconds + } + + if !store.ShouldApplyCapacityDecrease(*prior, readWindowSeconds, now) { log.Debug().Str("chain", chain).Str("endpoint", endpointID).Msg("Skipping capacity estimate decrease in standalone health checker, within cooldown") return } - observedCount, err := valkeyClient.GetCapacityCount(ctx, chain, endpointID, params.WindowSeconds) + observedCount, err := valkeyClient.GetCapacityCount(ctx, chain, endpointID, readWindowSeconds) if err != nil { log.Debug().Err(err).Str("chain", chain).Str("endpoint", endpointID).Msg("Failed to get capacity count in standalone health checker") return From 77d2200c079057cfc0d05fd7a7fee411214ccb3c Mon Sep 17 00:00:00 2001 From: Jacob Sussmilch Date: Wed, 1 Jul 2026 15:47:29 +1000 Subject: [PATCH 04/12] fix: KEEP-905 Fix premature backoff escalation for signal-seeded recovery performRecoveryCheck used CurrentBackoff == 0 as its only signal for "this is the first recovery attempt, don't multiply yet." Before this feature, CurrentBackoff always started at exactly 0, so that check worked. Now handleRateLimit can seed CurrentBackoff to a precise nonzero value from a Retry-After header or Infura's daily-quota MaxBackoff, which collided with that sentinel: the very first probe after a signaled rate limit already multiplied the backoff, instead of repeating the precise value once more first - exactly the same treatment a guessed InitialBackoff already gets. Gate on RecoveryAttempts instead, so a signal-seeded value gets the same "one repeat wait, then escalate" shape as the no-signal path, just with a more precise baseline. --- internal/server/rate_limit_scheduler.go | 17 ++++- internal/server/rate_limit_scheduler_test.go | 77 ++++++++++++++++++++ 2 files changed, 91 insertions(+), 3 deletions(-) diff --git a/internal/server/rate_limit_scheduler.go b/internal/server/rate_limit_scheduler.go index 63e0d62..0a01ef8 100644 --- a/internal/server/rate_limit_scheduler.go +++ b/internal/server/rate_limit_scheduler.go @@ -238,9 +238,20 @@ func (rls *RateLimitScheduler) performRecoveryCheck(ctx context.Context, chain, state.RecoveryAttempts++ state.LastRecoveryCheck = time.Now() - // Update current backoff for next iteration - if state.CurrentBackoff == 0 { - state.CurrentBackoff = rateLimitConfig.InitialBackoff + // Update current backoff for next iteration. Gate on RecoveryAttempts, not + // CurrentBackoff == 0: handleRateLimit can seed CurrentBackoff to a precise nonzero + // value (from a Retry-After header or Infura's daily-quota MaxBackoff), and a bare + // "== 0" check would then look like escalation had already happened, multiplying the + // backoff after only a single probe instead of giving the seeded value one repeat + // wait first - the same treatment a guessed InitialBackoff gets below. + if state.RecoveryAttempts == 1 { + // First attempt for this episode: only fall back to InitialBackoff if nothing + // more precise was seeded: a signal-seeded value is left untouched, so the + // second attempt still waits the exact seeded duration once more before this + // same "else" branch below starts multiplying it. + if state.CurrentBackoff == 0 { + state.CurrentBackoff = rateLimitConfig.InitialBackoff + } } else { newBackoff := int(float64(state.CurrentBackoff) * rateLimitConfig.BackoffMultiplier) state.CurrentBackoff = min(newBackoff, rateLimitConfig.MaxBackoff) diff --git a/internal/server/rate_limit_scheduler_test.go b/internal/server/rate_limit_scheduler_test.go index 4a0f6db..de645c2 100644 --- a/internal/server/rate_limit_scheduler_test.go +++ b/internal/server/rate_limit_scheduler_test.go @@ -383,3 +383,80 @@ func TestShouldResetBackoff(t *testing.T) { t.Error("Expected no reset for zero time") } } + +// TestPerformRecoveryCheckPreservesSeededBackoffForFirstAttempt guards against a +// regression where a CurrentBackoff seeded to a precise nonzero value (e.g. from a +// Retry-After header, via handleRateLimit's initialBackoffForSignal) collided with this +// function's "== 0" sentinel for "first attempt," causing the backoff to escalate one +// generation earlier than it would for a guessed InitialBackoff. The seeded value must +// be given one repeat wait, exactly like InitialBackoff is, before multiplying starts. +func TestPerformRecoveryCheckPreservesSeededBackoffForFirstAttempt(t *testing.T) { + // Endpoint that always fails the recovery health check, so CurrentBackoff's value + // after each call reflects only the backoff-update logic, not a success-reset. + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusTooManyRequests) + })) + defer server.Close() + + cfg := &config.Config{ + Endpoints: map[string]config.ChainEndpoints{ + "ethereum": { + "test-endpoint": config.Endpoint{ + Provider: "test-provider", + Role: "primary", + Type: "full", + HTTPURL: server.URL, + }, + }, + }, + } + + mockValkey := store.NewMockValkeyClient() + // Simulate handleRateLimit having just seeded CurrentBackoff=5 from a Retry-After + // header (RecoveryAttempts=0, a fresh episode - never decremented/reset elsewhere). + mockValkey.SetRateLimitState(context.Background(), "ethereum", "test-endpoint", store.RateLimitState{ + RateLimited: true, + CurrentBackoff: 5, + RecoveryAttempts: 0, + FirstRateLimited: time.Now(), + }) + + scheduler := NewRateLimitScheduler(cfg, mockValkey) + endpoint := cfg.Endpoints["ethereum"]["test-endpoint"] + rateLimitConfig := config.RateLimitRecovery{ + BackoffMultiplier: 2.0, + InitialBackoff: 300, + MaxBackoff: 7200, + MaxRetries: 10, + RequiredSuccesses: 2, + ResetAfter: 86400, + } + + // First recovery attempt: the seeded precise value must survive untouched, not be + // multiplied as if a prior attempt had already happened. + shouldContinue, err := scheduler.performRecoveryCheck(context.Background(), "ethereum", "test-endpoint", endpoint, rateLimitConfig) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + if !shouldContinue { + t.Fatal("Expected monitoring to continue after the first failed attempt") + } + state, _ := mockValkey.GetRateLimitState(context.Background(), "ethereum", "test-endpoint") + if state.CurrentBackoff != 5 { + t.Errorf("Expected CurrentBackoff to remain 5 after the first attempt, got %d", state.CurrentBackoff) + } + + // Second recovery attempt: now escalation is expected, exactly as it would be for a + // guessed InitialBackoff after its own second attempt. + shouldContinue, err = scheduler.performRecoveryCheck(context.Background(), "ethereum", "test-endpoint", endpoint, rateLimitConfig) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + if !shouldContinue { + t.Fatal("Expected monitoring to continue after the second failed attempt") + } + state, _ = mockValkey.GetRateLimitState(context.Background(), "ethereum", "test-endpoint") + if state.CurrentBackoff != 10 { + t.Errorf("Expected CurrentBackoff to escalate to 10 (5*2.0) on the second attempt, got %d", state.CurrentBackoff) + } +} From 288b4d25cfb79ee85713894376ef6d793bcf011a Mon Sep 17 00:00:00 2001 From: Jacob Sussmilch Date: Wed, 1 Jul 2026 15:50:00 +1000 Subject: [PATCH 05/12] fix: KEEP-905 Stop duplicate response headers leaking across retries Response headers were copied onto the client's http.ResponseWriter unconditionally right after the non-2xx check, before the embedded rate-limit body scan decided whether to abort. handleRequestHTTP's retry loop reuses the same ResponseWriter across attempts, so an aborted attempt (2xx response whose body embeds a rate-limit signal) left its headers sitting in the writer's header map uncommitted, and a subsequent successful retry's headers were then added on top rather than replacing them - producing duplicate header values (e.g. two Content-Length) in the response actually sent to the client. Headers are now copied only immediately before each point the function commits to writing a response (extracted into copyResponseHeaders, called right before each WriteHeader), never before an abort path. --- internal/server/server.go | 38 ++++++++++++++++------- internal/server/server_test.go | 57 ++++++++++++++++++++++++++++++++++ 2 files changed, 84 insertions(+), 11 deletions(-) diff --git a/internal/server/server.go b/internal/server/server.go index bde3777..c724f2d 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -1255,22 +1255,20 @@ func (s *Server) defaultForwardRequestWithBodyFunc(w http.ResponseWriter, ctx co return fmt.Errorf("HTTP %d: %s", resp.StatusCode, resp.Status) } - // Copy response headers - for key, values := range resp.Header { - // Skip CORS headers to avoid duplication - if strings.HasPrefix(key, "Access-Control-") { - continue - } - for _, value := range values { - w.Header().Add(key, value) - } - } - // For JSON responses within a bounded size, buffer and scan the body for an embedded // rate-limit error before forwarding. Some providers (e.g. Alchemy on batch requests) // return HTTP 200 with the rate-limit signal only inside the JSON-RPC body - invisible // to a status-code-only check. Larger bodies are streamed through unmodified and // uninspected, to bound memory use on large responses (e.g. eth_getLogs). + // + // Headers are copied onto w only once we've committed to writing this specific + // response (immediately before each w.WriteHeader below), not unconditionally up + // front: handleRequestHTTP's retry loop reuses the same w across attempts, so copying + // headers before knowing whether this attempt aborts (the embedded-signal branch + // below returns an error without ever calling WriteHeader) would leave them sitting + // in w's header map, where a subsequent successful retry's headers would be added on + // top of rather than replacing them - e.g. two Content-Length values in the response + // actually sent to the client. if isJSONContentType(resp.Header.Get("Content-Type")) { buffered, readErr := io.ReadAll(io.LimitReader(resp.Body, maxRateLimitScanBodyBytes+1)) if readErr == nil && int64(len(buffered)) <= maxRateLimitScanBodyBytes { @@ -1284,18 +1282,21 @@ func (s *Server) defaultForwardRequestWithBodyFunc(w http.ResponseWriter, ctx co return fmt.Errorf("rate limited: embedded JSON-RPC rate-limit error in 2xx response") } } + copyResponseHeaders(w, resp.Header) w.WriteHeader(resp.StatusCode) _, err = w.Write(buffered) return err } // Exceeded the scan cap (or a read error occurred): stream the buffered prefix // plus whatever remains, unmodified and uninspected. + copyResponseHeaders(w, resp.Header) w.WriteHeader(resp.StatusCode) _, err = io.Copy(w, io.MultiReader(bytes.NewReader(buffered), resp.Body)) return err } // Set response status + copyResponseHeaders(w, resp.Header) w.WriteHeader(resp.StatusCode) // Copy response body @@ -1303,6 +1304,21 @@ func (s *Server) defaultForwardRequestWithBodyFunc(w http.ResponseWriter, ctx co return err } +// copyResponseHeaders copies resp's headers onto w, skipping CORS headers (already set +// by the CORS middleware). Callers must only invoke this once they've committed to +// writing this specific response - see the comment at defaultForwardRequestWithBodyFunc's +// call sites for why copying headers before that commitment is unsafe. +func copyResponseHeaders(w http.ResponseWriter, respHeader http.Header) { + for key, values := range respHeader { + if strings.HasPrefix(key, "Access-Control-") { + continue + } + for _, value := range values { + w.Header().Add(key, value) + } + } +} + // maxRateLimitScanBodyBytes bounds how much of a 2xx response body is buffered to scan // for an embedded JSON-RPC rate-limit error. Single-request rate-limit bodies are tiny; // this cap exists so large legitimate responses aren't fully buffered in memory. diff --git a/internal/server/server_test.go b/internal/server/server_test.go index 54f97a5..f1f08f7 100644 --- a/internal/server/server_test.go +++ b/internal/server/server_test.go @@ -1057,6 +1057,63 @@ func TestDefaultForwardRequestWithBodyFuncDetectsEmbeddedBatchRateLimit(t *testi } } +// TestDefaultForwardRequestWithBodyFuncNoDuplicateHeadersAcrossRetries guards against a +// regression where response headers were copied onto w unconditionally, before the +// embedded-rate-limit body scan decided whether to abort. Since handleRequestHTTP's +// retry loop reuses the same http.ResponseWriter across attempts, an aborted attempt's +// headers were left sitting in w's header map and then added to (not replaced by) the +// next successful attempt's headers - producing duplicate header values in the response +// actually sent to the client. +func TestDefaultForwardRequestWithBodyFuncNoDuplicateHeadersAcrossRetries(t *testing.T) { + rateLimitedServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Header().Set("X-Test-Header", "from-rate-limited") + w.WriteHeader(http.StatusOK) + w.Write([]byte(`[{"jsonrpc":"2.0","id":1,"result":"0x1"},{"jsonrpc":"2.0","id":2,"error":{"code":-32005,"message":"Request limit exceeded"}}]`)) + })) + defer rateLimitedServer.Close() + + healthyServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Header().Set("X-Test-Header", "from-healthy") + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"jsonrpc":"2.0","id":1,"result":"0x2"}`)) + })) + defer healthyServer.Close() + + cfg := &config.Config{ + Endpoints: map[string]config.ChainEndpoints{ + "ethereum": { + "flaky": config.Endpoint{Provider: "alchemy", HTTPURL: rateLimitedServer.URL, Role: "primary", Type: "full"}, + "reliable": config.Endpoint{Provider: "drpc", HTTPURL: healthyServer.URL, Role: "primary", Type: "full"}, + }, + }, + } + valkeyClient := store.NewMockValkeyClient() + server := NewServer(cfg, valkeyClient, createTestConfig()) + + // Simulate exactly what handleRequestHTTP's retry loop does: the same + // http.ResponseWriter passed to a first attempt that aborts (embedded rate-limit + // signal, no write), then to a second attempt that succeeds and writes for real. + w := httptest.NewRecorder() + + if err := server.defaultForwardRequestWithBodyFunc(w, context.Background(), "POST", rateLimitedServer.URL, nil, http.Header{}); err == nil { + t.Fatal("Expected the first attempt (embedded rate-limit signal) to return an error") + } + + if err := server.defaultForwardRequestWithBodyFunc(w, context.Background(), "POST", healthyServer.URL, nil, http.Header{}); err != nil { + t.Fatalf("Expected the second attempt to succeed, got error: %v", err) + } + + values := w.Header().Values("X-Test-Header") + if len(values) != 1 { + t.Fatalf("Expected exactly one X-Test-Header value in the final response, got %v (duplicate headers leaked across retries)", values) + } + if values[0] != "from-healthy" { + t.Errorf("Expected the surviving header to be from the endpoint that actually succeeded, got %q", values[0]) + } +} + // TestDefaultProxyWebSocketHandshake429ParsesRetryAfter tests that a 429 during the WS // handshake dial (before any client upgrade) carries a Retry-After header through to the // returned RateLimitError's Signal field. From 3133ba69137dd7bfb3a5013fb40c65e087c6a8e4 Mon Sep 17 00:00:00 2001 From: Jacob Sussmilch Date: Wed, 1 Jul 2026 15:51:30 +1000 Subject: [PATCH 06/12] fix: KEEP-905 Guard against non-positive capacity max_requests An endpoint's static capacity.max_requests of 0 (e.g. omitted from the JSON) or negative made the "count >= max_requests" gating check in getEndpointsByRole true on the very first request, silently and permanently excluding the endpoint from selection with no distinguishing error from a real outage. LoadConfig now validates max_requests alongside window_seconds, warning and disabling capacity throttling for that endpoint when either is non-positive. --- internal/config/config.go | 26 +++++++++++---- internal/config/config_test.go | 58 ++++++++++++++++++++++++++++++++++ 2 files changed, 77 insertions(+), 7 deletions(-) diff --git a/internal/config/config.go b/internal/config/config.go index 9705b7a..e99f1b1 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -102,13 +102,16 @@ func LoadConfig(path string) (*Config, error) { } // validateEndpointCapacity catches an endpoint's static Capacity being configured with a -// non-positive WindowSeconds (e.g. omitted entirely, which JSON-decodes to the int zero -// value). WindowSeconds ends up as a divisor when computing the Valkey bucket key for -// capacity counting (see capacityBucketKey in internal/store/valkey.go), so a zero or -// negative value would panic on the endpoint's very first request. Rather than crash, -// disable proactive capacity throttling for that endpoint (matching today's behavior for -// an endpoint with no capacity configured at all) and warn loudly so the operator notices -// the misconfiguration. +// non-positive WindowSeconds or MaxRequests (e.g. either field omitted entirely, which +// JSON-decodes to the int zero value). WindowSeconds ends up as a divisor when computing +// the Valkey bucket key for capacity counting (see capacityBucketKey in +// internal/store/valkey.go), so a zero or negative value would panic on the endpoint's +// very first request. A MaxRequests of zero (or negative) makes the "count >= max" +// gating check in getEndpointsByRole true on the very first request, silently and +// permanently excluding the endpoint with no distinguishing error. Rather than crash or +// silently starve the endpoint, disable proactive capacity throttling for it (matching +// today's behavior for an endpoint with no capacity configured at all) and warn loudly +// so the operator notices the misconfiguration. func validateEndpointCapacity(chain, endpointID string, endpoint *Endpoint) { if endpoint.Capacity == nil { return @@ -120,6 +123,15 @@ func validateEndpointCapacity(chain, endpointID string, endpoint *Endpoint) { Int("window_seconds", endpoint.Capacity.WindowSeconds). Msg("Endpoint's capacity.window_seconds must be positive - disabling proactive capacity throttling for this endpoint") endpoint.Capacity = nil + return + } + if endpoint.Capacity.MaxRequests <= 0 { + log.Warn(). + Str("chain", chain). + Str("endpoint", endpointID). + Int("max_requests", endpoint.Capacity.MaxRequests). + Msg("Endpoint's capacity.max_requests must be positive - disabling proactive capacity throttling for this endpoint") + endpoint.Capacity = nil } } diff --git a/internal/config/config_test.go b/internal/config/config_test.go index f50a8cb..595c28b 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -210,6 +210,64 @@ func TestLoadConfigDisablesCapacityWithNonPositiveWindowSeconds(t *testing.T) { } } +// TestLoadConfigDisablesCapacityWithNonPositiveMaxRequests guards against a regression +// where a max_requests of 0 (e.g. omitted from the JSON) or negative would make the +// "count >= max_requests" gating check in getEndpointsByRole true on the endpoint's very +// first request, silently and permanently excluding it with no distinguishing error - +// LoadConfig should catch it and disable capacity for that endpoint instead. +func TestLoadConfigDisablesCapacityWithNonPositiveMaxRequests(t *testing.T) { + tmpFile := "test_zero_max_requests.json" + content := `{ + "ethereum": { + "zero-max": { + "provider": "alchemy", + "role": "primary", + "type": "full", + "http_url": "http://test.com", + "capacity": {"window_seconds": 10} + }, + "negative-max": { + "provider": "alchemy", + "role": "primary", + "type": "full", + "http_url": "http://test2.com", + "capacity": {"max_requests": -5, "window_seconds": 10} + }, + "valid-max": { + "provider": "alchemy", + "role": "primary", + "type": "full", + "http_url": "http://test3.com", + "capacity": {"max_requests": 100, "window_seconds": 10} + } + } + }` + if err := os.WriteFile(tmpFile, []byte(content), 0644); err != nil { + t.Fatalf("Failed to create test file: %v", err) + } + defer os.Remove(tmpFile) + + cfg, err := LoadConfig(tmpFile) + if err != nil { + t.Fatalf("LoadConfig failed: %v", err) + } + + endpoints := cfg.Endpoints["ethereum"] + + if endpoints["zero-max"].Capacity != nil { + t.Error("Expected Capacity to be disabled (nil) for an endpoint with omitted (zero) max_requests") + } + if endpoints["negative-max"].Capacity != nil { + t.Error("Expected Capacity to be disabled (nil) for an endpoint with negative max_requests") + } + if endpoints["valid-max"].Capacity == nil { + t.Fatal("Expected Capacity to remain set for an endpoint with a valid max_requests") + } + if endpoints["valid-max"].Capacity.MaxRequests != 100 { + t.Errorf("Expected valid endpoint's MaxRequests to be untouched (100), got %d", endpoints["valid-max"].Capacity.MaxRequests) + } +} + func TestDefaultRateLimitRecovery(t *testing.T) { config := DefaultRateLimitRecovery() From 8db395436fcd2cfb2bcf4e5438e4f94ef7f86a24 Mon Sep 17 00:00:00 2001 From: Jacob Sussmilch Date: Wed, 1 Jul 2026 15:59:47 +1000 Subject: [PATCH 07/12] fix: KEEP-905 dedupe capacity-decrease and backoff-seeding logic The AIMD capacity-decrease sequence and the recovery-backoff-seeding logic were duplicated verbatim between the load balancer (internal/server) and the standalone health checker (services/health-checker) - two separate processes that both mutate the same Valkey-persisted capacity estimate. Hand-duplicated copies are a correctness risk: any future change to the math or eligibility rules would need to be applied identically in both places, and a missed spot would make the two processes silently disagree about what the learned ceiling means for a given endpoint. Extract the decrease sequence into store.ApplyLearnedCapacityDecreaseIfEligible (internal/store/capacity_estimate.go) and the backoff-seeding logic into health.InitialBackoffForSignal (internal/health/ratelimit.go). Both server.Server.applyLearnedCapacityDecrease/initialBackoffForSignal and services/health-checker/main.go's applyStandaloneLearnedCapacityDecrease/ standaloneInitialBackoff are now thin wrappers delegating to these shared functions, preserving their existing signatures so all existing tests keep passing unmodified. As a byproduct of sharing the decrease sequence, the standalone health checker's wrapper now applies the same 2-second context timeout the load balancer already used, instead of an unbounded context.Background() - a Valkey call hanging in the standalone checker no longer risks blocking indefinitely. --- internal/health/ratelimit.go | 33 +++++++++++ internal/server/server.go | 85 +++-------------------------- internal/store/capacity_estimate.go | 63 +++++++++++++++++++++ services/health-checker/main.go | 85 ++++------------------------- 4 files changed, 116 insertions(+), 150 deletions(-) diff --git a/internal/health/ratelimit.go b/internal/health/ratelimit.go index 7e87bcc..0592718 100644 --- a/internal/health/ratelimit.go +++ b/internal/health/ratelimit.go @@ -5,6 +5,8 @@ import ( "strconv" "strings" "time" + + "aetherlay/internal/config" ) // RateLimitSignal describes a detected rate-limit condition and any recovery timing @@ -59,3 +61,34 @@ func DetectRateLimit(provider string, statusCode int, headers http.Header, rpcRe return sig } + +// InitialBackoffForSignal computes the first recovery-check wait time from whatever the +// provider actually told us, instead of always guessing from the configured +// RateLimitRecovery.InitialBackoff: +// - A parsed Retry-After is the most precise signal available and is used directly. +// - Infura's daily credit cap (402) can't be sped up by probing sooner, since it only +// resets once the day rolls over - Infura's docs don't guarantee an exact reset +// boundary, so rather than assume one, this seeds at the endpoint's own configured +// (or default) MaxBackoff, so a known-exhausted daily quota isn't re-probed on a +// short cycle. +// - Otherwise 0, which leaves the caller's recovery scheduler to fall back to +// InitialBackoff, exactly as it did before this signal existed. +// +// This is shared verbatim between the load balancer (internal/server) and the standalone +// health checker (services/health-checker) - both seed the same recovery backoff from the +// same signal, so they must apply identical logic. +func InitialBackoffForSignal(cfg *config.Config, chain, endpointID string, signal RateLimitSignal) int { + if signal.RetryAfter > 0 { + return int(signal.RetryAfter.Seconds()) + } + if signal.IsDailyQuota { + rlc := config.DefaultRateLimitRecovery() + if chainEndpoints, ok := cfg.GetEndpointsForChain(chain); ok { + if ep, ok := chainEndpoints[endpointID]; ok && ep.RateLimitRecovery != nil && ep.RateLimitRecovery.MaxBackoff != 0 { + rlc.MaxBackoff = ep.RateLimitRecovery.MaxBackoff + } + } + return rlc.MaxBackoff + } + return 0 +} diff --git a/internal/server/server.go b/internal/server/server.go index c724f2d..0aa0875 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -1385,63 +1385,13 @@ func (s *Server) recordCapacityUsage(chain string, endpoint config.Endpoint, end } // applyLearnedCapacityDecrease is called from handleRateLimit on every confirmed -// rate-limit signal. It only ever engages for endpoints with no static Capacity -// configured (static config, if present, is left completely untouched - see -// effectiveCapacityCeiling). A cooldown (the learning window itself) collapses several -// near-simultaneous hits from one episode into a single decrease. +// rate-limit signal. It delegates to store.ApplyLearnedCapacityDecreaseIfEligible, the +// same shared implementation the standalone health checker calls, so the two processes - +// both mutating the same Valkey-persisted estimate - never diverge. func (s *Server) applyLearnedCapacityDecrease(chain, endpointID string, endpoint config.Endpoint) { - if !s.appConfig.CapacityThrottlingEnabled || !s.appConfig.CapacityLearningEnabled || endpoint.Capacity != nil { - return - } - - params := config.ResolveCapacityLearning(endpoint.CapacityLearning) - now := time.Now() ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - - prior, err := s.valkeyClient.GetCapacityEstimate(ctx, chain, endpointID) - if err != nil { - log.Debug().Err(err).Str("chain", chain).Str("endpoint", endpointID).Msg("Failed to get capacity estimate") - return - } - - // Read against whichever window the usage counter is actually being written to right - // now (see capacityWindowSeconds): the estimate's own frozen window once one exists, - // not the live-resolved config - otherwise this read targets a different Valkey - // bucket key than recordCapacityUsage's writes, and observedCount below would always - // be stale/zero. Only before any estimate has been seeded is there no frozen value to - // match, so the live-resolved window is the correct (and only) choice at that point. - readWindowSeconds := params.WindowSeconds - if prior.HasEstimate { - readWindowSeconds = prior.WindowSeconds - } - - if !store.ShouldApplyCapacityDecrease(*prior, readWindowSeconds, now) { - log.Debug().Str("chain", chain).Str("endpoint", endpointID).Msg("Skipping capacity estimate decrease, within cooldown of the last decrease") - return - } - - observedCount, err := s.valkeyClient.GetCapacityCount(ctx, chain, endpointID, readWindowSeconds) - if err != nil { - log.Debug().Err(err).Str("chain", chain).Str("endpoint", endpointID).Msg("Failed to get capacity count for estimate decrease") - return - } - - effectiveNow := store.EffectiveMaxRequests(*prior, params, now) - newEstimate := store.ApplyCapacityDecrease(*prior, effectiveNow, observedCount, params.WindowSeconds, params, now) - - if err := s.valkeyClient.SetCapacityEstimate(ctx, chain, endpointID, newEstimate); err != nil { - log.Debug().Err(err).Str("chain", chain).Str("endpoint", endpointID).Msg("Failed to set capacity estimate") - return - } - - log.Info().Str("chain", chain).Str("endpoint", endpointID).Int64("new_estimate", newEstimate.MaxRequests).Int("window_seconds", newEstimate.WindowSeconds).Msg("Decreased learned capacity estimate after a rate-limit hit") - if metrics.EndpointCapacityEstimatedCeiling != nil { - metrics.EndpointCapacityEstimatedCeiling.WithLabelValues(chain, endpointID).Set(float64(newEstimate.MaxRequests)) - } - if metrics.EndpointCapacityEstimateDecreasedTotal != nil { - metrics.EndpointCapacityEstimateDecreasedTotal.WithLabelValues(chain, endpointID).Inc() - } + store.ApplyLearnedCapacityDecreaseIfEligible(ctx, s.valkeyClient, chain, endpointID, endpoint, s.appConfig.CapacityThrottlingEnabled, s.appConfig.CapacityLearningEnabled) } // bodyCarriesRateLimitSignal reports whether a 2xx JSON-RPC response body (a single @@ -1746,28 +1696,9 @@ func (s *Server) handleRateLimit(chain, endpointID, protocol string, signal heal s.rateLimitScheduler.StartMonitoring(chain, endpointID) } -// initialBackoffForSignal computes the first recovery-check wait time from whatever the -// provider actually told us, instead of always guessing via the scheduler's InitialBackoff: -// - A parsed Retry-After is the most precise signal available and is used directly. -// - Infura's daily credit cap (402) can't be sped up by probing sooner, since it only -// resets once the day rolls over - Infura's docs don't guarantee an exact reset -// boundary, so rather than assume one, this seeds at the endpoint's own configured -// (or default) MaxBackoff, so a known-exhausted daily quota isn't re-probed on a -// short cycle. -// - Otherwise 0, which leaves rate_limit_scheduler.go's calculateNextBackoff to fall -// back to InitialBackoff, exactly as it did before this signal existed. +// initialBackoffForSignal delegates to health.InitialBackoffForSignal, the same shared +// implementation the standalone health checker calls, so the two processes seed recovery +// backoff identically from the same signal. func (s *Server) initialBackoffForSignal(chain, endpointID string, signal health.RateLimitSignal) int { - if signal.RetryAfter > 0 { - return int(signal.RetryAfter.Seconds()) - } - if signal.IsDailyQuota { - rlc := config.DefaultRateLimitRecovery() - if chainEndpoints, ok := s.config.GetEndpointsForChain(chain); ok { - if ep, ok := chainEndpoints[endpointID]; ok && ep.RateLimitRecovery != nil && ep.RateLimitRecovery.MaxBackoff != 0 { - rlc.MaxBackoff = ep.RateLimitRecovery.MaxBackoff - } - } - return rlc.MaxBackoff - } - return 0 + return health.InitialBackoffForSignal(s.config, chain, endpointID, signal) } diff --git a/internal/store/capacity_estimate.go b/internal/store/capacity_estimate.go index ddc3f15..065f357 100644 --- a/internal/store/capacity_estimate.go +++ b/internal/store/capacity_estimate.go @@ -6,7 +6,9 @@ import ( "time" "aetherlay/internal/config" + "aetherlay/internal/metrics" + "github.com/rs/zerolog/log" "github.com/valkey-io/valkey-go" ) @@ -133,3 +135,64 @@ func (r *ValkeyClient) SetCapacityEstimate(ctx context.Context, chain, endpoint cmd := r.client.B().Set().Key(key).Value(string(jsonBytes)).Build() return r.client.Do(ctx, cmd).Error() } + +// ApplyLearnedCapacityDecreaseIfEligible is called on every confirmed rate-limit signal. +// It only ever engages for endpoints with no static Capacity configured - static config, +// if present, is left completely untouched. A cooldown (the learning window itself) +// collapses several near-simultaneous hits from one episode into a single decrease. +// +// This is shared verbatim between the load balancer (internal/server) and the standalone +// health checker (services/health-checker) - both mutate the same Valkey-persisted +// estimate for a given endpoint, so they must apply identical math or silently disagree +// about what the learned ceiling means for that endpoint. +func ApplyLearnedCapacityDecreaseIfEligible(ctx context.Context, valkeyClient ValkeyClientIface, chain, endpointID string, ep config.Endpoint, capacityThrottlingEnabled, capacityLearningEnabled bool) { + if !capacityThrottlingEnabled || !capacityLearningEnabled || ep.Capacity != nil { + return + } + + params := config.ResolveCapacityLearning(ep.CapacityLearning) + now := time.Now() + + prior, err := valkeyClient.GetCapacityEstimate(ctx, chain, endpointID) + if err != nil { + log.Debug().Err(err).Str("chain", chain).Str("endpoint", endpointID).Msg("Failed to get capacity estimate") + return + } + + // Read against whichever window the usage counter is actually being written to right + // now: the estimate's own frozen window once one exists, not the live-resolved + // config - otherwise this read targets a different Valkey bucket key than the + // dispatch-time usage writes, and observedCount below would always be stale/zero. + // Only before any estimate has been seeded is there no frozen value to match. + readWindowSeconds := params.WindowSeconds + if prior.HasEstimate { + readWindowSeconds = prior.WindowSeconds + } + + if !ShouldApplyCapacityDecrease(*prior, readWindowSeconds, now) { + log.Debug().Str("chain", chain).Str("endpoint", endpointID).Msg("Skipping capacity estimate decrease, within cooldown of the last decrease") + return + } + + observedCount, err := valkeyClient.GetCapacityCount(ctx, chain, endpointID, readWindowSeconds) + if err != nil { + log.Debug().Err(err).Str("chain", chain).Str("endpoint", endpointID).Msg("Failed to get capacity count for estimate decrease") + return + } + + effectiveNow := EffectiveMaxRequests(*prior, params, now) + newEstimate := ApplyCapacityDecrease(*prior, effectiveNow, observedCount, params.WindowSeconds, params, now) + + if err := valkeyClient.SetCapacityEstimate(ctx, chain, endpointID, newEstimate); err != nil { + log.Debug().Err(err).Str("chain", chain).Str("endpoint", endpointID).Msg("Failed to set capacity estimate") + return + } + + log.Info().Str("chain", chain).Str("endpoint", endpointID).Int64("new_estimate", newEstimate.MaxRequests).Int("window_seconds", newEstimate.WindowSeconds).Msg("Decreased learned capacity estimate after a rate-limit hit") + if metrics.EndpointCapacityEstimatedCeiling != nil { + metrics.EndpointCapacityEstimatedCeiling.WithLabelValues(chain, endpointID).Set(float64(newEstimate.MaxRequests)) + } + if metrics.EndpointCapacityEstimateDecreasedTotal != nil { + metrics.EndpointCapacityEstimateDecreasedTotal.WithLabelValues(chain, endpointID).Inc() + } +} diff --git a/services/health-checker/main.go b/services/health-checker/main.go index ef0b7ac..bba6ea5 100644 --- a/services/health-checker/main.go +++ b/services/health-checker/main.go @@ -36,33 +36,18 @@ var testExitAfterSetup bool // exitCode is used to track the exit code for the process var exitCode int -// standaloneInitialBackoff mirrors server.Server.initialBackoffForSignal: seed the first -// recovery-check wait from whatever the provider actually told us, instead of always -// guessing. A parsed Retry-After is used directly; Infura's daily credit cap (402) seeds -// from the endpoint's own (or default) MaxBackoff, since the docs don't guarantee an -// exact reset boundary. Otherwise 0, leaving the load balancer's scheduler to fall back -// to InitialBackoff once it picks up recovery monitoring. +// standaloneInitialBackoff delegates to health.InitialBackoffForSignal, the same shared +// implementation server.Server.initialBackoffForSignal calls, so the load balancer and +// the standalone health checker seed recovery backoff identically from the same signal. func standaloneInitialBackoff(cfg *config.Config, chain, endpointID string, signal health.RateLimitSignal) int { - if signal.RetryAfter > 0 { - return int(signal.RetryAfter.Seconds()) - } - if signal.IsDailyQuota { - rlc := config.DefaultRateLimitRecovery() - if chainEndpoints, ok := cfg.GetEndpointsForChain(chain); ok { - if ep, ok := chainEndpoints[endpointID]; ok && ep.RateLimitRecovery != nil && ep.RateLimitRecovery.MaxBackoff != 0 { - rlc.MaxBackoff = ep.RateLimitRecovery.MaxBackoff - } - } - return rlc.MaxBackoff - } - return 0 + return health.InitialBackoffForSignal(cfg, chain, endpointID, signal) } -// applyStandaloneLearnedCapacityDecrease mirrors server.Server.applyLearnedCapacityDecrease -// exactly, using the same shared store.ApplyCapacityDecrease/EffectiveMaxRequests math, so -// the load balancer and the standalone health checker - two separate processes mutating -// the same Valkey-persisted estimate - never diverge. Only engages for endpoints with no -// static Capacity configured, mirroring effectiveCapacityCeiling's rule in server.go. +// applyStandaloneLearnedCapacityDecrease resolves chain/endpointID to a config.Endpoint +// and delegates to store.ApplyLearnedCapacityDecreaseIfEligible, the same shared +// implementation server.Server.applyLearnedCapacityDecrease calls, so the load balancer +// and the standalone health checker - two separate processes mutating the same +// Valkey-persisted estimate - never diverge. func applyStandaloneLearnedCapacityDecrease(cfg *config.Config, valkeyClient store.ValkeyClientIface, capacityThrottlingEnabled, capacityLearningEnabled bool, chain, endpointID string) { chainEndpoints, ok := cfg.GetEndpointsForChain(chain) if !ok { @@ -72,56 +57,10 @@ func applyStandaloneLearnedCapacityDecrease(cfg *config.Config, valkeyClient sto if !ok { return } - if !capacityThrottlingEnabled || !capacityLearningEnabled || ep.Capacity != nil { - return - } - - params := config.ResolveCapacityLearning(ep.CapacityLearning) - now := time.Now() - ctx := context.Background() - - prior, err := valkeyClient.GetCapacityEstimate(ctx, chain, endpointID) - if err != nil { - log.Debug().Err(err).Str("chain", chain).Str("endpoint", endpointID).Msg("Failed to get capacity estimate in standalone health checker") - return - } - - // Read against the estimate's own frozen window once one exists (matching - // server.Server.capacityWindowSeconds/effectiveCapacityCeiling exactly), not the - // live-resolved config - otherwise this read targets a different Valkey bucket key - // than the load balancer's writes. Only before any estimate has been seeded is there - // no frozen value to match. - readWindowSeconds := params.WindowSeconds - if prior.HasEstimate { - readWindowSeconds = prior.WindowSeconds - } - - if !store.ShouldApplyCapacityDecrease(*prior, readWindowSeconds, now) { - log.Debug().Str("chain", chain).Str("endpoint", endpointID).Msg("Skipping capacity estimate decrease in standalone health checker, within cooldown") - return - } - - observedCount, err := valkeyClient.GetCapacityCount(ctx, chain, endpointID, readWindowSeconds) - if err != nil { - log.Debug().Err(err).Str("chain", chain).Str("endpoint", endpointID).Msg("Failed to get capacity count in standalone health checker") - return - } - - effectiveNow := store.EffectiveMaxRequests(*prior, params, now) - newEstimate := store.ApplyCapacityDecrease(*prior, effectiveNow, observedCount, params.WindowSeconds, params, now) - if err := valkeyClient.SetCapacityEstimate(ctx, chain, endpointID, newEstimate); err != nil { - log.Debug().Err(err).Str("chain", chain).Str("endpoint", endpointID).Msg("Failed to set capacity estimate in standalone health checker") - return - } - - log.Info().Str("chain", chain).Str("endpoint", endpointID).Int64("new_estimate", newEstimate.MaxRequests).Msg("Standalone health checker decreased learned capacity estimate after a rate-limit hit") - if metrics.EndpointCapacityEstimatedCeiling != nil { - metrics.EndpointCapacityEstimatedCeiling.WithLabelValues(chain, endpointID).Set(float64(newEstimate.MaxRequests)) - } - if metrics.EndpointCapacityEstimateDecreasedTotal != nil { - metrics.EndpointCapacityEstimateDecreasedTotal.WithLabelValues(chain, endpointID).Inc() - } + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + store.ApplyLearnedCapacityDecreaseIfEligible(ctx, valkeyClient, chain, endpointID, ep, capacityThrottlingEnabled, capacityLearningEnabled) } // createStandaloneRateLimitHandler creates a simple rate limit handler for the standalone health checker From 665731bcd576d37261893405057763732ac2ac8b Mon Sep 17 00:00:00 2001 From: Jacob Sussmilch Date: Wed, 1 Jul 2026 16:01:41 +1000 Subject: [PATCH 08/12] fix: KEEP-905 keep learned capacity window frozen across decreases ApplyLearnedCapacityDecreaseIfEligible correctly read against an existing estimate's frozen WindowSeconds (fixed previously for the read/observedCount path), but the decrease it persisted still passed the live-resolved params.WindowSeconds to ApplyCapacityDecrease instead of the frozen value. So each rate-limit hit after an operator changed an endpoint's capacity_learning.window_seconds would silently re-freeze the estimate to the new value, contradicting the documented invariant that the window is frozen for the estimate's lifetime once seeded, not just stable between reads. Reuse the same frozen-if-exists windowSeconds value for both the read (cooldown check, GetCapacityCount) and the write (ApplyCapacityDecrease), so the persisted window can only ever be set once, at first seed. Found by CodeRabbit's review of PR #37 (README.md:510's documented invariant). --- internal/server/capacity_learning_test.go | 60 +++++++++++++++++++++++ internal/store/capacity_estimate.go | 23 +++++---- 2 files changed, 73 insertions(+), 10 deletions(-) diff --git a/internal/server/capacity_learning_test.go b/internal/server/capacity_learning_test.go index a4dfa80..0f7fc9f 100644 --- a/internal/server/capacity_learning_test.go +++ b/internal/server/capacity_learning_test.go @@ -453,3 +453,63 @@ func TestApplyLearnedCapacityDecreaseReadsFrozenWindowNotLiveConfig(t *testing.T t.Errorf("Expected decrease to be grounded in the frozen window's observed count (8 -> 4), got MaxRequests=%d", final.MaxRequests) } } + +func TestApplyLearnedCapacityDecreasePersistsFrozenWindowNotLiveConfig(t *testing.T) { + cfg := &config.Config{ + Endpoints: map[string]config.ChainEndpoints{ + "ethereum": { + "ep1": config.Endpoint{ + Provider: "alchemy", + Role: "primary", + Type: "full", + HTTPURL: "http://ep1", + CapacityLearning: &config.CapacityLearning{WindowSeconds: 30}, // live value, differs from frozen + }, + }, + }, + } + valkeyClient := store.NewMockValkeyClient() + ctx := context.Background() + ep := cfg.Endpoints["ethereum"]["ep1"] + + // Seed an estimate frozen at 60s, with its cooldown already elapsed. + if err := valkeyClient.SetCapacityEstimate(ctx, "ethereum", "ep1", store.CapacityEstimate{ + HasEstimate: true, MaxRequests: 10, IncreaseStep: 1, WindowSeconds: 60, + LastDecreaseAt: time.Now().Add(-120 * time.Second), + }); err != nil { + t.Fatalf("SetCapacityEstimate failed: %v", err) + } + valkeyClient.IncrementCapacityCount(ctx, "ethereum", "ep1", 60) + + server := NewServer(cfg, valkeyClient, learningTestConfig()) + server.applyLearnedCapacityDecrease("ethereum", "ep1", ep) + + final, err := valkeyClient.GetCapacityEstimate(ctx, "ethereum", "ep1") + if err != nil { + t.Fatalf("GetCapacityEstimate failed: %v", err) + } + // The persisted estimate's WindowSeconds must stay frozen at the value it was seeded + // with (60), not get re-synced to the live config value (30) on this decrease - the + // window is documented as permanently frozen for the estimate's lifetime, not just + // stable between reads. + if final.WindowSeconds != 60 { + t.Errorf("Expected the persisted estimate to keep its frozen WindowSeconds (60), got %d", final.WindowSeconds) + } + + // A second decrease after another elapsed cooldown must still persist the original + // frozen window, not the live one - the freeze must survive across repeated decreases. + valkeyClient.SetCapacityEstimate(ctx, "ethereum", "ep1", store.CapacityEstimate{ + HasEstimate: true, MaxRequests: final.MaxRequests, IncreaseStep: final.IncreaseStep, WindowSeconds: final.WindowSeconds, + LastDecreaseAt: time.Now().Add(-120 * time.Second), + }) + valkeyClient.IncrementCapacityCount(ctx, "ethereum", "ep1", 60) + server.applyLearnedCapacityDecrease("ethereum", "ep1", ep) + + final2, err := valkeyClient.GetCapacityEstimate(ctx, "ethereum", "ep1") + if err != nil { + t.Fatalf("GetCapacityEstimate failed: %v", err) + } + if final2.WindowSeconds != 60 { + t.Errorf("Expected the frozen window to survive a second decrease (60), got %d", final2.WindowSeconds) + } +} diff --git a/internal/store/capacity_estimate.go b/internal/store/capacity_estimate.go index 065f357..0b17dd8 100644 --- a/internal/store/capacity_estimate.go +++ b/internal/store/capacity_estimate.go @@ -159,29 +159,32 @@ func ApplyLearnedCapacityDecreaseIfEligible(ctx context.Context, valkeyClient Va return } - // Read against whichever window the usage counter is actually being written to right - // now: the estimate's own frozen window once one exists, not the live-resolved - // config - otherwise this read targets a different Valkey bucket key than the - // dispatch-time usage writes, and observedCount below would always be stale/zero. - // Only before any estimate has been seeded is there no frozen value to match. - readWindowSeconds := params.WindowSeconds + // Once an estimate exists, its own WindowSeconds is permanently authoritative - not + // just for this read, but for the estimate this decrease persists below. Otherwise a + // later config change to the learning window would both target a different Valkey + // bucket key than the dispatch-time usage writes (making observedCount stale/zero) + // and silently re-freeze the estimate to the new value, contradicting the documented + // invariant that the window is frozen for the lifetime of the estimate. Only before + // any estimate has been seeded is there no frozen value yet, so the live-resolved + // config is used - and that first decrease is what freezes it from then on. + windowSeconds := params.WindowSeconds if prior.HasEstimate { - readWindowSeconds = prior.WindowSeconds + windowSeconds = prior.WindowSeconds } - if !ShouldApplyCapacityDecrease(*prior, readWindowSeconds, now) { + if !ShouldApplyCapacityDecrease(*prior, windowSeconds, now) { log.Debug().Str("chain", chain).Str("endpoint", endpointID).Msg("Skipping capacity estimate decrease, within cooldown of the last decrease") return } - observedCount, err := valkeyClient.GetCapacityCount(ctx, chain, endpointID, readWindowSeconds) + observedCount, err := valkeyClient.GetCapacityCount(ctx, chain, endpointID, windowSeconds) if err != nil { log.Debug().Err(err).Str("chain", chain).Str("endpoint", endpointID).Msg("Failed to get capacity count for estimate decrease") return } effectiveNow := EffectiveMaxRequests(*prior, params, now) - newEstimate := ApplyCapacityDecrease(*prior, effectiveNow, observedCount, params.WindowSeconds, params, now) + newEstimate := ApplyCapacityDecrease(*prior, effectiveNow, observedCount, windowSeconds, params, now) if err := valkeyClient.SetCapacityEstimate(ctx, chain, endpointID, newEstimate); err != nil { log.Debug().Err(err).Str("chain", chain).Str("endpoint", endpointID).Msg("Failed to set capacity estimate") From 53e44a7b16fa25f5f54f04500cb8affde62fc8e8 Mon Sep 17 00:00:00 2001 From: Jacob Sussmilch Date: Wed, 1 Jul 2026 16:06:31 +1000 Subject: [PATCH 09/12] fix: KEEP-905 bound capacity Valkey lookups to a request-scoped timeout capacityWindowSeconds, effectiveCapacityCeiling, and the capacity-ceiling gating check in getEndpointsByRole all called GetCapacityEstimate/ GetCapacityCount with context.Background(), so a slow or unresponsive Valkey could stall endpoint selection indefinitely instead of failing fast within the request's own timeout budget. Wrap each of these three call sites in a 2-second context.WithTimeout, matching the bound already used for the capacity-decrease and usage-recording paths (recordCapacityUsage, applyLearnedCapacityDecrease). Also removed a stale duplicate doc comment on capacityWindowSeconds left over from an earlier commit in this branch. Found by CodeRabbit's review of PR #37 (internal/server/server.go:888,1024). --- internal/server/capacity_learning_test.go | 55 +++++++++++++++++++++++ internal/server/server.go | 18 ++++---- 2 files changed, 65 insertions(+), 8 deletions(-) diff --git a/internal/server/capacity_learning_test.go b/internal/server/capacity_learning_test.go index 0f7fc9f..2595390 100644 --- a/internal/server/capacity_learning_test.go +++ b/internal/server/capacity_learning_test.go @@ -513,3 +513,58 @@ func TestApplyLearnedCapacityDecreasePersistsFrozenWindowNotLiveConfig(t *testin t.Errorf("Expected the frozen window to survive a second decrease (60), got %d", final2.WindowSeconds) } } + +// deadlineCheckingValkeyClient wraps MockValkeyClient to record whether any capacity +// lookup arrived with a context that has no deadline - i.e. context.Background(), which +// can let a slow Valkey stall endpoint selection indefinitely instead of bounding it to +// the request's timeout budget. +type deadlineCheckingValkeyClient struct { + *store.MockValkeyClient + sawNoDeadline bool +} + +func (c *deadlineCheckingValkeyClient) GetCapacityCount(ctx context.Context, chain, endpoint string, windowSeconds int) (int64, error) { + if _, ok := ctx.Deadline(); !ok { + c.sawNoDeadline = true + } + return c.MockValkeyClient.GetCapacityCount(ctx, chain, endpoint, windowSeconds) +} + +func (c *deadlineCheckingValkeyClient) GetCapacityEstimate(ctx context.Context, chain, endpoint string) (*store.CapacityEstimate, error) { + if _, ok := ctx.Deadline(); !ok { + c.sawNoDeadline = true + } + return c.MockValkeyClient.GetCapacityEstimate(ctx, chain, endpoint) +} + +func TestCapacityValkeyLookupsUseRequestScopedTimeouts(t *testing.T) { + cfg := &config.Config{ + Endpoints: map[string]config.ChainEndpoints{ + "ethereum": { + "ep1": config.Endpoint{Provider: "alchemy", Role: "primary", Type: "full", HTTPURL: "http://ep1"}, + }, + }, + } + inner := store.NewMockValkeyClient() + inner.PopulateStatuses(map[string]*store.EndpointStatus{ + "ethereum:ep1": {HasHTTP: true, HealthyHTTP: true}, + }) + ctx := context.Background() + if err := inner.SetCapacityEstimate(ctx, "ethereum", "ep1", store.CapacityEstimate{ + HasEstimate: true, MaxRequests: 5, IncreaseStep: 1, WindowSeconds: 60, LastDecreaseAt: time.Now(), + }); err != nil { + t.Fatalf("SetCapacityEstimate failed: %v", err) + } + + client := &deadlineCheckingValkeyClient{MockValkeyClient: inner} + server := NewServer(cfg, client, learningTestConfig()) + + // Exercises effectiveCapacityCeiling + GetCapacityCount via getEndpointsByRole. + server.getAvailableEndpoints("ethereum", false, false) + // Exercises capacityWindowSeconds directly. + server.capacityWindowSeconds("ethereum", "ep1", cfg.Endpoints["ethereum"]["ep1"]) + + if client.sawNoDeadline { + t.Error("Expected all capacity Valkey lookups in the selection path to use a context with a deadline, not context.Background()") + } +} diff --git a/internal/server/server.go b/internal/server/server.go index 0aa0875..f6b5a41 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -885,7 +885,9 @@ func (s *Server) getEndpointsByRole(chainEndpoints config.ChainEndpoints, role s // a provider signal. if s.appConfig.CapacityThrottlingEnabled { if maxRequests, windowSeconds, hasCeiling := s.effectiveCapacityCeiling(chain, endpointID, endpoint); hasCeiling { - count, err := s.valkeyClient.GetCapacityCount(context.Background(), chain, endpointID, windowSeconds) + capCtx, capCancel := context.WithTimeout(context.Background(), 2*time.Second) + count, err := s.valkeyClient.GetCapacityCount(capCtx, chain, endpointID, windowSeconds) + capCancel() if err == nil { if metrics.EndpointCapacityUtilization != nil { metrics.EndpointCapacityUtilization.WithLabelValues(chain, endpointID).Set(float64(count) / float64(maxRequests)) @@ -1021,7 +1023,9 @@ func (s *Server) effectiveCapacityCeiling(chain, endpointID string, ep config.En if !s.appConfig.CapacityLearningEnabled { return 0, 0, false } - estimate, err := s.valkeyClient.GetCapacityEstimate(context.Background(), chain, endpointID) + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + estimate, err := s.valkeyClient.GetCapacityEstimate(ctx, chain, endpointID) if err != nil || !estimate.HasEstimate { return 0, 0, false } @@ -1339,11 +1343,6 @@ func (s *Server) providerForEndpoint(chain, endpointID string) string { return chainEndpoints[endpointID].Provider } -// capacityWindowSeconds resolves the window width to track usage against: the static -// CapacityLimit's window if configured, else the (possibly overridden) adaptive -// learning window - so an endpoint with no static Capacity still accumulates real -// GetCapacityCount evidence before it's ever rate limited, for applyLearnedCapacityDecrease -// to seed an estimate from on the first hit. // capacityWindowSeconds resolves the window width to track usage against for the WRITE // path (recordCapacityUsage's usage counter). It must agree with whatever window // effectiveCapacityCeiling's READ path is watching, or gating silently stops working - @@ -1358,7 +1357,10 @@ func (s *Server) capacityWindowSeconds(chain, endpointID string, endpoint config return endpoint.Capacity.WindowSeconds } if s.appConfig.CapacityLearningEnabled { - if estimate, err := s.valkeyClient.GetCapacityEstimate(context.Background(), chain, endpointID); err == nil && estimate.HasEstimate { + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + estimate, err := s.valkeyClient.GetCapacityEstimate(ctx, chain, endpointID) + cancel() + if err == nil && estimate.HasEstimate { return estimate.WindowSeconds } } From 2564bd46acd6e7df994e41fe654bef2850417b7c Mon Sep 17 00:00:00 2001 From: Jacob Sussmilch Date: Wed, 1 Jul 2026 16:09:21 +1000 Subject: [PATCH 10/12] fix: KEEP-905 scope embedded rate-limit body scan to Alchemy only Every 2xx JSON-RPC response through the proxy was being fully buffered (up to maxRateLimitScanBodyBytes = 10MB) to scan for a rate-limit error embedded in the body - a quirk specific to Alchemy on batch requests, which can return HTTP 200 with the error only inside the JSON-RPC payload. Applying this to every provider regressed streaming for large legitimate responses (e.g. eth_getLogs), delaying time-to-first-byte until the full body was read, and held up to 10MB in memory per in-flight response regardless of provider. Add providerNeedsEmbeddedRateLimitScan, resolving the endpoint's provider before deciding whether to buffer at all. Only Alchemy responses are buffered and scanned; every other provider's 2xx JSON response streams through unmodified and uninspected, exactly as it did before this feature was added. Found by CodeRabbit's review of PR #37 (internal/server/server.go:1297). --- internal/server/server.go | 43 +++++++++++++++-------- internal/server/server_test.go | 62 ++++++++++++++++++++++++++++++++++ 2 files changed, 90 insertions(+), 15 deletions(-) diff --git a/internal/server/server.go b/internal/server/server.go index f6b5a41..748cd60 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -1259,11 +1259,15 @@ func (s *Server) defaultForwardRequestWithBodyFunc(w http.ResponseWriter, ctx co return fmt.Errorf("HTTP %d: %s", resp.StatusCode, resp.Status) } - // For JSON responses within a bounded size, buffer and scan the body for an embedded - // rate-limit error before forwarding. Some providers (e.g. Alchemy on batch requests) - // return HTTP 200 with the rate-limit signal only inside the JSON-RPC body - invisible - // to a status-code-only check. Larger bodies are streamed through unmodified and - // uninspected, to bound memory use on large responses (e.g. eth_getLogs). + // For JSON responses from a provider known to need it, buffer (within a bounded size) + // and scan the body for an embedded rate-limit error before forwarding. Some providers + // (currently Alchemy, on batch requests) return HTTP 200 with the rate-limit signal + // only inside the JSON-RPC body - invisible to a status-code-only check. This is + // scoped to those providers specifically: buffering every 2xx JSON response + // regardless of provider would delay time-to-first-byte and hold up to + // maxRateLimitScanBodyBytes in memory per in-flight request, for providers that never + // exhibit this behavior. Responses from other providers, and oversized bodies from + // providers that do need the scan, are streamed through unmodified and uninspected. // // Headers are copied onto w only once we've committed to writing this specific // response (immediately before each w.WriteHeader below), not unconditionally up @@ -1273,18 +1277,17 @@ func (s *Server) defaultForwardRequestWithBodyFunc(w http.ResponseWriter, ctx co // in w's header map, where a subsequent successful retry's headers would be added on // top of rather than replacing them - e.g. two Content-Length values in the response // actually sent to the client. - if isJSONContentType(resp.Header.Get("Content-Type")) { + chain, endpointID, endpointFound := s.findChainAndEndpointByURL(targetURL) + if endpointFound && isJSONContentType(resp.Header.Get("Content-Type")) && providerNeedsEmbeddedRateLimitScan(s.providerForEndpoint(chain, endpointID)) { buffered, readErr := io.ReadAll(io.LimitReader(resp.Body, maxRateLimitScanBodyBytes+1)) if readErr == nil && int64(len(buffered)) <= maxRateLimitScanBodyBytes { - if chain, endpointID, found := s.findChainAndEndpointByURL(targetURL); found { - if sig := bodyCarriesRateLimitSignal(s.providerForEndpoint(chain, endpointID), buffered, resp.Header); sig.IsRateLimited { - // The HTTP transaction itself succeeded (2xx) - don't mark the endpoint - // unhealthy, just flag it as rate limited so selection avoids it, and - // retry the same buffered request body against a different endpoint. - s.handleRateLimit(chain, endpointID, "http", sig) - log.Debug().Str("url", helpers.RedactAPIKey(targetURL)).Bool("daily_quota", sig.IsDailyQuota).Msg("2xx response carried an embedded rate-limit signal, retrying with a different endpoint") - return fmt.Errorf("rate limited: embedded JSON-RPC rate-limit error in 2xx response") - } + if sig := bodyCarriesRateLimitSignal(s.providerForEndpoint(chain, endpointID), buffered, resp.Header); sig.IsRateLimited { + // The HTTP transaction itself succeeded (2xx) - don't mark the endpoint + // unhealthy, just flag it as rate limited so selection avoids it, and + // retry the same buffered request body against a different endpoint. + s.handleRateLimit(chain, endpointID, "http", sig) + log.Debug().Str("url", helpers.RedactAPIKey(targetURL)).Bool("daily_quota", sig.IsDailyQuota).Msg("2xx response carried an embedded rate-limit signal, retrying with a different endpoint") + return fmt.Errorf("rate limited: embedded JSON-RPC rate-limit error in 2xx response") } copyResponseHeaders(w, resp.Header) w.WriteHeader(resp.StatusCode) @@ -1401,6 +1404,16 @@ func (s *Server) applyLearnedCapacityDecrease(chain, endpointID string, endpoint // Alchemy's HTTP-200-with-one-batch-item-429 blind spot visible, since HTTP status alone // can't see it. Returns a zero-value signal if the body isn't a recognizable JSON-RPC // response or carries no rate-limit error. +// providerNeedsEmbeddedRateLimitScan reports whether a provider is known to sometimes +// return a 2xx HTTP response with a rate-limit error embedded in the JSON-RPC body +// instead of a non-2xx status code. Currently only Alchemy, on batch requests - the +// specific quirk bodyCarriesRateLimitSignal exists to detect. Gating on this list keeps +// the buffer-and-scan cost (up to maxRateLimitScanBodyBytes held in memory, delaying +// time-to-first-byte) off every other provider's 2xx JSON responses. +func providerNeedsEmbeddedRateLimitScan(provider string) bool { + return strings.EqualFold(provider, "alchemy") +} + func bodyCarriesRateLimitSignal(provider string, body []byte, headers http.Header) health.RateLimitSignal { var single health.RpcResponse if err := json.Unmarshal(body, &single); err == nil && single.Error != nil { diff --git a/internal/server/server_test.go b/internal/server/server_test.go index f1f08f7..14ca930 100644 --- a/internal/server/server_test.go +++ b/internal/server/server_test.go @@ -1057,6 +1057,68 @@ func TestDefaultForwardRequestWithBodyFuncDetectsEmbeddedBatchRateLimit(t *testi } } +// TestDefaultForwardRequestWithBodyFuncSkipsEmbeddedScanForNonAlchemyProvider guards +// against buffering-and-scanning every provider's 2xx JSON responses: since only Alchemy +// is documented to embed a rate-limit error in an HTTP-200 batch response, a non-Alchemy +// provider's response carrying the same shape must be forwarded to the client verbatim +// and unscanned, not buffered into memory and treated as a rate-limit signal. +func TestDefaultForwardRequestWithBodyFuncSkipsEmbeddedScanForNonAlchemyProvider(t *testing.T) { + batchBody := `[{"jsonrpc":"2.0","id":1,"result":"0x1"},{"jsonrpc":"2.0","id":2,"error":{"code":-32005,"message":"Request limit exceeded"}}]` + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + w.Write([]byte(batchBody)) + })) + defer ts.Close() + + cfg := &config.Config{ + Endpoints: map[string]config.ChainEndpoints{ + "chainA": { + "ep1": config.Endpoint{Provider: "drpc", HTTPURL: ts.URL, Role: "primary", Type: "full"}, + }, + }, + } + valkeyClient := store.NewMockValkeyClient() + server := NewServer(cfg, valkeyClient, createTestConfig()) + + rec := httptest.NewRecorder() + err := server.defaultForwardRequestWithBodyFunc(rec, context.Background(), "POST", ts.URL, nil, http.Header{}) + if err != nil { + t.Fatalf("Expected no error - a non-Alchemy provider's body should never be scanned, got: %v", err) + } + if rec.Body.String() != batchBody { + t.Errorf("Expected the batch body to be forwarded verbatim, got %q", rec.Body.String()) + } + + state, stateErr := valkeyClient.GetRateLimitState(context.Background(), "chainA", "ep1") + if stateErr != nil { + t.Fatalf("Failed to get rate limit state: %v", stateErr) + } + if state.RateLimited { + t.Error("Expected a non-Alchemy provider's embedded error to NOT be treated as a rate-limit signal") + } +} + +func TestProviderNeedsEmbeddedRateLimitScan(t *testing.T) { + tests := []struct { + provider string + expected bool + }{ + {"alchemy", true}, + {"Alchemy", true}, + {"infura", false}, + {"drpc", false}, + {"", false}, + } + for _, tt := range tests { + t.Run(tt.provider, func(t *testing.T) { + if got := providerNeedsEmbeddedRateLimitScan(tt.provider); got != tt.expected { + t.Errorf("providerNeedsEmbeddedRateLimitScan(%q) = %v, want %v", tt.provider, got, tt.expected) + } + }) + } +} + // TestDefaultForwardRequestWithBodyFuncNoDuplicateHeadersAcrossRetries guards against a // regression where response headers were copied onto w unconditionally, before the // embedded-rate-limit body scan decided whether to abort. Since handleRequestHTTP's From a740aa067bcb4d88897b150a45d77d86f8c1d496 Mon Sep 17 00:00:00 2001 From: Jacob Sussmilch Date: Thu, 2 Jul 2026 12:27:22 +1000 Subject: [PATCH 11/12] fix: KEEP-905 exclude daily-quota signals from adaptive capacity decrease handleRateLimit called applyLearnedCapacityDecrease unconditionally for every confirmed rate-limit signal, including Infura's HTTP 402 daily-credit-cap exhaustion. That signal says nothing about an endpoint's short-term RPS capacity, so folding it into the AIMD estimator taught it the wrong lesson - halving a burst-capacity ceiling in response to a daily quota running out, contradicting the PR's own rationale for keeping the two signal types apart. Add an isDailyQuota parameter to the shared store.ApplyLearnedCapacityDecreaseIfEligible (used by both the load balancer and the standalone health checker) that short-circuits the decrease entirely for a daily-quota signal, leaving the long-cooldown recovery backoff already seeded for it (health.InitialBackoffForSignal) as the sole response. Covered independently for both call paths - the load balancer (TestApplyLearnedCapacityDecreaseSkipsForDailyQuotaSignal) and the standalone health checker (TestCreateStandaloneRateLimitHandlerSkipsCapacityEstimateForDailyQuota) - since the two are separate processes mutating the same persisted estimate. Raised in a PR #37 review comment questioning whether 402/daily-quota exhaustion was kept separate from burst 429s in the adaptive estimator. --- internal/server/capacity_learning_test.go | 49 ++++++++++++++++++----- internal/server/server.go | 6 +-- internal/store/capacity_estimate.go | 11 ++++- services/health-checker/main.go | 6 +-- services/health-checker/main_test.go | 31 +++++++++++++- 5 files changed, 83 insertions(+), 20 deletions(-) diff --git a/internal/server/capacity_learning_test.go b/internal/server/capacity_learning_test.go index 2595390..054027a 100644 --- a/internal/server/capacity_learning_test.go +++ b/internal/server/capacity_learning_test.go @@ -6,6 +6,7 @@ import ( "time" "aetherlay/internal/config" + "aetherlay/internal/health" "aetherlay/internal/helpers" "aetherlay/internal/store" ) @@ -40,7 +41,7 @@ func TestApplyLearnedCapacityDecreaseSeedsFromObservedCount(t *testing.T) { } server := NewServer(cfg, valkeyClient, learningTestConfig()) - server.applyLearnedCapacityDecrease("ethereum", "ep1", cfg.Endpoints["ethereum"]["ep1"]) + server.applyLearnedCapacityDecrease("ethereum", "ep1", cfg.Endpoints["ethereum"]["ep1"], health.RateLimitSignal{}) estimate, err := valkeyClient.GetCapacityEstimate(ctx, "ethereum", "ep1") if err != nil { @@ -75,7 +76,7 @@ func TestApplyLearnedCapacityDecreaseSkipsWithinCooldown(t *testing.T) { server := NewServer(cfg, valkeyClient, learningTestConfig()) ep := cfg.Endpoints["ethereum"]["ep1"] - server.applyLearnedCapacityDecrease("ethereum", "ep1", ep) + server.applyLearnedCapacityDecrease("ethereum", "ep1", ep, health.RateLimitSignal{}) first, _ := valkeyClient.GetCapacityEstimate(ctx, "ethereum", "ep1") if first.MaxRequests != 5 { t.Fatalf("Expected first decrease to seed 5, got %d", first.MaxRequests) @@ -83,7 +84,7 @@ func TestApplyLearnedCapacityDecreaseSkipsWithinCooldown(t *testing.T) { // Second hit immediately after - within the cooldown (the learning window itself) - // must not decrease again even though the observed count hasn't changed. - server.applyLearnedCapacityDecrease("ethereum", "ep1", ep) + server.applyLearnedCapacityDecrease("ethereum", "ep1", ep, health.RateLimitSignal{}) second, _ := valkeyClient.GetCapacityEstimate(ctx, "ethereum", "ep1") if second.MaxRequests != 5 { t.Errorf("Expected MaxRequests to remain 5 (no double decrease within cooldown), got %d", second.MaxRequests) @@ -107,7 +108,7 @@ func TestApplyLearnedCapacityDecreaseAppliesAgainAfterCooldownElapsed(t *testing } server := NewServer(cfg, valkeyClient, learningTestConfig()) - server.applyLearnedCapacityDecrease("ethereum", "ep1", ep) + server.applyLearnedCapacityDecrease("ethereum", "ep1", ep, health.RateLimitSignal{}) seeded, _ := valkeyClient.GetCapacityEstimate(ctx, "ethereum", "ep1") if seeded.MaxRequests != 5 { @@ -128,7 +129,7 @@ func TestApplyLearnedCapacityDecreaseAppliesAgainAfterCooldownElapsed(t *testing valkeyClient.IncrementCapacityCount(ctx, "ethereum", "ep1", 60) } - server.applyLearnedCapacityDecrease("ethereum", "ep1", ep) + server.applyLearnedCapacityDecrease("ethereum", "ep1", ep, health.RateLimitSignal{}) final, err := valkeyClient.GetCapacityEstimate(ctx, "ethereum", "ep1") if err != nil { @@ -157,7 +158,7 @@ func TestApplyLearnedCapacityDecreaseSkipsWhenStaticCapacityConfigured(t *testin ctx := context.Background() server := NewServer(cfg, valkeyClient, learningTestConfig()) - server.applyLearnedCapacityDecrease("ethereum", "ep1", cfg.Endpoints["ethereum"]["ep1"]) + server.applyLearnedCapacityDecrease("ethereum", "ep1", cfg.Endpoints["ethereum"]["ep1"], health.RateLimitSignal{}) estimate, _ := valkeyClient.GetCapacityEstimate(ctx, "ethereum", "ep1") if estimate.HasEstimate { @@ -165,6 +166,34 @@ func TestApplyLearnedCapacityDecreaseSkipsWhenStaticCapacityConfigured(t *testin } } +// TestApplyLearnedCapacityDecreaseSkipsForDailyQuotaSignal confirms an Infura-style +// daily-credit-cap exhaustion (HTTP 402) never feeds the AIMD estimator: that signal +// says nothing about the endpoint's short-term RPS capacity, so folding it in would +// incorrectly halve a burst-capacity ceiling in response to a daily quota running out. +func TestApplyLearnedCapacityDecreaseSkipsForDailyQuotaSignal(t *testing.T) { + cfg := &config.Config{ + Endpoints: map[string]config.ChainEndpoints{ + "ethereum": { + "ep1": config.Endpoint{Provider: "infura", Role: "primary", Type: "full", HTTPURL: "http://ep1"}, + }, + }, + } + valkeyClient := store.NewMockValkeyClient() + ctx := context.Background() + + for i := 0; i < 10; i++ { + valkeyClient.IncrementCapacityCount(ctx, "ethereum", "ep1", 60) + } + + server := NewServer(cfg, valkeyClient, learningTestConfig()) + server.applyLearnedCapacityDecrease("ethereum", "ep1", cfg.Endpoints["ethereum"]["ep1"], health.RateLimitSignal{IsDailyQuota: true}) + + estimate, _ := valkeyClient.GetCapacityEstimate(ctx, "ethereum", "ep1") + if estimate.HasEstimate { + t.Error("Expected no learned estimate to be seeded from a daily-quota (402) signal") + } +} + func TestApplyLearnedCapacityDecreaseSkipsWhenLearningDisabled(t *testing.T) { cfg := &config.Config{ Endpoints: map[string]config.ChainEndpoints{ @@ -179,7 +208,7 @@ func TestApplyLearnedCapacityDecreaseSkipsWhenLearningDisabled(t *testing.T) { appConfig := learningTestConfig() appConfig.CapacityLearningEnabled = false server := NewServer(cfg, valkeyClient, appConfig) - server.applyLearnedCapacityDecrease("ethereum", "ep1", cfg.Endpoints["ethereum"]["ep1"]) + server.applyLearnedCapacityDecrease("ethereum", "ep1", cfg.Endpoints["ethereum"]["ep1"], health.RateLimitSignal{}) estimate, _ := valkeyClient.GetCapacityEstimate(ctx, "ethereum", "ep1") if estimate.HasEstimate { @@ -440,7 +469,7 @@ func TestApplyLearnedCapacityDecreaseReadsFrozenWindowNotLiveConfig(t *testing.T } server := NewServer(cfg, valkeyClient, learningTestConfig()) - server.applyLearnedCapacityDecrease("ethereum", "ep1", ep) + server.applyLearnedCapacityDecrease("ethereum", "ep1", ep, health.RateLimitSignal{}) final, err := valkeyClient.GetCapacityEstimate(ctx, "ethereum", "ep1") if err != nil { @@ -482,7 +511,7 @@ func TestApplyLearnedCapacityDecreasePersistsFrozenWindowNotLiveConfig(t *testin valkeyClient.IncrementCapacityCount(ctx, "ethereum", "ep1", 60) server := NewServer(cfg, valkeyClient, learningTestConfig()) - server.applyLearnedCapacityDecrease("ethereum", "ep1", ep) + server.applyLearnedCapacityDecrease("ethereum", "ep1", ep, health.RateLimitSignal{}) final, err := valkeyClient.GetCapacityEstimate(ctx, "ethereum", "ep1") if err != nil { @@ -503,7 +532,7 @@ func TestApplyLearnedCapacityDecreasePersistsFrozenWindowNotLiveConfig(t *testin LastDecreaseAt: time.Now().Add(-120 * time.Second), }) valkeyClient.IncrementCapacityCount(ctx, "ethereum", "ep1", 60) - server.applyLearnedCapacityDecrease("ethereum", "ep1", ep) + server.applyLearnedCapacityDecrease("ethereum", "ep1", ep, health.RateLimitSignal{}) final2, err := valkeyClient.GetCapacityEstimate(ctx, "ethereum", "ep1") if err != nil { diff --git a/internal/server/server.go b/internal/server/server.go index 748cd60..70c7838 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -1393,10 +1393,10 @@ func (s *Server) recordCapacityUsage(chain string, endpoint config.Endpoint, end // rate-limit signal. It delegates to store.ApplyLearnedCapacityDecreaseIfEligible, the // same shared implementation the standalone health checker calls, so the two processes - // both mutating the same Valkey-persisted estimate - never diverge. -func (s *Server) applyLearnedCapacityDecrease(chain, endpointID string, endpoint config.Endpoint) { +func (s *Server) applyLearnedCapacityDecrease(chain, endpointID string, endpoint config.Endpoint, signal health.RateLimitSignal) { ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - store.ApplyLearnedCapacityDecreaseIfEligible(ctx, s.valkeyClient, chain, endpointID, endpoint, s.appConfig.CapacityThrottlingEnabled, s.appConfig.CapacityLearningEnabled) + store.ApplyLearnedCapacityDecreaseIfEligible(ctx, s.valkeyClient, chain, endpointID, endpoint, s.appConfig.CapacityThrottlingEnabled, s.appConfig.CapacityLearningEnabled, signal.IsDailyQuota) } // bodyCarriesRateLimitSignal reports whether a 2xx JSON-RPC response body (a single @@ -1703,7 +1703,7 @@ func (s *Server) handleRateLimit(chain, endpointID, protocol string, signal heal // throughput ceiling from it - only engages for endpoints with no static Capacity. if chainEndpoints, ok := s.config.GetEndpointsForChain(chain); ok { if ep, ok := chainEndpoints[endpointID]; ok { - s.applyLearnedCapacityDecrease(chain, endpointID, ep) + s.applyLearnedCapacityDecrease(chain, endpointID, ep, signal) } } diff --git a/internal/store/capacity_estimate.go b/internal/store/capacity_estimate.go index 0b17dd8..32d6c73 100644 --- a/internal/store/capacity_estimate.go +++ b/internal/store/capacity_estimate.go @@ -141,12 +141,19 @@ func (r *ValkeyClient) SetCapacityEstimate(ctx context.Context, chain, endpoint // if present, is left completely untouched. A cooldown (the learning window itself) // collapses several near-simultaneous hits from one episode into a single decrease. // +// isDailyQuota excludes Infura-style daily-credit-cap exhaustion (HTTP 402) from this +// math entirely: that signal says nothing about the endpoint's short-term RPS capacity, +// so folding it into the AIMD estimator would teach it the wrong lesson - halving a +// burst-capacity ceiling in response to a daily quota running out. The long-cooldown +// recovery backoff already seeded for this signal (see health.InitialBackoffForSignal) +// is the correct - and separate - response. +// // This is shared verbatim between the load balancer (internal/server) and the standalone // health checker (services/health-checker) - both mutate the same Valkey-persisted // estimate for a given endpoint, so they must apply identical math or silently disagree // about what the learned ceiling means for that endpoint. -func ApplyLearnedCapacityDecreaseIfEligible(ctx context.Context, valkeyClient ValkeyClientIface, chain, endpointID string, ep config.Endpoint, capacityThrottlingEnabled, capacityLearningEnabled bool) { - if !capacityThrottlingEnabled || !capacityLearningEnabled || ep.Capacity != nil { +func ApplyLearnedCapacityDecreaseIfEligible(ctx context.Context, valkeyClient ValkeyClientIface, chain, endpointID string, ep config.Endpoint, capacityThrottlingEnabled, capacityLearningEnabled, isDailyQuota bool) { + if !capacityThrottlingEnabled || !capacityLearningEnabled || ep.Capacity != nil || isDailyQuota { return } diff --git a/services/health-checker/main.go b/services/health-checker/main.go index bba6ea5..b92240a 100644 --- a/services/health-checker/main.go +++ b/services/health-checker/main.go @@ -48,7 +48,7 @@ func standaloneInitialBackoff(cfg *config.Config, chain, endpointID string, sign // implementation server.Server.applyLearnedCapacityDecrease calls, so the load balancer // and the standalone health checker - two separate processes mutating the same // Valkey-persisted estimate - never diverge. -func applyStandaloneLearnedCapacityDecrease(cfg *config.Config, valkeyClient store.ValkeyClientIface, capacityThrottlingEnabled, capacityLearningEnabled bool, chain, endpointID string) { +func applyStandaloneLearnedCapacityDecrease(cfg *config.Config, valkeyClient store.ValkeyClientIface, capacityThrottlingEnabled, capacityLearningEnabled bool, chain, endpointID string, signal health.RateLimitSignal) { chainEndpoints, ok := cfg.GetEndpointsForChain(chain) if !ok { return @@ -60,7 +60,7 @@ func applyStandaloneLearnedCapacityDecrease(cfg *config.Config, valkeyClient sto ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - store.ApplyLearnedCapacityDecreaseIfEligible(ctx, valkeyClient, chain, endpointID, ep, capacityThrottlingEnabled, capacityLearningEnabled) + store.ApplyLearnedCapacityDecreaseIfEligible(ctx, valkeyClient, chain, endpointID, ep, capacityThrottlingEnabled, capacityLearningEnabled, signal.IsDailyQuota) } // createStandaloneRateLimitHandler creates a simple rate limit handler for the standalone health checker @@ -96,7 +96,7 @@ func createStandaloneRateLimitHandler(cfg *config.Config, valkeyClient store.Val // Check for rate limits first (this signal), then approximate the endpoint's // safe throughput ceiling from it - only engages for endpoints with no static // Capacity, and shares the exact same math as the load balancer's own path. - applyStandaloneLearnedCapacityDecrease(cfg, valkeyClient, capacityThrottlingEnabled, capacityLearningEnabled, chain, endpointID) + applyStandaloneLearnedCapacityDecrease(cfg, valkeyClient, capacityThrottlingEnabled, capacityLearningEnabled, chain, endpointID, signal) } } diff --git a/services/health-checker/main_test.go b/services/health-checker/main_test.go index 18270e2..6275e5a 100644 --- a/services/health-checker/main_test.go +++ b/services/health-checker/main_test.go @@ -295,7 +295,7 @@ func TestApplyStandaloneLearnedCapacityDecreaseSeedsFromObservedCount(t *testing valkeyClient.IncrementCapacityCount(ctx, "mainnet", "mock", 60) } - applyStandaloneLearnedCapacityDecrease(cfg, valkeyClient, true, true, "mainnet", "mock") + applyStandaloneLearnedCapacityDecrease(cfg, valkeyClient, true, true, "mainnet", "mock", health.RateLimitSignal{}) estimate, err := valkeyClient.GetCapacityEstimate(ctx, "mainnet", "mock") if err != nil { @@ -323,7 +323,7 @@ func TestApplyStandaloneLearnedCapacityDecreaseSkipsWhenStaticCapacityConfigured valkeyClient := store.NewMockValkeyClient() ctx := context.Background() - applyStandaloneLearnedCapacityDecrease(cfg, valkeyClient, true, true, "mainnet", "mock") + applyStandaloneLearnedCapacityDecrease(cfg, valkeyClient, true, true, "mainnet", "mock", health.RateLimitSignal{}) estimate, _ := valkeyClient.GetCapacityEstimate(ctx, "mainnet", "mock") if estimate.HasEstimate { @@ -353,3 +353,30 @@ func TestCreateStandaloneRateLimitHandlerAlsoSeedsCapacityEstimate(t *testing.T) t.Errorf("Expected a learned estimate of 4 (8 observed * 0.5), got %+v", estimate) } } + +// TestCreateStandaloneRateLimitHandlerSkipsCapacityEstimateForDailyQuota mirrors +// TestApplyLearnedCapacityDecreaseSkipsForDailyQuotaSignal in the server package, but +// through the standalone health checker's own production entry point +// (createStandaloneRateLimitHandler), not just the shared store function directly - this +// is what actually confirms the daily-quota exclusion is wired through this process's +// handler, not only proven correct in the load balancer's call path. +func TestCreateStandaloneRateLimitHandlerSkipsCapacityEstimateForDailyQuota(t *testing.T) { + cfg := mockConfig() + valkeyClient := store.NewMockValkeyClient() + ctx := context.Background() + + for i := 0; i < 8; i++ { + valkeyClient.IncrementCapacityCount(ctx, "mainnet", "mock", 60) + } + + handler := createStandaloneRateLimitHandler(cfg, valkeyClient, true, true) + handler("mainnet", "mock", "http", health.RateLimitSignal{IsRateLimited: true, IsDailyQuota: true}) + + estimate, err := valkeyClient.GetCapacityEstimate(ctx, "mainnet", "mock") + if err != nil { + t.Fatalf("GetCapacityEstimate failed: %v", err) + } + if estimate.HasEstimate { + t.Errorf("Expected no learned estimate to be seeded from a daily-quota (402) signal, got %+v", estimate) + } +} From 56a5a2ea333bb1a8ad2259f2b93e8fbfd2f3a1cc Mon Sep 17 00:00:00 2001 From: Jacob Sussmilch Date: Thu, 2 Jul 2026 12:27:35 +1000 Subject: [PATCH 12/12] test: KEEP-905 cover body-scan-cap fail-open behavior for embedded rate-limit scan defaultForwardRequestWithBodyFunc's embedded-rate-limit body scan already fails open when a 2xx JSON body exceeds maxRateLimitScanBodyBytes - streaming it through unmodified and unscanned rather than buffering it fully - but nothing asserted that behavior. Add a regression test with an oversized Alchemy body carrying an embedded rate-limit error beyond the scan cap, confirming it streams through verbatim and the endpoint is not marked rate limited. Raised in a PR #37 review comment asking for coverage of the scan-cap boundary. --- internal/server/server_test.go | 53 ++++++++++++++++++++++++++++++++++ 1 file changed, 53 insertions(+) diff --git a/internal/server/server_test.go b/internal/server/server_test.go index 14ca930..71d29a8 100644 --- a/internal/server/server_test.go +++ b/internal/server/server_test.go @@ -1099,6 +1099,59 @@ func TestDefaultForwardRequestWithBodyFuncSkipsEmbeddedScanForNonAlchemyProvider } } +// TestDefaultForwardRequestWithBodyFuncFailsOpenWhenBodyExceedsScanCap tests the +// intentional fail-open behavior when an Alchemy 2xx JSON body exceeds +// maxRateLimitScanBodyBytes: rather than fully buffering an unbounded body in memory, +// the scan is skipped entirely and the body is streamed through unmodified - even if it +// carries an embedded rate-limit error the scan would otherwise have caught. This +// guards the tradeoff described at maxRateLimitScanBodyBytes's declaration against a +// silent regression (e.g. accidentally treating a truncated scan as proof of no +// embedded error, or losing bytes from the streamed tail). +func TestDefaultForwardRequestWithBodyFuncFailsOpenWhenBodyExceedsScanCap(t *testing.T) { + // An embedded rate-limit error placed first, followed by enough padding to push the + // total body past maxRateLimitScanBodyBytes - so the scan never gets far enough to see + // it, and the whole oversized body must stream through untouched. + padding := strings.Repeat("a", int(maxRateLimitScanBodyBytes)) + oversizedBody := fmt.Sprintf( + `[{"jsonrpc":"2.0","id":1,"error":{"code":-32005,"message":"Request limit exceeded"}},{"jsonrpc":"2.0","id":2,"result":"0x%s"}]`, + padding, + ) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + w.Write([]byte(oversizedBody)) + })) + defer ts.Close() + + cfg := &config.Config{ + Endpoints: map[string]config.ChainEndpoints{ + "chainA": { + "ep1": config.Endpoint{Provider: "alchemy", HTTPURL: ts.URL, Role: "primary", Type: "full"}, + }, + }, + } + valkeyClient := store.NewMockValkeyClient() + server := NewServer(cfg, valkeyClient, createTestConfig()) + + rec := httptest.NewRecorder() + err := server.defaultForwardRequestWithBodyFunc(rec, context.Background(), "POST", ts.URL, nil, http.Header{}) + if err != nil { + t.Fatalf("Expected no error - an oversized body must fail open and stream through, got: %v", err) + } + if rec.Body.String() != oversizedBody { + t.Errorf("Expected the oversized body to be forwarded verbatim (got length %d, want %d)", rec.Body.Len(), len(oversizedBody)) + } + + state, stateErr := valkeyClient.GetRateLimitState(context.Background(), "chainA", "ep1") + if stateErr != nil { + t.Fatalf("Failed to get rate limit state: %v", stateErr) + } + if state.RateLimited { + t.Error("Expected the embedded rate-limit error beyond the scan cap to go undetected, not mark the endpoint rate limited") + } +} + func TestProviderNeedsEmbeddedRateLimitScan(t *testing.T) { tests := []struct { provider string