Skip to content

feat(global-router): add DynamoDB-backed session affinity - #17675

Open
sachalmalick wants to merge 10 commits into
mainfrom
sachalmalick/global-router-dynamodb-affinity
Open

sachalmalick wants to merge 10 commits into
mainfrom
sachalmalick/global-router-dynamodb-affinity

Conversation

@sachalmalick

@sachalmalick sachalmalick commented Oct 30, 2026 •

Copy link
Copy Markdown
Contributor

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.

  • Affinity contracts: separate SessionAffinity behavior from the typed AffinityStore, with idle expiry, revision-based conditional writes, and bounded fail-open selection.
  • DynamoDB adapter: a dedicated crate isolates AWS access and item encoding from routing logic; scope/session partition keys and model/policy sort keys isolate bindings.
  • Frontend configuration: select memory or DynamoDB, configure scope, table, region, endpoint, consistent reads, and timeout; reuse one SDK client.

PoolSelector::select is 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 expiresAt are deployment configuration.

Validation

57 tests passed with Rust 1.96.1 across the router, adapter, and frontend, including these behavior checks:

Test Evidence
independent_frontends_share_session_affinity HTTP requests through two frontends reuse one pool binding via shared memory storage.
independent_selectors_and_restarts_reuse_persisted_bindings Fresh selectors reuse DynamoDB Local bindings.
conditional_writes_expiration_and_schema_match_the_store_contract DynamoDB Local enforces competing writes, expired-record replacement, key isolation, and schema validation.
Failure/timeout contract tests Failed reads do not write; stalls respect the total budget; ambiguous writes are not rolled back; conflict retries are bounded.

The 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
  • I used Dynamo's /pr-checklist skill.

Summary by CodeRabbit

  • New Features
    • Added session affinity, allowing requests with the same session to prefer the same eligible pool across frontend instances.
    • Added in-memory and DynamoDB-backed affinity storage options. DynamoDB storage can preserve session bindings across restarts.
    • Added configurable affinity scope, routing policy, timeout, and storage settings.
  • Bug Fixes
    • Affinity lookup or storage issues no longer prevent routing; requests fall back to the best available candidate.
    • Requests without a session continue to route without using affinity storage.

Signed-off-by: Sachal Malick <s@chal.ai>
@sachalmalick
sachalmalick requested review from a team as code owners October 30, 2026 08:17
@copy-pr-bot

copy-pr-bot Bot commented Oct 30, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@coderabbitai

coderabbitai Bot commented Oct 30, 2026

Copy link
Copy Markdown
Contributor

Review in Change Stack →

No actionable comments were generated in the recent review.

Recent review info
Run configuration
Commits

Reviewing files that changed from the base of the PR and between a416113 and 501eb21.

Files ignored due to path filters: 1
  • Cargo.lock is excluded by !**/*.lock
Files selected for processing: 19
  • Cargo.toml
  • 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
  • lib/global-router-store-dynamodb/Cargo.toml
  • lib/global-router-store-dynamodb/src/item.rs
  • lib/global-router-store-dynamodb/src/lib.rs
  • lib/global-router-store-dynamodb/tests/store.rs
  • lib/global-router/Cargo.toml
  • lib/global-router/src/affinity/memory.rs
  • lib/global-router/src/affinity/mod.rs
  • lib/global-router/src/affinity/selector.rs
  • lib/global-router/src/affinity/store.rs
  • lib/global-router/src/affinity/tracker.rs
  • lib/global-router/src/lib.rs
  • lib/global-router/src/selector.rs
  • lib/global-router/tests/session_affinity.rs

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

Walkthrough

The Global Router now supports asynchronous session affinity backed by interchangeable stores. The pull request adds an in-memory store, a DynamoDB implementation, and Global Frontend options to configure and use either backend.

Changes

Session affinity

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: Normal

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: Passed 4 · Failed 1

Failed Failed checks: 1 warning

Check name Status Explanation Resolution
Docstring Coverage Warning 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. Documentat… Write docstrings for the functions missing them to satisfy the coverage threshold.
Passed Passed checks: 4 passed
Check name Status Explanation
Title check Passed The title clearly and concisely identifies the main change: adding DynamoDB-backed session affinity to the Global Router.
Description check Passed 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, limitati…
Linked Issues check Passed The changes satisfy the coding requirements in #17674. The router adds replaceable AffinityStore and SessionAffinity contracts, retains InMemoryAffinityStore, and adds the DynamoDB adapter. Keys…
Out of Scope Changes check Passed 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. …
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 replaceable AffinityStore and SessionAffinity contracts, retains InMemoryAffinityStore, 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 #17599 provide 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.

@dynamo-review-agent dynamo-review-agent Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Questions for the author:

  • lib/global-router/src/selector.rs: Making the object-safe PoolSelector async through async_trait adds 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.

Comment thread lib/global-frontend/src/proxy.rs Resolved
Comment thread lib/global-router/src/affinity/mod.rs Resolved
Comment thread lib/global-router/src/selector.rs Resolved
Comment thread lib/global-frontend/src/affinity.rs Resolved

@nnshah1 nnshah1 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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):

  1. Write the binding only when the pin changes or is half way to expiry (tracker.rs).
  2. Document that pins can split between regions on a multi-region table (store-dynamodb/src/lib.rs).
  3. The default in-memory path has more cost per request than main (selector.rs).
  4. 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.

sachalmalick and others added 5 commits November 3, 2026 11:36
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>
@github-actions github-actions Bot added the documentation Improvements or additions to documentation label Nov 4, 2026
sachalmalick and others added 2 commits November 4, 2026 13:59
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>
@sachalmalick
sachalmalick requested a review from a team as a code owner November 5, 2026 07:17
@sachalmalick

Copy link
Copy Markdown
Contributor Author

/ok to test f75a6fc

@github-actions

github-actions Bot commented Nov 5, 2026 •

Copy link
Copy Markdown
Contributor

@sachalmalick
sachalmalick enabled auto-merge (squash) November 5, 2026 10:17
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>
@sachalmalick

Copy link
Copy Markdown
Contributor Author

/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>
@sachalmalick

Copy link
Copy Markdown
Contributor Author

/ok to test

@copy-pr-bot

copy-pr-bot Bot commented Nov 5, 2026

Copy link
Copy Markdown

/ok to test

@sachalmalick, there was an error processing your request: E1

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/1/

@sachalmalick

Copy link
Copy Markdown
Contributor Author

/ok to test 02dc35a

@sachalmalick
sachalmalick disabled auto-merge November 5, 2026 11:39
@sachalmalick
sachalmalick enabled auto-merge (squash) November 5, 2026 11:39

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation feat size/XXL

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[FEATURE]: Add pluggable persistent session affinity to the Global Router

5 participants