[Store] OffsetAllocator: survive restarts with configurable persistence

[LujhCoconut]

Description

Add configurable crash-safe persistence to OffsetAllocatorStorageBackend, allowing cached KV data to survive restarts.

Background

Currently OffsetAllocatorStorageBackend::Init() opens kv_cache.data with O_TRUNC, destroying all cached data on every process restart. While this is acceptable for LLM inference (KV cache is ephemeral — loss only causes cache miss, not correctness errors), some deployments benefit from warm restarts.

This PR adds an optional persistence layer with three configurable modes, defaulting to existing O_TRUNC behavior (zero overhead, zero change for existing users). Also includes graceful-shutdown checkpoint (destructor forces SaveMetadata on clean exit) and eviction-tombstone deduplication (unordered_set + consecutive-failure counter with rate-limited logging).

Persistence Modes

Mode	Behavior	Recovery	Write Overhead	Default
kDisabled	O_TRUNC on Init, no metadata written	Empty	Zero	✅
kRelaxed	Periodic checkpoint every N seconds (default 60s). datasync() data + atomic-write metadata on each checkpoint. Destructor forces final checkpoint on clean exit.	Latest checkpoint	Low (data fsync + metadata write only once per interval)	—
kStrict	Every BatchOffload is durable before returning	Complete	High	—

Usage / Configuration

All persistence settings are controlled via environment variables. No configuration changes are required for existing deployments — the default kDisabled mode is identical to current behavior.

Enabling persistence (set before starting the storage node):

# Enable periodic checkpoint mode (recommended)
export MOONCAKE_OFFSET_PERSIST_MODE=relaxed

# Optional: change checkpoint interval from default 60s to 120s
export MOONCAKE_OFFSET_PERSIST_INTERVAL_SECONDS=120

# Start the storage node as usual
./mooncake_store_server ...

Full list of persistence environment variables:

Variable	Values	Default	Description
MOONCAKE_OFFSET_PERSIST_MODE	disabled, relaxed, strict	disabled	Persistence strategy
MOONCAKE_OFFSET_PERSIST_INTERVAL_SECONDS	integer ≥ 5	60	Checkpoint interval for kRelaxed mode

Use cases:

Scenario	Recommended Mode	Why
Default / existing deployment	disabled (or unset)	Zero overhead, behavior unchanged
Warm restart after rolling upgrade	relaxed	~84 MB checkpoint every 60s; lost data on crash ≤ 60s; destructor forces final checkpoint on clean exit
Low-latency inference, can't tolerate recomputation after crash	strict	Every batch is durable; high write overhead — not recommended for production
Development / testing	relaxed with persist_interval_seconds=10	Frequent checkpoints for fast iteration

Verifying persistence is working:

# 1. Write data, then check that kv_cache.meta exists
ls -lh /data/file_storage/kv_cache.meta

# 2. Restart the server — recovery log should show:
#    "OffsetAllocatorStorageBackend recovered: N keys, M bytes"

# 3. Check checkpoint health via metrics (if exposed):
#    GetMetadataSaveFailures() == 0  → all checkpoints succeeded
#    GetMetadataLoadFallbacks() == 0 → no recovery fallbacks to fresh start

Rollback safety: if you need to revert to an older binary without persistence support, the old binary ignores kv_cache.meta and kv_cache.meta.tmp files (treats them as unknown on-disk artifacts). Data will be lost on restart (same as pre-PR behavior), but no crash or corruption occurs. To clean up leftover files after rollback, delete kv_cache.meta and kv_cache.meta.tmp from the storage directory.

Monitoring (via public getters on OffsetAllocatorStorageBackend):

Metric	Getter	Meaning
Save failures	GetMetadataSaveFailures()	Cumulative count of SaveMetadata failures (disk full, permission errors, etc.)
Load fallbacks	GetMetadataLoadFallbacks()	Cumulative count of recovery attempts that fell back to fresh start
Consecutive failures	(logged at WARNING every 10th)	How many consecutive checkpoints have failed — resets to 0 on success
Last save latency	(internal last_save_metadata_cost_us_)	Microseconds spent in last successful SaveMetadata call

Architecture

On-disk layout: two files per arena

File	Contents	Format
kv_cache.data	Record data (unchanged)	[u32 key_len][u32 value_len][key bytes][value bytes]
kv_cache.meta	Full allocator state + FIFO index + eviction tombstones	struct_pb serialized OffsetAllocatorPersistedMetadata

