kvcache-ai / kvcache-ai/Mooncake
[RFC]: RDMA Endpoint and Slice Completion Ownership
- Dominant language
- C++
- Stars
- 6.6k
- Forks
- 1.2k
- Avg merge
- 3d 5h
- Merged PRs (30d)
- 312
Description
### Changes proposed
## Summary
This RFC proposes a clearer ownership model for RDMA endpoint lifecycle,
posted work requests, completion handling, and slice retry. The immediate
trigger is the concurrency risk exposed by the dedicated RDMA poller model,
where one worker drains all completion queues while other workers continue
posting sends to the same endpoints.
The short-term mitigation is to keep the historical shared worker model as the
default and make the dedicated poller opt-in. The long-term fix should define
per-post completion ownership so stale completions cannot mutate a slice that
has already been retried on another endpoint.
## Background
RDMA transfer work is represented by `Transport::Slice`. A slice currently
serves several roles at once:
- It is the user-visible transfer fragment.
- It is the `wr_id` token returned by CQ polling.
- It carries retry state such as `retry_cnt`, peer NIC path, destination rkey,
source lkey, endpoint pointer, QP depth pointer, and status.
Historically, each worker both posted its assigned send shards and polled a
disjoint shard of CQs. The dedicated poller change separated those roles:
worker 0 polls all CQs, while workers 1..N post send shards. That can improve
polling latency and CQ drain throughput, but it also concentrates completion
error handling and endpoint deletion on one thread while multiple posting
threads continue to use the same endpoint store and endpoint objects.
## Problem Statement
The current model lacks a precise ownership boundary between:
- A post worker that prepares and posts a slice.
- The CQ poller that receives the completion for that post.
- Endpoint lifecycle paths that retire, delete, or reclaim endpoints.
- Retry paths that rewrite the same `Slice` object for a new attempt.
The main risk is not only concurrent polling of the same CQ. It is that a
single `Slice` object may be reused across retry attempts while an old WR
completion can still arrive. Because the completion token is the mutable
`Slice*` itself, the poller cannot tell whether the completion belongs to the
current attempt or a stale attempt.
This can lead to incorrect behavior such as:
- A stale completion marking a retried slice successful or failed.
- A stale completion deleting an endpoint using an old raw endpoint pointer.
- WR/QP depth counters being decremented through a pointer that belongs to an
endpoint being retired or reclaimed.
- Retry count and final slice status being updated more than once for the same
logical attempt.
- Endpoint retire paths racing with new posts because deletion entry points are
spread across completion errors, async events, passive handshakes, and
endpoint-store eviction.
## Current Worker Models
### Historical Shared Model
Each worker posts and polls a disjoint shard:
```text
worker 0: post shard 0/N/... and poll CQ 0/N/...
worker 1: post shard 1/N/... and poll CQ 1/N/...
worker 2: post shard 2/N/... and poll CQ 2/N/...
```
This model avoids concurrent polling of the same CQ and keeps post/poll work
more local to each worker, but it does not fully define endpoint or slice
ownership.
### Dedicated Poller Model
Worker roles are separated:
```text
worker 0: poll all CQs
worker 1..N-1: post all send shards
```
This model can improve CQ drain behavior, but it exposes more cross-thread
interleavings between posting, completion error handling, retry, and endpoint
retirement.
## Goals
- Define when a completion is allowed to mutate a `Slice`.
- Ensure stale completions only release resources owned by that posted WR.
- Avoid use-after-free and stale raw pointer decisions during endpoint retire
and reclaim.
- Preserve retry behavior across peer rail failover and local RNIC handoff.
- Make dedicated polling safe enough to be enabled deliberately by operators.
## Non-goals
- Redesign all RDMA transport path-selection logic.
- Remove endpoint pooling or the SIEVE/FIFO endpoint stores.
- Require dedicated polling to be enabled by default.
- Optimize every hot path before the ownership model is correct.
## Proposed Direction
### 1. Keep Dedicated Polling Opt-in #3090
Until the ownership model is fixed, the default worker mode should remain the
historical shared model. Dedicated polling should be an explicit opt-in for
deployments that understand the risk and benefit from the performance behavior.
This is a mitigation, not a complete correctness fix.
### 2. Introduce a Per-post Completion Token
`wr_id` should identify a posted attempt, not the reusable `Slice` object alone.
For example:
```cpp
struct PostedSlice {
Transport::Slice* slice;
std::atomic* qp_depth;
std::atomic* cq_outstanding;
RdmaEndPoint* endpoint;
uint64_t endpoint_generation;
uint32_t slice_attempt;
};
```
The exact allocation strategy is open for discussion. The important property is
that each posted WR owns a stable token until its completion has been processed.
Completion handling should:
1. Release the resources associated with that posted attempt, such as QP depth
and CQ outstanding counters.
2. Check whether the token still matches the slice's current attempt.
3. Mutate the slice status or retry state only when the completion matches the
current attempt.
4. Treat stale completions as resource-release events, not as current slice
completion events.
### 3. Add Slice Attempt Identity
Each slice should carry a monotonically increasing attempt id. Retry or handoff
increments the attempt before the slice is re-enqueued. A completion token
records the attempt id used at post time.
If a completion arrives for an old attempt, the poller must not mark the slice
successful, mark it failed, retry it again, or delete the current endpoint based
on the stale token.
### 4. Add Endpoint Generation
Each endpoint instance should have a generation or retire epoch. A post token
records the generation used at post time. Endpoint retirement should invalidate
future posts on that generation, but allow completions for already-posted WRs
to drain and release their counters.
Endpoint generation also makes it explicit when a raw endpoint pointer observed
from an old completion no longer refers to the active endpoint for that peer.
### 5. Unify Endpoint Retirement
All endpoint deletion paths should converge on one retire API:
- CQ completion error handling.
- `IBV_EVENT_QP_FATAL`.
- Context fatal events and GID changes.
- Passive handshake stale-endpoint handling.
- Endpoint-store eviction.
The retire API should remove the endpoint from the active map, mark it
non-postable, and defer physical QP destruction until all posted attempts have
completed or a bounded timeout policy is applied.
## Alternatives Considered
### Worker-level Post/Poll Mutex
A striped mutex around `submitPostSend()` and completion handling can reduce a
small race window, but it does not cover async events, passive handshakes,
endpoint-store eviction, `disconnectAllEndpoints()`, or reclaim. It is therefore
not sufficient as a correctness boundary.
### Shared Worker Mode Only
Keeping shared worker mode permanently avoids the dedicated-poller exposure, but
it leaves endpoint and slice ownership ambiguous. It is useful as a mitigation,
not as the long-term model.
### Endpoint Locks Only
Endpoint locks protect endpoint-local state, but they do not identify whether a
completion belongs to the current slice attempt. Slice attempt identity is still
needed.
## Compatibility and Migration
The short-term mitigation is backward-compatible:
- Keep shared worker scheduling as the default.
- Keep dedicated polling available behind `MC_RDMA_DEDICATED_POLLER`.
The long-term token model can be introduced internally without changing the
public Transfer Engine API. It may change memory allocation behavior on the RDMA
hot path, so the implementation should include benchmarks for high-throughput
small-slice workloads.
## Testing Plan
The correctness tests should include fake or controllable RDMA paths that can
force specific interleavings:
- Post succeeds, slice is retried, then old completion arrives.
- Post partially fails and returns a `bad_wr` chain.
- Completion error retires an endpoint while another worker is posting.
- Async QP fatal races with completion handling.
- Endpoint reclaim attempts to destroy QPs while completions are still pending.
- Local RNIC handoff rewrites a slice while old flush completions are draining.
The tests should verify:
- No slice reaches both success and failed states.
- Retry count advances only for the owning attempt.
- QP depth and CQ outstanding counters never go negative or leak.
- Stale completions do not delete or mutate the active endpoint generation.
- Waiting-list endpoints are eventually reclaimed after completions drain.
## Open Questions
- Should `PostedSlice` be allocated per WR, per batch, or from an endpoint-local
pool?
- Should `wr_id` point to a token object, or encode an index into a pool?
- Should stale completions be logged at trace level for diagnosis, or remain
silent on the hot path?
- What timeout policy is acceptable when a retired endpoint never drains all WRs?
- Can dedicated polling become the default again after this ownership model is
implemented and tested?
### Before submitting a new issue...
- [ ] Make sure you already searched for relevant issues and read the [documentation](https://kvcache-ai.github.io/Mooncake/)
Contributor guide
Research direction
Start by tracing submitPostSend(), CQ completion handling, endpoint retirement paths, and the MC_RDMA_DEDICATED_POLLER configuration described here. Use controllable RDMA interleavings to assess the proposed per-post token, slice-attempt, and endpoint-generation model. Done means stale completions release only their own resources, current slices are not incorrectly mutated, and the listed counters and endpoint reclamation tests pass.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- networking
- Issue type
- Refactor
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100