Skip to content

Fix/store offload main#3

Open
pingzhuu wants to merge 18 commits into
mainfrom
fix/store-offload-main
Open

Fix/store offload main#3
pingzhuu wants to merge 18 commits into
mainfrom
fix/store-offload-main

Conversation

@pingzhuu

@pingzhuu pingzhuu commented Jul 3, 2026

Copy link
Copy Markdown
Owner

Description

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

How Has This Been Tested?

Test commands:

# Example: bash scripts/run_ci_test.sh

Test results:

  • Unit tests pass
  • Integration tests pass (if applicable)
  • Manual testing done (describe below)

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

AI Assistance Disclosure

  • No AI tools were used
  • AI tools were used (specify below)

pingzhuu added 13 commits July 3, 2026 12:50
GroupOffloadingKeysByBucket previously returned early when the input
exhausted mid-bucket, dumping the partial keys back into
ungrouped_offloading_objects_. Under low inflow (e.g. after a pressure
test rampdown, only 151-491 keys/call), the pool never reached
bucket_keys_limit (1000), so keys sat indefinitely — their master-side
offloading_task expired before ever being written to SSD.

Fix: emit the partial bucket to buckets_keys immediately at input end
instead of re-buffering. Trade a small per-bucket metadata overhead for
correctness: no indefinite backlog.
…ug H)

When a store-server restarts with a new client_id, ScanMeta re-registers
its LOCAL_DISK replicas via AddReplica. The previous code's VisitReplicas
predicate required matching client_id, so the stale replica (from the
dead client) blocked the update: the new replica was silently dropped
and reads followed the dead transport_endpoint (INVALID_KEY / RDMA error).

Fix: before checking for an existing LOCAL_DISK replica, erase any
LOCAL_DISK replicas whose client_id differs from the caller's. This
clears stale entries from restarted store-servers so the new replica
can be added or the same-client_id one updated.
The previous PutStart check used metadata.HasReplica(&Replica::fn_is_completed)
which matches ANY completed replica type, including LOCAL_DISK. After a
store-server restart, ScanMeta re-registers LOCAL_DISK replicas, leaving
metadata with only a completed LOCAL_DISK replica. The next PutStart for
that key was then rejected with OBJECT_ALREADY_EXISTS, blocking writes
to keys whose cache lives on SSD.

Fix: narrow the check to completed MEMORY replicas only. A LOCAL_DISK-
only metadata object should not block a fresh PutStart — the disk
replica's purpose is to serve cache reads, not to gate writes.
…xits on first pin (Bug K)

PutEnd's offload-on-write path (!offload_on_evict_) lacked two safeguards:

  1. No 'if (task_created) return' early-exit inside the VisitReplicas
     visitor. With replica_num>1 (MOONCAKE_STORE_REPLICA_NUM=2), the
     second MEMORY replica in a different segment would also succeed at
     PushOffloadingQueue + inc_refcnt, but the second emplace would
     silently fail (key already in offloading_tasks), leaking the refcnt.

  2. No offloading_tasks.count(key)==0 pre-check. If an offloading task
     was already in flight (rare race), the visitor's first iteration
     would inc_refcnt + emplace-fail, also leaking.

Fix: add the early-exit + pre-check + emplace return-value check with
dec_refcnt rollback. This closes the refcnt leak that manifested as
OFFLOAD-COMPLETE gap (admitted != completed) under replica_num=2.
…pin refcnt leak (Bug J part 1)

BatchEvict's try_evict_or_offload lambda iterates candidate metadata and
pins one MEMORY replica per key for SSD offload. Two issues caused a
permanent refcnt leak with replica_num>1:

  1. No offloading_tasks.count(key)==0 pre-check. After cycle N pins
     replica A, FetchOffloadingTasks on the store-server side clears the
     master's offloading_objects map; cycle N+1 can then pin replica B
     for the same key. The emplace silently no-ops (key already in
     offloading_tasks from cycle N), but B's inc_refcnt is never rolled
     back — the replica is pinned forever, never evictable.

  2. emplace return value unchecked. Even within a single cycle, a rare
     race between BatchEvict and another Path (PutEnd offload, tenant
     quota eviction) can pre-insert the key, making emplace fail.

Fix: add the count>0 pre-check and the emplace return-value check with
dec_refcnt rollback. Defense-in-depth: even if the count pre-check
passes, the emplace check catches races.
… (Bug J part 2)