Write ordering (crash consistency):

1. data_file_->datasync() — ensure all written records are durable
2. struct_pb::to_pb(meta) → write to .meta.tmp → fsync(tmp) → rename(tmp → meta) → fsync(parent_dir) — atomic metadata update
3. complete_handler() — notify master of new keys

Recovery flow (Init()):

1. Check if kv_cache.meta exists
2. If yes (and persist_mode != kDisabled): LoadMetadata() → deserialize_from<OffsetAllocator>(blob) → visit_used_nodes() to walk all used allocator nodes → read each key from data file → rebuild shard maps + total_size_/total_keys_ → restore FIFO index → Phase C delete evicted-but-resurrected keys
3. If no (or recovery fails): FallbackGuard increments metadata_load_fallbacks_ counter, fall through to fresh-start path (O_TRUNC)

Key RAII guards:

• DirtyGuard: wraps SaveMetadata() body. On any failure (including std::bad_alloc), restores metadata_dirty_ = true and increments metadata_save_failures_. On success, disarm() keeps dirty=false.
• FallbackGuard: wraps TryRecoverFromMetadata(). On any return false, increments metadata_load_fallbacks_. On success, disarm().
• FdGuard: anonymous namespace RAII fd closer, shared by Init(), TryRecoverFromMetadata(), LoadMetadata(), and SaveMetadata()'s directory fsync.

Graceful-shutdown checkpoint: ~OffsetAllocatorStorageBackend() checks persist_mode, initialized_, data_file_, and metadata_dirty_, then forces datasync() + SaveMetadata() + clear() on clean exit (SIGTERM, rolling restart). Wrapped in try/catch(...) — destuctor never throws. Move-ctor/default-copy-delete explicitly declared to avoid implicit suppression.

Eviction-tombstone dedup: all_evicted_this_batch_ is std::unordered_set<std::string> — natural dedup, O(1) lookup for stale-tombstone filtering. metadata_consecutive_failures_ counter logs every 10th consecutive checkpoint failure at WARNING level (using bare LOG(WARNING), not LOG_FIRST_N — the latter counts by call-site, not content).

Tombstone filtering: runs on every persist-enabled batch (not just checkpoint-triggering ones), removing keys that were successfully re-written in the current batch from all_evicted_this_batch_. Prevents stale-tombstone data-loss bug where a re-written key would be deleted on the next checkpoint's recovery.

Corrupt-node recovery: RebuildShardMapsFromAllocator uses free_leaked_node() lambda (constructs temporary OffsetAllocationHandle with requested_size=0 then lets RAII free the node) at every corrupt-record skip branch. Duplicate-key nodes are freed via the normal allocation_ptr RAII path (keep-first, free-duplicate). createHandleAtNode cross-validates that real_offset matches the node's stored dataOffset << multiplier_bits — mismatch returns nullopt for defense-in-depth.

Compatibility

• Default kDisabled: zero behavior change. All 56 existing tests pass.
• kRelaxed/kStrict: data files backward-compatible (record format unchanged). Metadata has version field — unrecognized versions trigger fresh-start fallback.
• FileStorage / MasterService: no changes needed.
• all_evicted_this_batch_ type change: internal only — no on-disk format impact (serialized via assign(begin,end) to vector field in struct_pb).

Known Limitations (V1)

Limitation	Impact	Mitigation
Crash between eviction and checkpoint	Evicted key may resurrect on restart	kRelaxed accepted trade-off; destructor checkpoint reduces window for clean exits; kStrict eliminates it
Full metadata serialization per checkpoint	~84 MB per checkpoint (1M nodes)	metadata_dirty_ skips idle checkpoints; future: incremental WAL
Duplicate key on recovery (pinned reader + overwrite)	May keep stale version	V1 limitation; cumulative tombstone set planned for V2

Post-Merge Optimizations

• Optimization 1 ✅ Implemented: graceful-shutdown checkpoint (this PR)
• Optimization 2 ✅ Implemented: unordered_set dedup + consecutive-failure counter (this PR)
• Optimization 3 (V2): key→offset index in separate kv_cache.index file for fast recovery (no data-file I/O)

Module

• Mooncake Store (mooncake-store)

Type of Change

• New feature

How Has This Been Tested?

Test commands:

cd build && cmake --build . --target storage_backend_test
./mooncake-store/tests/storage_backend_test

Test results:

