[Store] Prevent drain success with local disk replicas

[feichai0017]

Description

Fixes #2817.

Prevent drain jobs from reporting success while LOCAL_DISK replicas remain on source clients:

• Capture source segment ownership when the drain job is created.
• Use one shared predicate for scheduling and completion checks.
• Fail unsupported local-disk moves, restore source segments to OK, and keep memory migration behavior unchanged.
• Cover both local-disk-only and mixed MEMORY + LOCAL_DISK objects.

This change fails closed; it does not add cross-node local-disk migration.

Module

• Mooncake Store (mooncake-store)

Type of Change

• Bug fix

How Has This Been Tested?

Test commands:

./scripts/code_format.sh --base upstream/main --check
SKIP=mooncake-code-format pre-commit run --files mooncake-store/include/master_service.h mooncake-store/src/master_service.cpp mooncake-store/tests/master_service_test.cpp
cmake --build build --target master_service_test -j4
./build/mooncake-store/tests/master_service_test

Test results:

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

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 on the touched files and all applicable 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; under 500 LOC)

AI Assistance Disclosure

• AI tools were used: Codex assisted with implementation, review, formatting, and test execution.

[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 handling for unsupported LOCAL_DISK replicas during drain jobs. Since LOCAL_DISK replicas are client-scoped and cannot currently be migrated, the system now identifies them, associates them with the source segments owned by the client, and fails the drain job. A new unit test has been added to verify this failure behavior. I have no additional feedback to provide as there are no review comments.

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

Files with missing lines	Patch %	Lines
mooncake-store/src/master_service.cpp	88.88%	5 Missing ⚠️

📢 Thoughts on this report? Let us know!

[he-yufeng]

Reviewed as a maintainer — I triaged the originating issue (#2817), so I checked this against that root cause: a LOCAL_DISK-only object contributes no segment name (get_segment_names() returns empty for it), so the old drain scheduler couldn't see it and a segment could reach DRAINED with on-disk data still present. This is exactly the fail-closed first step the issue called for, and the implementation is sound.

The LOCAL_DISK association is the right handle. Since a LOCAL_DISK replica has no segment name, keying it by its owning client via get_local_disk_client_id() and matching against source_segments_by_client is correct — it fails the drain only for LOCAL_DISK replicas whose holder actually owns a source segment being drained, and leaves replicas on unrelated clients untouched. MakeDrainLocalDiskUnitKey (the local-disk:<client>:... prefix) keeps these units distinct from the segment-based ones, and the terminal_failed_unit_keys.insert(...).second guard dedups so an object reachable under several source segments is only counted once.

Completion is the authoritative gate. MaybeCompleteDrainJob re-scans LOCAL_DISK replicas and fails the job with an explicit "LOCAL_DISK replicas remain on source clients" message whenever any still resolve to a source client. So even if the set shifts between creation and completion (e.g. an eviction removes the object — in which case there's no data left to lose), the final decision reflects the actual state at completion rather than a stale creation-time count. Counting them as failed_units at creation plus this re-check is a sound belt-and-suspenders.

Rollback is handled. The new GetClientIdBySegmentName failure path in CreateDrainJob restores every already-updated draining segment back to OK before returning the error, so a mid-creation failure doesn't strand segments in DRAINING.

The regression test pins a local-disk-only object failing the drain, and the change correctly fails closed without attempting cross-node migration — matching the "surface it, don't silently drop, defer actual migration to an RFC" split from the issue discussion. LGTM.

[he-yufeng]

Read through the drain path (CreateDrainJob, ScheduleDrainJobTasks, MaybeCompleteDrainJob) and issue #2817. This is a real correctness fix: a drain that reports success while un-migratable LOCAL_DISK replicas still sit on the source-owning clients is silently lying about durability, and failing closed (no cross-node local-disk migration yet) is the right conservative call. A few things are done well: capturing source_segments_by_client at create time to scope local-disk replicas to their owner, the CreateDrainJob rollback that does GetClientIdBySegmentName first and restores the already-draining segments to OK before returning on error, and reusing the existing failure path in MaybeCompleteDrainJob (which restores job.request.segments to OK when all_remaining_terminal_failed) so the source segments don't get stranded in DRAINING.

Two points, neither blocking:

1. Accounting consistency for the local-disk terminal failure. The local-disk unit increments failed_units and inserts into terminal_failed_unit_keys, and the whole thing then relies on all_remaining_terminal_failed flipping true to trigger the restore + failure message. Could you confirm the local-disk unit is counted consistently on both sides - i.e. it's included in whatever total/remaining-unit denominator all_remaining_terminal_failed evaluates, not only the failed side? The case I want to be sure about is an object with both a memory replica and a local-disk replica on the source client: the memory replica migrates, the local-disk one is terminal-failed, and the job should deterministically reach all_remaining_terminal_failed (and restore) rather than hang as "running" or miscount.

2. The local-disk detection (VisitReplicas(&Replica::fn_is_local_disk_replica, ...) plus the source_segments_by_client membership check) now appears in both ScheduleDrainJobTasks and MaybeCompleteDrainJob. Since both encode the same "fail closed on a source-owned local-disk replica" invariant, a small shared helper would keep them from drifting if the predicate or the ownership rule changes later.

[feichai0017]

Addressed both follow-ups in c298dc5c. The completion gate is set-based rather than counter-based: it rebuilds the same local-disk unit key that scheduling inserts into terminal_failed_unit_keys. The new mixed-replica test verifies that MEMORY migrates to the target, LOCAL_DISK remains, the job reaches FAILED, and the source segment returns to OK. I also moved the shared source-owned LOCAL_DISK predicate into one helper. All 143 master_service_test cases pass on Ubuntu.