[Bugfix] Roll back failed OffsetAllocator evictions

[feichai0017]

Description

Fixes #2963.

OffsetAllocatorStorageBackend previously removed FIFO victims before notifying the master. If notification failed, the local replicas were lost while the master could retain stale LOCAL_DISK metadata.

This PR keeps victims and allocation handles pending until notification succeeds, restores the shard map, FIFO index, and counters on failure, and adds a regression test covering rollback plus retry order.

#2932 also touches OffsetAllocator for persistence, but does not address notification-failure rollback; this PR is limited to that correctness issue.

Module

• Mooncake Store (mooncake-store)

Type of Change

• Bug fix

How Has This Been Tested?

Test commands:

pre-commit run --files mooncake-store/include/storage_backend.h mooncake-store/src/storage_backend.cpp mooncake-store/tests/storage_backend_test.cpp
./scripts/code_format.sh --check --base upstream/main
cmake --build build --target storage_backend_test -j 16
./build/mooncake-store/tests/storage_backend_test

Test results:

• Unit tests pass (72/72 on Ubuntu)
• Integration tests pass (not applicable)
• Manual testing done (not applicable)

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 (repository-wide hooks currently fail on unrelated existing files; touched-file hooks pass)
• I have updated the documentation (not applicable for this internal bug fix)
• I have added tests to prove my changes are effective
• For changes >500 LOC: I have filed an RFC issue (not applicable; implementation is under 500 LOC)

AI Assistance Disclosure

• AI tools were used

Codex assisted with diagnosis, implementation, testing, and review.

[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 transactional eviction mechanism in the OffsetAllocatorStorageBackend. It adds a PendingEviction structure to hold evicted metadata and allocation handles, preventing their reuse until the master accepts the replica-removal notification. If the notification fails, the eviction is rolled back using the new RestorePreparedEviction method. A unit test has been added to verify this rollback behavior on notification failure. There are no review comments, so no feedback is provided.

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.

[codecov-commenter]

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

Codecov Report

❌ Patch coverage is 90.47619% with 10 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
mooncake-store/src/storage_backend.cpp	80.48%	8 Missing ⚠️
mooncake-store/tests/storage_backend_test.cpp	96.87%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!

[he-yufeng]

Thanks for tightening this up — the core change is correct and closes a real consistency hole. Before this, if eviction_handler rejected the removal we'd already have erased the victim from shard.map and dropped its AllocationPtr, so the extent could be reused while the master still believed the local replica existed. Pinning the evicted ObjectEntry (and with it the refcounted AllocationPtr) in PendingEviction until the notify succeeds, then committing by dropping the last ref or restoring on rejection, is the right shape. The emplace_back(vkey, it->second) copy is safe precisely because AllocationPtr is a shared_ptr, so the subsequent shard.map.erase(it) doesn't free the extent. Totals stay symmetric (sub at prepare, add on restore, untouched on commit), and the fallback loop's switch to pending_eviction.objects.size() is fine since pending is always drained at the top of each turn.

Two things I'd like to resolve before approving.

1. The refactor drops the eviction_handler null check that both notify sites used to carry (if (eviction_on && eviction_handler && ...)), and NotifyAndCommitPreparedEviction invokes eviction_handler(evicted_keys) unconditionally. The interface declares EvictionHandler eviction_handler = nullptr and documents it as optional, so an empty std::function is a legal argument. The sole production caller in file_storage.cpp always passes a live handler, so this isn't reachable today — but a test or future caller that runs with eviction_policy != NONE and no handler would now hit std::bad_function_call instead of the old "evict locally, skip master notify" behavior. Could you early-return inside the helper to preserve that, e.g. if (!eviction_handler) { pending.objects.clear(); return {}; } at the top? That keeps the null-handler path a local-only commit like before.

2. RestorePreparedEviction uses CHECK(map_inserted) / CHECK(fifo_inserted), which abort the process if the key or FIFO slot is already present. The notify call runs with eviction_mutex_ released (EvictToMakeRoom unlocks on return, and the handler is invoked before restore re-locks), so there's a window where another writer could re-insert an evicted key or reuse its fifo_seq. Is BatchOffload serialized against a concurrent Put/offload of the same key such that this can't happen? If the invariant holds, a one-line comment on the CHECKs would make it self-documenting; if it doesn't, I'd rather reconcile (keep the live entry / re-seat under a fresh seq) than turn a recoverable master rejection into a hard abort.

Neither is a big change — happy to approve once the null-handler guard is back and the restore-under-concurrency question is settled.

[feichai0017]

@he-yufeng Addressed both points in a31db23:

• Added the defensive null-handler guard in NotifyAndCommitPreparedEviction.
• Documented that the restore CHECKs rely on the existing single-writer precondition for BatchOffload; production offloads are serialized by heartbeat_thread_, and readers do not mutate either index.

[he-yufeng]

Re-reviewed after the rollback-assumptions commit. The core is correct and the guard commit closes the robustness gap cleanly.

The pin-until-notify shape is right: EvictToMakeRoom moves each victim's ObjectEntry into PendingEviction before erasing it from shard.map, and since AllocationPtr is a shared_ptr the extent stays alive (not returned to the allocator) until NotifyAndCommitPreparedEviction either drops the last reference on success or hands the entry back to RestorePreparedEviction on rejection. Totals stay symmetric (sub at prepare, add on restore, untouched on commit), and RestorePreparedEviction copies total_size/fifo_seq into locals before the std::move, so there's no use-after-move.

I also checked the sequencing concern: every NotifyAndCommitPreparedEviction runs before the following allocate() (the (B) block ahead of (C), and again inside the fallback loop before each retry), so committed extents are freed in time and holding the pin doesn't starve the allocation.

On the CHECK in RestorePreparedEviction: it's safe because the offload path is genuinely single-writer. BatchOffload is only reached through FileStorage::OffloadObjects, which is only called from Heartbeat(), which only runs on the single heartbeat_thread_ (nothing else calls either, and the client-buffer GC thread never offloads). So no writer can recreate the key or reuse its FIFO sequence during the notify window and the CHECK can't fire in correct operation. Documenting the precondition inline and warning future maintainers to revisit the DCHECK_GE guards if concurrent offload is ever added is the right call.

LGTM.

[LujhCoconut]

A clean and correct fix. LGTM. OffsetAllocatorStorageBackend is still not production ready. Thanks for making this further.