[TransferEngine] Reject concurrent overlapping memory registrations

[feichai0017]

Description

registerLocalMemory() and registerLocalMemoryBatch() checked local_memory_regions_ before transport registration, then inserted the regions after transport calls completed. Concurrent callers could both pass the overlap check and register the same or overlapping memory.

This PR adds an in-flight region map protected by the existing mutex. Classic single and batch registrations now reserve their complete interval set atomically, promote it after transport success, and release it on failure. A deterministic blocking transport test covers both APIs and verifies failed registrations release their reservations.

This change is independent of #2869, which handles transport error propagation.

Module

• Transfer Engine (mooncake-transfer-engine)

Type of Change

• Bug fix

How Has This Been Tested?

Test commands:

cmake -S . -B build -G Ninja -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=OFF -DBUILD_BENCHMARK=OFF -DWITH_TE=ON -DWITH_STORE=OFF -DWITH_STORE_RUST=OFF -DWITH_EP=OFF -DWITH_P2P_STORE=OFF -DUSE_TCP=ON -DENABLE_DEBUG_SYMBOLS=OFF
cmake --build build --target transport_uint_test -j 12
./build/mooncake-transfer-engine/tests/transport_uint_test
./build/mooncake-transfer-engine/tests/transport_uint_test --gtest_filter=TransportTest.ConcurrentRegisterLocalMemoryRejectsOverlap:TransportTest.ConcurrentRegisterLocalMemoryBatchRejectsOverlap --gtest_repeat=20

Test results:

• Unit tests pass (14/14 on Ubuntu; concurrency tests pass 20 repetitions)
• Integration tests pass (not applicable)
• Manual testing done: with reservation calls removed, both deterministic concurrency tests return success for the second overlapping registration and call the transport twice

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 (touched-file hooks pass)
• I have updated the documentation (not applicable; no user-facing API change)
• I have added tests to prove my changes are effective
• For changes >500 LOC: I have filed an RFC issue (not applicable)

AI Assistance Disclosure

• AI tools were used: Codex assisted with test preparation.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[gemini-code-assist]

Code Review

This pull request introduces a mechanism to prevent concurrent registration of overlapping memory regions in the TransferEngineImpl class. It adds a new registering_memory_regions_ map to track regions currently undergoing registration, along with helper methods tryReserveMemoryRegions, commitMemoryRegions, and releaseMemoryRegions to manage this state safely. Unit tests are also added to verify concurrent registration behavior and proper cleanup on failure. There are no review comments, so I have no additional feedback to provide.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[he-yufeng]

Reviewed as a maintainer: traced the reservation lifecycle and the overlap check rather than just the diff. This correctly closes the TOCTOU window and holds up.

What I checked:

The race is actually closed. The old path checked checkOverlap, dropped the lock for the transport calls, then inserted — so two callers could both pass the check and both register overlapping memory. Now tryReserveMemoryRegions does the whole check-all-then-reserve-all under one unique_lock, and crucially hasOverlapLocked consults both local_memory_regions_ and the new registering_memory_regions_. So while a registration is mid-flight (lock released during the transport calls) its interval is visible as an in-flight reservation, and a concurrent overlapping caller is rejected. That is the window that was open before.

Reservation is all-or-nothing. If any interval in a batch overlaps, tryReserveMemoryRegions rolls back the ones it already inserted (the reserved vector) and returns false, so a partially-reserved batch is never left behind. Intra-batch overlaps fall out of the same mechanism for free: once interval A is inserted into registering_memory_regions_, checking B runs hasOverlapLocked against the map that now contains A, so a batch that overlaps itself (including a duplicate address) is rejected at the second interval.

Promote/release are consistent. commitMemoryRegions erases from the in-flight map and inserts into local_memory_regions_ under the lock; releaseMemoryRegions erases on the failure path. No nested locking, so no deadlock, and the shared_mutex write side stays exclusive throughout. The blocking-transport test covering release-on-failure for both APIs is the right thing to pin.

One non-blocking note: the reserve → transport → commit/release sequence is not exception-safe. releaseMemoryRegions only runs on the ret < 0 return path, so if a transport's registerLocalMemory/registerLocalMemoryBatch ever threw instead of returning an error code, the reservation would leak in registering_memory_regions_ and block that address permanently. Transports return error codes by convention today so this is latent, but a small RAII guard that releases unless commit was reached would make it robust against a future throwing transport.

LGTM.

[codecov-commenter]

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 94.40559% with 8 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...cake-transfer-engine/tests/transport_uint_test.cpp	94.50%	5 Missing ⚠️
...ncake-transfer-engine/src/transfer_engine_impl.cpp	94.23%	3 Missing ⚠️

📢 Thoughts on this report? Let us know!