EvictTenantMemoryForQuota has its own try_evict_or_offload lambda with
the same shape as BatchEvict's. Apply the same two safeguards:

  1. offloading_tasks.count(key)==0 pre-check before VisitReplicas
  2. emplace return-value check with dec_refcnt rollback

Without these, tenant-quota-driven eviction under replica_num>1 can pin
a second replica for a key whose offloading_tasks entry was inserted by
BatchEvict, leaking the refcnt and making the replica unevictable.
BatchOffload inserted new buckets into lru_index_ with timestamp 0,
making them the smallest (i.e. 'oldest') entry. SelectEvictionCandidate
picks lru_index_.begin(), so once the SSD fills up the newly written
bucket — whose last_access_ns_ is still 0 — is selected before any
older bucket whose BatchLoad read updated its timestamp.

Confirmed via diagnostic logs on a small-SSD test store:
  [LRU-BUG] BatchOffload commit: bucket_id=X lru_ts=0 keys=5
  [LRU-BUG] EvictCandidate: bucket_id=X ts=0 (matches actual)
Same bucket_id was committed and immediately selected for eviction.

After fix:
  [LRU-FIX] BatchOffload commit: bucket_id=X lru_ts=<steady_clock ns>
  EvictCandidate picks oldest bucket by real timestamp, not newest.

Production impact: hit rate on a full SSD recovered from ~27% to ~80%.

