kvcache-ai / kvcache-ai/Mooncake
[RFC]: LOCAL_DISK data migration during drain — Options (a) and (c)
- Dominant language
- C++
- Stars
- 6.6k
- Forks
- 1.2k
- Avg merge
- 3d 5h
- Merged PRs (30d)
- 312
Description
### Changes proposed
## Background
Issue [#2817](https://github.com/kvcache-ai/Mooncake/issues/2817) identified that `ScheduleDrainJobTasks` silently ignores LOCAL_DISK replicas because `Replica::get_segment_names()` returns an empty vector for `LOCAL_DISK`. PR [#2834](https://github.com/kvcache-ai/Mooncake/pull/2834) implements a fail-closed safety net (Step 1 in he-yufeng's two-step plan), but the actual migration capability — Step 2 — remains unimplemented.
This RFC follows up on [#2817](https://github.com/kvcache-ai/Mooncake/issues/2817) and the fail-closed fix in [#2834](https://github.com/kvcache-ai/Mooncake/pull/2834), proposing a concrete plan for Step 2 — actual LOCAL_DISK migration during drain.
It evaluates two viable approaches:
- **Option (a)**: Targeted file transfer — `BatchLoad` → `TransferWrite` → `BatchOffload`
- **Option (c)**: Model LOCAL_DISK as a `NoFSegment`-style allocator-managed resource
Option (b) (drop LOCAL_DISK and re-offload on target) is excluded because LOCAL_DISK data is **not recomputable** for some use cases.
---
## Current State: LOCAL_DISK in the Drain Flow
```
Drain Job
└─> ScheduleDrainJobTasks()
└─> GetReplicaSegmentNames()
└─> LOCAL_DISK → empty vector ← structurally invisible
└─> CreateMoveTask() — only for MEMORY replicas
└─> Client::Move()
└─> ExecuteReplicaTransfer()
└─> if (!source.is_memory_replica()) → INVALID_PARAMS ← hard rejection
```
Both the scheduling and execution paths structurally exclude LOCAL_DISK replicas.
---
## Option (a): Targeted File Transfer
### Approach
```
Data: source SSD ────→ source RAM ────→ target RAM ────→ target SSD
Action: BatchLoad TransferWrite BatchOffload
```
Three building blocks already exist in the codebase:
- **`BatchLoad`**: Reads data from local SSD files into staging buffers (all four storage backends).
- **`TransferWrite`** (TransferEngine): RDMA/TCP transfer from a local buffer to a remote MEMORY replica.
- **`BatchOffload`**: Writes data from memory buffers to local SSD on the target node.
The Promotion path (`ProcessPromotionTasks`) already proves the first two steps work together:
`AllocateBatch() → BatchLoad() → TransferWrite()`.
### What Already Exists
| Component | Location | Status |
| --- | --- | --- |
| BatchLoad (SSD → RAM) | `storage_backend.cpp` | ✅ 4 backends |
| TransferWrite (RDMA/TCP) | `transfer_task.cpp:981` | ✅ generic buffer support |
| BatchOffload (RAM → SSD) | `storage_backend.cpp` | ✅ 3 backends |
| Promotion: BatchLoad → TransferWrite | `file_storage.cpp:728` | ✅ proven pattern |
### What Needs to Be Built
1. **`Client::ExecuteReplicaTransfer()`** — lift the `is_memory_replica()` restriction (client_service.cpp:3182). Add a LOCAL_DISK path:
```
if (source.is_local_disk_replica()) {
allocate staging buffer;
BatchLoad(slices); // SSD → staging buffer
TransferWrite(target, slices); // staging buffer → target MEMORY replica
release staging buffer;
}
```
2. **`MasterService::ScheduleDrainJobTasks()`** — recognize LOCAL_DISK replicas as drainable (currently excluded via `get_segment_names()` returning empty). Generate move units for keys that have only LOCAL_DISK replicas on the draining segment.
3. **Staging buffer management** — reuse the existing `client_buffer_allocator_` (same pool used by Promotion). With `local_buffer_size` as hard cap, failures trigger GC retry, then skip. Concurrency is naturally bounded by `max_concurrency` (default 4).
4. **Source LOCAL_DISK cleanup** — after successful migration, delete the source `.bucket`/`.meta` entries. This requires a new cleanup path (MEMORY replicas use `discarded_replicas_` with delayed release; LOCAL_DISK has no equivalent).
5. **Deduplication** — a key may have both MEMORY and LOCAL_DISK replicas on the draining node. When both need to move, transferring both wastes bandwidth (LOCAL_DISK data is a superset of MEMORY data). Optimization: if a MEMORY replica exists for the same key on the same node, use the existing zero-copy TransferWrite path and skip the BatchLoad for LOCAL_DISK.
### Considerations
| Concern | Analysis |
| --- | --- |
| **Staging buffer OOM** | Buffer pool has O(n) hard cap. GC + retry on failure. Same risk profile as Promotion. |
| **Duplicate transfer** | Can be addressed by preferring MEMORY replica over LOCAL_DISK when both exist on same node. |
| **Source cleanup** | New path required. MEMORY replicas use `discarded_replicas_` (10-min delayed release). LOCAL_DISK cleanup is simpler — just delete the files. |
| **Concurrent drain pressure** | `max_concurrency=4` per drain job + 500ms scheduling interval provides natural throttling. |
### Estimation
- **Core change**: ~200–300 lines in `client_service.cpp` (new LOCAL_DISK branch in `ExecuteReplicaTransfer`, following the Promotion pattern)
- **Drain scheduling**: ~50–100 lines in `master_service.cpp` (recognize LOCAL_DISK replicas)
- **Cleanup**: ~50–100 lines (delete LOCAL_DISK entries after successful move)
- **Total**: moderate effort, low architectural risk
---
## Option (c): Model LOCAL_DISK as NoFSegment
### Approach
Treat LOCAL_DISK as a block-addressable resource managed by a `BufferAllocator`, similar to how NoF (NVMe-oF) segments work. This would allow LOCAL_DISK replicas to participate in the existing `REPLICA_MOVE` primitive (allocator-to-allocator copy).
### Why This Requires Architectural Changes
**1. Storage model mismatch — files vs blocks**
```
LOCAL_DISK (current): NoF (what option c requires):
┌─────────────────────┐ ┌─────────────────────┐
│ File system │ │ NVMe namespace │
│ bucket_1.bucket │ │ base=0 │
│ bucket_2.bucket │ │ offset=4096 (LBA 8) │
│ .meta │ │ offset=8192 │
│ Addressing: path │ │ Addressing: LBA │
│ I/O: open/read/write│ │ I/O: SPDK NVMe-oF │
└─────────────────────┘ └─────────────────────┘
```
NoF uses block-level LBA addressing via SPDK. LOCAL_DISK uses file-level path addressing. Bridging these requires a **file-backed block abstraction layer**.
**2. Replica data structure incompatibility**
| Aspect | NoF (`NoFReplicaData`) | LOCAL_DISK (`LocalDiskReplicaData`) |
| --- | --- | --- |
| Core field | `unique_ptr` (offset + endpoint) | `UUID client_id` + `transport_endpoint` |
| Allocator | CachelibBufferAllocator / OffsetBufferAllocator | None (client manages files) |
| get_segment_names() | Returns segment name | Returns empty vector |
To model LOCAL_DISK as NoF, `LocalDiskReplicaData` must carry an `AllocatedBuffer`, which means a new allocator type that manages file offsets rather than memory offsets.
**3. Requires a new `FileBackedBufferAllocator`**
Existing allocators (`CachelibBufferAllocator`, `OffsetBufferAllocator`) manage memory regions. LOCAL_DISK data is in files managed by `BucketStorageBackend` where multiple keys share a single bucket file. The allocator must:
- Track available space across bucket files
- Handle bucket fragmentation (keys of different sizes packed together)
- Coordinate with the bucket's LRU eviction policy
- Present a flat offset space that maps to `(file_path, file_offset)` pairs
This is effectively building a **file-system-backed block allocator** — a significant subsystem.
**4. Transfer path incompatibility**
```
NoF transfer path: LOCAL_DISK transfer path:
SPDK → NVMe-oF → remote NVMe open/read → user buffer → TransferEngine → remote
```
NoF's `submitSpdkNofOperation` issues NVMe commands directly. LOCAL_DISK reads must go through file I/O. To unify them, LOCAL_DISK would need an SPDK backend that emulates NVMe on top of files — unrealistic.
**5. Drain infrastructure doesn't support NoF segments**
| Drain step | NoF support | Needs |
| --- | --- | --- |
| `ValidateDrainRequest` | ❌ Only queries `segment_manager_` | Add `nof_segment_manager_` |
| `MaybeCompleteDrainJob` | ❌ Only sets `segment_manager_` status | Add NoF segment status handling |
| `ExecuteReplicaTransfer` | ❌ Hardcoded `is_memory_replica()` only | Add NoF source support |
| `CreateMoveTask` / `MoveStart` | ❌ Only allocates MEMORY replicas | Support NoF→NoF or NoF→MEMORY |
### Estimation
- **New allocator**: major new subsystem (thousands of lines)
- **SPDK/file bridge**: large effort, architecturally questionable
- **Drain full-chain NoF support**: hundreds of lines across validation, scheduling, and execution
- **Total**: high effort, high architectural risk, uncertain benefit
---
## Comparison
| | Option (a): File Transfer | Option (c): Model as NoF |
| --- | --- | --- |
| **Reuses existing code** | ✅ BatchLoad, TransferWrite, BatchOffload all exist | ❌ Requires new allocator, new transfer path |
| **New code required** | ~400–500 lines | Several thousand lines |
| **Architectural risk** | Low — pipelining existing primitives | High — abstracting files as blocks |
| **Drain scheduling change** | Recognize LOCAL_DISK in `ScheduleDrainJobTasks` | Full NoF drain support chain |
| **Staging buffer** | Reuse existing `client_buffer_allocator_` | N/A (block-level access) |
| **Transfer protocol** | File I/O + RDMA/TCP (existing) | Must bridge file I/O to SPDK NVMe-oF |
| **Source cleanup** | New delete path (simple) | Via allocator free (needs to be built) |
| **Deduplication** | Solvable — prefer MEMORY over LOCAL_DISK | Same |
| **Fits architecture** | ✅ Composes existing layers | ❌ Forces square peg into round hole |
---
## Open Questions
1. **Target replica type**: After moving LOCAL_DISK data, should the target store it as MEMORY (simpler, handles drain immediately) or trigger an offload to target's LOCAL_DISK (preserves storage tier semantics)?
2. **Partial drain semantics**: If some LOCAL_DISK objects are successfully moved and others fail, should the drain report PARTIAL rather than FAILED? This provides better visibility than the current all-or-nothing approach.
3. **Staging buffer sizing**: Should LOCAL_DISK drain use the same `local_buffer_size` pool as Promotion, or have a separate configurable limit?
### Before submitting a new issue...
- [ ] Make sure you already searched for relevant issues and read the [documentation](https://kvcache-ai.github.io/Mooncake/)
Contributor guide
Research direction
Start by reading Client::ExecuteReplicaTransfer() in client_service.cpp:3182 and MasterService::ScheduleDrainJobTasks(), then compare the Promotion path in file_storage.cpp:728 and TransferWrite in transfer_task.cpp:981. Resolve the open questions about target replica type, partial drains, and staging limits before choosing an approach. Done means LOCAL_DISK replicas can be migrated during drain with source cleanup and defined failure semantics.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- backend, distributed-systems
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100