Fix/store offload bugs

[pingzhuu]

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)

[gemini-code-assist]

Code Review

This pull request introduces background worker threads to execute L2-to-L1 promotion-on-hit tasks asynchronously in FileStorage, along with corresponding configuration options, unit tests, and benchmarks. It also updates PosixFile to perform chunked vector reads and writes with enhanced error logging. The review feedback highlights critical issues: potential data corruption in PosixFile::vector_write and PosixFile::vector_read due to incorrect offset handling on partial writes/reads, a concurrency race condition during shutdown in FileStorage that could leak enqueued tasks, and a potential crash during shutdown caused by a strict CHECK on task enqueuing.

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.

[gemini-code-assist]

Issue: Potential Data Corruption on Partial Write

In PosixFile::vector_write, if ::pwritev performs a partial write (i.e., returns fewer bytes than the chunk's total size), the loop still advances idx by UIO_MAXIOV in the next iteration. This means the remaining unwritten bytes of the current chunk are skipped, and subsequent chunks are written at incorrect offsets, leading to silent data corruption.

Since PosixFile expects exact writes, we should fail immediately if ret is less than the expected chunk size.

    for (int idx = 0; idx < iovcnt; idx += UIO_MAXIOV) {
        int chunk_cnt = std::min(iovcnt - idx, UIO_MAXIOV);
        size_t chunk_bytes = 0;
        for (int i = 0; i < chunk_cnt; ++i) {
            chunk_bytes += iov[idx + i].iov_len;
        }
        ssize_t ret = ::pwritev(fd_, iov + idx, chunk_cnt, cur_offset);
        if (ret < 0) {
            int saved_errno = errno;
            LOG(ERROR) << "pwritev failed for file: " << filename_
                       << ", errno=" << saved_errno
                       << " (" << strerror(saved_errno) << ")"
                       << ", fd=" << fd_
                       << ", iovcnt=" << iovcnt
                       << ", total_bytes=" << total_bytes
                       << ", offset=" << offset
                       << ", chunk_start=" << idx
                       << ", chunk_cnt=" << chunk_cnt;
            return make_error<size_t>(ErrorCode::FILE_WRITE_FAIL);
        }
        if (static_cast<size_t>(ret) < chunk_bytes) {
            LOG(ERROR) << "pwritev partial write for file: " << filename_
                       << ", expected=" << chunk_bytes
                       << ", written=" << ret;
            return make_error<size_t>(ErrorCode::FILE_WRITE_FAIL);
        }
        written_total += ret;
        cur_offset += ret;
    }

[gemini-code-assist]

Issue: Potential Data Corruption on Partial Read

In PosixFile::vector_read, if ::preadv performs a partial read (returns fewer bytes than the chunk's total size), the loop still advances idx by UIO_MAXIOV in the next iteration. This shifts the offsets for subsequent chunks, reading data into the wrong buffers and corrupting the read result.

We should break early if ret is less than the expected chunk size (which indicates EOF or a short read), stopping further reads.

    for (int idx = 0; idx < iovcnt; idx += UIO_MAXIOV) {
        int chunk_cnt = std::min(iovcnt - idx, UIO_MAXIOV);
        size_t chunk_bytes = 0;
        for (int i = 0; i < chunk_cnt; ++i) {
            chunk_bytes += iov[idx + i].iov_len;
        }
        ssize_t ret = ::preadv(fd_, iov + idx, chunk_cnt, cur_offset);
        if (ret < 0) {
            int saved_errno = errno;
            LOG(ERROR) << "preadv failed for file: " << filename_
                       << ", errno=" << saved_errno
                       << " (" << strerror(saved_errno) << ")"
                       << ", fd=" << fd_
                       << ", iovcnt=" << iovcnt
                       << ", total_bytes=" << total_bytes
                       << ", offset=" << offset
                       << ", chunk_start=" << idx
                       << ", chunk_cnt=" << chunk_cnt;
            return make_error<size_t>(ErrorCode::FILE_READ_FAIL);
        }
        read_total += ret;
        cur_offset += ret;
        if (static_cast<size_t>(ret) < chunk_bytes) {
            break;  // EOF or partial read, stop reading further
        }
    }

[gemini-code-assist]

Issue: Concurrency Race Condition during Shutdown

In FileStorage::~FileStorage(), the lock on promotion_queue_mutex_ is released at line 251, and then promotion_workers_running_ is set to false at line 252.

During this tiny window, another thread (such as the heartbeat thread or a concurrent test runner) could call EnqueuePromotionTask(), which acquires the lock, sees promotion_workers_running_ as true, and successfully pushes a task to promotion_task_queue_. Since the queue was already cleared, this newly enqueued task will be leaked, leaving its master slot permanently pinned.

To prevent this race condition, promotion_workers_running_ should be set to false inside the lock block before clearing the queue.

Suggested change

	std::vector<PromotionTaskItem> queued_tasks_to_release;
	{
	std::lock_guard<std::mutex> lock(promotion_queue_mutex_);
	while (!promotion_task_queue_.empty()) {
	queued_tasks_to_release.push_back(
	std::move(promotion_task_queue_.front()));
	promotion_task_queue_.pop_front();
	}
	}
	promotion_workers_running_.store(false);
	std::vector<PromotionTaskItem> queued_tasks_to_release;
	{
	std::lock_guard<std::mutex> lock(promotion_queue_mutex_);
	promotion_workers_running_.store(false);
	while (!promotion_task_queue_.empty()) {
	queued_tasks_to_release.push_back(
	std::move(promotion_task_queue_.front()));
	promotion_task_queue_.pop_front();
	}
	}

[gemini-code-assist]

Issue: Potential Crash on Shutdown / Enqueue Failure

Using CHECK(enqueued) will abruptly terminate the program if EnqueuePromotionTask returns false. This can easily happen during shutdown when promotion_workers_running_ is set to false while the heartbeat thread is still processing or pulling tasks.

Instead of crashing the program, we should handle the failure gracefully by logging a warning and calling ReleasePromotionTask to ensure the master slot is released immediately.

    auto enqueue_batch = [this, &enqueued_tasks](const auto& tasks) {
        for (const auto& task : tasks) {
            const bool enqueued = EnqueuePromotionTask(
                task, /*allow_over_capacity_for_pulled_task=*/
                config_.promotion_queue_capacity > 0);
            if (!enqueued) {
                LOG(WARNING) << "Failed to enqueue promotion task for key=" << task.key
                             << "; releasing master slot";
                ReleasePromotionTask(task.key, task.tenant_id);
            } else {
                ++enqueued_tasks;
            }
        }
    };