Only affects deployments using MOONCAKE_OFFLOAD_BUCKET_EVICTION_POLICY=lru
(the default fifo policy is keyed by bucket_id insertion order and is
not affected). ScanMeta recovery path (line 1592) intentionally keeps
timestamp 0 — recovered buckets are genuinely older than new writes.
… (PR kvcache-ai#2529)

ProcessPromotionTasks previously ran SSD read + RDMA write for each
candidate inline on the store-server's heartbeat thread. With
promotion_max_per_heartbeat=64, a single heartbeat tick spent hundreds
of milliseconds in batch_get + TransferWrite, stalling the next
heartbeat tick. Under sustained cache churn this caused promotion task
TTL expiry (default put_start_release_timeout_sec) and a sustained
outage of the promotion queue.

Fix: split the per-key body into ProcessPromotionTask and hand
candidates to a pool of background worker threads via a mutex/cv
guarded queue. ProcessPromotionTasks keeps pulling from the master
until the local queue reaches its soft budget, then returns; workers
drain the queue in batches, decoupling heartbeat latency from
promotion work.

Env vars (default in parens):
  MOONCAKE_OFFLOAD_PROMOTION_WORKER_THREADS (1)
    - 0 disables workers and falls back to inline execution
    - 4 recommended for production under load
  MOONCAKE_OFFLOAD_PROMOTION_QUEUE_CAPACITY (1024)
    - 0 = unbounded
  MOONCAKE_OFFLOAD_PROMOTION_DRAIN_BATCH_SIZE (64)

When workers are not running (Init not yet called, or
promotion_worker_threads=0), ProcessPromotionTasks falls back to the
pre-worker synchronous loop so existing tests and standalone callers
keep working.

File changes:
  - include/storage_backend.h: add promotion_worker_threads /
    promotion_queue_capacity / promotion_drain_batch_size to
    FileStorageConfig
  - include/file_storage.h: declare ProcessPromotionTask /
    EnqueuePromotionTask / ReleasePromotionTask /
    PromotionWorkerThreadFunc + add worker member fields
  - src/file_storage.cpp: read env vars; spawn workers in Init; join
    workers in destructor; rewrite ProcessPromotionTasks to pull+enqueue;
    factor per-key body into ProcessPromotionTask. Heartbeat() now
    always calls ProcessPromotionTasks even when offloading_objects is
    empty, so promotion backlog can drain independently of offload.
…r_threads 1->4

After PR kvcache-ai#2529 decoupled promotion execution from the synchronous
heartbeat path into a worker pool, the old conservative defaults
(promotion_max_per_heartbeat=1, promotion_worker_threads=1) became the
bottleneck. Under 64-concurrency pressure testing, promotion task
admission could not keep up: 25259 of ~70000 admitted tasks expired
their TTL because the master admitted only 2 tasks/sec while the
client demanded 70 tasks/sec.

Raising defaults:
  promotion_max_per_heartbeat: 1 -> 64  (master_config.h, 4 config classes)
    Admission rate at heartbeat=5s, 4 servers: 64*4/5 ~= 51/s per server,
    204/s aggregate — well above the 70/s demand.
  promotion_worker_threads: 1 -> 4  (storage_backend.h FileStorageConfig)
    With drain_batch_size=64, 4 workers process 256 tasks in parallel,
    matching the admission rate.

After this change promotion task expired drops to 0 with no extra env
config. Tunable via MOONCAKE_OFFLOAD_PROMOTION_WORKER_THREADS env var
(0 disables async workers and falls back to inline execution).
…r threads (Bug F)

OffloadObjects previously iterated buckets_keys serially, issuing one
BatchOffload (-> WriteBucket -> pwritev) per bucket per heartbeat tick.
Each bucket is a separate file with no shared write state, so this
serialization was purely artificial; on a single NVMe device it limited
write throughput to ~1.6 GB/s while the device is capable of ~4.7 GB/s.

Fix: capture the per-bucket body as a lambda returning std::optional<
ErrorCode> (nullopt on success/recoverable failure, error code on
unrecoverable failure) and dispatch the buckets across a fixed set of
worker threads. Workers pull bucket indices from a shared atomic
counter and stop as soon as one bucket returns an unrecoverable error,
so the caller joins all workers before returning. The serial path is
retained when offload_write_threads<=1 or there is only one bucket,
preserving the legacy behavior with no thread overhead.

Data safety:
  - failed_tasks and all_bucket_keys are merged under failed_tasks_mutex
    instead of being shared raw across threads.
  - staging_bufs (PinnedBufferPool::Buffer RAII handles) are local per
    bucket, so no cross-thread pool interaction.
  - bucket_complete_handler captures [this, offload_start,
    complete_handler] — ssd_metric_ counters are designed for
    concurrent inc/observe; complete_handler client_->NotifyOffload-
    Success RPCs are independent.
  - The caller joins all worker threads before OffloadObjects returns,
    so captures stay valid for the lifetime of the parallel region.

Env var (added): MOONCAKE_OFFLOAD_WRITE_THREADS (default 4). Set to 1
to fall back to the legacy serial loop.
Add INFO/WARNING tags at three locations where offloading tasks can
silently become invisible to operators:

  [OFFLOAD-NOTASK]  NotifyOffloadSuccess: store worker reports a
                    completed offload for a key whose offloading_task
                    was already expired or cleaned up by BatchEvict.
                    Previously this was silent; AddReplica would still
                    register the LOCAL_DISK replica, but the orphan
                    path was invisible.

  [CLEANUP-ORPHAN]  EraseMetadata: metadata is being erased while an
                    offloading_task is still registered for the key
                    (BatchEvict's second pass racing the store worker
                    in-flight offload). Previously dec_refcnt + erase
                    happened silently.

  [PUSH-EMPTY]      PushOffloadingQueue: caller tried to offload a
                    replica whose segment_names is empty — silent
                    no-op previously, masked as 'nothing to offload'.

These tags make it possible to grep master logs for orphan-task
signals during diagnosis, matching the [OFFLOAD-SKIP] /
[OFFLOAD-PARTIAL] tags already added by Bug D and the [BATCHEVICT-*]
/ [PUTEND-OFFLOAD-EMPLACE-FAIL] / [TENANT-EVICT-OFFLOAD-EMPLACE-FAIL]
tags added by Bug J/K.
vector_write and vector_read previously called ::pwritev/::preadv with
the full iovec array without clamping to UIO_MAXIOV (1024). When a
bucket contained more than ~513 keys (each k/v pair contributing 2
iovecs), pwritev returned EINVAL (error 'Invalid argument') which the
caller translated to FILE_WRITE_FAIL.

Reproduced on the new main-based branch: launching production SSD offload
with BUCKET_KEYS_LIMIT=1000 immediately printed
  vector_write failed for: <bucket_id>, error: FILE_WRITE_FAIL
for every bucket write. Files were deleted by PosixFile's destructor
('Deleted corrupted file: ...bucket'). The cluster stayed at 0 SSD usage
and never served a single cache hit.

Fix (re-applied from the old fix/store-offload-clean branch):
  1. Split the iovec array into chunks of UIO_MAXIOV before calling
     pwritev/preadv, advancing the file offset by bytes written/read
     in each chunk, so buckets with > 1024 iovecs work correctly.
  2. Verify the total bytes match the expected byte count across all
     chunks. Partial writes are surfaced as FILE_WRITE_FAIL.
  3. Log errno(strerror) + fd + iovcnt + total_bytes + offset +
     chunk_start/chunk_cnt on read/write/pwritev/preadv failures so
     future FILE_WRITE_FAIL reports point at the actual syscall
     failure rather than just the ErrorCode.

Upstream main is missing both the chunking and the diagnostics.
Replace the previous A-K doc (which targeted the old fix/store-offload-clean
branch based on b57d18a) with a new record covering the 12 commits on
fix/store-offload-main (origin/main @ 85e0ed6 + 12 patches).

Each section now maps one-to-one to a commit and includes:
- the root cause and reproduction symptoms
- the code-path details that caused silent failures
- verification results from test_lru_bug_repro.py,
  test_store_restart_cache.py, and AI_ERA_PRESSURE pressure testing

A new 'Already covered upstream' table records the three bugs (A/B/C)
that upstream PRs kvcache-ai#2286, kvcache-ai#2599, kvcache-ai#2658 fixed since our old base, plus the
SSD capacity re-report that upstream file_storage.cpp already implements.
This is why those fixes are not reapplied on this branch.

Also notes the posix_file UIO_MAXIOV fix (commit 7936a3b) which is
mandatory on upstream main — without it, every bucket write with > ~513
keys returns FILE_WRITE_FAIL (EINVAL), and the cluster stays at 0 SSD
usage / 0 cache hits.
@github-actions github-actions Bot added documentation Improvements or additions to documentation run-ci store labels Jul 3, 2026
pingzhuu added 5 commits July 3, 2026 15:35
strerror() returns a pointer to a static internal buffer that is shared
across all threads; concurrent calls (now possible with parallel
WriteBucket and promotion worker threads) corrupt the returned message.

Switch all 4 call sites (write/read/pwritev/preadv failure logs) to
strerror_r, the POSIX thread-safe variant that writes into a
caller-provided stack buffer.
WriteBucket previously used the shared member aligned_io_buffer_ (inside
the USE_URING / O_DIRECT path) to stage data before write_aligned. With
parallel WriteBucket dispatch (introduced by Bug F's worker threads),
multiple concurrent calls overwrite each other's stage buffer, causing
SSD data corruption.

Fix: always posix_memalign a per-call buffer and free it when the
function exits via the temp_buffer unique_ptr RAII guard. This removes
the shared-buffer race entirely at the cost of one posix_memalign +
free per bucket write — negligible compared to the pwritev itself.
ProcessPromotionTasks' master-pull loop used:
  std::max<size_t>(1, config_.promotion_queue_capacity)
When promotion_queue_capacity is 0 (meaning 'unbounded'), std::max(1, 0)
returns 1, so the loop breaks as soon as a single task is queued —
defeating the unbounded behavior entirely.

Fix: only enforce the cap when capacity > 0, matching the guard in
EnqueuePromotionTask (line ~1066).
Init used std::max<uint32_t>(1, config_.promotion_worker_threads) which
clamped 0 to 1, always spawning at least one worker and setting
promotion_workers_running_ = true. This prevented the synchronous
(inline-in-heartbeat) fallback path in ProcessPromotionTasks from ever
being used, contradicting the documented 0 = disable-workers behavior.

Fix: only spawn workers and set the flag when worker_threads > 0.
When worker_threads == 0, promotion_workers_running_ stays false and
ProcessPromotionTasks falls through to the inline loop.
…isk_object_count sync

EvictTenantMemoryForQuota's try_evict_group_or_object lambda and its
'pass' loop previously called the legacy 3-arg EraseMetadata overload,
which forwards to the 5-arg implementation with shard=nullptr. This
skipped OnDiskReplicaRemoved, so disk_object_count was never
decremented when LOCAL_DISK replicas were erased during tenant-quota
eviction. Over time disk_object_count drifted upward, causing
BatchEvict's eviction_base = metadata.size() - disk_object_count to
underflow and degrade eviction decisions.

Fix: add a MetadataShardAccessorRW* parameter to the try_evict_group_
or_object lambda and forward &shard from the pass loop. Both EraseMetadata
call sites in EvictTenantMemoryForQuota now use the 5-arg overload,
matching the BatchEvict path which already passes &shard.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation run-ci store

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant