kvcache-ai / kvcache-ai/Mooncake

[RFC]: Refactor Mooncake Store segment management around SegmentPool

Open
#3,360 1 comment 2 reactions 0 assignees View on GitHub
RFC
Dominant language
C++
Stars
6.6k
Forks
1.2k
Avg merge
3d 5h
Merged PRs (30d)
312

Description

### Changes proposed

## Summary

Mooncake Store currently splits mounted-region ownership between `SegmentManager` and `NoFSegmentManager`, while client-local SSD state is embedded in `SegmentManager` even though it is not an addressable allocation region.

This RFC proposes a behavior-preserving refactor:

- `SegmentPool` owns client-mounted, addressable regions that the Master can allocate from: Memory, CXL, and NoF.
- `LocalSsdManager` owns per-client local SSD registration, capacity, usage, and existing local tier state.
- Existing promotion and offload flows coordinate the two components; they are not owned by `SegmentPool`.

```mermaid
flowchart LR
subgraph Current
M1[MasterService] --> SM[SegmentManager]
SM --> MEM1[Memory]
SM --> CXL1[CXL]
SM --> LSSD1[LocalSSD state]
M1 --> NM[NoFSegmentManager]
NM --> NOF1[NoF]
end

subgraph Proposed
M2[MasterService] --> SP[SegmentPool]
SP --> MEM2[Memory regions]
SP --> CXL2[CXL regions]
SP --> NOF2[NoF regions]
M2 --> LS[LocalSsdManager]
LS --> LSSD2[Client-local object stores]
LS -. existing offload / promotion .-> SP
end
```

## Motivation

The current layout duplicates mounted-segment catalogs, indexes, allocators, locking, and lifecycle code across Memory/CXL and NoF. It also couples LocalSSD registration and task state to the memory segment manager.

A resource should belong to `SegmentPool` because it is an addressable allocation region, not because of its physical medium. This keeps NoF in the pool while separating LocalSSD, whose data is object/file based and identified by client.

## Target architecture

For addressable Memory/CXL regions, the conceptual dependency remains `object catalog / request orchestration -> SegmentPool -> RegionDriver`. Replica placement is a stateless execution seam between the object layer and the pool: it borrows a lock-scoped placement view but owns no catalog, target, allocator, or physical resource.

```mermaid
flowchart TB
subgraph ObjectLayer["Object layer"]
MS["MasterService
request orchestration"]
OC["Object catalog
MetadataStore and object shards"]
MS --> OC
end

subgraph PlacementLayer["Stateless replica placement"]
RP["ReplicaPlacement for SegmentPool
static source binding"]
RA["ReplicaAllocator
one concrete-policy dispatch"]
RP --> RA
end

subgraph PoolLayer["SegmentPool: mounted-region state owner"]
SP["SegmentPool
pool mutex and scoped access"]
CAT["MountedRegionCatalog
client, name, host, and status indexes"]
PI["PlacementIndex
logical groups and target pointers"]
REG["RegionDriverRegistry"]
SP -->|owns| CAT
SP -->|owns| PI
SP -->|owns| REG
end

subgraph DriverLayer["Physical resource ownership"]
MEM["MemoryRegionDriver
per-region allocators"]
CXL["CxlRegionDriver
global allocator and logical bindings"]
RR["RegionResource
stable AllocationTarget and allocator handle"]
BA["BufferAllocatorBase"]
REG --> MEM
REG --> CXL
MEM -->|owns| RR
CXL -->|owns| RR
RR -->|target calls| BA
end

MS -->|object lifecycle requests replicas| RP
OC -. logical replica descriptors resolve through .-> SP
RP -->|AcquirePlacementAccess| SP
RA -->|lock-scoped group selection| PI
PI -. stable non-owning pointer .-> RR
SP -->|mount, remount, and unmount| REG

subgraph SeparateOwners["Separate owners"]
LS["LocalSsdManager
client-local object storage"]
NRP["ReplicaPlacement for NoF"]
NF["NoFSegmentManager
catalog and resources until #3364"]
NRP --> NF
end

MS --> LS
MS --> NRP
RA -. LocalSSDMetricsView only .-> LS
NF -. planned NoF ownership move in issue 3364 .-> SP
```

The runtime paths use those ownership boundaries as follows:

```mermaid
flowchart LR
subgraph PutStart["PutStart hot path"]
P1["MasterService
metadata shard locked"] --> P2["ReplicaPlacement"]
P2 --> P3["ScopedPlacementReadAccess
holds pool shared lock"]
P3 --> P4["ReplicaAllocator
concrete policy"]
P4 --> P5["PlacementIndex
group and target lookup"]
P5 --> P6["AllocationTarget"]
P6 --> P7["BufferAllocatorBase::allocate"]
end

subgraph Lifecycle["Mount and unmount lifecycle"]
L1["MasterService"] --> L2["SegmentPool write transaction"]
L2 --> L3["RegionDriver PrepareOpen or PrepareAdopt"]
L3 --> L4["Validate catalog and indexes"]
L4 --> L5["Commit driver resource"]
L5 --> L6["Publish catalog and placement
no failing steps remain"]
L2 --> L7["Unmount deactivate or reactivate or erase"]
end

subgraph Snapshot["Snapshot and restore"]
S1["MasterSnapshotCodec"] --> S2["SegmentPoolSnapshotView
fork-safe read"]
S2 --> S3["Catalog view and RegionResourceReadView"]
S4["Fully validated allocator payload"] --> S5["RegionDriver PrepareAdopt"]
S5 --> S6["SegmentPool restore transaction"]
end
```

Key invariants:

