[Store] Support LOCAL_DISK replica migration during drain job

[liangxu2000]

Description

Problem

When a Mooncake Store node is drained (decommission / scale-in), the drain
scheduler (ScheduleDrainJobTasks) migrates all replicas to a target node.
However, LOCAL_DISK replicas (client-side SSD offload files written by
FileStorage) were silently skipped because Replica::get_segment_names()
returned an empty vector for the LOCAL_DISK type. Any data that had been
evicted from memory and offloaded to the local SSD was left behind on the
decommissioned node and effectively lost to the cluster.

Related issues / PRs (chronological):

1. #2817 — identified
   the problem: drain jobs do not migrate LOCAL_DISK data.
2. #2834 — Step 1:
   fail-closed safety net that prevents a segment from being marked DRAINED
   while LOCAL_DISK residuals exist (merged).
3. #2885 — RFC
   analyzing implementation options for Step 2 (Option (a): targeted file
   transfer vs Option (c): model LOCAL_DISK as NoF-style block resource).
   Concluded Option (a) is preferred: low architectural risk, ~400–500
   lines, reuses existing BatchLoad / TransferWrite primitives.
4. This PR — implements Step 2 based on the Option (a) design from
   [RFC]: LOCAL_DISK data migration during drain — Options (a) and (c) #2885: actual LOCAL_DISK migration during drain, so residuals are moved
   rather than merely blocking completion.

Solution — Targeted File Transfer (reuse the existing Move flow)

Core insight.
A LOCAL_DISK replica is just a file on the source client's
local SSD. To migrate it, we need to (a) read that file into a memory buffer,
then (b) RDMA-write the buffer into a MEMORY segment on the target node.
Mooncake already has battle-tested primitives for both steps:
BatchLoad (O_DIRECT file → aligned memory) and TransferWrite (registered
memory → remote segment). The promotion path (ProcessPromotionTasks) uses
exactly this pattern within a single node; we simply redirect the
TransferWrite destination to a remote target allocated by the drain
scheduler.

Why reuse the Move flow instead of a new mechanism?
The existing drain
infrastructure (ScheduleDrainJobTasks → CreateMoveTask → client executes
transfer → MoveEnd finalizes metadata) already provides:

• per-key retry with configurable kMaxDrainUnitRetries
• concurrency control (max_concurrency)
• completion tracking (MaybeCompleteDrainJob)
• segment state machine (OK → DRAINING → DRAINED)

Building a separate “LOCAL_DISK drain” path would duplicate all of this. By
plugging into the existing Move flow, we get it for free with minimal blast
radius.

Data flow (per key):

Source client                              Target client
───────────                              ─────────────
SSD file (bucket storage)
    │
    │ BatchLoad (O_DIRECT, align_down/align_up)
    ▼
