[runzhech]

Description

During RdmaContext initialization, the context remained in the DEVICE_DISABLED state while RDMA resources were being allocated. If initialization failed after partially creating resources, the error path called disable() for rollback. However, disable() returned immediately for a DEVICE_DISABLED context, leaving resources such as the device context, protection domain, epoll file descriptor, completion channels, and completion queues unreleased.

Because the state machine had no explicit initialization state, it could not distinguish a fully disabled context from a partially initialized one. Cleanup therefore relied on an inaccurate logical state instead of the resources actually owned by the context, potentially causing resource leaks and inconsistent lifecycle behavior after initialization failures.

To fix it, I've made the following changes:

• add an explicit DEVICE_INITIALIZING state
• roll back partially initialized RDMA resources based on actual handles
• make context cleanup idempotent and independent of disabled state
• prevent inbound connections and MR operations during initialization
• clean up residual resources during context destruction

Module

• Transfer Engine (mooncake-transfer-engine)
• Mooncake Store (mooncake-store)
• Mooncake EP (mooncake-ep)
• Mooncake PG (mooncake-pg)
• Integration (mooncake-integration)
• P2P Store (mooncake-p2p-store)
• Python Wheel (mooncake-wheel)
• Common (mooncake-common)
• Mooncake RL (mooncake-rl)
• CI/CD
• Docs
• Other

Type of Change

• Bug fix
• New feature
• Refactor
• Breaking change
• Documentation update
• Performance improvement
• Other

How Has This Been Tested?

Test results:

• Unit tests pass
• Integration tests pass (if applicable)
• Manual testing done (describe below)

The test I've performed manually is described below:

Mock or wrap the ibv_* calls used by RdmaContext::enable() and inject failures at each initialization stage, including PD, epoll, completion channel, CQ, notification CQ, and port query creation.

For each failure, verify that:

initialization returns an error;
context status returns to DEVICE_DISABLED;
all previously created resources are released exactly once;
all handles, pointers, and containers are reset.
Call disable() again to verify idempotent cleanup, then remove the failure and call enable() to verify successful retry.

Checklist

• I have performed a self-review of my own code
• I have formatted my code using ./scripts/code_format.sh
• I have run pre-commit run --all-files and all hooks pass
• I have updated the documentation (if applicable)
• I have added tests to prove my changes are effective
• For changes >500 LOC: I have filed an RFC issue

AI Assistance Disclosure

• No AI tools were used
• AI tools were used (specify below)

VSCode Copilot assisted with reading and editing the source files. I have reviewed every changed line and can defend the change end-to-end.

[Copilot]

Pull request overview

This PR fixes RDMA resource leaks when RdmaContext::enable() fails part-way through initialization by introducing an explicit DEVICE_INITIALIZING state and making cleanup independent of the prior logical “disabled” state.

Changes:

• Add DEVICE_INITIALIZING and switch enable() to transition DISABLED -> INITIALIZING -> (ENABLED|PAUSED) with explicit rollback on failure.
• Refactor teardown into cleanupResources() / rollbackEnable() and make destruction clean up any residual partially-created RDMA handles.
• Block inbound RDMA connection setup and MR operations while a context is initializing.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.

File	Description
mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp	Tightens inbound connection gating to only allow ENABLED/PAUSED contexts.
mooncake-transfer-engine/tent/src/transport/rdma/context.cpp	Adds DEVICE_INITIALIZING, introduces rollback/cleanup helpers, and updates enable/disable + MR guards.
mooncake-transfer-engine/tent/include/tent/transport/rdma/context.h	Exposes the new status and declares cleanup/rollback helpers.

Suppressed comments (2)

mooncake-transfer-engine/tent/src/transport/rdma/context.cpp:620

• The added DEVICE_INITIALIZING guard now trips this LOG(FATAL) path too, but the message says "not constructed". During initialization the context is constructed, it’s just not ready/enabled yet, which makes debugging harder.

    if (status_ == DEVICE_DISABLED || status_ == DEVICE_INITIALIZING ||
        status_ == DEVICE_UNINIT) {
        LOG(FATAL) << "RDMA context " << name() << " not constructed";
        return -1;
    }

mooncake-transfer-engine/tent/src/transport/rdma/context.cpp:645

• The added DEVICE_INITIALIZING guard now trips this LOG(FATAL) path too, but the message says "not constructed". During initialization the context is constructed, it’s just not ready/enabled yet, which makes debugging harder.

    if (status_ == DEVICE_DISABLED || status_ == DEVICE_INITIALIZING ||
        status_ == DEVICE_UNINIT) {
        LOG(FATAL) << "RDMA context " << name() << " not constructed";
        return -1;
    }

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated no new comments.

Suppressed comments (2)

mooncake-transfer-engine/tent/src/transport/rdma/context.cpp:488

• cleanupResources() iterates and clears mr_set_ without holding mr_set_mutex_, while registerMemReg()/unregisterMemReg() mutate the same set under that mutex. This breaks the locking discipline and can lead to a data race/UB if cleanup runs concurrently with MR registration/unregistration.