• Unit tests pass: 56/56 (24 OffsetAllocator tests + 6 eviction-specific)
• Builds clean with zero warnings
• Default kDisabled behavior verified unchanged
• Persistence-specific tests to follow in a separate PR

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
  • Design doc: docs/source/deployment/offset-evict.md
  • Implementation plan: /home/ljh/develop/todo/plan_offset_presist.md (v22)
  • Optimization plan: /home/ljh/develop/todo/opt_persist.md
• I have added tests to prove my changes are effective (persistence-specific tests to follow)
• For changes >500 LOC: I have filed an RFC issue

AI Assistance Disclosure

• AI tools were used

AI tools (Claude Code) assisted with: design iteration (~22 rounds adversarial review), code generation (all 7 files), iterative bug fixing, and optimization design. The human submitter reviewed, understood, and approved all changes.

[gemini-code-assist]

Code Review

This pull request introduces a persistence and recovery mechanism for the OffsetAllocatorStorageBackend. Key additions include data synchronization interfaces, metadata serialization, and recovery routines to rebuild shard maps and restore the FIFO index from persisted checkpoints. The review feedback identifies a critical bug where stale tombstones are not filtered out when checkpointing is skipped, potentially leading to data loss on recovery. Additionally, the reviewer suggests defensive checks to prevent recovering corrupted empty keys and recommends using FdGuard for directory file descriptors to ensure resource safety.

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.

[LujhCoconut]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces persistence and recovery capabilities to the OffsetAllocatorStorageBackend. It adds support for periodic or strict metadata checkpointing, including serialization of the allocator state and FIFO index, as well as recovery logic to rebuild shard maps from the data file. Additionally, it updates PosixFile and UringFile with datasync operations and options to prevent unlinking files on write failures. The review feedback highlights two important issues: a potential space leak of allocator nodes during recovery when corrupt records are skipped, and a failure in SaveMetadata when using relative paths due to an empty parent path being passed to open.

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.

[he-yufeng]

I focused this pass on the crash-consistency of the persistence path — the part most likely to have a subtle durability bug — and the ordering is correct at every level, which is great to see:

• Data-before-metadata: data_file_->datasync() runs before SaveMetadata(), so the arena bytes an offset refers to are durable before the metadata that claims that offset is used. That's the ordering that matters for recovery correctness.
• Atomic metadata swap: SaveMetadata writes to .tmp, fdatasyncs the tmp fd before the rename, renames, then fsyncs the parent directory. The directory fsync is the step that's usually missing — without it a crash right after rename can lose the swap even though the tmp bytes were durable — so it's good to see it here.
• Fault coverage: the test_metadata_write_failure_step_ injection at the pre-fdatasync (step==2) and post-rename (step==3) points, plus SetDeleteOnWriteFail(false) guarding the data file across the metadata write, means the failure windows are actually exercised rather than assumed.

Defaulting to kDisabled keeps this a no-op for existing deployments, and gating kRelaxed on persist_interval_seconds >= 5 is a sane floor.

I haven't done a full pass on the recovery reconstruction (visit_used_nodes / createHandleAtNode / RebuildShardMapsFromAllocator) or the per-mode BatchOffload semantics, so I'm leaving this as a comment rather than an approval — but the durability discipline is solid.

[LujhCoconut]

/gemini review

[gemini-code-assist]

Code Review

This pull request implements a persistence and recovery mechanism for the OffsetAllocatorStorageBackend. It introduces a datasync interface for storage files, persistence configuration modes, and metadata serialization to enable state recovery upon restart. The review feedback suggests improving the robustness of metadata loading by handling EINTR and short reads in pread, and using fstat instead of lseek to query the data file size safely.

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

Files with missing lines	Patch %	Lines
mooncake-store/src/storage_backend.cpp	58.91%	279 Missing ⚠️
...-store/include/offset_allocator/offset_allocator.h	66.66%	10 Missing ⚠️
mooncake-store/src/offset_allocator.cpp	78.26%	5 Missing ⚠️
mooncake-store/src/uring_file.cpp	0.00%	3 Missing ⚠️
mooncake-store/src/posix_file.cpp	71.42%	2 Missing ⚠️
mooncake-store/include/file_interface.h	50.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[LujhCoconut]

Please consider merging #2964 first, and I will rebase afterwards to adapt to its changes, since #2964 has a smaller change surface. This PR is completely orthogonal to #2964 — one addresses runtime correctness, the other addresses
cross-restart data survival.