Staging buffer (4096-aligned,
  in registered memory region)
    │
    │ TransferWrite (RDMA)
    └───────────────────────────▶  Newly allocated MEMORY replica
                                   (in target's mounted segment)

Implementation steps:

1. Visibility. Replica::get_segment_names() now returns the
   transport_endpoint for LOCAL_DISK replicas. Previously it returned an
   empty vector, making them invisible to the scheduler.

2. Scheduling with MEMORY-first dedup. ScheduleDrainJobTasks scans all
   replicas on the draining segment. If a key has both a MEMORY and a
   LOCAL_DISK replica on the same client, only the MEMORY replica is
   scheduled for migration (it is already in RAM — cheaper to transfer). The
   redundant LOCAL_DISK metadata is cleaned up in MoveEnd (see step 5).
   Only keys that are LOCAL_DISK-only (memory was already evicted) go
   through the file-transfer path.

3. File transfer branch. ExecuteReplicaTransfer gains a LOCAL_DISK
   source branch:

   • LoadBatchFromLocalDisk calls AllocateBatch (allocates an
     O_DIRECT-aligned staging buffer in registered memory) + BatchLoad
     (reads the SSD file via align_down/align_up offset correction).
   • TransferWrite RDMA-writes the staging buffer into the target's freshly
     allocated MEMORY replica.
   • A replica_move_local_disk_transfer_success log is emitted on success,
     making the path observable and testable.

4. is_drain flag (see Design notes). A boolean is threaded from the
   scheduler through ReplicaMovePayload → Client::Move → MoveEnd RPC.

5. Conditional metadata cleanup. When a drain move relocates a MEMORY
   replica, the same key's LOCAL_DISK replica on the source client is now
   redundant (it is a disk backup of the memory data we just moved, and the
   target already holds that data). MoveEnd drops the LOCAL_DISK master
   metadata only when is_drain=true. The physical SSD file is
   intentionally left on disk as a safety net (no destructive deletion);
   physical cleanup is a separate future step.

6. Completion guard. MaybeCompleteDrainJob checks for LOCAL_DISK
   residuals before marking a segment DRAINED. If any LOCAL_DISK-only
   keys remain (e.g. because they failed), the segment stays in DRAINING
   and the job is marked FAILED.

What is NOT introduced: no new RPC, no new master state machine, no new
task type. The entire feature plugs into the existing Move pipeline.

Testing

5 end-to-end tests drive real eviction-driven offload, manual moves, and
failure injection, verifying the full per-key data route with byte-for-byte
comparison and direct log evidence that the LOCAL_DISK transfer path
(LoadBatchFromLocalDisk → staging buffer → TransferWrite) is exercised.

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

Changes

File	Change
include/replica.h	get_segment_names() returns the transport_endpoint for LOCAL_DISK so the drain scheduler can locate them.
src/client_service.cpp	ExecuteReplicaTransfer: new LOCAL_DISK source branch (SSD -> staging buffer via LoadBatchFromLocalDisk -> TransferWrite -> target MEMORY). Adds a ..._local_disk_transfer_success log so the path is observable/testable.
src/file_storage.cpp	LoadBatchFromLocalDisk (AllocateBatch + BatchLoad) exposed for the client.
src/master_service.cpp	CreateMoveTask: client fallback for LOCAL_DISK sources, plus an is_drain flag (default false; drain sets it true). ScheduleDrainJobTasks: scans LOCAL_DISK replicas with MEMORY-first dedup. MaybeCompleteDrainJob: checks LOCAL_DISK residue before marking a segment DRAINED. MoveEnd: only for drain moves (is_drain=true) drops the redundant LOCAL_DISK metadata after a MEMORY move, with a scope-narrowing fix for an EDEADLK on segment_mutex_ (unique -> shared on the same std::shared_mutex).
include/task_manager.h, include/client_service.h / src/client_service.cpp, include/master_client.h / src/master_client.cpp, include/rpc_service.h / src/rpc_service.cpp	Thread an is_drain flag from the drain scheduler through the move task payload (ReplicaMovePayload), the client (Client::Move / ExecuteTask), and the MoveEnd RPC chain, so MoveEnd can distinguish drain moves from the public create_move_task API. Manual moves default to is_drain=false.
include/real_client.h / src/real_client.cpp	mountLocalDiskSegment() + getOffloadRpcAddr() (used by the e2e test to register a LOCAL_DISK segment).
include/master_config.h / tests/test_server_helpers.h	InProcMasterConfig / builder gain offload_on_evict, passed through to WrappedMasterServiceConfig (test infra).
tests/e2e/CMakeLists.txt	Registers local_disk_drain_test.
tests/e2e/local_disk_drain_test.cpp	New end-to-end tests (see below).

Design notes

• is_drain flag — distinguishing drain moves from manual moves. The
  public create_move_task API (used for load-balancing) and the internal
  drain scheduler both go through the same Move → MoveEnd code path.
  Without a distinguishing flag, MoveEnd would unconditionally drop the
  source client's LOCAL_DISK metadata after any MEMORY move — silently
  discarding a durable SSD backup on a still-live node during a manual move.
  To prevent this, a boolean is_drain is threaded through:
  ScheduleDrainJobTasks sets it true in ReplicaMovePayload →
  task_manager.h carries it in the task → Client::Move / ExecuteTask
  forwards it → MasterClient::MoveEnd passes it in the RPC →
  MasterService::MoveEnd gates the LOCAL_DISK metadata removal on
  is_drain == true. Manual moves default to is_drain=false, preserving
  the source's LOCAL_DISK replica.
• MEMORY-first dedup (drain moves only). When a key has both a MEMORY
  and a LOCAL_DISK replica on the draining client, only the MEMORY replica
  is migrated; the LOCAL_DISK metadata is dropped in MoveEnd (the SSD file
  is intentionally retained on disk as a safety net — physical cleanup is
  future work).
• Eviction watermark is GLOBAL. get_global_mem_used_ratio() is
  allocated / capacity summed across all mounted segments. Mounting a
  large target segment therefore dilutes the ratio and can stop eviction. The
  test mounts the drain target only after eviction has freed memory (see
  test scenarios).
• offload_on_evict (easy to confuse with enable_offload). It is gated
  by enable_offload: offload_on_evict_ = enable_offload_ && config.offload_on_evict, so it can only take effect when offload is enabled.
  It changes when offload happens, not whether:
  • offload_on_evict=false (default): every PutEnd offloads to SSD, so a key
    already has a LOCAL_DISK replica by the time eviction runs.
  • offload_on_evict=true: PutEnd does not offload; a key is queued for
    offload only when eviction picks it and finds no LOCAL_DISK replica yet,
    which yields a two-step eviction (queue offload + pin, then drop MEMORY in
    the next cycle once the disk backup exists).

How Has This Been Tested?

Test commands:

cd build
cmake --build . --target local_disk_drain_test -j$(nproc)
GLOG_logtostderr=1 ./mooncake-store/tests/e2e/local_disk_drain_test

Test scenarios (5 tests, all pass):

1. LocalDiskDrainTest.DrainLocalDiskReplicasToPeer
   Purpose: verify the baseline happy path — drain correctly migrates data
   when every key has both MEMORY + LOCAL_DISK replicas, using the MEMORY-first
   dedup policy.
   Setup: 4 keys x 4 KB on client1 (MEMORY + LOCAL_DISK); drain client1 →
   client2. Asserts: drain SUCCEEDED, MEMORY replica migrated (dedup),
   LOCAL_DISK metadata dropped, all keys byte-for-byte readable on client2.

2. LocalDiskDrainTest.ManualMovePreservesLocalDiskBackup
   Purpose: regression guard for the is_drain gating — prove that a
   manual move (is_drain=false) does NOT drop the source's LOCAL_DISK
   replicas, and that a subsequent drain correctly uses the LOCAL_DISK transfer
   path (including non-aligned O_DIRECT reads).
   Setup: manually move each key's MEMORY replica client1 → client2 via the
   public create_move_task API, then drain client1 → client3. Keys use a mix
   of 4096-aligned and deliberately non-aligned sizes ({4096, 7001, 12000,
   5000}) filled with a position-dependent pattern. Asserts: all 4 keys
   migrated via LOCAL_DISK path (local_disk_transfer_success count),
   byte-for-byte correct on client3 (guards the O_DIRECT aligned-read offset
   fix in BatchLoad). If the manual move had wrongly dropped LOCAL_DISK
   metadata, succeeded_units would be 0.

3. LocalDiskDrainTest.DrainFailsGracefullyWhenLocalDiskDataMissing
   Purpose: verify graceful failure handling — when the source SSD files
   are missing/corrupt, the drain job must reach a terminal FAILED state
   (no hang), correctly account retries, and NOT corrupt surviving replicas.
   Setup: put 4 keys on client1 (MEMORY + LOCAL_DISK), manually move MEMORY
   to client2 (leaving client1 LOCAL_DISK-only), delete client1's SSD files,
   then drain client1 → client3. Asserts: job FAILED (not stuck),
   succeeded_units == 0, local_disk_load_failed error branch hit,
   source segment NOT marked DRAINED, client2's MEMORY copy byte-for-byte
   intact. Exercises the retry / terminal-failure accounting
   (kMaxDrainUnitRetries = 3).

4. LocalDiskDrainEvictTest.DrainAfterEviction_OffloadOnEvict
   Purpose: verify drain works end-to-end in a realistic eviction scenario
   with offload_on_evict=true (two-step eviction), where some keys are
   MEMORY-only and some are LOCAL_DISK-only at drain time.
   Setup: 16 keys x 4 MB filling a 64 MB segment; eviction offloads to SSD
   and frees MEMORY for ~8 keys. Asserts: drain migrates ~8 keys via MEMORY
   path and ~8 keys via LOCAL_DISK path; all 16 keys byte-for-byte readable
   on target.

5. LocalDiskDrainEvictTest.DrainAfterEviction_ImmediateOffload
   Purpose: same as above but with offload_on_evict=false (offload at
   PutEnd), verifying that both offload modes produce a correct drain outcome.
   Setup: identical topology; eviction frees MEMORY for ~8 keys. Asserts:
   drain again exercises both MEMORY and LOCAL_DISK paths successfully.

Test results:

• Unit tests pass
• Integration tests pass (e2e local_disk_drain_test, 5/5 PASSED)

[==========] 5 tests from 2 test suites ran.
[  PASSED  ] 5 tests.

Log evidence (what you see when running the tests):

The tests emit structured glog output (GLOG_logtostderr=1) and use a
CapturingLogSink to programmatically assert that specific code paths were
taken. Here is what the key evidence looks like:

Test 2 — LOCAL_DISK transfer path exercised (success):

I0723 ... action=replica_move_local_disk_transfer_success, key=manual_move_key_0, size=4096
I0723 ... action=replica_move_local_disk_transfer_success, key=manual_move_key_1, size=7001
I0723 ... action=replica_move_local_disk_transfer_success, key=manual_move_key_2, size=12000
I0723 ... action=replica_move_local_disk_transfer_success, key=manual_move_key_3, size=5000

This proves the full SSD → staging buffer → TransferWrite → target MEMORY
path ran for every key, including non-aligned sizes (7001, 12000, 5000 bytes)
that exercise the O_DIRECT align_down/align_up offset correction in
BatchLoad. The test asserts this count >= 4.

Test 3 — graceful failure when SSD files are deleted (4 keys × 3 retries):

W0723 ... LoadBatchFromLocalDisk: BatchLoad failed for key=fail_key_0, error=FILE_OPEN_FAIL
E0723 ... error=local_disk_load_failed, error_code=FILE_OPEN_FAIL, key=fail_key_0
W0723 ... LoadBatchFromLocalDisk: BatchLoad failed for key=fail_key_0, error=FILE_OPEN_FAIL
E0723 ... error=local_disk_load_failed, error_code=FILE_OPEN_FAIL, key=fail_key_0
W0723 ... LoadBatchFromLocalDisk: BatchLoad failed for key=fail_key_0, error=FILE_OPEN_FAIL
E0723 ... error=local_disk_load_failed, error_code=FILE_OPEN_FAIL, key=fail_key_0
... (same pattern for fail_key_1, fail_key_2, fail_key_3)

Each key fails exactly 3 times (kMaxDrainUnitRetries = 3), producing 12
local_disk_load_failed lines total. The test asserts:

• status_name == "FAILED" (job reached terminal state, no hang)
• succeeded_units == 0 (nothing was migrated)
• source segment NOT marked DRAINED
• client2's surviving MEMORY copy is byte-for-byte intact (the failed drain
  did not corrupt or drop the live replica)

