[KastanDay]

Description

Client::MountSegmentAndGetId registers the buffer with the transfer engine, then asks the master to MountSegment it. When that RPC fails (the master restarting or briefly unreachable), it returned the error without unregistering the buffer. Its callers (allocate_and_mount_segment, MountSegment) treat a failed mount as nothing mounted and free the buffer, so the registration outlives its memory:

• Every later mount at that address fails. The next buffer the allocator places there is rejected as Transfer Engine does not support overlapped memory region (ERR_ADDRESS_OVERLAPPED), so a store that adds segments at runtime can never grow again until the process restarts.
• On RDMA the freed region stays pinned (64 GiB per segment in our case), invisible to the store's own accounting.

The fix unregisters the buffer on that error path. The unregistration uses the default update_metadata=true, mirroring the registration just above it and UnmountSegmentImpl. ERR_ADDRESS_NOT_REGISTERED is tolerated there, as elsewhere in this file.

#2336 contains the same rollback, as part of a larger rework of the destructor and unmount paths that is marked "DO NOT Merge Now". This PR is only the mount-failure fix, with a regression test.

Module

• Transfer Engine (mooncake-transfer-engine)
• Mooncake Store (mooncake-store)
• Mooncake Conductor (mooncake-conductor)
• Reshard (mooncake-reshard)
• 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?

Unit test. The new NonHAReconnectTest.FailedMountLeavesBufferUnregistered stops the in-process master, mounts a buffer (registration succeeds, MountSegment fails), restarts the master on the same ports, and mounts the same buffer again.

Test commands (Ubuntu 22.04 with dependencies.sh, as in CI):

cmake -S . -B build-test -DBUILD_UNIT_TESTS=ON -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=OFF -DWITH_EP=OFF -DSTORE_USE_ETCD=ON -DCMAKE_BUILD_TYPE=Release
cmake --build build-test --target non_ha_reconnect_test
./build-test/mooncake-store/tests/non_ha_reconnect_test

Test results: on this branch and on a backport to v0.3.13.post1:

	With this fix	With client_service.cpp reverted
FailedMountLeavesBufferUnregistered	passes (3.0 s)	fails (12 s): every remount gets register_local_memory_failed ... error=-7 ("overlapped memory region")
the other 3 NonHAReconnectTest tests	pass	pass

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

End to end on RDMA. Our setup: stores that add and free 64 GiB segments at runtime with allocate_and_mount_segment / unmount_and_free_segment, on 8×H200 nodes with 8 RDMA NICs each.

• Wheel: mooncake-transfer-engine-cuda13 for CPython 3.12, built from v0.3.13.post1 plus this commit with the release workflow's cuda13 x86_64 profile in pytorch/manylinux2_28-builder:cuda13.0. It has the same 85 files as the PyPI release.
• Master restarted 4 s into a round of 64 GiB allocate_and_mount_segment calls on 3 stores: each mount failed (mount_segment_to_master_failed ... error=RPC_FAIL). Every store kept running. VmPin stayed at exactly the live segments' registrations, with no extra 64 GiB per NIC. After the master returned, all 3 grew from 2 to 8 segments with no failures and no overlapped-region errors.
• Master scaled to 0 while all 3 stores tried to add a segment: the same result. With the master back, all reached 10 segments.
• Same outage, stock wheel next to the fix: one store on the stock 0.3.13.post1 wheel and two on the fix. The stock store's pinned memory grew by 64 GiB × 8 NICs (the stranded registration), and it had to restart. The two stores with the fix were unaffected.

Checklist

• I have performed a self-review of my own code
• I have formatted my code using ./scripts/code_format.sh (the changed files pass clang-format 20.1.8, the version the script requires)
• I have run pre-commit on the files changed in this PR 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)

[he-yufeng]

Verified the leak path against current main (245e710) before writing this:

• Client::MountSegmentAndGetId registers with the transfer engine at client_service.cpp:3901 and then calls master_client_.MountSegment at :3922. On RPC failure it returns without unregistering, and since the segment never entered mounted_segments_ (insertion happens only after a successful master mount), UnmountSegment can't find it later (INVALID_PARAMS), so no in-process cleanup path exists.
• The failure is permanent from the process's point of view: registerLocalMemory rejects any later buffer overlapping that address with ERR_ADDRESS_OVERLAPPED (transfer_engine_impl.cpp:766 and friends), which matches the "store can never grow again" symptom in the description.
• The in-repo caller I checked (Client::MountSegment at :3816) just propagates the error, so on the RDMA path the pinned pages indeed outlive the freed buffer.

The fix is placed correctly: same mutex scope and same ERR_ADDRESS_NOT_REGISTERED tolerance as the existing unregister in UnmountSegmentImpl, and there is no double-unregister risk because the failed segment never made it into mounted_segments_. The test is the right regression shape too: master down, mount fails, master back, the same buffer must mount, which fails deterministically if the registration leaked.

One note, not a request: if the master actually applied the mount but the response was lost, the client now unregisters and frees while the master keeps the segment record. That ambiguity predates this change (the old code ended in the same state, plus the leak), so it shouldn't block anything here, but it may be worth a follow-up look in the reconnect/remount path.

Scope of this review: code-path read against current main; I did not rebuild locally.

[he-yufeng]

Read this against current main and traced the engine paths it touches. The bug and the fix both check out.

Mechanism, for the record:

• TransferEngineImpl::registerLocalMemory reserves the region in tryReserveMemoryRegions before any transport sees it, and that bookkeeping is protocol-agnostic, so a stale reservation also breaks later TCP mounts, not just RDMA. The overlap error is exactly where a retry at the same address dies.
• The client's own overlap guard in MountSegmentAndGetId can't catch this because mounted_segments_ only records successful mounts. After a failed mount the engine reservation is the sole artifact, which is why every later buffer landing at that address gets refused.
• unregisterLocalMemory(void*, bool update_metadata = true) matches the registration's update_metadata=true, and tolerating ERR_ADDRESS_NOT_REGISTERED follows the same pattern as the other unregister call sites in this file (UnmountSegmentImpl, the hot-cache unregister).

The test is a real red/green: pre-fix the second MountSegment never reaches the master, it dies inside registerLocalMemory with OVERLAPPED, so the retry loop exhausts and the assert fires. CI confirms the binary runs green here (CTest leg, non_ha_reconnect_test 27s pass). The two red legs on this run are the repo-wide wheel-test torch gap that #4440 fixed after this CI ran; a merge with current main should clear them.

Two non-blocking observations:

• unregisterLocalMemory only erases the engine bookkeeping when every transport returns 0 (first-error early return). A partial transport failure would leave the reservation despite the new cleanup, with only the log line to show for it. Mirrors the existing behavior everywhere else, nothing to do in this PR.
• Good call keeping this scoped to the mount-failure path instead of pulling it out of #2336; that one is parked and this stands on its own.

I did not rerun the test locally; verification here is code-level tracing plus the CI run above.