Repository navigation
feat(global-router): add DynamoDB-backed session affinity - #17675
sachalmalick wants to merge 10 commits into
Conversation
Signed-off-by: Sachal Malick <s@chal.ai>
|
No actionable comments were generated in the recent review.
|
| Layer / File(s) | Summary |
|---|---|
Affinity contracts and in-memory store lib/global-router/Cargo.toml, lib/global-router/src/affinity/*, lib/global-router/src/lib.rs |
Defines affinity keys, bindings, storage and session-affinity interfaces, and an in-memory store with conditional writes and expiry sweeping. |
Asynchronous affinity selection lib/global-router/src/selector.rs, lib/global-router/src/affinity/selector.rs, lib/global-router/tests/session_affinity.rs |
Makes pool selection asynchronous. Affinity selection reads and conditionally updates bindings, retries conflicts once, and falls back to the best candidate on storage errors or timeout. Tests cover shared bindings, key isolation, expiry, conflicts, and failure handling. |
DynamoDB persistence Cargo.toml, lib/global-router-store-dynamodb/* |
Adds DynamoDB key and record encoding, reads, conditional writes, and integration tests for conflicts and persistence across selector restarts. |
Frontend configuration and routing lib/global-frontend/Cargo.toml, lib/global-frontend/src/affinity.rs, lib/global-frontend/src/main.rs, lib/global-frontend/src/proxy.rs, lib/global-frontend/tests/forwarding.rs |
Adds memory or DynamoDB affinity configuration and startup wiring. Request forwarding awaits pool selection. Tests cover configuration validation and shared in-memory affinity across frontend instances. |
Priority:
Estimated code review effort: 4 (Complex) | ~50 minutes
Merge Risk: 
Minimal · up to 501eb
-
No identified issue prevents merging after normal checks.
Pre-merge checks: 
4 · 
1
-
Failed checks: 1 warning Passed checks: 4 passed-
Full details: Description check
-
Explanation
The description includes all required sections: Overview, Key Changes, Validation, Issues and Ongoing Work, and the pull request checklist. It explains the implementation, validation results, limitations, and linked issues in sufficient detail.
Full details: Linked Issues check
-
Explanation
The changes satisfy the coding requirements in
#17674. The router adds replaceableAffinityStoreandSessionAffinitycontracts, retainsInMemoryAffinityStore, and adds the DynamoDB adapter. Keys include scope, policy, model, and session. The implementation supports conditional writes, idle expiry, configurable region and endpoint access, and bounded fail-open selection with one conflict retry and a shared timeout. Tests cover persisted reuse across selectors and restarts, key isolation, expiry, competing writes, backend failures, timeouts, and ambiguous writes. The DynamoDB adapter uses the same key and access paths for regional and Global Tables behavior. The reported cross-region divergence is consistent with the issue's accepted soft-affinity behavior.#11403,#17639, and#17599provide related context and do not add coding requirements for this PR.
Full details: Out of Scope Changes check
-
Explanation
The changes stay within
#17674. The new async selection path, frontend configuration, DynamoDB crate, item encoding, in-memory store, and tests directly support pluggable persistent session affinity. Supporting dependency changes and forwarding updates enable the new async selector contract. The PR does not implement the excluded provisional dispatch/acceptance lifecycle, peer replication of bookings or pins, affinity metrics, or forwarding failover.
-
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 4.35% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 23 functions across 15 files. (4 skipped: 4 unsupported.)
Resolution
Write docstrings for the functions missing them to satisfy the coverage threshold.
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
There was a problem hiding this comment.
Questions for the author:
- lib/global-router/src/selector.rs: Making the object-safe
PoolSelectorasync throughasync_traitadds a boxed future allocation to every request, including the supported round-robin and least-in-flight paths that remain synchronous; the memory-affinity path additionally boxes the nested affinity/store calls. Please provide before/after allocation counts and latency or throughput measurements for round-robin, least-in-flight, and in-memory affinity so the inference-path regression can be assessed.
nnshah1
left a comment
There was a problem hiding this comment.
Four items about the cost of affinity per request, in addition to the comments from @hhzhang16 and the review bot (I did not repeat them):
- Write the binding only when the pin changes or is half way to expiry (
tracker.rs). - Document that pins can split between regions on a multi-region table (
store-dynamodb/src/lib.rs). - The default in-memory path has more cost per request than
main(selector.rs). - Export the store cost as metrics, with the same names as the local router's
dynamo_router_overhead_*(selector.rs).
🤖 Generated with Claude Code, reviewed and posted by @nnshah1.
Review follow-ups on #17675, plus a read-through cache and a tighter storage budget. - Cache: `CachedAffinityStore` wraps a remote store with a TTL'd, capacity-bounded read-through cache, so a chatty session costs one lookup per TTL rather than one per request. A cached entry carries the expectation its write must match, so a binding another replica changed loses its conditional write, invalidates the entry and reads through: staleness costs a conflict, never a wrong pin. Configured with --affinity-cache-ttl-ms (default 1000) and --affinity-cache-entries. - Budget: --affinity-timeout-ms and AffinityConfig default from 100ms to 20ms, since the cache now absorbs most lookups. - Tenancy: AffinityKey gains `tenant`, a BLAKE3 digest of the request's `authorization` header derived at the edge, so two tenants reusing one session ID never share a binding and no credential reaches a key or a stored row. Selection still precedes the pool's own authorization; the README says so and says what that does and does not protect. - Degraded path: repeated storage failures now open a breaker for --affinity-breaker-cooldown-ms instead of only suppressing a log, so an outage costs one deadline per cooldown rather than one per request. A probe is claimed by pushing the deadline forward, so a probe whose client left cannot wedge the breaker open. The frontend exports affinity_lookups_total (hit/miss/conflict/error/timeout/skipped/repair) and affinity_cache_total. - Repair: `Expected` replaces the bare expected-revision option and adds `Unreadable`, the one expectation that can replace a row no reader can decode. Without it such a row is unreachable by any conditional write and TTL cannot sweep it when `expiresAt` is what is corrupt, stranding the session until the table is fixed by hand. Covered against DynamoDB Local, including the mixed-version case. - Expiry: the idle deadline is computed at full resolution and rounded up to the next Unix second, so a binding written near a second boundary is no longer cut short by up to a second. - CI: a new `rust-affinity-store-tests` job runs the DynamoDB Local tests against an `amazon/dynamodb-local` service container, so the adapter's conditional writes, TTL encoding and repair path are covered by the green checks rather than only by hand. - Docs: the Global Frontend README documents the supported consumer, table schema, TTL, IAM policy, scope and tenancy semantics, and startup configuration, and no longer says pins are per process. - `PoolSelector` and the affinity contracts are documented as internal to the Global Router, so the async `select` is not read as a stable extension point. - Test and comment fixes from the review bot: a non-discriminating policy-isolation assertion now builds a selector on a distinct policy and asserts an independent binding, and two comments that only narrated the next statement are gone. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Sachal Malick <s@chal.ai>
One conflict: main changed `lib/global-router/src/affinity.rs`, which this branch had split into `lib/global-router/src/affinity/`, so git saw a modify/delete. Main's change was confined to the `set_occupancy` test helper, which `PoolSignals::record_load` and `record_kv` now take a `RelaySource` for; that was ported into `affinity/selector.rs`, where the helper lives, and main's leftover copy of the old file was deleted. `lib/global-router/src/lib.rs` auto-merged with both sides kept: main's `RelaySource` export and this branch's `pub mod affinity`. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Sachal Malick <s@chal.ai>
Two resolutions: `Cargo.toml` had a textual conflict where both sides added a workspace path dependency on adjacent lines. Both are kept, with `dynamo-global-router-store-dynamodb` staying beside the other `global-*` crates and main's `dynamo-http-server` after it. `Cargo.lock` merged textually but not correctly. Main's new `dynamo-http-server` brings a second `rustls` (0.23.45 beside 0.21.12) and `tokio-rustls` (0.26.4 beside 0.24.1), so those dependency entries have to be version-qualified; the merged file left them bare, which `cargo --locked` rejects. Regenerated, and `cargo metadata --locked` now passes. Main's rewrite of `serve` onto per-connection shutdown signals (#17717) auto-merged with this branch's tenant extraction in `forward`, since they touch different parts of proxy.rs; `main.rs` picked up `serve` no longer returning a Result alongside this branch's affinity wiring. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Sachal Malick <s@chal.ai>
One conflict, in the `needs` list of `pre-merge-status-check` in `.github/workflows/pre-merge.yml`: main added `rust-feature-builds` (#17958) and this branch had added `rust-affinity-store-tests`, both before `rust-clippy`. Both are kept here, in job order; the job blocks themselves merged cleanly. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Sachal Malick <s@chal.ai>
`rust-affinity-store-tests` would have been a blocking check on every PR that touches Rust, pulling a Docker Hub image onto a hosted runner, to cover one crate that few PRs change. That cost falls on the whole repository rather than on this feature, so the job goes. The tests stay `#[ignore]`d. The header of tests/store.rs now says why, and gives the commands to start DynamoDB Local and run them. pre-merge.yml is back to main's exactly: the job and its entry in the status check's `needs` were this branch's only changes to it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Sachal Malick <s@chal.ai>
Four conflicts. lib/global-frontend/src/main.rs: main (#17716) now creates `in_flight` and `PoolSignals` before choosing a selector, so Relay subscriptions can gate the pool catalog under every policy, and it moved `follow_relays` out of the affinity arm. This branch's re-declaration of `in_flight` before the selector is dropped: kept, it would shadow main's tracker, so the proxy would count requests in one tracker while affinity's signals read another. Main's early `Metrics::new` is dropped in favour of this branch's, which runs after the selector so it can attach the affinity counters. The "--pool-relay is used only by affinity" warning stays deleted, as main intended, since relays now gate every policy; the DynamoDB-requires-affinity check stays. lib/global-frontend/Cargo.toml: both sides added dependencies; all are kept. lib/global-frontend/README.md: both sides added a section at the same place. This branch's "Session affinity storage" subsection stays under "Pool selection", followed by main's new "Relay-initiated state". Cargo.lock: regenerated from main's, rather than trusting a textual merge. Every locked version now matches either main or this branch; four AWS SDK crates that re-resolved to newer releases were pinned back to this branch's versions, so the merge brings in no upgrades of its own. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Sachal Malick <s@chal.ai>
Addresses the second review round on this PR.
`bind` wrote a new revision on every session request, even when the
session stayed on its pool: one conditional write per turn, conflicts
between parallel requests of one session, and on an MREC table a pin
that could flip between regions for as long as both saw traffic. The
snapshot now carries the pin's `expires_at`, and `bind` writes only for
a session with no live pin, a move, an undecodable record, or a pin past
half its idle window. Most turns only read; behind the cache, a kept pin
is usually no store call at all.
Since a kept pin no longer writes, a replica's cached entry for a
session another replica moved is used until it ages out, at most the
cache TTL. It still cannot overwrite: any write from it carries the old
revision and loses. The cache docs and README say so.
The MREC docs now say pins are soft, last writer wins between regions.
`AffinityKey` shares `scope` and `policy` as `Arc<str>`, so the memory
store's per-request cost drops along with the skipped write.
Metrics follow the local router's routing-overhead names:
`overhead_affinity_store_ms`, a histogram of time in storage bucketed
like `dynamo_router_overhead_shared_cache_query_ms`;
`affinity_store_errors_total{kind}`, one per request that routed
without affinity; and `affinity_store_writes_total{reason}`. The
metrics catalog lists every affinity family.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Sachal Malick <s@chal.ai>
|
/ok to test f75a6fc |
Two conflicts, both from main's DynamoDB registration store for the Global Discovery Service (#17691), which added the same AWS SDK workspace dependencies this branch had. Cargo.toml: main's declarations are taken. They pin lower minimum versions than this branch did and turn off aws-sdk-dynamodb's default features, which pulled in a legacy hyper 0.14 and rustls 0.21 client. Both DynamoDB stores now share one SDK and one HTTPS client. Cargo.lock: main's lock, with Cargo adding only this branch's crate and its dependencies; no locked version differs from main. The AWS SDK crates therefore move to main's versions, and the legacy client's packages leave the lock. The router, DynamoDB store and frontend tests pass against these versions, and clippy is clean. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Sachal Malick <s@chal.ai>
|
/ok to test 73192ac |
…tacked PR The metrics catalog is under /docs/, so its entries for the affinity metrics put this PR on the docs team's review as well. They move to a PR stacked on this one, so the code can be reviewed by the global-router owners alone. The crate README stays here: the global-router owners own it, and it documents how to deploy the store. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Sachal Malick <s@chal.ai>
|
/ok to test |
@sachalmalick, there was an error processing your request: See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/1/ |
|
/ok to test 02dc35a |
Overview
Add pluggable session affinity with an optional DynamoDB store to the Global Router. Frontend replicas and restarts can reuse persisted session-to-pool bindings while retaining the existing headroom placement policy and in-memory default.
Key Changes
574 added production Rust lines, excluding tests, whole-line comments, and blanks.
SessionAffinitybehavior from the typedAffinityStore, with idle expiry, revision-based conditional writes, and bounded fail-open selection.PoolSelector::selectis now async and the affinity constructor accepts an injected tracker. AWS dependencies compile into the frontend even with memory selected, but AWS initialization occurs only for DynamoDB.Selection and consistency
A session request reads its binding, applies the existing pool eligibility/headroom rules, then conditionally writes or refreshes the binding. One conflict retry and all storage work share a configurable deadline (100 ms by default); failures retain the selected candidate and ambiguous writes are not rolled back.
DynamoDB normally adds one sequential read and conditional write before forwarding; a conflict can repeat both. Async trait dispatch also adds boxed futures, including on other selection policies. Throughput and latency impact have not been benchmarked.
Regional tables and Global Tables MREC use the same schema and accessors. Conditional writes serialize against the local replica; cross-region concurrent writes can select different pools before replication converges. This is soft affinity, not global mutual exclusion. Table provisioning, replicas, and TTL on
expiresAtare deployment configuration.Validation
57 tests passed with Rust 1.96.1 across the router, adapter, and frontend, including these behavior checks:
independent_frontends_share_session_affinityindependent_selectors_and_restarts_reuse_persisted_bindingsconditional_writes_expiration_and_schema_match_the_store_contractThe two DynamoDB Local tests are explicitly ignored by default and were run manually with
--include-ignored; current CI covers the other workspace tests. MREC replication was not exercised.Frontend validation on macOS required a temporary external patch guarding an existing Linux-only call in
velo 0.20.0. That patch is absent from this PR; repository source and the registry dependency remain unchanged.Issues and Ongoing Work
Pull Request Checklist
/pr-checklistskill.Summary by CodeRabbit