[Store] Fix SSD offload publish-before-commit race (#2799)

[LujhCoconut]

Description

Fixed #2799 .

In BucketStorageBackend::BatchOffload, commit the local index before calling NotifyOffloadSuccess. This closes the race where Master redirects reads to this node before object_bucket_map_ contains the key, causing INVALID_KEY errors.

If NotifyOffloadSuccess fails after local commit, roll back the local index via RollbackCommittedBucket: remove metadata entries, wait for in-flight reads, and clean up on-disk files. On timeout we leave the orphan files rather than risk I/O errors for active readers.

Also evict cached file handles in CleanupOrphanedBucket before deleting files to avoid stale handles.

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

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

[gemini-code-assist]

Code Review

This pull request modifies the BatchOffload process to commit metadata to the local index before notifying the master. If the notification fails, a new RollbackCommittedBucket method is invoked to safely revert the local commit, wait for inflight reads to drain, and clean up the on-disk files. Additionally, CleanupOrphanedBucket is updated to evict cached file handles before deletion. Feedback points out a use-after-move issue in BatchOffload where metadatas[i] is moved into object_bucket_map_ but later passed to complete_handler; copying the metadata instead of moving it is recommended to prevent fragile behavior.

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 81.13208% with 20 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
mooncake-store/src/storage_backend.cpp	71.92%	16 Missing ⚠️
mooncake-store/tests/storage_backend_test.cpp	91.83%	4 Missing ⚠️

📢 Thoughts on this report? Let us know!

[he-yufeng]

Thanks for tracking this down, @LujhCoconut — I walked the fix through the offload and read paths and the logic is correct. Committing the local index ahead of NotifyOffloadSuccess is the right way to close the #2799 window: a concurrent BatchLoad that arrives after the Master redirects reads here now finds the key in object_bucket_map_ instead of returning INVALID_KEY.

A few things I verified while going through it, for the record:

• Empty transport_endpoint in the local index is fine. The copy inserted into object_bucket_map_ carries an empty transport_endpoint (it only gets populated inside complete_handler before the RPC). The local read path (BatchLoad → ReadPlan) only consumes bucket_id/offset/key_size/data_size, and storage_backend.cpp never reads transport_endpoint off a map entry — that field only matters on the copy that travels to the Master. Switching the insert from std::move to a copy is the correct call here, and it also resolves the use-after-move the bot flagged.

• Size accounting stays balanced. Commit adds bucket->data_size + bucket->meta_size, and since BuildBucket sets bucket->data_size = Σ(object_total_size + key_size), the rollback's per-key Σ(data_size + key_size) plus meta_size subtracts exactly what was added. It also matches the existing eviction path, so total_size_ doesn't drift on rollback.

• No TOCTOU between rollback and in-flight reads. BatchLoad acquires the BucketReadGuard (bumping inflight_reads_) while still holding the shared mutex_, in the same critical section as the lookup. Rollback removes the entries under the exclusive lock, so a reader either grabbed its guard first (Phase 2 drains it) or can't find the key after removal. The spin-then-sleep drain mirrors DeleteBucket, and leaving orphans on a drain timeout instead of deleting under an active reader is the safe choice.

One follow-up I'd like to see: the rollback path (complete_handler failing after the local commit) has no test yet — this PR only touches the header and the .cpp. storage_backend_test.cpp already exercises BatchOffload, and you can drive the rollback deterministically by passing a complete_handler that returns a non-OK ErrorCode, then asserting the keys are gone from object_bucket_map_, total_size_ is back to its pre-offload value, and the on-disk bucket file is removed. This is data-integrity recovery code, so a regression here would be silent — worth locking down. The logic checks out so I'm approving, but I'd strongly encourage landing that test alongside it.

Minor, non-blocking: on the 10s drain timeout, Phase 1 has already decremented total_size_ while the files stay on disk until the next Init() orphan scan, so physical usage briefly exceeds the accounted size. That's the same trade-off DeleteBucket makes and the path should be effectively unreachable in practice — just flagging it for awareness.