feat: KEEP-905 Rate-limit-aware load balancing with adaptive capacity estimation#37
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (6)
🚧 Files skipped from review as they are similar to previous changes (6)
WalkthroughThis PR adds capacity configuration and documentation, structured rate-limit signal detection, Valkey-backed capacity counters and learned estimates, server-side capacity-aware dispatch and recovery, and standalone health-checker wiring. It also updates tests and metrics to cover capacity throttling, adaptive learning, provider-specific rate-limit handling, and the new configuration flags. ChangesRelated issues: None identified. Sequence Diagram(s)sequenceDiagram
participant Client
participant Server
participant Valkey
participant Backend
Client->>Server: request
Server->>Server: recordCapacityUsage
Server->>Backend: forward request
Backend-->>Server: 2xx/429/402
Server->>Server: DetectRateLimit / body scan
Server->>Valkey: update rate-limit or capacity state
Poem: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
services/health-checker/main.go (1)
128-138: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
CurrentBackoffis reset on every rate-limit detection, not just the first, undoing the scheduler's exponential-backoff progress.Unlike
FirstRateLimited, which is guarded to only be set once (if state.FirstRateLimited.IsZero()),state.CurrentBackoffis unconditionally recomputed fromstandaloneInitialBackoffon every call to this handler. Since the standalone health checker periodically re-checks already-rate-limited endpoints (that's its job — detecting recovery), each subsequent detection will overwrite whatever exponential backoff the main load balancer's recovery scheduler has already built up (via multiplyingCurrentBackoffon failed recovery attempts) back down to the freshly-seeded value (often0for a plain 429 withoutRetry-After). The line-range change notes confirm the prior behavior intentionally left this field "unset/unchanged" — this change removes that protection.This can cause the recovery scheduler to probe/retry the endpoint far more aggressively than the configured backoff schedule intends, since the two processes race on the same shared Valkey key.
🐛 Proposed fix: only seed CurrentBackoff on first detection
// Mark as rate limited but don't start recovery scheduling (that's for the main load balancer) now := time.Now() + wasAlreadyRateLimited := state.RateLimited state.RateLimited = true state.LastRecoveryCheck = now - state.CurrentBackoff = standaloneInitialBackoff(cfg, chain, endpointID, signal) + if !wasAlreadyRateLimited { + state.CurrentBackoff = standaloneInitialBackoff(cfg, chain, endpointID, signal) + }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@services/health-checker/main.go` around lines 128 - 138, The rate-limit handler in health-checker is overwriting the shared recovery backoff on every detection instead of only seeding it once. Update the logic around `state.CurrentBackoff` and `standaloneInitialBackoff` so it is only initialized when the endpoint is first marked rate-limited, similar to the existing `state.FirstRateLimited` guard, and leave any existing backoff value unchanged on later checks. Keep the rest of the state updates in the same handler (`state.RateLimited`, `state.LastRecoveryCheck`, `FirstRateLimited`) intact.
🧹 Nitpick comments (3)
services/health-checker/main_test.go (1)
266-283: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider adding a regression test for repeated rate-limit detections.
This test only exercises a single call to the handler. Given the
CurrentBackoffoverwrite issue flagged inmain.go(line 133), a follow-up test asserting that a second call for an already-rate-limited endpoint does not reset an in-progress backoff would guard against regressions once that fix lands.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@services/health-checker/main_test.go` around lines 266 - 283, Extend TestCreateStandaloneRateLimitHandlerSeedsBackoffFromSignal to cover repeated rate-limit detections: after the first call to createStandaloneRateLimitHandler seeds the backoff for the same endpoint, invoke the returned handler again with another RateLimitSignal and verify the existing rate-limit state in valkeyClient.GetRateLimitState is not reset or overwritten. Use the createStandaloneRateLimitHandler and GetRateLimitState symbols to locate the behavior and assert CurrentBackoff remains the in-progress value after the second detection.internal/store/capacity_estimate.go (1)
76-94: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low value
EffectiveMaxRequestsgrows the ceiling unbounded over elapsed wall-clock time.The additive-increase step compounds indefinitely (
steps*estimate.IncreaseStepwith no cap), so after a long idle/healthy stretch the effective ceiling can grow far beyond anything ever observed, effectively neutralizing the learned throttle until the next real rate-limit hit corrects it. This may be an intentional trade-off (avoids introducing new config per the PR's stated goal), but worth confirming it's not an oversight, since a very long clean period could let a traffic burst through unchecked.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/store/capacity_estimate.go` around lines 76 - 94, `EffectiveMaxRequests` currently increases the ceiling without any upper bound based on elapsed time, so cap the additive growth to a safe maximum instead of letting `steps*estimate.IncreaseStep` grow indefinitely. Update the logic in `EffectiveMaxRequests` to clamp the returned value using an existing learned limit or another explicit ceiling derived from `CapacityEstimate`/`CapacityLearning`, while preserving the current stateless behavior and using `estimate.MaxRequests`, `estimate.LastDecreaseAt`, and `params.IncreaseStep` as the key symbols to adjust.internal/server/server.go (1)
945-1031: 🚀 Performance & Scalability | 🔵 Trivial | 🏗️ Heavy lift
effectiveCapacityCeilingis re-resolved per candidate on every retry.
selectBestEndpointByRoleis invoked once per retry attempt (from the outer loop inhandleRequestHTTP/handleRequestWS), and for each retry it re-resolveseffectiveCapacityCeilingfor every candidate endpoint again — on top of the resolution already done once per candidate ingetEndpointsByRole's proactive skip check. For a request needing several retries across several candidates with learned estimates, this multiplies Valkey round-trips (GetCapacityEstimate) well beyond what's structurally necessary for a single incoming request. The in-functionceilingscache only avoids the double-resolution within one call toselectBestEndpointByRole, not across the getEndpointsByRole call or across retries.Consider resolving each candidate's ceiling once per incoming request (e.g., cache on
EndpointWithIDor a per-request map) and reusing it acrossgetEndpointsByRoleand all retry iterations ofselectBestEndpointByRole.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/server/server.go` around lines 945 - 1031, `selectBestEndpointByRole` is re-resolving `effectiveCapacityCeiling` for the same candidates on every retry, causing repeated Valkey reads. Move ceiling resolution to a per-request cache so `getEndpointsByRole` and `selectBestEndpointByRole` reuse the same `endpointCeiling`/resolved capacity data instead of calling `effectiveCapacityCeiling` repeatedly. Use the existing `endpointCeiling` cache pattern and wire it through the request flow in `handleRequestHTTP`/`handleRequestWS` so each endpoint’s ceiling is computed once per incoming request.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@internal/server/server.go`:
- Around line 1269-1297: The JSON-response rate-limit scan in the proxy path is
buffering every 2xx JSON body up to maxRateLimitScanBodyBytes before any bytes
are forwarded, which regresses streaming and increases memory use. Restrict this
behavior in the response handling block around bodyCarriesRateLimitSignal,
providerForEndpoint, and findChainAndEndpointByURL to only the
providers/endpoints that actually need embedded 200-rate-limit detection (for
example the known Alchemy batch case), or switch to a streaming-safe inspection
approach instead of fully buffering generic JSON responses. Keep the existing
non-JSON and oversized-body streaming behavior unchanged.
- Around line 881-903: The Valkey capacity lookup in the endpoint selection path
still uses a non-request context, and capacity limits are not validated before
use. Update the calls in the selection logic and the related capacity lookup
path in `effectiveCapacityCeiling` to use the request-scoped context or a short
timeout instead of `context.Background()`, and add validation so
`CapacityLimit.WindowSeconds` and `CapacityLimit.MaxRequests` must be positive
before they are applied. Ensure `GetCapacityCount`, `capacityBucketKey`, and the
ceiling check reject invalid values rather than allowing panics or making every
endpoint skippable.
In `@internal/store/valkey.go`:
- Around line 392-441: Guard zero WindowSeconds in the capacity bucket path:
capacityBucketKey currently divides by windowSeconds, so a 0 value from
IncrementCapacityCount/GetCapacityCount will panic. Add a zero-value fallback or
validation before calling capacityBucketKey, and keep the behavior consistent
with CapacityLearning.WindowSeconds by defaulting to a safe window (for example
60) or rejecting invalid config earlier in the load path.
In `@README.md`:
- Around line 507-510: The `window_seconds` behavior described in the README is
inconsistent with the implementation in `ApplyCapacityDecrease` and
`effectiveCapacityCeiling`; the stored window is overwritten on every decrease
instead of staying fixed after the first estimate. Either update the README to
match the current behavior or change `ApplyCapacityDecrease` to preserve
`prior.WindowSeconds` when persisting the estimate, so existing endpoints keep
their original window size.
---
Outside diff comments:
In `@services/health-checker/main.go`:
- Around line 128-138: The rate-limit handler in health-checker is overwriting
the shared recovery backoff on every detection instead of only seeding it once.
Update the logic around `state.CurrentBackoff` and `standaloneInitialBackoff` so
it is only initialized when the endpoint is first marked rate-limited, similar
to the existing `state.FirstRateLimited` guard, and leave any existing backoff
value unchanged on later checks. Keep the rest of the state updates in the same
handler (`state.RateLimited`, `state.LastRecoveryCheck`, `FirstRateLimited`)
intact.
---
Nitpick comments:
In `@internal/server/server.go`:
- Around line 945-1031: `selectBestEndpointByRole` is re-resolving
`effectiveCapacityCeiling` for the same candidates on every retry, causing
repeated Valkey reads. Move ceiling resolution to a per-request cache so
`getEndpointsByRole` and `selectBestEndpointByRole` reuse the same
`endpointCeiling`/resolved capacity data instead of calling
`effectiveCapacityCeiling` repeatedly. Use the existing `endpointCeiling` cache
pattern and wire it through the request flow in
`handleRequestHTTP`/`handleRequestWS` so each endpoint’s ceiling is computed
once per incoming request.
In `@internal/store/capacity_estimate.go`:
- Around line 76-94: `EffectiveMaxRequests` currently increases the ceiling
without any upper bound based on elapsed time, so cap the additive growth to a
safe maximum instead of letting `steps*estimate.IncreaseStep` grow indefinitely.
Update the logic in `EffectiveMaxRequests` to clamp the returned value using an
existing learned limit or another explicit ceiling derived from
`CapacityEstimate`/`CapacityLearning`, while preserving the current stateless
behavior and using `estimate.MaxRequests`, `estimate.LastDecreaseAt`, and
`params.IncreaseStep` as the key symbols to adjust.
In `@services/health-checker/main_test.go`:
- Around line 266-283: Extend
TestCreateStandaloneRateLimitHandlerSeedsBackoffFromSignal to cover repeated
rate-limit detections: after the first call to createStandaloneRateLimitHandler
seeds the backoff for the same endpoint, invoke the returned handler again with
another RateLimitSignal and verify the existing rate-limit state in
valkeyClient.GetRateLimitState is not reset or overwritten. Use the
createStandaloneRateLimitHandler and GetRateLimitState symbols to locate the
behavior and assert CurrentBackoff remains the in-progress value after the
second detection.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 90b898f9-771f-43fa-b429-9e3f1026bff8
📒 Files selected for processing (23)
README.mdconfigs/endpoints-example.jsoninternal/config/config.gointernal/config/config_test.gointernal/health/checker.gointernal/health/health_test.gointernal/health/ratelimit.gointernal/health/ratelimit_test.gointernal/helpers/helpers.gointernal/metrics/definitions.gointernal/server/capacity_learning_test.gointernal/server/capacity_test.gointernal/server/failover_soak_test.gointernal/server/rate_limit_scheduler.gointernal/server/server.gointernal/server/server_test.gointernal/store/capacity_estimate.gointernal/store/capacity_estimate_test.gointernal/store/testutils.gointernal/store/valkey.gointernal/store/valkey_test.goservices/health-checker/main.goservices/health-checker/main_test.go
| // Proactively skip an endpoint that has hit its capacity ceiling (static | ||
| // or learned) for the current window, so selection routes to another | ||
| // endpoint before the provider's own rate limiter would trigger. | ||
| // Independent of RateLimited above: this is a self-imposed budget, not | ||
| // a provider signal. | ||
| if s.appConfig.CapacityThrottlingEnabled { | ||
| if maxRequests, windowSeconds, hasCeiling := s.effectiveCapacityCeiling(chain, endpointID, endpoint); hasCeiling { | ||
| count, err := s.valkeyClient.GetCapacityCount(context.Background(), chain, endpointID, windowSeconds) | ||
| if err == nil { | ||
| if metrics.EndpointCapacityUtilization != nil { | ||
| metrics.EndpointCapacityUtilization.WithLabelValues(chain, endpointID).Set(float64(count) / float64(maxRequests)) | ||
| } | ||
| if count >= maxRequests { | ||
| log.Debug().Str("chain", chain).Str("endpoint", endpointID).Str("role", role).Msg("Skipping endpoint at its capacity ceiling") | ||
| if metrics.EndpointCapacitySkippedTotal != nil { | ||
| metrics.EndpointCapacitySkippedTotal.WithLabelValues(chain, endpointID).Inc() | ||
| } | ||
| continue | ||
| } | ||
| } | ||
| } | ||
| } | ||
|
|
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map relevant files and locate the symbols mentioned in the review comment.
git ls-files | rg '(^|/)(internal/server/server\.go|.*config.*\.go|.*valkey.*\.go|.*capacity.*\.go|.*learning.*\.go)$'
printf '\n--- server.go outline ---\n'
ast-grep outline internal/server/server.go --view expanded | sed -n '1,260p'
printf '\n--- search for capacity-related call sites ---\n'
rg -n "GetCapacityCount|GetCapacityEstimate|recordCapacityUsage|applyLearnedCapacityDecrease|effectiveCapacityCeiling|getEndpointsByRole|getAvailableEndpoints|selectBestEndpointByRole|CapacityThrottlingEnabled|CapacityLearningEnabled|WindowSeconds|MaxRequests" internal/server/server.go internal
printf '\n--- config validation search ---\n'
rg -n "WindowSeconds|MaxRequests|CapacityLimit|validate.*capacity|capacity.*validate|WindowSeconds > 0|MaxRequests > 0|must be.*positive|positive.*capacity" internalRepository: project-aethermesh/aetherlay
Length of output: 47690
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n--- internal/server/server.go around relevant sections ---\n'
sed -n '870,1045p' internal/server/server.go
printf '\n--- internal/server/server.go around Valkey timeout-bearing helpers ---\n'
sed -n '1338,1398p' internal/server/server.go
printf '\n--- internal/config/config.go around CapacityLimit and validation ---\n'
sed -n '1,260p' internal/config/config.go
printf '\n--- search for any positive/non-zero capacity validation elsewhere ---\n'
rg -n "positive|non-zero|nonzero|must be > 0|must be positive|validate.*Capacity|CapacityLimit" internal/config internal/helpers internal/serverRepository: project-aethermesh/aetherlay
Length of output: 22068
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n--- internal/store/valkey.go around client setup and capacity methods ---\n'
sed -n '1,220p' internal/store/valkey.go
printf '\n--- internal/store/valkey.go around capacity methods ---\n'
sed -n '400,520p' internal/store/valkey.go
printf '\n--- internal/store/capacity_estimate.go ---\n'
sed -n '1,220p' internal/store/capacity_estimate.goRepository: project-aethermesh/aetherlay
Length of output: 16595
🌐 Web query:
valkey-go Do(ctx context.Context) command timeout context.Background deadline client timeout
💡 Result:
In valkey-go, the Do(ctx, cmd) method respects the provided context's deadline and cancellation [1][2]. If the context expires (e.g., via context.WithTimeout) before a response is received, the operation will return early with an error, typically context.DeadlineExceeded [1][3][2]. Key points regarding timeout and context behavior: 1. Context Deadline: When you pass a context with a timeout to Do, the client monitors this deadline [1][2]. Note that returning early due to a deadline does not guarantee the command was not sent to or processed by the server; it only means the client stopped waiting for the response [1][3][2]. 2. Client vs. Context Timeouts: - Context Timeout: Controlled by the context passed to individual Do calls, affecting how long the client waits for that specific operation [1][2]. - ConnWriteTimeout: A client-level configuration (set in ClientOption) that controls the maximum duration the client waits for responses to pipeline commands [4][5]. It acts as a safety mechanism to detect unresponsive connections or local congestion [4]. If not explicitly set, it defaults to the Dialer's KeepAlive interval multiplied by 10 [4][5]. 3. Cancellation Behavior: Manually canceling a context is reliably respected in "pipeline mode," which is used for concurrent requests on the same connection [1][2]. For dedicated connections (often used for blocking requests), you may need to enable ClientOption.AlwaysPipelining to ensure context cancellation is immediately effective [1][2]. 4. Automatic Retries: The client automatically retries read-only commands by default if they fail, provided the context deadline has not been exceeded [1][2]. If the context is canceled or expired, the retry logic terminates [6][7]. You can disable this behavior using the DisableRetry option or configure the retry delay [1][2].
Citations:
- 1: https://pkg.go.dev/github.com/valkey-io/valkey-go/valkeyprob
- 2: https://pkg.go.dev/github.com/valkey-io/valkey-go
- 3: https://github.com/valkey-io/valkey-go/blob/f68a678f7c413b13e6579e3ec30e554bf2cb224b/README.md
- 4: https://github.com/valkey-io/valkey-go/blob/e384f2521e5b0e9dc17efade48d3a7f92302c859/valkey.go
- 5: https://deepwiki.com/valkey-io/valkey-go/2.2-client-configuration
- 6: https://github.com/valkey-io/valkey-go/blob/main/client.go
- 7: https://deepwiki.com/valkey-io/valkey-go/4.3-command-execution
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n--- capacity bucket key helper ---\n'
rg -n "func capacityBucketKey|capacityBucketKey\\(" internal/store/valkey.go internal/store
printf '\n--- surrounding helper implementation ---\n'
sed -n '520,620p' internal/store/valkey.go
printf '\n--- config tests around capacity defaults/validation ---\n'
sed -n '340,420p' internal/config/config_test.goRepository: project-aethermesh/aetherlay
Length of output: 3206
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n--- handleRequestHTTP / handleRequestWS around getAvailableEndpoints ---\n'
sed -n '407,670p' internal/server/server.go
printf '\n--- capacityBucketKey implementation ---\n'
sed -n '390,445p' internal/store/valkey.goRepository: project-aethermesh/aetherlay
Length of output: 13816
Add request-scoped timeouts to the Valkey capacity lookups and validate positive capacity limits.
internal/server/server.go:888,1024still usecontext.Background(), so these Valkey reads can outlive the request timeout and stall endpoint selection when Valkey is slow. Thread the handler context through these calls or wrap them in a short timeout.CapacityLimit.WindowSecondsandMaxRequestsstill need positive-value validation;windowSeconds == 0panics incapacityBucketKey, andmaxRequests <= 0makes the endpoint always skippable.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@internal/server/server.go` around lines 881 - 903, The Valkey capacity lookup
in the endpoint selection path still uses a non-request context, and capacity
limits are not validated before use. Update the calls in the selection logic and
the related capacity lookup path in `effectiveCapacityCeiling` to use the
request-scoped context or a short timeout instead of `context.Background()`, and
add validation so `CapacityLimit.WindowSeconds` and `CapacityLimit.MaxRequests`
must be positive before they are applied. Ensure `GetCapacityCount`,
`capacityBucketKey`, and the ceiling check reject invalid values rather than
allowing panics or making every endpoint skippable.
…indow_seconds An endpoint's static capacity.window_seconds is used as a divisor when computing the Valkey bucket key for capacity counting. Omitting the field (JSON zero value) or setting it negative caused a panic on the endpoint's first request, and since capacity gating evaluates every endpoint in a role before selection, this took down the entire chain, not just the misconfigured endpoint. LoadConfig now validates window_seconds at load time, warning and disabling capacity throttling for that endpoint rather than crashing. capacityBucketKey and MockValkeyClient's capacity methods also clamp defensively, since they're reachable through the public ValkeyClientIface beyond just the config path.
recordCapacityUsage and applyLearnedCapacityDecrease resolved an endpoint's learning window from the live config on every call, while effectiveCapacityCeiling's gating/weighting read path used the window frozen into the endpoint's CapacityEstimate at last-seed time. Since the Valkey bucket key is derived from the window value itself, a config change to an endpoint's capacity_learning.window_seconds after it already had a learned estimate caused writes and reads to target different buckets - the gate would then read a permanently stale/zero count and never trip again, silently defeating the proactive throttle until another real rate-limit hit happened to re-sync the frozen value. Both the load balancer and the standalone health checker now consult the estimate's frozen window (once one exists) for the cooldown check and observed-count read, falling back to the live-resolved default only before any estimate has been seeded.
…very performRecoveryCheck used CurrentBackoff == 0 as its only signal for "this is the first recovery attempt, don't multiply yet." Before this feature, CurrentBackoff always started at exactly 0, so that check worked. Now handleRateLimit can seed CurrentBackoff to a precise nonzero value from a Retry-After header or Infura's daily-quota MaxBackoff, which collided with that sentinel: the very first probe after a signaled rate limit already multiplied the backoff, instead of repeating the precise value once more first - exactly the same treatment a guessed InitialBackoff already gets. Gate on RecoveryAttempts instead, so a signal-seeded value gets the same "one repeat wait, then escalate" shape as the no-signal path, just with a more precise baseline.
Response headers were copied onto the client's http.ResponseWriter unconditionally right after the non-2xx check, before the embedded rate-limit body scan decided whether to abort. handleRequestHTTP's retry loop reuses the same ResponseWriter across attempts, so an aborted attempt (2xx response whose body embeds a rate-limit signal) left its headers sitting in the writer's header map uncommitted, and a subsequent successful retry's headers were then added on top rather than replacing them - producing duplicate header values (e.g. two Content-Length) in the response actually sent to the client. Headers are now copied only immediately before each point the function commits to writing a response (extracted into copyResponseHeaders, called right before each WriteHeader), never before an abort path.
An endpoint's static capacity.max_requests of 0 (e.g. omitted from the JSON) or negative made the "count >= max_requests" gating check in getEndpointsByRole true on the very first request, silently and permanently excluding the endpoint from selection with no distinguishing error from a real outage. LoadConfig now validates max_requests alongside window_seconds, warning and disabling capacity throttling for that endpoint when either is non-positive.
The AIMD capacity-decrease sequence and the recovery-backoff-seeding logic were duplicated verbatim between the load balancer (internal/server) and the standalone health checker (services/health-checker) - two separate processes that both mutate the same Valkey-persisted capacity estimate. Hand-duplicated copies are a correctness risk: any future change to the math or eligibility rules would need to be applied identically in both places, and a missed spot would make the two processes silently disagree about what the learned ceiling means for a given endpoint. Extract the decrease sequence into store.ApplyLearnedCapacityDecreaseIfEligible (internal/store/capacity_estimate.go) and the backoff-seeding logic into health.InitialBackoffForSignal (internal/health/ratelimit.go). Both server.Server.applyLearnedCapacityDecrease/initialBackoffForSignal and services/health-checker/main.go's applyStandaloneLearnedCapacityDecrease/ standaloneInitialBackoff are now thin wrappers delegating to these shared functions, preserving their existing signatures so all existing tests keep passing unmodified. As a byproduct of sharing the decrease sequence, the standalone health checker's wrapper now applies the same 2-second context timeout the load balancer already used, instead of an unbounded context.Background() - a Valkey call hanging in the standalone checker no longer risks blocking indefinitely.
ApplyLearnedCapacityDecreaseIfEligible correctly read against an existing estimate's frozen WindowSeconds (fixed previously for the read/observedCount path), but the decrease it persisted still passed the live-resolved params.WindowSeconds to ApplyCapacityDecrease instead of the frozen value. So each rate-limit hit after an operator changed an endpoint's capacity_learning.window_seconds would silently re-freeze the estimate to the new value, contradicting the documented invariant that the window is frozen for the estimate's lifetime once seeded, not just stable between reads. Reuse the same frozen-if-exists windowSeconds value for both the read (cooldown check, GetCapacityCount) and the write (ApplyCapacityDecrease), so the persisted window can only ever be set once, at first seed. Found by CodeRabbit's review of PR project-aethermesh#37 (README.md:510's documented invariant).
capacityWindowSeconds, effectiveCapacityCeiling, and the capacity-ceiling gating check in getEndpointsByRole all called GetCapacityEstimate/ GetCapacityCount with context.Background(), so a slow or unresponsive Valkey could stall endpoint selection indefinitely instead of failing fast within the request's own timeout budget. Wrap each of these three call sites in a 2-second context.WithTimeout, matching the bound already used for the capacity-decrease and usage-recording paths (recordCapacityUsage, applyLearnedCapacityDecrease). Also removed a stale duplicate doc comment on capacityWindowSeconds left over from an earlier commit in this branch. Found by CodeRabbit's review of PR project-aethermesh#37 (internal/server/server.go:888,1024).
Every 2xx JSON-RPC response through the proxy was being fully buffered (up to maxRateLimitScanBodyBytes = 10MB) to scan for a rate-limit error embedded in the body - a quirk specific to Alchemy on batch requests, which can return HTTP 200 with the error only inside the JSON-RPC payload. Applying this to every provider regressed streaming for large legitimate responses (e.g. eth_getLogs), delaying time-to-first-byte until the full body was read, and held up to 10MB in memory per in-flight response regardless of provider. Add providerNeedsEmbeddedRateLimitScan, resolving the endpoint's provider before deciding whether to buffer at all. Only Alchemy responses are buffered and scanned; every other provider's 2xx JSON response streams through unmodified and uninspected, exactly as it did before this feature was added. Found by CodeRabbit's review of PR project-aethermesh#37 (internal/server/server.go:1297).
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
internal/server/capacity_learning_test.go (1)
521-568: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert that capacity lookups were actually exercised.
TestCapacityValkeyLookupsUseRequestScopedTimeoutsonly fails when a no-deadline call is observed. If the selection path stops callingGetCapacityCount/GetCapacityEstimate,sawNoDeadlineremains false and the test still passes.Proposed test hardening
type deadlineCheckingValkeyClient struct { *store.MockValkeyClient sawNoDeadline bool + lookupCount int } func (c *deadlineCheckingValkeyClient) GetCapacityCount(ctx context.Context, chain, endpoint string, windowSeconds int) (int64, error) { + c.lookupCount++ 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) { + c.lookupCount++ if _, ok := ctx.Deadline(); !ok { c.sawNoDeadline = true } return c.MockValkeyClient.GetCapacityEstimate(ctx, chain, endpoint) } @@ + if client.lookupCount == 0 { + t.Fatal("Expected the selection path to perform at least one capacity Valkey lookup") + } if client.sawNoDeadline { t.Error("Expected all capacity Valkey lookups in the selection path to use a context with a deadline, not context.Background()") }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/server/capacity_learning_test.go` around lines 521 - 568, The test only detects missing deadlines when Valkey lookups happen, so it can still pass if getAvailableEndpoints, effectiveCapacityCeiling, or capacityWindowSeconds stop calling GetCapacityCount/GetCapacityEstimate. Harden TestCapacityValkeyLookupsUseRequestScopedTimeouts by asserting the deadlineCheckingValkeyClient was actually invoked at least once (for example by tracking call counts or a separate “wasCalled” flag) in addition to sawNoDeadline, so the test fails if the capacity lookup path is no longer exercised.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@internal/server/rate_limit_scheduler.go`:
- Around line 247-257: The retry backoff logic in rateLimitRecovery scheduling
must handle persisted state where RecoveryAttempts is greater than zero but
CurrentBackoff is still zero. Update the branch in the scheduler logic around
state.RecoveryAttempts and state.CurrentBackoff so later attempts initialize a
non-zero backoff (for example by seeding InitialBackoff when CurrentBackoff is
zero) before applying BackoffMultiplier, ensuring the recovery loop does not
keep scheduling zero-delay checks.
In `@internal/store/valkey_test.go`:
- Around line 332-342: The test only covers MockValkeyClient, so it misses the
real ValkeyClient path where raw windowSeconds is still used. Add a shared
normalization helper for capacity windows and use it in both ValkeyClient and
MockValkeyClient inside IncrementCapacityCount and GetCapacityCount so
non-positive values are clamped consistently. Then expand the regression
coverage to exercise the production ValkeyClient path, not just the mock, to
verify 0 and negative windowSeconds no longer reach bucket-key or EXPIRE logic.
---
Nitpick comments:
In `@internal/server/capacity_learning_test.go`:
- Around line 521-568: The test only detects missing deadlines when Valkey
lookups happen, so it can still pass if getAvailableEndpoints,
effectiveCapacityCeiling, or capacityWindowSeconds stop calling
GetCapacityCount/GetCapacityEstimate. Harden
TestCapacityValkeyLookupsUseRequestScopedTimeouts by asserting the
deadlineCheckingValkeyClient was actually invoked at least once (for example by
tracking call counts or a separate “wasCalled” flag) in addition to
sawNoDeadline, so the test fails if the capacity lookup path is no longer
exercised.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: d49b942c-9ad3-43c4-ac0d-029f149d805b
📒 Files selected for processing (13)
internal/config/config.gointernal/config/config_test.gointernal/health/ratelimit.gointernal/server/capacity_learning_test.gointernal/server/rate_limit_scheduler.gointernal/server/rate_limit_scheduler_test.gointernal/server/server.gointernal/server/server_test.gointernal/store/capacity_estimate.gointernal/store/testutils.gointernal/store/valkey.gointernal/store/valkey_test.goservices/health-checker/main.go
🚧 Files skipped from review as they are similar to previous changes (4)
- internal/store/testutils.go
- internal/store/valkey.go
- services/health-checker/main.go
- internal/server/server.go
| 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) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Handle zero backoff on later attempts.
If an existing persisted state has RecoveryAttempts > 0 and CurrentBackoff == 0, the new else branch multiplies zero by the multiplier and keeps scheduling zero-delay recovery checks indefinitely.
Proposed fix
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 {
+ if state.CurrentBackoff == 0 {
+ state.CurrentBackoff = rateLimitConfig.InitialBackoff
+ }
newBackoff := int(float64(state.CurrentBackoff) * rateLimitConfig.BackoffMultiplier)
state.CurrentBackoff = min(newBackoff, rateLimitConfig.MaxBackoff)
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| 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) | |
| 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 { | |
| if state.CurrentBackoff == 0 { | |
| state.CurrentBackoff = rateLimitConfig.InitialBackoff | |
| } | |
| newBackoff := int(float64(state.CurrentBackoff) * rateLimitConfig.BackoffMultiplier) | |
| state.CurrentBackoff = min(newBackoff, rateLimitConfig.MaxBackoff) | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@internal/server/rate_limit_scheduler.go` around lines 247 - 257, The retry
backoff logic in rateLimitRecovery scheduling must handle persisted state where
RecoveryAttempts is greater than zero but CurrentBackoff is still zero. Update
the branch in the scheduler logic around state.RecoveryAttempts and
state.CurrentBackoff so later attempts initialize a non-zero backoff (for
example by seeding InitialBackoff when CurrentBackoff is zero) before applying
BackoffMultiplier, ensuring the recovery loop does not keep scheduling
zero-delay checks.
| 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) | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Cover the production Valkey path, not just the mock.
This regression test only verifies MockValkeyClient, which already normalizes windowSeconds <= 0 to 1. The real ValkeyClient still uses the raw value for bucket-key calculation and EXPIRE, so 0/negative values can still panic or issue invalid Valkey commands despite this test passing.
Suggested direction
func normalizeCapacityWindowSeconds(windowSeconds int) int {
if windowSeconds <= 0 {
return 1
}
return windowSeconds
}Use that shared helper in both mock and production IncrementCapacityCount / GetCapacityCount, then add coverage for the production helper/path.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@internal/store/valkey_test.go` around lines 332 - 342, The test only covers
MockValkeyClient, so it misses the real ValkeyClient path where raw
windowSeconds is still used. Add a shared normalization helper for capacity
windows and use it in both ValkeyClient and MockValkeyClient inside
IncrementCapacityCount and GetCapacityCount so non-positive values are clamped
consistently. Then expand the regression coverage to exercise the production
ValkeyClient path, not just the mock, to verify 0 and negative windowSeconds no
longer reach bucket-key or EXPIRE logic.
Test coverage summaryThe automated suite covers the rate-limiting system at three levels: Pure math ( Integration ( End-to-end failover ( Manual verification (not committed - documented here for the record)To specifically confirm the adaptive/learned capacity approximation converges sensibly rather than diverging or collapsing, I ran two scratch simulations directly against the AIMD functions (deleted afterward since they duplicate coverage already in the tests above):
Both confirm the approximation behaves as designed under repeated cycles, not just in a single isolated decrease. |
…ease handleRateLimit called applyLearnedCapacityDecrease unconditionally for every confirmed rate-limit signal, including Infura's HTTP 402 daily-credit-cap exhaustion. That signal says nothing about an endpoint's short-term RPS capacity, so folding it into the AIMD estimator taught it the wrong lesson - halving a burst-capacity ceiling in response to a daily quota running out, contradicting the PR's own rationale for keeping the two signal types apart. Add an isDailyQuota parameter to the shared store.ApplyLearnedCapacityDecreaseIfEligible (used by both the load balancer and the standalone health checker) that short-circuits the decrease entirely for a daily-quota signal, leaving the long-cooldown recovery backoff already seeded for it (health.InitialBackoffForSignal) as the sole response. Covered independently for both call paths - the load balancer (TestApplyLearnedCapacityDecreaseSkipsForDailyQuotaSignal) and the standalone health checker (TestCreateStandaloneRateLimitHandlerSkipsCapacityEstimateForDailyQuota) - since the two are separate processes mutating the same persisted estimate. Raised in a PR project-aethermesh#37 review comment questioning whether 402/daily-quota exhaustion was kept separate from burst 429s in the adaptive estimator.
…te-limit scan defaultForwardRequestWithBodyFunc's embedded-rate-limit body scan already fails open when a 2xx JSON body exceeds maxRateLimitScanBodyBytes - streaming it through unmodified and unscanned rather than buffering it fully - but nothing asserted that behavior. Add a regression test with an oversized Alchemy body carrying an embedded rate-limit error beyond the scan cap, confirming it streams through verbatim and the endpoint is not marked rate limited. Raised in a PR project-aethermesh#37 review comment asking for coverage of the scan-cap boundary.
|
Creating a preprod image in order to test it out. |
2 participants
Summary
Aetherlay previously reacted to rate limiting with a single, blunt signal: an HTTP 429 on
the live proxy path, at three call sites, with no way to avoid triggering a provider's
limit in the first place. This PR adds three layers that work together - reactive
detection, proactive static throttling, and adaptive learned throttling - so the load
balancer can serve every request even when one or more configured providers are rate
limited or otherwise failing, without requiring any new configuration.
Linear: KEEP-905
How rate-limit detection works (reactive)
All rate-limit detection funnels through one new function,
health.DetectRateLimit(
internal/health/ratelimit.go), called at every request path (HTTP proxy, WebSockethandshake, and both the integrated and standalone health checkers):
exhausted," which can't be sped up by probing sooner (unlike a 429 burst limit), so it's
classified separately (
IsDailyQuota) and seeds a much longer recovery wait from theendpoint's own
max_backoff, not a short guess.Alchemy on batch requests) only signal rate limiting this way, invisible to a
status-code-only check. The HTTP proxy path now buffers and scans 2xx JSON bodies
(bounded to 10MB to protect memory on large legitimate responses like
eth_getLogs) forthis case, and retries the same request against a different endpoint instead of
forwarding a partial/rate-limited response to the client.
Retry-Afterheader, when present, is parsed and used to seed the endpoint'srecovery backoff precisely instead of guessing.
Once detected,
handleRateLimitmarks the endpoint'sRateLimitState.RateLimited = truein Valkey and kicks off the existing
RateLimitScheduler, which probes the endpoint onexponential backoff (real
eth_blockNumbercalls) until it's seen enough consecutiveclean responses to flip back to healthy.
How load balancing works
Every request goes through two stages, then a retry loop:
getEndpointsByRole) - for each endpoint in a role (public/primary/fallback), three independent gates must all pass: health status, thereactive
RateLimitedflag above, and a proactive capacity gate (below).selectBestEndpointByRole) - within the highest-priority non-emptyrole, picks the endpoint with the lowest utilization score. If every candidate has a
resolvable capacity ceiling (static or learned), scores are utilization ratios
(
24h request count / endpoint's own ceiling projected to a 24h-equivalent budget) soa higher-capacity endpoint isn't penalized for carrying proportionally more traffic. If
any candidate has no ceiling at all, this falls back to raw 24h request count (today's
original least-used heuristic) to avoid comparing endpoints on incompatible units.
the failing endpoint is dropped from that request's remaining candidates and retried
against the next one, up to
PROXY_MAX_RETRIES. A client-facing request only fails ifevery available endpoint in every role is exhausted.
Proactive capacity throttling (static, pre-existing shape extended here)
An endpoint can optionally declare
capacity: {max_requests, window_seconds}- aself-imposed ceiling checked via a Valkey window counter (
IncrementCapacityCount/GetCapacityCount) before every dispatch. Once an endpoint hits its configured ceilingfor the current window, it's excluded from candidates for that window - independent of
whether the provider has actually rejected anything yet.
Adaptive capacity estimation (new - no configuration required)
Since none of the major providers (Infura, Alchemy, dRPC) expose a live "remaining quota"
signal, guessing a static ceiling up front is unreliable. This PR adds a learned
alternative that requires zero config by default (
CAPACITY_LEARNING_ENABLED=true):capacityis neverproactively throttled until it has actually been rate limited once. The first hit seeds
an estimate from the real observed request count in that window (
GetCapacityCount),halved for safety margin - not a guess.
grounded in whichever is more conservative - the currently-effective (grown) estimate,
or the freshly observed count. A cooldown (the endpoint's own learning window) collapses
several near-simultaneous hits from one retry storm into a single decrease.
(recomputed per decrease, scaled to the endpoint's own magnitude) for every elapsed
interval - derived purely from elapsed wall-clock time at read time, with zero
background goroutines and zero counters to race across replicas.
effectiveCapacityCeilingfeeds the exact samegating and weighting logic as static
capacity- static config still wins when presentand is left completely untouched; learned estimates only ever engage for endpoints
without it.
services/health-checker/main.go), whichshares the identical
store.ApplyCapacityDecreasemath with the load balancer, sinceboth processes mutate the same persisted estimate.
Test plan
internal/server/failover_soak_test.go(new): drives the real HTTP forwarding path(not a stub) against local
httptestbackends and fires 100-150 requests per scenario,asserting every single request succeeds (
200) while one or more endpointsmisbehave simultaneously - always-429, always-500 ("or something else"), a mix of both,
and a statically-capped endpoint exhausting its ceiling mid-run.
internal/health/ratelimit_test.go,internal/store/capacity_estimate_test.go: pure,table-driven tests for detection and the AIMD math (decrease/cooldown/growth).
internal/server/capacity_test.go,capacity_learning_test.go: gating/weightingintegration tests for both static and learned ceilings, including regression coverage
confirming existing static-only behavior is unchanged.
go build ./...,go vet ./...,go test ./...,gofmt -l .all pass.Summary by CodeRabbit
Retry-After, daily-quota hints, and embedded JSON-RPC rate-limit errors inside HTTP 200).