Tests 4/5 — eviction-driven drain (final job query):

drain job status: SUCCEEDED
succeeded=16  failed=0  blocked=0  migrated_bytes=67108864

All 16 keys (8 via MEMORY path + 8 via LOCAL_DISK path) migrated successfully;
migrated_bytes = 16 × 4 MB = 64 MiB confirms no data was silently dropped.

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 added tests to prove my changes are effective
• For changes >500 LOC: I have filed an RFC issue

RFC note: the design is documented in drain-local-disk-migration-rfc.md
(Option A vs Option C, follow-up to #2817 / #2834). The non-test source change
is under 500 LOC; please confirm whether a separate GitHub RFC issue is still
required before merge.

AI Assistance Disclosure

• AI tools were used

AI coding assistants were used to help implement the feature, write the tests,
diagnose test behavior, and review the change. The human submitter has reviewed
the changes and is responsible for defending them end-to-end.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[codecov-commenter]

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

Codecov Report

❌ Patch coverage is 89.47368% with 90 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
mooncake-store/src/master_service.cpp	58.82%	42 Missing ⚠️
mooncake-store/tests/e2e/local_disk_drain_test.cpp	95.99%	26 Missing ⚠️
mooncake-store/src/client_service.cpp	69.56%	14 Missing ⚠️
mooncake-store/src/file_storage.cpp	80.00%	4 Missing ⚠️
mooncake-store/src/real_client.cpp	55.55%	4 Missing ⚠️

📢 Thoughts on this report? Let us know!