Skip to content

[TransferEngine] Make single unregisterLocalMemory best-effort #2962

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
he-yufeng wants to merge 1 commit into
base: main
Choose a base branch
from
+5 −1
Open

[TransferEngine] Make single unregisterLocalMemory best-effort #2962

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion mooncake-transfer-engine/src/transfer_engine_impl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -616,10 +616,14 @@ int TransferEngineImpl::registerLocalMemory(void* addr, size_t length,

int TransferEngineImpl::unregisterLocalMemory(void* addr,
bool update_metadata) {
// Best-effort: try every transport so one failure can't leave the region
// registered on the others; mirrors unregisterLocalMemoryBatch (#2869).
int first_error = 0;
for (auto& transport : multi_transports_->listTransports()) {
int ret = transport->unregisterLocalMemory(addr, update_metadata);
if (ret) return ret;
if (ret && !first_error) first_error = ret;
}
if (first_error) return first_error;

std::unique_lock<std::shared_mutex> lock(mutex_);
eraseMemoryRegionLocked(addr);
Comment on lines +626 to 629

@gemini-code-assist gemini-code-assist Bot Jul 16, 2026

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.

high

If any transport fails to unregister, first_error is set to a non-zero value. In the current implementation, if first_error is non-zero, the function returns early on line 626, skipping the local bookkeeping cleanup (eraseMemoryRegionLocked(addr)).

This has two severe consequences:

  1. Permanent Address Block: The memory region remains in the local_memory_regions_ map. Any future attempt to register a memory region at the same address (or an overlapping address) will fail with ERR_ADDRESS_OVERLAPPED in tryReserveMemoryRegions, even if the physical memory was freed and reallocated.
  2. Impossible Recovery: If the user retries unregisterLocalMemory, the transports that did succeed in the first attempt might now return an error (such as ERR_ADDRESS_NOT_REGISTERED), which will again set first_error and prevent the local bookkeeping from ever being cleaned up.

By performing the local eraseMemoryRegionLocked(addr) regardless of first_error (while still returning the error to the caller), we ensure the local engine's state remains clean and future registrations are not permanently blocked.

Suggested change
if (first_error) return first_error;
std::unique_lock<std::shared_mutex> lock(mutex_);
eraseMemoryRegionLocked(addr);
std::unique_lock<std::shared_mutex> lock(mutex_);
eraseMemoryRegionLocked(addr);
if (first_error) return first_error;

Expand Down
Loading