void RdmaContext::cleanupResources() {
    for (auto& entry : mr_set_) {
        int ret = verbs_.ibv_dereg_mr(entry);
        if (ret) PLOG(ERROR) << "ibv_dereg_mr";
    }
    mr_set_.clear();

mooncake-transfer-engine/tent/include/tent/transport/rdma/context.h:130

• The comment above cleanupResources() says resources are released in “reverse construction order”, but the implementation does not strictly follow that order (e.g., event_fd_ is closed before completion channels are destroyed). Either update the comment or reorder cleanup to match, to avoid misleading future maintainers.

    // Release every resource currently owned by this context in reverse
    // construction order. This is intentionally state-independent so it can
    // roll back a partially completed enable() and can safely be retried.
    void cleanupResources();

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

Suppressed comments (3)

mooncake-transfer-engine/tent/src/transport/rdma/context.cpp:540

• The error message says "not constructed" even when the context is constructed but merely DISABLED/INITIALIZING. This makes failures harder to diagnose (especially now that INITIALIZING is a distinct state). Consider updating the message to reflect the actual invalid state (e.g., "not enabled" / "not ready" and ideally include the current status) and apply the same fix to the similar LOG(FATAL) messages in warmup/unregister paths.

    if (status_ == DEVICE_DISABLED || status_ == DEVICE_INITIALIZING ||
        status_ == DEVICE_UNINIT) {
        LOG(FATAL) << "RDMA context " << name() << " not constructed";
        return nullptr;

mooncake-transfer-engine/tent/src/transport/rdma/context.cpp:315

• The logic for "do we currently own any resources that need cleanup" is embedded as a long predicate here (and a similar but not identical predicate exists in disable()). Consider extracting a small helper (e.g., hasOwnedResources()) and reusing it in both places to reduce the chance of future resource fields being missed or the predicates drifting apart.

    if (status_ != DEVICE_UNINIT || native_context_ || native_pd_ ||
        event_fd_ >= 0 || !comp_channel_.empty() || !cq_list_.empty() ||
        notify_cq_) {
        disable();
    }

mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp:662

• This error is now emitted for states beyond "down" (e.g., DEVICE_INITIALIZING and DEVICE_UNINIT). To improve debuggability, consider including the actual context status in the log message (or using a status-specific message) so callers can distinguish "initializing" vs "disabled" vs "uninitialized".

    auto ctx_status = context->status();
    if (ctx_status != RdmaContext::DEVICE_ENABLED &&
        ctx_status != RdmaContext::DEVICE_PAUSED) {
        std::stringstream ss;
        ss << "Device is down: " << peer_desc.local_nic_path;
        LOG(ERROR) << ss.str();

[staryxchen]

Thanks for the contribution. The core leak fix is correct, but I think the state-machine extension is more than the current call graph requires.

cleanupResources() (nullity-based, state-independent cleanup) is the real fix for the PD / epoll / comp_channel / CQ / device_context leaks — that part belongs in disable() and a simpler disable() early-return on resource-nullity would suffice. The std::make_unique() change is also correct and necessary: a CQ
that fails construct() never enters cq_list_, so disable() can't reclaim it — this must be fixed at the allocation site.

What I'd question is DEVICE_INITIALIZING + rollbackEnable(). In the TENT call graph, enable() is only called from construct() (single-threaded install()), and disable() only from ~RdmaContext() and enable()'s own failure path — both same-thread, externally serialized by the transport lifecycle. There's no concurrent
caller the new state coordinates against today. Once disable()'s early-return is resource-based, enable() could keep calling disable() directly; rollbackEnable() is needed only because disable() now rejects DEVICE_INITIALIZING.

The new state + CAS are defensible as future-proofing, but for the present code they add machinery without addressing a current race. I'd lean toward trimming them unless there's a concurrent caller I'm missing.

[runzhech]

@staryxchen As you noted, there's no concurrent caller today that would need the CAS/state machinery. I've made the following changes:

• Removed DEVICE_INITIALIZING and rollbackEnable().
• Made disable() resource-driven: early-return on empty resources, do actual cleanup via cleanupResources().
• enable() now calls disable() directly on failure.

Please take another look.

[staryxchen]

mr_set_ is iterated and cleared here without holding mr_set_mutex_, while every other access site takes the lock. In the current call graph disable() runs during teardown (same thread, after workers are stopped), so this isn't a live race today — but the broken locking discipline could bite if teardown ordering ever changes. Consider holding mr_set_mutex_ around the iterate-and-clear to keep the invariant, or add a comment documenting the relies-on-teardown-ordering assumption

[staryxchen]

the comment says "reverse construction order", but the actual cleanup isn't strictly reverse — cq_list_ is cleaned before notify_cq_ (constructed after it), and event_fd_ before comp_channel_ (also constructed after). Functionally the order is fine (no interdependency requiring strict reverse), but the comment misleads future readers. It's better to drop "reverse construction order" from the comment.

[runzhech]

@staryxchen
The latest commit made the following 2 changes:
Held mr_set_mutex_ around the mr_set_ iterate-and-clear in cleanupResources() to match the locking discipline of the other access sites.
Removed "reverse construction order" from the cleanupResources() comment, since the actual cleanup order is not strictly reverse.
Thanks for catching those.

[staryxchen]

Thanks for the contribution! LGTM