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..e99f1b1 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 @@ -15,11 +17,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" @@ -70,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 } } @@ -77,6 +101,40 @@ func LoadConfig(path string) (*Config, error) { return &config, nil } +// validateEndpointCapacity catches an endpoint's static Capacity being configured with a +// 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 + } + 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 + 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 + } +} + // 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) { @@ -149,3 +207,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..595c28b 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -1,6 +1,7 @@ package config import ( + "encoding/json" "os" "testing" ) @@ -151,6 +152,122 @@ 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) + } +} + +// 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() @@ -261,3 +378,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..0592718 --- /dev/null +++ b/internal/health/ratelimit.go @@ -0,0 +1,94 @@ +package health + +import ( + "net/http" + "strconv" + "strings" + "time" + + "aetherlay/internal/config" +) + +// 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 +} + +// 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/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..054027a --- /dev/null +++ b/internal/server/capacity_learning_test.go @@ -0,0 +1,599 @@ +package server + +import ( + "context" + "testing" + "time" + + "aetherlay/internal/config" + "aetherlay/internal/health" + "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"], health.RateLimitSignal{}) + + 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, health.RateLimitSignal{}) + 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, 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) + } +} + +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, health.RateLimitSignal{}) + + 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, health.RateLimitSignal{}) + + 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"], health.RateLimitSignal{}) + + estimate, _ := valkeyClient.GetCapacityEstimate(ctx, "ethereum", "ep1") + if estimate.HasEstimate { + t.Error("Expected no learned estimate when a static Capacity is configured") + } +} + +// 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{ + "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"], health.RateLimitSignal{}) + + 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") + } +} + +// 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, health.RateLimitSignal{}) + + 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) + } +} + +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, health.RateLimitSignal{}) + + 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, health.RateLimitSignal{}) + + 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) + } +} + +// 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/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..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) @@ -375,7 +386,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 +405,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/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) + } +} diff --git a/internal/server/server.go b/internal/server/server.go index 7353295..70c7838 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,31 @@ 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 { + 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)) + } + 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 +944,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 +1007,32 @@ 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 + } + 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 + } + 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 +1244,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") @@ -1166,18 +1259,51 @@ 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 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 + // 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. + 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 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) + _, 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 @@ -1185,6 +1311,130 @@ 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. +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 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 { + 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 + } + } + 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, s.capacityWindowSeconds(chain, endpointID, 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 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, 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, signal.IsDailyQuota) +} + +// 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. +// 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 { + 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 +1509,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 +1662,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 +1685,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 +1697,23 @@ 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, signal) + } + } // Start rate limit recovery monitoring s.rateLimitScheduler.StartMonitoring(chain, endpointID) } + +// 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 { + return health.InitialBackoffForSignal(s.config, chain, endpointID, signal) +} diff --git a/internal/server/server_test.go b/internal/server/server_test.go index 0568233..71d29a8 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,453 @@ 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") + } +} + +// 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") + } +} + +// 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 + 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 +// 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. +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..32d6c73 --- /dev/null +++ b/internal/store/capacity_estimate.go @@ -0,0 +1,208 @@ +package store + +import ( + "context" + "encoding/json" + "time" + + "aetherlay/internal/config" + "aetherlay/internal/metrics" + + "github.com/rs/zerolog/log" + "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() +} + +// 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. +// +// 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, isDailyQuota bool) { + if !capacityThrottlingEnabled || !capacityLearningEnabled || ep.Capacity != nil || isDailyQuota { + 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 + } + + // 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 { + windowSeconds = prior.WindowSeconds + } + + 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, 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, 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/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..6d05449 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,62 @@ 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() + if windowSeconds <= 0 { + windowSeconds = 1 + } + 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() + if windowSeconds <= 0 { + windowSeconds = 1 + } + 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..ee07403 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,60 @@ 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 { + // 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) +} + +// 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..4c6843c 100644 --- a/internal/store/valkey_test.go +++ b/internal/store/valkey_test.go @@ -286,6 +286,97 @@ 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) + } +} + +// 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() + 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..b92240a 100644 --- a/services/health-checker/main.go +++ b/services/health-checker/main.go @@ -36,9 +36,36 @@ var testExitAfterSetup bool // exitCode is used to track the exit code for the process var exitCode int +// 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 { + return health.InitialBackoffForSignal(cfg, chain, endpointID, signal) +} + +// 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, signal health.RateLimitSignal) { + chainEndpoints, ok := cfg.GetEndpointsForChain(chain) + if !ok { + return + } + ep, ok := chainEndpoints[endpointID] + if !ok { + return + } + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + store.ApplyLearnedCapacityDecreaseIfEligible(ctx, valkeyClient, chain, endpointID, ep, capacityThrottlingEnabled, capacityLearningEnabled, signal.IsDailyQuota) +} + // 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 +79,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 +91,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, signal) } } @@ -73,6 +106,8 @@ func RunHealthChecker( corsHeaders string, corsMethods string, corsOrigin string, + capacityLearningEnabled bool, + capacityThrottlingEnabled bool, ephemeralChecksEnabled bool, ephemeralChecksHealthyThreshold int, ephemeralChecksInterval int, @@ -170,7 +205,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 +272,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..6275e5a 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,169 @@ 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", health.RateLimitSignal{}) + + 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", health.RateLimitSignal{}) + + 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) + } +} + +// 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) + } +}