- `SegmentPool` is the only owner of the Memory/CXL mounted catalog, indexes, placement index, and driver registry.
- The object catalog owns logical objects and replica descriptors only; mounted-region lookup, liveness, capacity, and physical resources resolve through `SegmentPool`.
- `RegionDriver` owns physical resources; `PlacementIndex` contains only stable non-owning target pointers.
- The pool shared lock stays held across policy selection and allocator calls, preserving allocation-versus-unmount ordering.
- `ReplicaAllocator` owns policy configuration and scratch state, not resources; adding a policy does not expand `SegmentPool`.
- `LocalSsdManager` remains independent. NoF uses the same placement executor but retains its own catalog/resources until #3364.
- Snapshot codecs consume catalog/resource views and never move serializer responsibilities into drivers or the pool lifecycle API.

## Proposed boundaries

| Component | Owns | Does not own |
| --- | --- | --- |
| `SegmentPool` | Mount/remount, lookup, status, allocation, and two-phase unmount for Memory, CXL, and NoF regions | LocalSSD object state; promotion/offload workflows |
| `LocalSsdManager` | Per-client registration, capacity, usage, enablement, and current local SSD state | Address-range allocation; segment-name indexes |
| `MasterService` | Request and object-lifecycle orchestration | Duplicate resource catalogs |

The external Memory, CXL, NoF, and LocalSSD RPC contracts remain unchanged. Type-specific RPC requests are converted at the Master boundary.

## Compatibility and invariants

This is a refactor, not a traffic migration or placement-policy change.

- Keep exactly one authoritative owner for every catalog, allocator, and LocalSSD record; no shadow state or dual writes.
- Preserve RPC signatures, error codes, mount/remount behavior, allocation order, and two-phase unmount semantics.
- Preserve the existing snapshot wire format, including the `ld` field and legacy LocalSSD decoding.
- Keep CXL represented as a MEMORY replica.
- Keep promotion targeting MEMORY replicas; NoF promotion is a separate feature.
- Preserve the current absence of NoF state from snapshots unless a separate RFC changes it.

```mermaid
flowchart TD
A[Extract LocalSSD model and snapshot seam] --> B[Move ownership to LocalSsdManager]
B --> C[Refocus SegmentManager as Memory/CXL SegmentPool]
C --> D[Move NoF catalog, allocator, and lifecycle as one unit]
D --> E[Remove compatibility adapters and run full validation]
```

## Implementation and tracking

This RFC is tracked through its native GitHub sub-issues. They should land in dependency order so each merged revision has one state owner and remains independently testable.

### Active Memory/CXL implementation stack (#3363)

The implementation is split so reusable mechanisms can be reviewed and merged before the final one-time `MasterService` cutover. The first four PRs do not switch `MasterService`; the last PR switches callers and deletes the obsolete paths in the same revision.

| Order | PR | Scope | Base | Status |
| ---: | --- | --- | --- | --- |
| 1 | #3703 | Stateful Memory/CXL region resource drivers | `main` | Draft |
| 2 | #3704 | Pointer-based placement index and concrete replica allocator | #3703 | Draft |
| 3 | #3705 | SegmentPool catalog, indexes, and transactional lifecycle | #3704 | Draft |
| 4 | #3706 | Fork-safe SegmentPool and allocator snapshot codecs | #3705 | Draft |
| 5 | #3707 | MasterService cutover and removal of legacy Memory/CXL paths | #3706 | Draft |

Review and merge in table order. #3364 remains responsible for moving NoF catalog/resource ownership into `SegmentPool`; #3365 remains responsible for cleanup that depends on that NoF move.

## Alternatives considered

- **Keep separate managers:** rejected because catalog, allocator, lifecycle, and locking behavior continue to diverge.
- **Put LocalSSD in `SegmentPool`:** rejected because LocalSSD is a client-local object store rather than a Master-allocatable address region.
- **Generalize promotion in the same change:** rejected because adding CXL/NoF target behavior would make regressions difficult to separate from ownership refactoring.

## Risks

- Snapshot regressions while moving LocalSSD ownership.
- Lock-order changes across metadata, segment, and LocalSSD state.
- Accidental placement or metrics changes during allocator migration.

Each tracking issue will carry explicit compatibility constraints and targeted tests.

## Related work

- #3158 removes business decisions from observability metrics and discusses the current `SsdMetricsProvider` naming. This RFC focuses on resource ownership and lifecycle. Any LocalSSD usage-view changes should coordinate with #3158 rather than duplicate its metrics work.

## Call for contributors

Contributions are welcome. Please comment on an unassigned tracking issue before starting and include a short implementation outline. The issues are dependency ordered; parallel work is welcome only where the issue says it is safe.

Contributors should:

1. keep the change behavior-preserving and maintain a single source of truth;
2. avoid unrelated promotion, placement, snapshot-format, or metrics changes;
3. use a `[Store]` PR title and link both this RFC and the tracking issue;
4. add or update focused tests and run the relevant Mooncake Store test targets;
5. use the repository PR template, run formatting/pre-commit on touched files, and disclose AI assistance when applicable.

The first snapshot/model extraction issue is the best entry point for contributors unfamiliar with the allocation path. The NoF integration issue is better suited to contributors familiar with allocator and unmount semantics.

### Before submitting a new issue...

- [x] Make sure you already searched for relevant issues and read the [documentation](https://kvcache-ai.github.io/Mooncake/)

Contributor guide

Open the contributing guide

Research direction

Start with the dependency-ordered tracking issues and read the existing MasterService, SegmentManager, NoFSegmentManager, and LocalSSD paths before choosing a sub-issue. Review the #3703–#3707 implementation stack and run the relevant Mooncake Store test targets. Done means behavior, RPCs, snapshots, locking, and ownership invariants remain unchanged, with each resource having one owner.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
backend, distributed-systems
Issue type
Refactor
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.