diff --git a/.dockerignore b/.dockerignore index d9bdcd4e01..3b0bfaba0d 100644 --- a/.dockerignore +++ b/.dockerignore @@ -1,2 +1,5 @@ **/.dockerignore **/*.Dockerfile +extern/yalantinglibs +build/ +dist/ diff --git a/docs/store-offload-bug-fixes.md b/docs/store-offload-bug-fixes.md new file mode 100644 index 0000000000..8b292ee959 --- /dev/null +++ b/docs/store-offload-bug-fixes.md @@ -0,0 +1,341 @@ +# Mooncake Store SSD Offload Bug 修复记录 + +> 分支: `fix/store-offload-main` (基于 upstream `main` @ `85e0ed69`) +> 仓库: `pingzhuu/Mooncake` +> 最后更新: 2026-07-03 + +## 背景 + +store-nvme 部署（.15 master + 2 store-server, .18 2 store-server, 4× NVMe SSD） +在 AI_ERA_PRESSURE 压测中出现缓存命中率从 ~80% 断崖式下跌到 ~27% 的问题。 +经多轮诊断后确认根因为 LRU eviction 在提交新 bucket 时使用了时间戳 0， +并伴有一组并发性能 / 元数据正确性 / 串行 I/O 瓶颈等问题。 + +本分支将旧分支 (`fix/store-offload-clean` 基于 `b57d18a4`) 的修复方案 +逐 commit 重新应用到 upstream main，且每个修复独立成 commit 以便审阅与 +cherry-pick。本文件按 commit 顺序记录每个修复。 + +## 已被上游覆盖（无需在本分支修改） + +| Bug | 上游 PR / 说明 | +|-----|---------------| +| A — BatchEvict 重复迭代 / 过量淘汰 | PR #2286 (`be1ceba6`) — Phase 1/2 重构并列扫描候选 | +| B — offload cap 绑定 offload_force_evict_ | PR #2599 (`077b696f`) — `--offloading_queue_limit` / `--offload_cap_ratio` 已为 gflag | +| C — 跳过的 offload key 静默丢弃 → orphan task | PR #2658 (`70ad6c7d`) — `data_size=-1` NACK + NotifyOffloadSuccess 处理 | +| 0ca9f054 等价的 SSD 容量 re-report | upstream `file_storage.cpp` Heartbeat 中已有 MountLocalDiskSegment + ReportSsdCapacity 逻辑 | + +## 本分支提交列表 + +### Bug D: partial bucket 永久滞留于 ungrouped_offloading_objects_ + +**commit**: `81ed7c08` +**文件**: `mooncake-store/src/storage_backend.cpp` `GroupOffloadingKeysByBucket` + +**根因**: 当输入 offloading_objects 在凑满 `bucket_keys_limit` 之前耗尽， +旧逻辑把剩余 key 倒回 `ungrouped_offloading_objects_` 并 `return {}`。压测 +high-watermark 之后每个 heartbeat 只有 ~151-491 key 到达，永远凑不满 +1000 key/bucket → key 永久滞留 → master 端 offloading_task TTL 超 +`put_start_release_timeout_sec` 而 expired。 + +**修复**: 输入耗尽时立即把已收集的 partial bucket 推入 `buckets_keys` +（带 `[OFFLOAD-PARTIAL]` 日志），不再回灌 ungrouped 池。代价是少量额外 +bucket 元数据开销，换得正确性：不会无限积压。 + +--- + +### Bug H: AddReplica 不替换 stale LOCAL_DISK replica + +**commit**: `7f57409c` +**文件**: `mooncake-store/src/master_service.cpp` `AddReplica` + +**根因**: store-server 重启后 ScanMeta 回告时，`AddReplica` 的 `VisitReplicas` +predicate 要求 `client_id` 匹配；旧 client_id 的 replica 不匹配 → 新 replica +被静默丢弃，master 中残留 dead transport_endpoint。后续 GET 取不到新 transport +→ `INVALID_KEY` / RDMA 错误。 + +**修复**: 在检查是否已有 LOCAL_DISK replica 之前，erase 掉所有 +`client_id != caller` 的 LOCAL_DISK replica（`[ADDREPLICA-STALE]` 日志）， +并通过 `EraseReplicasWithCacheTotalAccounting` + `OnDiskReplicaRemoved` 保持 +shard-level disk_object_count 与本地 SSD usage 计数一致。剩余同 client_id +的 replica 走原路径 update endpoint / size。 + +--- + +### Bug I: PutStart OBJECT_ALREADY_EXISTS 阻塞 LOCAL_DISK-only metadata + +**commit**: `05835124` +**文件**: `mooncake-store/src/master_service.cpp` `PutStart` + +**根因**: PutStart 用 `metadata.HasReplica(&Replica::fn_is_completed)` 判断是否 +已存在写入——`fn_is_completed` 匹配任意已完成 replica（含 LOCAL_DISK）。 +store 重启 + ScanMeta 后 metadata 仅含一个 completed LOCAL_DISK replica，导致 +每次相同 key 的 PutStart 都返回 `OBJECT_ALREADY_EXISTS`，把已写到 SSD 上的 +数据永久 gate 掉。 + +**修复**: 改为 `metadata.HasReplica([](r){ r.is_memory_replica() && r.is_completed(); })`。 +LOCAL_DISK-only metadata 不应阻塞新的 PutStart——它的作用是服务 cache 读而不是 +gate 写。 + +**验证**: `tests/test_store_restart_cache.py` Step 6 PUT 阶段 300000 key +`ok=300000 fail=0`。 + +--- + +### Bug K: PutEnd offload 路径 double-pin + emplace 失败不回退 + +**commit**: `d2048ae6` +**文件**: `mooncake-store/src/master_service.cpp` `PutEnd` + +**根因**: PutEnd 的 `!offload_on_evict_` 分支内 VisitReplicas 缺少 +`if (task_created) return` early-exit，且 `emplace` 返回值未检查。`replica_num=2` +时第二个 replica 仍然 `PushOffloadingQueue + inc_refcnt`，但第二份 emplace +被同 key 已存在的 task 排斥 → refcnt 永久泄漏。即使 `replica_num=1`， +若 `offloading_tasks[key]` 已存在（与 BatchEvict 路径的 race），也会泄漏。 + +**修复**: 进 lambda 前加 `offloading_tasks.count(key)==0` 前置检查； +lambda 内 `if (task_created) return` 早退；`emplace` 返回 `inserted=false` +时立即 `dec_refcnt` 回退（`[PUTEND-OFFLOAD-EMPLACE-FAIL]` 日志）。 + +**验证**: 旧分支在生产压测下 OFFLOAD-COMPLETE gap 从 87623 → 0； +新分支压测 5 分钟 OFFLOAD-COMPLETE gap=0。 + +--- + +### Bug J part 1: BatchEvict 跨周期重复 pin → refcnt 泄漏 + +**commit**: `77b43718` +**文件**: `mooncake-store/src/master_service.cpp` `BatchEvict::try_evict_or_offload` + +**根因**: `try_evict_or_offload` 在把 key 推入 offloading queue 之前没有检查 +`offloading_tasks.count(key) > 0`。BatchEvict 第 N 轮把 replica_A 推入队列并 pin； +FetchOffloadingTasks 把 master side `offloading_objects` 清空后，第 N+1 轮可以 +对同一 key 的 replica_B 再 pin 一次。`emplace(key, T2)` 静默失败（key 已有 T1）， +但 replica_B 的 `inc_refcnt` 不会回退 → replica 永久 pin，不可淘汰。 + +**修复**: VisitReplicas 前加 `count(key)==0` 前置检查；emplace 返回 `inserted=false` +时立即 `dec_refcnt` 回退（`[BATCHEVICT-OFFLOAD-EMPLACE-FAIL]` 日志）。双重防御。 + +--- + +### Bug J part 2: EvictTenantMemoryForQuota 同一泄漏路径 + +**commit**: `d3690852` +**文件**: `mooncake-store/src/master_service.cpp` `EvictTenantMemoryForQuota::try_evict_or_offload` + +**根因**: 与 Bug J part 1 同形状，tenant-quota 驱动 evict 路径同样存在 +`emplace` 失败不回退与缺少 count 前置检查的问题。在 BatchEvict 已插入 task 后 +紧接着触发 tenant-quota evict，会对同 key 的第二个 replica 造成 refcnt 泄漏。 + +**修复**: 完全对称的两项防护（count 前置 + emplace 返回值 + `dec_refcnt` 回退， +`[TENANT-EVICT-OFFLOAD-EMPLACE-FAIL]` 日志）。 + +--- + +### LRU bug: BatchOffload 提交新 bucket 时把 LRU 时间戳写为 0 + +**commit**: `40ab7c7d` +**文件**: `mooncake-store/src/storage_backend.cpp` `BatchOffload` + +**根因**: BatchOffload 提交新 bucket 时执行 +`buckets_.emplace(bucket_id, std::move(bucket)); lru_index_.emplace(0LL, bucket_id);`。 +`0` 严格小于任何真实 `steady_clock::now()` 时间戳，所以新写入的 bucket 在 +`SelectEvictionCandidate` 的 LRU 路径下永远是首选项。一旦 SSD 写满触发 +PrepareEviction，新写入的 cache 还没被读到就被淘汰掉，命中率从 ~80% 跌至 ~27%。 + +诊断日志确认： +``` +[LRU-BUG] BatchOffload commit: bucket_id=X lru_ts=0 keys=5 +[LRU-BUG] EvictCandidate: bucket_id=X ts=0 (matches actual) +``` +同一个 bucket_id 刚提交就被选为淘汰候选。 + +**修复**: 提交时 +```cpp +int64_t now_ns = std::chrono::steady_clock::now().time_since_epoch().count(); +bucket->last_access_ns_.store(now_ns, std::memory_order_relaxed); +buckets_.emplace(bucket_id, std::move(bucket)); +lru_index_.emplace(now_ns, bucket_id); +``` +ScanMeta 恢复路径（line ~1592）有意保持 `0`，因为重启后从磁盘扫出的 bucket +确实应当被视为最旧。 + +**影响**: 仅当 `MOONCAKE_OFFLOAD_BUCKET_EVICTION_POLICY=lru` 时触发；默认 `fifo` +按 bucket_id 单调递增淘汰，不受影响。本分支生产配置使用 LRU，必须修。 + +**验证**: `tests/test_lru_bug_repro.py` 在 30MB SSD + 100MB DRAM 的测试 store 上 +通过（`b15 survived — older data was evicted first`）。 + +--- + +### Promotion worker pool（PR #2529 commit 1/3，本分支独立重做） + +**commit**: `dc40e854` +**文件**: `mooncake-store/include/file_storage.h`、`include/storage_backend.h`、`src/file_storage.cpp` + +**根因**: 上游 `ProcessPromotionTasks` 在 store-server 的 heartbeat 线程内 +同步执行每个 promotion task（SSD read + RDMA write）。`promotion_max_per_heartbeat=1` +时每 tick 只 2/s 准入；调到 64 后同一 tick 内 batch_get + TransferWrite 64 个 task +会阻塞 heartbeat 线程几百毫秒，拖垮心跳节奏导致积压。 + +**修复**: 把 per-key body 拆出为 `ProcessPromotionTask(task, preferred_segments)` +返回 `PromotionExecutionResult`，并用 mutex/cv 守护的 `promotion_task_queue_` ++ 后台 worker pool 异步消费。`ProcessPromotionTasks` 持续从 master +`PromotionObjectHeartbeat` 拉 task 直到本地队列达到 soft budget。 + +新增 env vars: +- `MOONCAKE_OFFLOAD_PROMOTION_WORKER_THREADS` (默认 1，生产建议 4) + - `0` 走旧 inline 路径，便于测试与回退 +- `MOONCAKE_OFFLOAD_PROMOTION_QUEUE_CAPACITY` (默认 1024，0 = 无限) +- `MOONCAKE_OFFLOAD_PROMOTION_DRAIN_BATCH_SIZE` (默认 64) + +Init 时 spawn worker；~FileStorage join worker。 + +--- + +### Promotion 默认值调整 + +**commit**: `5cb97e1d` +**文件**: `mooncake-store/include/master_config.h`、`mooncake-store/include/storage_backend.h` + +**根因**: PR #2529 把 promotion 从 heartbeat 同步路径剥离后，老保守默认值 +(`promotion_max_per_heartbeat=1`、`promotion_worker_threads=1`) 仍然存在，导致 +64 并发压测下 25259/~70000 admitted task TTL expired——准入 2/s 满足不了 70/s +的实际需求。 + +**修复**: +- `master_config.h` 4 个 config struct 的 `promotion_max_per_heartbeat`: 1 → 64 +- `FileStorageConfig::promotion_worker_threads`: 1 → 4 +- 64×4/5s ≈ 51/s/server，4 server 聚合约 204/s 远超 70/s 需求 + +调优后 promotion task expired 从 25259 → 0。 + +--- + +### Bug F: OffloadObjects 串行 WriteBucket → SSD 写吞吐瓶颈 + +**commit**: `e738cfd9` +**文件**: `mooncake-store/src/file_storage.cpp`、`mooncake-store/include/storage_backend.h` + +**根因**: `OffloadObjects` 串行遍历每个 bucket 调 `BatchOffload` (→ WriteBucket → +pwritev)。每个 bucket 是独立文件，无共享写状态，串行化纯属人为瓶颈。单盘 NVMe +~1.6 GB/s vs 实测可达 ~4.7 GB/s，跟不上 high-watermark 触发后 cache 涌入速率。 + +**修复**: 把单个 bucket 的处理体捕获为返回 `std::optional` 的 lambda +（nullopt = 成功/可恢复；含值为不可恢复）。 +- `offload_write_threads <= 1` 或仅 1 bucket → 走旧串行路径，零开销 +- 否则 spawn `write_threads` 个 std::thread，从原子 `next_index` 抢 bucket， + 任何 bucket 返回错误即设 `first_error_set` 让其他 worker 立即退出； + caller join 全部 worker 后再返回，保证捕获的引用生命周期安全 + +数据安全： +- `failed_tasks` / `all_bucket_keys` 用 `failed_tasks_mutex` 合并 +- staging_bufs 是 per-bucket RAII 局部 +- complete_handler 的 ssd_metric_ counters 与 client NotifyOffloadSuccess RPC + 均为独立路径 + +新增 env: `MOONCAKE_OFFLOAD_WRITE_THREADS` (默认 4)；=1 退化到串行。 + +--- + +### Bug E: orphan task 诊断日志 + +**commit**: `64863cb0` +**文件**: `mooncake-store/src/master_service.cpp` + +修复 silent orphan-task 路径，加 3 个 tag： +- `[OFFLOAD-NOTASK]` — NotifyOffloadSuccess 收到 complete 但 master 没有 + offloading_task（已被 expired / 清理）。之前静默 AddReplica，现在可见 +- `[CLEANUP-ORPHAN]` — EraseMetadata 删 metadata 时该 key 仍挂 offloading_task + （BatchEvict 第二遍 vs store worker in-flight race） +- `[PUSH-EMPTY]` — PushOffloadingQueue 收到 segment_names 为空的 replica + （之前静默 no-op） + +与 Bug D 的 `[OFFLOAD-PARTIAL]`、Bug J/K 的 emplace-fail 日志一起构成 +orphan-task 诊断套件。 + +--- + +### pwritev/preadv 按 UIO_MAXIOV 分块 + errno 诊断 + +**commit**: `7936a3b4` +**文件**: `mooncake-store/src/posix_file.cpp` + +**根因**: `vector_write` / `vector_read` 直接把整个 iovec 数组传给 `pwritev` / +`preadv`，没有按 `UIO_MAXIOV` (1024) 分块。`BUCKET_KEYS_LIMIT=1000` × 每 key +2 iovec > 1024 → pwritev 返回 `EINVAL` → 转译成 `FILE_WRITE_FAIL`。 + +新分支部署时第一波压测立刻暴露： +``` +vector_write failed for: , error: FILE_WRITE_FAIL +Deleted corrupted file: .../X.bucket +``` +SSD 使用率保持 0，cache hit 为 0。 + +**修复**: +1. 分块循环 `for (int idx = 0; idx < iovcnt; idx += UIO_MAXIOV)`，每块 + `int chunk_cnt = std::min(iovcnt - idx, UIO_MAXIOV)`；累计 offset 与 totals +2. 总字节不匹配 → `FILE_WRITE_FAIL` +3. 失败时打 `errno(strerror) + fd + iovcnt + total_bytes + offset + + chunk_start/chunk_cnt`，便于后续 FILE_WRITE_FAIL 排查直接看到 syscall 错误 + +未在上游 main 中找到对应修复，故本 commit 是必需的。 + +## 配置参数（生产推荐） + +### store-server 端（`launch_store_server.sh`） + +| Env | 值 | 说明 | +|-----|----|------| +| `MOONCAKE_OFFLOAD_HEARTBEAT_INTERVAL_SECONDS` | 5 | 压测下更快处理积压 | +| `MOONCAKE_OFFLOAD_LOCAL_BUFFER_SIZE_BYTES` | 21474836480 (20GB) | batch-get SSD 默认 1.2GB 在并发下 BUFFER_OVERFLOW | +| `MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT` | 1000 | 每桶最多 1000 key | +| `MOONCAKE_OFFLOAD_BUCKET_SIZE_LIMIT_BYTES` | 2GB | | +| `MOONCAKE_OFFLOAD_BUCKET_EVICTION_POLICY` | lru | 需配合 LRU bug 修复 | +| `MOONCAKE_OFFLOAD_PROMOTION_WORKER_THREADS` | 4 (新默认) | promotion worker 池 | +| `MOONCAKE_OFFLOAD_PROMOTION_DRAIN_BATCH_SIZE` | 64 (新默认) | | +| `MOONCAKE_OFFLOAD_WRITE_THREADS` | 4 (新默认) | parallel WriteBucket | + +### master 端（`launch_master.sh`） + +| Flag | 值 | 说明 | +|------|----|------| +| `--eviction_high_watermark_ratio` | 0.7 | DRAM 触发 BatchEvict 的水位 | +| `--eviction_ratio` | 0.2 | 目标：DRAM 使用率降到 20% | +| `--put_start_release_timeout_sec` | 120 | 测试时缩短，加速过期诊断 | +| `--promotion_max_per_heartbeat` | 64 (新默认) | 每 tick 准入上限 | +| `--promotion_admission_threshold` | 2 | CountMinSketch 二次访问准入 | +| `--offloading_queue_limit` (upstream PR #2599) | 默认 50000 | 如调高需同步 offload_cap_ratio | +| `--offload_cap_ratio` (upstream PR #2599) | 默认 0.5 | offload_cap = queue_limit × ratio | + +## 验证 + +### test_lru_bug_repro.py +小容量 store (100MB DRAM + 30MB SSD) 上写 30 batch、READ 全部、再写 +b15/b16 触发 eviction。 +- buggy wheel: `b15 evicted` + `[LRU-BUG] lru_ts=0` 日志 +- fixed wheel: `b15 survived — older data was evicted first` + +### test_store_restart_cache.py +写 300000 key → 等 offload → baseline GET → restart store + master → ScanMeta → +再次 GET → PutStart 同批 key。 +- Step 1 WRITE: ok=299592 fail=408（DRAM 满，非 bug） +- Step 5 GET after restart: **hit=299592 miss=0**（验证 Bug H + I） +- Step 6 PUT after restart: **ok=300000 fail=0**（验证 Bug I） + +### AI_ERA_PRESSURE 压测 +4 store-server × 128GB DRAM / 128GB SSD = 512GB / 512GB 总量； +64 并发；glm-5.2 prefill + decoder；5 分钟采样。 +- hit_rate (window): **88.76%**（修复前 27%） +- FILE_WRITE_FAIL: 0 +- save queue full / BUFFER_OVERFLOW: 0 +- Promotion expired / OFFLOAD-EXPIRED: 0 +- OFFLOAD-COMPLETE gap: 0 +- req_err: 0 + +## 测试脚本 + +| 脚本 | 用途 | +|------|------| +| `tests/test_lru_bug_repro.py` | LRU ts=0 bug 复现 / 验证 | +| `tests/run_restart_cache_test.sh` | store restart cache 恢复完整流程 | +| `tests/test_store_restart_cache.py` | 上述脚本调用的 Python 测试本体 | diff --git a/mooncake-integration/store/store_py.cpp b/mooncake-integration/store/store_py.cpp index ec97258201..4e2439382d 100644 --- a/mooncake-integration/store/store_py.cpp +++ b/mooncake-integration/store/store_py.cpp @@ -2218,6 +2218,24 @@ PYBIND11_MODULE(store, m) { return real_client ? real_client->get_offload_rpc_read_count() : 0; }) + .def("get_metrics_summary", + [](MooncakeStorePyWrapper &self) -> py::object { + auto real_client = + std::dynamic_pointer_cast(self.store_); + if (!real_client || !real_client->client_) return py::none(); + auto result = real_client->client_->GetSummaryMetrics(); + if (!result) return py::none(); + return py::str(*result); + }) + .def("get_metrics", + [](MooncakeStorePyWrapper &self) -> py::object { + auto real_client = + std::dynamic_pointer_cast(self.store_); + if (!real_client || !real_client->client_) return py::none(); + auto result = real_client->client_->SerializeMetrics(); + if (!result) return py::none(); + return py::str(*result); + }) .def( "get_buffer", [](MooncakeStorePyWrapper &self, const std::string &key) { diff --git a/mooncake-store/include/file_interface.h b/mooncake-store/include/file_interface.h index c2ab05ce6e..6450ba8c7d 100644 --- a/mooncake-store/include/file_interface.h +++ b/mooncake-store/include/file_interface.h @@ -143,6 +143,12 @@ class StorageFile { return FileLockRAII(fd_, FileLockRAII::LockType::READ); } + virtual tl::expected datasync() { + return {}; + } + + int fd() const { return fd_; } + /** * @brief Gets the current error code * @return Current error code @@ -170,7 +176,8 @@ class PosixFile : public StorageFile { tl::expected vector_write(const iovec *iov, int iovcnt, off_t offset) override; tl::expected vector_read(const iovec *iov, int iovcnt, - off_t offset) override; + off_t offset) override; + tl::expected datasync() override; }; #ifdef USE_URING @@ -219,7 +226,7 @@ class UringFile : public StorageFile { // Flush data to stable storage via IORING_FSYNC_DATASYNC. // Must be called after write_aligned and before writing dependent metadata. - tl::expected datasync(); + tl::expected datasync() override; // Buffer registration — delegates to the shared ring (process-wide). // Static variant: no file instance needed. Must be called once from a diff --git a/mooncake-store/include/file_storage.h b/mooncake-store/include/file_storage.h index 6aa2ac3d56..d335675697 100644 --- a/mooncake-store/include/file_storage.h +++ b/mooncake-store/include/file_storage.h @@ -1,5 +1,10 @@ #pragma once +#include +#include +#include +#include + #include "client_service.h" #include "client_buffer.hpp" #include "storage_backend.h" @@ -101,6 +106,36 @@ class FileStorage { */ tl::expected ProcessPromotionTasks(); + // Single-key promotion body factored out of ProcessPromotionTasks so it + // can be called by both the inline fallback path (when workers are not + // running) and the background worker threads. Returns a structured result + // so callers can compose failure reporting. + struct PromotionExecutionResult { + ErrorCode terminal_error{ErrorCode::OK}; + bool alloc_attempted{false}; + bool write_attempted{false}; + bool notify_success_attempted{false}; + bool notify_failure_attempted{false}; + bool completed{false}; + }; + PromotionExecutionResult ProcessPromotionTask( + const PromotionTaskItem& task, + const std::vector& preferred_segments); + + // Hand a task to a background worker. Returns false (and the caller will + // fall back to ReleasePromotionTask) when workers are not running or the + // local queue is at capacity. + bool EnqueuePromotionTask(const PromotionTaskItem& task); + + // Best-effort NotifyPromotionFailure; swallows errors because the master + // reaper is the long-stop. + void ReleasePromotionTask(const std::string& key, + const std::string& tenant_id); + + // Background worker main loop. Drains promotion_task_queue_ in batches of + // promotion_drain_batch_size. + void PromotionWorkerThreadFunc(); + tl::expected IsEnableOffloading(); tl::expected BatchLoad( @@ -145,6 +180,17 @@ class FileStorage { std::thread client_buffer_gc_thread_; std::future rescan_future_; std::atomic metadata_resync_pending_{false}; + + // Background promotion workers. ProcessPromotionTasks pulls candidates + // from the master and hands them to a pool of worker threads via + // promotion_task_queue_, so the heartbeat path is not gated by per-key + // SSD read + RDMA write latency. Workers drain the queue in batches of + // config_.promotion_drain_batch_size. + std::atomic promotion_workers_running_{false}; + std::vector promotion_worker_threads_; + std::mutex promotion_queue_mutex_; + std::condition_variable promotion_queue_cv_; + std::deque promotion_task_queue_; }; } // namespace mooncake diff --git a/mooncake-store/include/master_config.h b/mooncake-store/include/master_config.h index 0275c57e0f..8e719e6d8e 100644 --- a/mooncake-store/include/master_config.h +++ b/mooncake-store/include/master_config.h @@ -131,7 +131,7 @@ struct MasterConfig { // on the client; serializing them avoids blocking past the client- // liveness window. Default 1 is conservative; small-object or RDMA- // rich clusters may safely raise it. - uint32_t promotion_max_per_heartbeat = 1; + uint32_t promotion_max_per_heartbeat = 64; }; class MasterServiceSupervisorConfig { @@ -210,7 +210,7 @@ class MasterServiceSupervisorConfig { bool promotion_on_hit = false; uint32_t promotion_admission_threshold = 2; uint32_t promotion_queue_limit = 50000; - uint32_t promotion_max_per_heartbeat = 1; + uint32_t promotion_max_per_heartbeat = 64; // Pod identity for K8s label-based routing std::string pod_name; @@ -402,7 +402,7 @@ class WrappedMasterServiceConfig { bool promotion_on_hit = false; uint32_t promotion_admission_threshold = 2; uint32_t promotion_queue_limit = 50000; - uint32_t promotion_max_per_heartbeat = 1; + uint32_t promotion_max_per_heartbeat = 64; std::string ha_backend_type = "etcd"; std::string ha_backend_connstring; std::string cluster_id = DEFAULT_CLUSTER_ID; @@ -982,7 +982,7 @@ class MasterServiceConfig { bool promotion_on_hit = false; uint32_t promotion_admission_threshold = 2; uint32_t promotion_queue_limit = 50000; - uint32_t promotion_max_per_heartbeat = 1; + uint32_t promotion_max_per_heartbeat = 64; std::string ha_backend_type = "etcd"; std::string ha_backend_connstring; std::string cluster_id = DEFAULT_CLUSTER_ID; diff --git a/mooncake-store/include/storage_backend.h b/mooncake-store/include/storage_backend.h index 5fd071b59b..78f932d757 100644 --- a/mooncake-store/include/storage_backend.h +++ b/mooncake-store/include/storage_backend.h @@ -226,6 +226,36 @@ struct FileStorageConfig { uint32_t client_buffer_gc_interval_seconds = 1; uint64_t client_buffer_gc_ttl_ms = 5000; + // Background worker settings for L2->L1 promotion-on-hit execution. + // Set promotion_worker_threads to 0 to fall back to the synchronous + // (inline-in-heartbeat) execution path used before PR #2529. + // 4 matches production load where each heartbeat pulls 64 tasks and each + // task does a ~1-2 ms SSD read + RDMA write; 4 workers drain 64 tasks in + // ~16 ms, comfortably inside one heartbeat tick. + uint32_t promotion_worker_threads = 4; + + // Parallel WriteBucket thread pool size for BatchOffload. Each heartbeat + // tick may dispatch many independent buckets; writing them in parallel + // raises single-NVMe write throughput from ~1.6 GB/s (serial) to ~4.7 GB/s + // (4 threads) so the SSD pipeline does not gate cache fill rate. 0 keeps + // the legacy serial loop. + uint32_t offload_write_threads = 4; + // Parallel bucket-read thread pool size for BatchLoad. Each BatchGet may + // fan out to N independent bucket files; reading them in parallel raises + // single-NVME read throughput from ~1.4 GB/s (serial, QD=1) to ~4-5 GB/s + // (QD=N). 0 or 1 keeps the legacy serial loop. Mirrors offload_write_threads. + uint32_t offload_read_threads = 4; + // Parallel endpoint thread pool size for LOCAL_DISK offload reads. Each + // batch_get_into[_multi_buffers] may fan out to N independent store pods + // (separate transport endpoints); fetching from them in parallel raises + // aggregate read bandwidth from single-pod (e.g. ~6 GB/s) to N-pod + // (e.g. ~48 GB/s across 8 pods). 0 or 1 keeps the legacy serial loop. + uint32_t offload_endpoint_threads = 8; + // Soft local backlog cap. 0 = unbounded. + uint32_t promotion_queue_capacity = 1024; + // Per-worker drain batch size. + uint32_t promotion_drain_batch_size = 64; + // Use io_uring for file I/O instead of POSIX pread/pwrite bool use_uring = false; diff --git a/mooncake-store/src/file_storage.cpp b/mooncake-store/src/file_storage.cpp index 00e7db4bcd..b726ff5cd7 100644 --- a/mooncake-store/src/file_storage.cpp +++ b/mooncake-store/src/file_storage.cpp @@ -1,6 +1,11 @@ #include "file_storage.h" +#include #include +#include +#include +#include +#include #include #include @@ -83,6 +88,25 @@ FileStorageConfig FileStorageConfig::FromEnvironment() { GetEnvOr("MOONCAKE_OFFLOAD_CLIENT_BUFFER_GC_TTL_MS", config.client_buffer_gc_ttl_ms); + config.promotion_worker_threads = + GetEnvOr("MOONCAKE_OFFLOAD_PROMOTION_WORKER_THREADS", + config.promotion_worker_threads); + config.promotion_queue_capacity = + GetEnvOr("MOONCAKE_OFFLOAD_PROMOTION_QUEUE_CAPACITY", + config.promotion_queue_capacity); + config.promotion_drain_batch_size = + GetEnvOr("MOONCAKE_OFFLOAD_PROMOTION_DRAIN_BATCH_SIZE", + config.promotion_drain_batch_size); + config.offload_write_threads = + GetEnvOr("MOONCAKE_OFFLOAD_WRITE_THREADS", + config.offload_write_threads); + config.offload_read_threads = + GetEnvOr("MOONCAKE_OFFLOAD_READ_THREADS", + config.offload_read_threads); + config.offload_endpoint_threads = + GetEnvOr("MOONCAKE_OFFLOAD_ENDPOINT_THREADS", + config.offload_endpoint_threads); + auto use_uring_str = GetEnvStringOr("MOONCAKE_OFFLOAD_USE_URING", GetEnvStringOr("MOONCAKE_USE_URING", "false")); @@ -226,6 +250,13 @@ FileStorage::~FileStorage() { if (client_buffer_gc_thread_.joinable()) { client_buffer_gc_thread_.join(); } + promotion_workers_running_.store(false); + promotion_queue_cv_.notify_all(); + for (auto& worker : promotion_worker_threads_) { + if (worker.joinable()) { + worker.join(); + } + } } tl::expected FileStorage::Init() { @@ -314,6 +345,15 @@ tl::expected FileStorage::Init() { client_buffer_gc_running_.store(true); client_buffer_gc_thread_ = std::thread(&FileStorage::ClientBufferGCThreadFunc, this); + if (config_.promotion_worker_threads > 0) { + promotion_workers_running_.store(true); + const uint32_t worker_count = config_.promotion_worker_threads; + promotion_worker_threads_.reserve(worker_count); + for (uint32_t i = 0; i < worker_count; ++i) { + promotion_worker_threads_.emplace_back( + &FileStorage::PromotionWorkerThreadFunc, this); + } + } return {}; } @@ -425,9 +465,17 @@ tl::expected FileStorage::OffloadObjects( // orphaned offloading_tasks and release source replica refcounts. std::vector failed_tasks; std::unordered_set all_bucket_keys; + std::mutex failed_tasks_mutex; - for (const auto& keys : buckets_keys) { - for (const auto& k : keys) all_bucket_keys.insert(k); + // The loop body is captured as a lambda so it can be dispatched either + // serially (offload_write_threads <= 1) or in parallel via a thread pool. + auto process_bucket = [&](const std::vector& keys) + -> std::optional { + std::vector local_failed; + { + std::lock_guard lk(failed_tasks_mutex); + for (const auto& k : keys) all_bucket_keys.insert(k); + } std::unordered_map> batch_object; std::unordered_map> storage_keys_by_tenant; @@ -457,12 +505,14 @@ tl::expected FileStorage::OffloadObjects( if (it != user_batch_object.end()) { batch_object.emplace(storage_key, std::move(it->second)); } else { - failed_tasks.push_back(task); + local_failed.push_back(task); } } } if (batch_object.empty()) { - continue; + std::lock_guard lk(failed_tasks_mutex); + for (auto& t : local_failed) failed_tasks.push_back(std::move(t)); + return std::nullopt; } auto eviction_handler = [this](const std::vector& @@ -513,7 +563,7 @@ tl::expected FileStorage::OffloadObjects( LOG(ERROR) << "D2H staging failed for key: " << obj_key; pinned_buffer_pool_->Release(std::move(buf)); obj_success = false; - failed_tasks.push_back(task_by_storage_key.at(obj_key)); + local_failed.push_back(task_by_storage_key.at(obj_key)); break; } host_slices.emplace_back(Slice{buf.data, slice.size}); @@ -530,9 +580,9 @@ tl::expected FileStorage::OffloadObjects( auto offload_start = std::chrono::steady_clock::now(); auto bucket_complete_handler = [this, offload_start, complete_handler]( - const std::vector& keys, + const std::vector& bucket_keys, std::vector& metadatas) -> ErrorCode { - auto res = complete_handler(keys, metadatas); + auto res = complete_handler(bucket_keys, metadatas); if (res == ErrorCode::OK && ssd_metric_) { auto elapsed_us = std::chrono::duration_cast( @@ -542,11 +592,11 @@ tl::expected FileStorage::OffloadObjects( for (const auto& metadata : metadatas) { total_bytes += metadata.data_size; } - ssd_metric_->ssd_write_ops.inc(keys.size()); + ssd_metric_->ssd_write_ops.inc(bucket_keys.size()); ssd_metric_->ssd_write_bytes.inc(total_bytes); ssd_metric_->ssd_write_latency_us.observe(elapsed_us); ssd_metric_->ssd_write_latency_summary.observe(elapsed_us); - ssd_metric_->ssd_total_ops.inc(keys.size()); + ssd_metric_->ssd_total_ops.inc(bucket_keys.size()); ssd_metric_->ssd_total_bytes.inc(total_bytes); ssd_metric_->ssd_total_latency_us.observe(elapsed_us); ssd_metric_->ssd_total_latency_summary.observe(elapsed_us); @@ -566,16 +616,71 @@ tl::expected FileStorage::OffloadObjects( if (offload_res.error() == ErrorCode::KEYS_ULTRA_LIMIT) { MutexLocker locker(&offloading_mutex_); enable_offloading_ = false; - return tl::make_unexpected(offload_res.error()); + return offload_res.error(); } if (offload_res.error() == ErrorCode::INVALID_READ) { for (const auto& [key, _] : host_batch_object) { - failed_tasks.push_back(task_by_storage_key.at(key)); + local_failed.push_back(task_by_storage_key.at(key)); } } else { - return tl::make_unexpected(offload_res.error()); + return offload_res.error(); + } + } + std::lock_guard lk(failed_tasks_mutex); + for (auto& t : local_failed) failed_tasks.push_back(std::move(t)); + return std::nullopt; + }; + + const uint32_t write_threads = + std::max(1, config_.offload_write_threads); + if (write_threads == 1 || buckets_keys.size() <= 1) { + // Serial path — preserves prior behavior and avoids thread-pool + // overhead for the single-bucket (or disabled) case. + for (const auto& keys : buckets_keys) { + auto err = process_bucket(keys); + if (err.has_value()) { + return tl::make_unexpected(err.value()); } } + } else { + // Parallel path — dispatch independent buckets concurrently. + // Each bucket is a separate file on SSD with no shared write state, + // so writes can safely overlap. A bounded set of worker threads pulls + // buckets from a shared index, capping memory pressure regardless of + // how many buckets GroupOffloadingKeysByBucket produced. + const size_t n = buckets_keys.size(); + std::atomic next_index{0}; + std::atomic first_error_set{false}; + std::optional first_error; // protected by result_mutex + std::mutex result_mutex; + + auto worker = [&]() { + for (;;) { + if (first_error_set.load()) return; + size_t i = next_index.fetch_add(1, std::memory_order_acq_rel); + if (i >= n) return; + auto err = process_bucket(buckets_keys[i]); + if (err.has_value()) { + std::lock_guard lk(result_mutex); + if (!first_error_set.exchange(true)) { + first_error = err.value(); + } + return; + } + } + }; + + std::vector workers; + workers.reserve(write_threads); + for (uint32_t i = 0; i < write_threads; ++i) { + workers.emplace_back(worker); + } + for (auto& w : workers) { + if (w.joinable()) w.join(); + } + if (first_error.has_value()) { + return tl::make_unexpected(first_error.value()); + } } // Keys skipped by GroupOffloadingKeysByBucket don't appear in any bucket, @@ -701,23 +806,24 @@ tl::expected FileStorage::Heartbeat() { } } - if (offloading_objects.empty()) { - return {}; - } - // === STEP 2: Persist offloaded objects (trigger actual data migration) === - auto offload_result = OffloadObjects(offloading_objects); - if (!offload_result) { - LOG(ERROR) << "Failed to persist objects with error: " - << offload_result.error(); - return offload_result; - } + if (!offloading_objects.empty()) { + // === STEP 2: Persist offloaded objects (trigger actual data + // migration) === + auto offload_result = OffloadObjects(offloading_objects); + if (!offload_result) { + LOG(ERROR) << "Failed to persist objects with error: " + << offload_result.error(); + return offload_result; + } - VLOG(1) << "Completed heartbeat with offloaded objects count: " - << offloading_objects.size(); + VLOG(1) << "Completed heartbeat with offloaded objects count: " + << offloading_objects.size(); + } - // Drive any pending L2->L1 promotion work for this client. Failures - // inside ProcessPromotionTasks are logged per-key and do not propagate; - // promotion is best-effort and must never break offload. + // Pull any pending L2->L1 promotion work for this client. Execution is + // delegated to background workers so the heartbeat path stays responsive. + // Always called, even when offloading_objects is empty, so promotion + // backlog can drain independently of the offload path. (void)ProcessPromotionTasks(); // TODO(eviction): Implement an LRU eviction mechanism to manage local @@ -753,125 +859,274 @@ tl::expected FileStorage::ProcessPromotionTasks() { // No segment preference from the client: let master pick from any // DRAM segment. - const std::vector preferred_segments; + const std::vector sync_preferred_segments; + + // Preserve pre-worker semantics for tests and pre-Init callers: execute + // inline. ProcessPromotionTask is the per-key body that was previously + // inlined here; see it for the failure-path documentation. + if (!promotion_workers_running_.load()) { + for (const auto& task : promotion_objects) { + (void)ProcessPromotionTask(task, sync_preferred_segments); + } + return {}; + } - // The master caps per-heartbeat work via PromotionObjectHeartbeat, - // returning at most one task per call so the heartbeat thread stays - // within the client-liveness window even for large objects. Leftover - // work stays queued in the master's promotion_objects map and is - // returned on subsequent heartbeats; we process whatever we received - // here without a second client-side cap. - for (const auto& task : promotion_objects) { - const auto& key = task.key; - const auto& tenant_id = task.tenant_id; - const int64_t size = task.size; - const auto storage_key = MakeTenantScopedStorageKey(tenant_id, key); - if (size <= 0) { - LOG(WARNING) << "Skipping promotion for key=" << key - << " with non-positive size=" << size; - continue; + // Hand off execution to background workers so the heartbeat thread stays + // responsive. Keep pulling from the master until the local queue reaches + // its soft budget, so backlog drain is no longer gated by one + // PromotionObjectHeartbeat call per client tick. + size_t queue_depth = 0; + { + std::lock_guard lock(promotion_queue_mutex_); + queue_depth = promotion_task_queue_.size(); + } + + const size_t worker_count = + std::max(1, config_.promotion_worker_threads); + const size_t drain_batch_size = + std::max(1, config_.promotion_drain_batch_size); + size_t remaining_pull_budget = worker_count * drain_batch_size; + + if (config_.promotion_queue_capacity > 0) { + if (queue_depth >= config_.promotion_queue_capacity) { + VLOG(1) << "ProcessPromotionTasks skipped master pull because " + << "local promotion queue is full: queue_depth=" + << queue_depth + << ", queue_capacity=" << config_.promotion_queue_capacity; + return {}; } + remaining_pull_budget = + std::min(remaining_pull_budget, + config_.promotion_queue_capacity - queue_depth); + } - auto alloc_result = client_->PromotionAllocStart( - key, tenant_id, static_cast(size), preferred_segments); - if (!alloc_result) { - // AllocStart failed (typically NO_AVAILABLE_HANDLE under - // DRAM pressure). No staged buffer to release, but the - // task entry already claimed a promotion_in_flight_ slot - // at admission. Notify the master to release it - // immediately; otherwise the slot stays pinned for the - // reaper TTL (~10 min default), turning transient DRAM - // pressure into a sustained outage of promotion_queue_limit_. - // Notify is idempotent and handles alloc_id == 0 correctly. - VLOG(1) << "PromotionAllocStart failed for key=" << key - << ", error=" << alloc_result.error() - << " (likely no free DRAM); releasing master slot"; - auto release = client_->NotifyPromotionFailure(key, tenant_id); - if (!release) { - VLOG(1) << "Promotion: NotifyPromotionFailure failed for key=" - << key << ", error=" << release.error() - << "; master reaper will reclaim on TTL expiry"; + size_t pulled_tasks = 0; + size_t enqueued_tasks = 0; + size_t released_tasks = 0; + size_t heartbeat_rounds = 0; + auto enqueue_batch = [this, &enqueued_tasks, + &released_tasks](const auto& tasks) { + for (const auto& task : tasks) { + if (EnqueuePromotionTask(task)) { + ++enqueued_tasks; + } else { + ReleasePromotionTask(task.key, task.tenant_id); + ++released_tasks; } - continue; } + }; - // Every failure path past this point has a master-side staged - // PROCESSING MEMORY buffer and an incremented in-flight slot. - // Eagerly notify the master on failure so the buffer is - // reclaimed and the slot is freed; otherwise transient SSD - // throttling or RDMA flakes saturate promotion_queue_limit_ - // for the full reaper TTL. NotifyPromotionFailure is - // idempotent and best-effort — the reaper is the long-stop. - auto release_master_state = [this, &key, &tenant_id]() { - auto release = client_->NotifyPromotionFailure(key, tenant_id); - if (!release) { - VLOG(1) << "Promotion: NotifyPromotionFailure failed for key=" - << key << ", error=" << release.error() - << "; master reaper will reclaim on TTL expiry"; + // The master caps per-heartbeat work via PromotionObjectHeartbeat, + // returning at most a bounded number of tasks per call so the heartbeat + // thread stays within the client-liveness window even for large objects. + // Leftover work stays queued in the master's promotion_objects map and is + // returned on subsequent heartbeats; we process whatever we received + // here without a second client-side cap, looping until the soft budget + // is filled or the master has nothing more to send. + enqueue_batch(promotion_objects); + pulled_tasks += promotion_objects.size(); + ++heartbeat_rounds; + while (pulled_tasks < remaining_pull_budget) { + std::vector next_batch; + client_->PromotionObjectHeartbeat(next_batch); + if (next_batch.empty()) break; + enqueue_batch(next_batch); + pulled_tasks += next_batch.size(); + ++heartbeat_rounds; + { + std::lock_guard lock(promotion_queue_mutex_); + if (config_.promotion_queue_capacity > 0 && + promotion_task_queue_.size() >= + config_.promotion_queue_capacity) { + break; } - }; - - // (a) Allocate an O_DIRECT-aligned staging buffer and read the bytes - // from the local SSD backend into it. AllocateBatch returns a - // shared_ptr whose BufferHandles RAII-release the - // staging space when the local goes out of scope. - std::vector single_key{storage_key}; - std::vector single_size{size}; - auto allocate_res = AllocateBatch(single_key, single_size); - if (!allocate_res) { - LOG(WARNING) << "Promotion: AllocateBatch failed for key=" << key - << ", error=" << allocate_res.error(); - release_master_state(); - continue; - } - auto staging = allocate_res.value(); - auto load_res = BatchLoad(staging->slices); - if (!load_res) { - LOG(WARNING) << "Promotion: BatchLoad failed for key=" << key - << ", error=" << load_res.error(); - release_master_state(); - continue; } + } - // (b) TE-write from the staging slice into the freshly-allocated - // MEMORY replica. Slice ptr may have been bumped by O_DIRECT offset - // correction in BatchLoad, so re-read it from the slice map. - auto slice_it = staging->slices.find(storage_key); - if (slice_it == staging->slices.end()) { - LOG(WARNING) << "Promotion: staging slice missing for key=" << key; - release_master_state(); - continue; - } - std::vector tx_slices{slice_it->second}; - ErrorCode write_err = client_->PromotionWrite( - alloc_result.value().memory_descriptor, tx_slices); - if (write_err != ErrorCode::OK) { - LOG(WARNING) << "Promotion: TransferWrite failed for key=" << key - << ", error=" << write_err; - release_master_state(); - continue; - } + VLOG(1) << "ProcessPromotionTasks pulled " << pulled_tasks + << " promotion candidate(s) from master in " << heartbeat_rounds + << " heartbeat round(s), enqueued " << enqueued_tasks + << ", released " << released_tasks; + return {}; +} + +FileStorage::PromotionExecutionResult FileStorage::ProcessPromotionTask( + const PromotionTaskItem& task, + const std::vector& preferred_segments) { + PromotionExecutionResult result; + const auto& key = task.key; + const auto& tenant_id = task.tenant_id; + const int64_t size = task.size; + const auto storage_key = MakeTenantScopedStorageKey(tenant_id, key); + + if (size <= 0) { + LOG(WARNING) << "Skipping promotion for key=" << key + << " with non-positive size=" << size; + result.terminal_error = ErrorCode::INVALID_PARAMS; + return result; + } + + result.alloc_attempted = true; + auto alloc_result = client_->PromotionAllocStart( + key, tenant_id, static_cast(size), preferred_segments); + if (!alloc_result) { + // AllocStart failed (typically NO_AVAILABLE_HANDLE under DRAM + // pressure). No staged buffer to release, but the task entry + // already claimed a promotion_in_flight_ slot at admission. + // Notify the master to release it immediately; otherwise the + // slot stays pinned for the reaper TTL (~10 min default), + // turning transient DRAM pressure into a sustained outage of + // promotion_queue_limit_. Notify is idempotent and handles + // alloc_id == 0 correctly. + VLOG(1) << "PromotionAllocStart failed for key=" << key + << ", error=" << alloc_result.error() + << " (likely no free DRAM); releasing master slot"; + ReleasePromotionTask(key, tenant_id); + result.notify_failure_attempted = true; + result.terminal_error = alloc_result.error(); + return result; + } - // (c) Commit. Master flips the PROCESSING replica to COMPLETE and it - // becomes visible to readers. - auto notify_res = client_->NotifyPromotionSuccess(key, tenant_id); - if (!notify_res) { - // The write landed but the commit failed. We can't retry the - // commit (the success path is one-shot via alloc_id), and we - // don't know whether the failure was transient or structural. - // Release the master-side state so the slot is reusable; the - // bytes we wrote become stranded under a soon-to-be-erased - // PROCESSING replica, which is harmless. - LOG(WARNING) << "Promotion: NotifyPromotionSuccess failed for key=" - << key << ", error=" << notify_res.error(); - release_master_state(); - continue; + // Every failure path past this point has a master-side staged PROCESSING + // MEMORY buffer and an incremented in-flight slot. Eagerly notify the + // master on failure so the buffer is reclaimed and the slot is freed; + // otherwise transient SSD throttling or RDMA flakes saturate + // promotion_queue_limit_ for the full reaper TTL. + // (a) Allocate an O_DIRECT-aligned staging buffer and read the bytes from + // the local SSD backend into it. + std::vector single_key{storage_key}; + std::vector single_size{size}; + auto allocate_res = AllocateBatch(single_key, single_size); + if (!allocate_res) { + LOG(WARNING) << "Promotion: AllocateBatch failed for key=" << key + << ", error=" << allocate_res.error(); + ReleasePromotionTask(key, tenant_id); + result.notify_failure_attempted = true; + result.terminal_error = allocate_res.error(); + return result; + } + auto staging = allocate_res.value(); + auto load_res = BatchLoad(staging->slices); + if (!load_res) { + LOG(WARNING) << "Promotion: BatchLoad failed for key=" << key + << ", error=" << load_res.error(); + ReleasePromotionTask(key, tenant_id); + result.notify_failure_attempted = true; + result.terminal_error = load_res.error(); + return result; + } + + // (b) TE-write from the staging slice into the freshly-allocated MEMORY + // replica. Slice ptr may have been bumped by O_DIRECT offset correction + // in BatchLoad, so re-read it from the slice map. + result.write_attempted = true; + auto slice_it = staging->slices.find(storage_key); + if (slice_it == staging->slices.end()) { + LOG(WARNING) << "Promotion: staging slice missing for key=" << key; + ReleasePromotionTask(key, tenant_id); + result.notify_failure_attempted = true; + result.terminal_error = ErrorCode::INVALID_KEY; + return result; + } + std::vector tx_slices{slice_it->second}; + ErrorCode write_err = client_->PromotionWrite( + alloc_result.value().memory_descriptor, tx_slices); + if (write_err != ErrorCode::OK) { + LOG(WARNING) << "Promotion: TransferWrite failed for key=" << key + << ", error=" << write_err; + ReleasePromotionTask(key, tenant_id); + result.notify_failure_attempted = true; + result.terminal_error = write_err; + return result; + } + + // (c) Commit. Master flips the PROCESSING replica to COMPLETE and it + // becomes visible to readers. + result.notify_success_attempted = true; + auto notify_res = client_->NotifyPromotionSuccess(key, tenant_id); + if (!notify_res) { + // The write landed but the commit failed. We can't retry the commit + // (the success path is one-shot via alloc_id), and we don't know + // whether the failure was transient or structural. Release the + // master-side state so the slot is reusable; the bytes we wrote + // become stranded under a soon-to-be-erased PROCESSING replica, + // which is harmless. + LOG(WARNING) << "Promotion: NotifyPromotionSuccess failed for key=" + << key << ", error=" << notify_res.error(); + ReleasePromotionTask(key, tenant_id); + result.notify_failure_attempted = true; + result.terminal_error = notify_res.error(); + return result; + } + + result.completed = true; + VLOG(1) << "Promotion completed for key=" << key << ", size=" << size; + return result; +} + +bool FileStorage::EnqueuePromotionTask(const PromotionTaskItem& task) { + std::lock_guard lock(promotion_queue_mutex_); + if (!promotion_workers_running_.load()) { + LOG(WARNING) << "Promotion workers are not running, rejecting key=" + << task.key; + return false; + } + if (config_.promotion_queue_capacity > 0 && + promotion_task_queue_.size() >= config_.promotion_queue_capacity) { + LOG(WARNING) << "Promotion queue full (" << promotion_task_queue_.size() + << "), rejecting key=" << task.key; + return false; + } + promotion_task_queue_.push_back(task); + promotion_queue_cv_.notify_one(); + return true; +} + +void FileStorage::ReleasePromotionTask(const std::string& key, + const std::string& tenant_id) { + auto release = client_->NotifyPromotionFailure(key, tenant_id); + if (!release) { + VLOG(1) << "Promotion: NotifyPromotionFailure failed for key=" << key + << ", error=" << release.error() + << "; master reaper will reclaim on TTL expiry"; + } +} + +void FileStorage::PromotionWorkerThreadFunc() { + VLOG(1) << "action=promotion_worker_started"; + const std::vector preferred_segments; + const size_t drain_batch_size = + std::max(1, config_.promotion_drain_batch_size); + + while (true) { + std::vector tasks; + tasks.reserve(drain_batch_size); + + { + std::unique_lock lock(promotion_queue_mutex_); + promotion_queue_cv_.wait(lock, [this]() { + return !promotion_workers_running_.load() || + !promotion_task_queue_.empty(); + }); + + if (!promotion_workers_running_.load() && + promotion_task_queue_.empty()) { + break; + } + + while (!promotion_task_queue_.empty() && + tasks.size() < drain_batch_size) { + tasks.push_back(std::move(promotion_task_queue_.front())); + promotion_task_queue_.pop_front(); + } } - VLOG(1) << "Promotion completed for key=" << key << ", size=" << size; + for (const auto& task : tasks) { + (void)ProcessPromotionTask(task, preferred_segments); + } } - return {}; + VLOG(1) << "action=promotion_worker_stopped"; } tl::expected FileStorage::BatchLoad( diff --git a/mooncake-store/src/master_service.cpp b/mooncake-store/src/master_service.cpp index c051fe1682..a83bd09352 100644 --- a/mooncake-store/src/master_service.cpp +++ b/mooncake-store/src/master_service.cpp @@ -1736,6 +1736,15 @@ MasterService::EraseMetadata( // becomes an orphan that only expires after 600s. auto offload_it = tenant_state.offloading_tasks.find(key); if (offload_it != tenant_state.offloading_tasks.end()) { + // BatchEvict is deleting metadata while the store worker still has an + // in-flight offload for this key. Without this cleanup the task would + // be an orphan that only expires after put_start_release_timeout_sec. + LOG(WARNING) + << "[CLEANUP-ORPHAN] erasing metadata for key=" << key + << " tenant=" << tenant_id + << " while offloading_task is still registered (source_id=" + << offload_it->second.source_id + << "); dec_refcnt applied, offloading_task erased"; auto source = metadata.GetReplicaByID(offload_it->second.source_id); if (source != nullptr) { source->dec_refcnt(); @@ -2998,7 +3007,19 @@ auto MasterService::PutStart(const UUID& client_id, const std::string& key, it = tenant_state.metadata.end(); } else { auto& metadata = it->second; - if (metadata.HasReplica(&Replica::fn_is_completed) || + // Block PutStart only when a completed MEMORY replica is + // present. A LOCAL_DISK-only metadata (e.g. left over after + // a store-server restart + ScanMeta recovery) must not + // block writes — the LOCAL_DISK replica is there to serve + // cache reads, not to gate new PutStart calls. Blocking on + // any completed replica type (the previous behavior) caused + // OBJECT_ALREADY_EXISTS for every key whose LOCAL_DISK + // replica was re-registered after restart. + bool has_completed_memory = metadata.HasReplica( + [](const Replica& r) { + return r.is_memory_replica() && r.is_completed(); + }); + if (has_completed_memory || metadata.put_start_time + put_start_discard_timeout_sec_ >= now) { @@ -3139,24 +3160,43 @@ auto MasterService::PutEnd(const UUID& client_id, const std::string& key, if (enable_offload_ && !offload_on_evict_) { auto& tenant_state = accessor.GetTenantState(); - bool task_created = false; - metadata.VisitReplicas( - [](const Replica& replica) { - return replica.is_completed() && replica.is_memory_replica(); - }, - [this, &object_id, &tenant_state, &task_created](Replica& replica) { - auto result = PushOffloadingQueue(object_id, replica); - if (result) { - if (!task_created) { + // Skip if an offloading task is already in flight for this key — + // emplace would silently fail, leaving the inc_refcnt below as a + // permanent leak. With replica_num>1, the visitor would keep iterating + // after the first emplace, double-pinning the second replica. + if (tenant_state.offloading_tasks.count(object_id.user_key) == 0) { + bool task_created = false; + metadata.VisitReplicas( + [](const Replica& replica) { + return replica.is_completed() && replica.is_memory_replica(); + }, + [this, &object_id, &tenant_state, + &task_created](Replica& replica) { + if (task_created) return; // only pin one replica per key + auto result = PushOffloadingQueue(object_id, replica); + if (result) { replica.inc_refcnt(); - tenant_state.offloading_tasks.emplace( - object_id.user_key, - OffloadingTask{replica.id(), - std::chrono::system_clock::now()}); - task_created = true; + auto [it, inserted] = + tenant_state.offloading_tasks.emplace( + object_id.user_key, + OffloadingTask{ + replica.id(), + std::chrono::system_clock::now()}); + if (!inserted) { + // Race: another path already queued this key. + // Roll back the refcnt so it doesn't leak. + replica.dec_refcnt(); + LOG(WARNING) + << "[PUTEND-OFFLOAD-EMPLACE-FAIL] key=" + << object_id.user_key + << " already has offloading_task; rolled " + "back refcnt"; + } else { + task_created = true; + } } - } - }); + }); + } } // If the object is completed, remove it from the processing set. @@ -3205,6 +3245,29 @@ auto MasterService::AddReplica(const UUID& client_id, const std::string& key, return tl::make_unexpected(ErrorCode::INVALID_PARAMS); } + // Erase stale LOCAL_DISK replicas from a previous (restarted) store-server + // incarnation. Without this, a stale transport_endpoint lingers in metadata + // after the store-server restarts with a new client_id; the visitor below + // requires matching client_id, so the new replica would be silently dropped + // and reads would follow the dead endpoint (INVALID_KEY / RDMA errors). + bool had_completed_disk = metadata.HasReplica([](const Replica& r) { + return r.is_local_disk_replica() && r.is_completed(); + }); + size_t stale_erased = EraseReplicasWithCacheTotalAccounting( + metadata, + [&client_id](const Replica& rep) { + return rep.is_local_disk_replica() && + rep.get_descriptor().get_local_disk_descriptor().client_id != + client_id; + }); + if (stale_erased > 0) { + auto& shard = accessor.GetShard(); + shard.OnDiskReplicaRemoved(had_completed_disk, metadata); + LOG(INFO) << "[ADDREPLICA-STALE] erased " << stale_erased + << " stale LOCAL_DISK replica(s) for key=" << key + << " tenant=" << normalized_tenant; + } + if (!metadata.HasReplica(&Replica::fn_is_local_disk_replica)) { std::vector replicas; replicas.emplace_back(std::move(replica)); @@ -4927,6 +4990,23 @@ auto MasterService::NotifyOffloadSuccess( << ". Expected ReplicaType::LOCAL_DISK."; return tl::make_unexpected(ErrorCode::INVALID_PARAMS); } + if (task_it == tenant_state.offloading_tasks.end()) { + // Worker reports offload success for a key the master does + // not remember queuing. Either the offloading_task already + // expired (put_start_release_timeout_sec) or the metadata + // was erased by BatchEvict's cleanup phase. Either way the + // caller has done the work; we accept the LOCAL_DISK + // replica via AddReplica below but log the mismatch so + // orphan-task sources are visible. + LOG(INFO) + << "[OFFLOAD-NOTASK] NotifyOffloadSuccess received a " + "complete for key=" + << request_object_id.user_key + << " tenant=" << request_object_id.tenant_id + << " but no offloading_task is registered (already " + "expired or cleaned up); will add LOCAL_DISK " + "replica via AddReplica"; + } // Existing orphan objects can only bypass tenant registration // for a master-admitted offload completion. Without this task @@ -5013,6 +5093,14 @@ tl::expected MasterService::PushOffloadingQueue( const ObjectIdentity& object_id, Replica& replica) { const auto& segment_names = replica.get_segment_names(); if (segment_names.empty()) { + // Replica has no transport segments — nothing to offload. This is + // unusual for a MEMORY replica that triggered BatchEvict; log it so + // silent no-ops show up in diagnostics instead of mysteriously + // skipping the offload path. + LOG(WARNING) << "[PUSH-EMPTY] PushOffloadingQueue called for key=" + << object_id.user_key + << " tenant=" << object_id.tenant_id + << " but replica has no segment_names; offload skipped"; return {}; } for (const auto& segment_name_it : segment_names) { @@ -6891,22 +6979,36 @@ MasterService::EvictTenantMemoryForQuota(const std::string& tenant_id, } bool queued = false; - metadata.VisitReplicas( - is_evictable_memory_replica, - [this, &key, &normalized_tenant, &tenant_state, &queued, - &now](Replica& replica) { - if (queued) { - return; - } - auto result = PushOffloadingQueue( - MakeObjectIdentity(key, normalized_tenant), replica); - if (result) { - replica.inc_refcnt(); - tenant_state.offloading_tasks.emplace( - key, OffloadingTask{replica.id(), now}); - queued = true; - } - }); + // Skip if an offloading task is already in flight for this key to + // avoid the cross-cycle re-pin refcnt leak (see Bug J part 1). + if (tenant_state.offloading_tasks.count(key) == 0) { + metadata.VisitReplicas( + is_evictable_memory_replica, + [this, &key, &normalized_tenant, &tenant_state, &queued, + &now](Replica& replica) { + if (queued) return; + auto result = PushOffloadingQueue( + MakeObjectIdentity(key, normalized_tenant), + replica); + if (result) { + replica.inc_refcnt(); + auto [it, inserted] = + tenant_state.offloading_tasks.emplace( + key, OffloadingTask{replica.id(), now}); + if (!inserted) { + replica.dec_refcnt(); + LOG(WARNING) + << "[TENANT-EVICT-OFFLOAD-EMPLACE-FAIL] " + "key=" + << key + << " already has offloading_task; rolled " + "back refcnt"; + } else { + queued = true; + } + } + }); + } if (queued) { ++offload_queued_this_call; @@ -6925,7 +7027,8 @@ MasterService::EvictTenantMemoryForQuota(const std::string& tenant_id, [&, this](const std::string& key, ObjectMetadata& metadata, TenantState& tenant_state, std::vector>& deferred_replicas, - bool allow_soft_pinned) -> TenantQuotaEvictionResult { + bool allow_soft_pinned, + MetadataShardAccessorRW* shard_ptr) -> TenantQuotaEvictionResult { if (!metadata.IsGrouped()) { uint64_t freed = try_evict_or_offload(key, metadata, tenant_state, deferred_replicas); @@ -6972,7 +7075,8 @@ MasterService::EvictTenantMemoryForQuota(const std::string& tenant_id, ++result.evicted_objects; } if (member_key != key && !member_metadata.IsValid()) { - EraseMetadata(tenant_state, member_it, normalized_tenant); + EraseMetadata(tenant_state, member_it, normalized_tenant, + QuotaEraseMode::kFull, shard_ptr); } } return result; @@ -7006,11 +7110,12 @@ MasterService::EvictTenantMemoryForQuota(const std::string& tenant_id, auto evict_result = try_evict_group_or_object( it->first, metadata, tenant_state, deferred_replicas, - allow_soft_pinned); + allow_soft_pinned, &shard); total.freed_bytes += evict_result.freed_bytes; total.evicted_objects += evict_result.evicted_objects; if (!metadata.IsValid()) { - it = EraseMetadata(tenant_state, it, normalized_tenant); + it = EraseMetadata(tenant_state, it, normalized_tenant, + QuotaEraseMode::kFull, &shard); } else { ++it; } @@ -7134,24 +7239,40 @@ void MasterService::BatchEvict(double evict_ratio_target, } // Queue one MEMORY replica for offload; others will be evicted below. + // Skip if an offloading task is already in flight for this key — + // otherwise the emplace below silently fails and inc_refcnt leaks. bool queued = false; - metadata.VisitReplicas( - [](const Replica& r) { - return r.is_memory_replica() && r.is_completed() && - r.get_refcnt() == 0; - }, - [this, &tenant_id, &key, &tenant_state, &queued, - &now](Replica& replica) { - if (queued) return; // only need to pin one replica for offload - auto result = PushOffloadingQueue( - MakeObjectIdentity(key, tenant_id), replica); - if (result) { - replica.inc_refcnt(); - tenant_state.offloading_tasks.emplace( - key, OffloadingTask{replica.id(), now}); - queued = true; - } - }); + if (tenant_state.offloading_tasks.count(key) == 0) { + metadata.VisitReplicas( + [](const Replica& r) { + return r.is_memory_replica() && r.is_completed() && + r.get_refcnt() == 0; + }, + [this, &tenant_id, &key, &tenant_state, &queued, + &now](Replica& replica) { + if (queued) return; // only pin one replica for offload + auto result = PushOffloadingQueue( + MakeObjectIdentity(key, tenant_id), replica); + if (result) { + replica.inc_refcnt(); + auto [it, inserted] = + tenant_state.offloading_tasks.emplace( + key, OffloadingTask{replica.id(), now}); + if (!inserted) { + // Cross-cycle re-pin: another iteration already + // queued this key. Roll back to avoid refcnt leak. + replica.dec_refcnt(); + LOG(WARNING) + << "[BATCHEVICT-OFFLOAD-EMPLACE-FAIL] key=" + << key + << " already has offloading_task; rolled " + "back refcnt"; + } else { + queued = true; + } + } + }); + } if (queued) { offload_queued_this_cycle++; diff --git a/mooncake-store/src/posix_file.cpp b/mooncake-store/src/posix_file.cpp index e5f0158276..746d6da571 100644 --- a/mooncake-store/src/posix_file.cpp +++ b/mooncake-store/src/posix_file.cpp @@ -1,3 +1,4 @@ +#include #include #include #include @@ -55,6 +56,14 @@ tl::expected PosixFile::write(std::span data, ssize_t written = ::write(fd_, ptr, remaining); if (written == -1) { if (errno == EINTR) continue; + int saved_errno = errno; + char errbuf[256]; + strerror_r(saved_errno, errbuf, sizeof(errbuf)); + LOG(ERROR) << "write failed for file: " << filename_ + << ", errno=" << saved_errno + << " (" << errbuf << ")" + << ", fd=" << fd_ + << ", remaining=" << remaining; return make_error(ErrorCode::FILE_WRITE_FAIL); } remaining -= written; @@ -86,6 +95,15 @@ tl::expected PosixFile::read(std::string &buffer, ssize_t n = ::read(fd_, ptr, length - read_bytes); if (n == -1) { if (errno == EINTR) continue; + int saved_errno = errno; + char errbuf[256]; + strerror_r(saved_errno, errbuf, sizeof(errbuf)); + LOG(ERROR) << "read failed for file: " << filename_ + << ", errno=" << saved_errno + << " (" << errbuf << ")" + << ", fd=" << fd_ + << ", length=" << length + << ", read_bytes=" << read_bytes; buffer.clear(); return make_error(ErrorCode::FILE_READ_FAIL); } @@ -108,12 +126,42 @@ tl::expected PosixFile::vector_write(const iovec *iov, return make_error(ErrorCode::FILE_NOT_FOUND); } - ssize_t ret = ::pwritev(fd_, iov, iovcnt, offset); - if (ret < 0) { + size_t total_bytes = 0; + for (int i = 0; i < iovcnt; ++i) total_bytes += iov[i].iov_len; + + size_t written_total = 0; + off_t cur_offset = offset; + + for (int idx = 0; idx < iovcnt; idx += UIO_MAXIOV) { + int chunk_cnt = std::min(iovcnt - idx, UIO_MAXIOV); + ssize_t ret = ::pwritev(fd_, iov + idx, chunk_cnt, cur_offset); + if (ret < 0) { + int saved_errno = errno; + char errbuf[256]; + strerror_r(saved_errno, errbuf, sizeof(errbuf)); + LOG(ERROR) << "pwritev failed for file: " << filename_ + << ", errno=" << saved_errno + << " (" << errbuf << ")" + << ", fd=" << fd_ + << ", iovcnt=" << iovcnt + << ", total_bytes=" << total_bytes + << ", offset=" << offset + << ", chunk_start=" << idx + << ", chunk_cnt=" << chunk_cnt; + return make_error(ErrorCode::FILE_WRITE_FAIL); + } + written_total += ret; + cur_offset += ret; + } + + if (written_total != total_bytes) { + LOG(ERROR) << "pwritev partial write for file: " << filename_ + << ", expected=" << total_bytes + << ", written=" << written_total; return make_error(ErrorCode::FILE_WRITE_FAIL); } - return ret; + return written_total; } tl::expected PosixFile::vector_read(const iovec *iov, @@ -123,12 +171,53 @@ tl::expected PosixFile::vector_read(const iovec *iov, return make_error(ErrorCode::FILE_NOT_FOUND); } - ssize_t ret = ::preadv(fd_, iov, iovcnt, offset); - if (ret < 0) { - return make_error(ErrorCode::FILE_READ_FAIL); + size_t total_bytes = 0; + for (int i = 0; i < iovcnt; ++i) total_bytes += iov[i].iov_len; + + size_t read_total = 0; + off_t cur_offset = offset; + + for (int idx = 0; idx < iovcnt; idx += UIO_MAXIOV) { + int chunk_cnt = std::min(iovcnt - idx, UIO_MAXIOV); + ssize_t ret = ::preadv(fd_, iov + idx, chunk_cnt, cur_offset); + if (ret < 0) { + int saved_errno = errno; + char errbuf[256]; + strerror_r(saved_errno, errbuf, sizeof(errbuf)); + LOG(ERROR) << "preadv failed for file: " << filename_ + << ", errno=" << saved_errno + << " (" << errbuf << ")" + << ", fd=" << fd_ + << ", iovcnt=" << iovcnt + << ", total_bytes=" << total_bytes + << ", offset=" << offset + << ", chunk_start=" << idx + << ", chunk_cnt=" << chunk_cnt; + return make_error(ErrorCode::FILE_READ_FAIL); + } + read_total += ret; + cur_offset += ret; + if (ret == 0) break; // EOF } - return ret; + return read_total; +} + +tl::expected PosixFile::datasync() { + if (fd_ < 0) { + return make_error(ErrorCode::FILE_NOT_FOUND); + } + if (::fdatasync(fd_) != 0) { + int saved_errno = errno; + char errbuf[256]; + strerror_r(saved_errno, errbuf, sizeof(errbuf)); + LOG(ERROR) << "fdatasync failed for file: " << filename_ + << ", errno=" << saved_errno + << " (" << errbuf << ")" + << ", fd=" << fd_; + return make_error(ErrorCode::FILE_WRITE_FAIL); + } + return {}; } } // namespace mooncake \ No newline at end of file diff --git a/mooncake-store/src/real_client.cpp b/mooncake-store/src/real_client.cpp index dc544b7872..fe9cc29011 100644 --- a/mooncake-store/src/real_client.cpp +++ b/mooncake-store/src/real_client.cpp @@ -11,9 +11,11 @@ #include // for dlsym (Python detection) #include // for atexit #include +#include #include #include #include +#include #include #include @@ -4667,18 +4669,92 @@ RealClient::batch_get_into_internal(const std::vector &keys, } size_t offload_object_count = 0; + for (const auto &oe : offload_objects) offload_object_count += oe.second.size(); auto start_read_store_time = std::chrono::steady_clock::now(); - for (auto &offload_objects_it : offload_objects) { - offload_object_count += offload_objects_it.second.size(); - auto batch_get_offload_result = batch_get_into_offload_object_internal( - offload_objects_it.first, offload_objects_it.second); - if (!batch_get_offload_result) { - LOG(ERROR) << "Batch get store object failed with error: " - << batch_get_offload_result.error(); - for (const auto &offload_object_it : offload_objects_it.second) { - results[valid_local_disk_operations.at(offload_object_it.first) - .original_index] = - tl::make_unexpected(batch_get_offload_result.error()); + { + std::ostringstream oss; + oss << "action=batch_get_into_endpoint_split"; + for (const auto &oe : offload_objects) { + oss << " ep=" << oe.first << " keys=" << oe.second.size(); + } + VLOG(1) << oss.str(); + } + // Each endpoint targets an independent store pod (separate NVMe, separate + // transport). Fetching from them in parallel raises aggregate read + // bandwidth from single-pod to N-pod. Different workers touch disjoint + // keys in `results` (a key belongs to exactly one endpoint), so the only + // shared mutable state is `results` itself, indexed by original_index. + { + const size_t n_endpoints = offload_objects.size(); + const uint32_t endpoint_threads = std::max( + 1, GetEnvOr("MOONCAKE_OFFLOAD_ENDPOINT_THREADS", 8)); + const bool use_parallel = + (endpoint_threads > 1 && n_endpoints > 1); + + VLOG(1) << "action=offload_endpoint_dispatch endpoints=" << n_endpoints + << " threads=" << endpoint_threads + << " parallel=" << (use_parallel ? "true" : "false"); + + // Snapshot endpoints into a vector so workers can index concurrently + // without touching the unordered_map. + std::vector>*>> + endpoint_snapshot; + endpoint_snapshot.reserve(n_endpoints); + for (auto &oe : offload_objects) { + endpoint_snapshot.emplace_back(oe.first, &oe.second); + } + + auto process_endpoint = [&](const std::string &endpoint, + std::unordered_map> &objects) + -> std::optional { + auto batch_get_offload_result = + batch_get_into_offload_object_internal(endpoint, objects); + if (!batch_get_offload_result) { + LOG(ERROR) << "Batch get store object failed with error: " + << batch_get_offload_result.error(); + for (const auto &offload_object_it : objects) { + results[valid_local_disk_operations + .at(offload_object_it.first) + .original_index] = + tl::make_unexpected( + batch_get_offload_result.error()); + } + return batch_get_offload_result.error(); + } + return std::nullopt; + }; + + if (!use_parallel) { + for (auto &[endpoint, objects_ptr] : endpoint_snapshot) { + if (auto err = process_endpoint(endpoint, *objects_ptr); + err.has_value()) { + // Continue other endpoints on error to maximize partial + // success; first error is logged above. + } + } + } else { + std::atomic next_index{0}; + std::vector workers; + workers.reserve(endpoint_threads); + for (uint32_t i = 0; i < endpoint_threads; ++i) { + workers.emplace_back([&]() { + for (;;) { + size_t i = next_index.fetch_add( + 1, std::memory_order_acq_rel); + if (i >= n_endpoints) return; + auto &[endpoint, objects_ptr] = + endpoint_snapshot[i]; + (void)process_endpoint(endpoint, *objects_ptr); + // Errors are logged + results marked inside + // process_endpoint; we do not early-terminate so + // other endpoints can still succeed. + } + }); + } + for (auto &w : workers) { + if (w.joinable()) w.join(); } } } @@ -4691,10 +4767,14 @@ RealClient::batch_get_into_internal(const std::vector &keys, std::chrono::duration_cast( end_time - start_read_store_time) .count(); - // LOG(INFO) << "Time taken for batch_get_into: " << elapsed_time - // << "us, read store: " << read_store_time - // << "us, with memory key count: " << valid_operations.size() - // << ", offload key count: " << offload_object_count; + VLOG(1) << "action=batch_get_into_summary" + << " total_keys=" << num_keys + << " memory_keys=" << valid_operations.size() + << " local_disk_keys=" << offload_object_count + << " disk_keys=" << disk_operations.size() + << " endpoints_with_offload=" << offload_objects.size() + << " total_us=" << elapsed_time + << " store_read_us=" << read_store_time; return results; } @@ -5093,13 +5173,45 @@ RealClient::batch_get_into_multi_buffers_internal( .emplace(key, std::move(user_slices)); } - for (auto &[endpoint, objects] : offload_objects) { - if (objects.empty()) continue; + // Each endpoint targets an independent store pod. Fetching in + // parallel raises aggregate read bandwidth across all store pods + // (mirrors the offload_endpoint_dispatch in batch_get_into_internal + // above). Different workers write disjoint indices in `results` + // because a key belongs to exactly one endpoint. + const size_t n_endpoints = offload_objects.size(); + const uint32_t endpoint_threads = std::max( + 1, GetEnvOr("MOONCAKE_OFFLOAD_ENDPOINT_THREADS", 8)); + const bool use_parallel = + (endpoint_threads > 1 && n_endpoints > 1); + + VLOG(1) << "action=offload_multi_buffers_endpoint_dispatch" + << " endpoints=" << n_endpoints + << " threads=" << endpoint_threads + << " parallel=" << (use_parallel ? "true" : "false"); + + // Snapshot endpoints so workers can index without touching the + // unordered_map. Pointers remain valid because offload_objects + // is local and outlives the workers. + std::vector>*>> + endpoint_snapshot; + endpoint_snapshot.reserve(n_endpoints); + for (auto &oe : offload_objects) { + if (oe.second.empty()) continue; + endpoint_snapshot.emplace_back(oe.first, &oe.second); + } + const size_t n_active = endpoint_snapshot.size(); + + auto process_endpoint = + [&](const std::string &endpoint, + std::unordered_map> + &objects) -> void { auto read_result = batch_get_into_offload_object_internal(endpoint, objects); // On success: results[original_index] was already pre-filled - // with total_size when valid_local_disk_ops was built; nothing - // to update. Only overwrite on failure. + // with total_size when valid_local_disk_ops was built; + // nothing to update. Only overwrite on failure. if (!read_result) { for (auto &[key, slices] : objects) { auto disk_it = valid_local_disk_ops.find(key); @@ -5110,6 +5222,31 @@ RealClient::batch_get_into_multi_buffers_internal( tl::make_unexpected(read_result.error()); } } + }; + + if (!use_parallel) { + for (auto &[endpoint, objects_ptr] : endpoint_snapshot) { + process_endpoint(endpoint, *objects_ptr); + } + } else { + std::atomic next_index{0}; + std::vector workers; + workers.reserve(endpoint_threads); + for (uint32_t i = 0; i < endpoint_threads; ++i) { + workers.emplace_back([&]() { + for (;;) { + size_t i = next_index.fetch_add( + 1, std::memory_order_acq_rel); + if (i >= n_active) return; + auto &[endpoint, objects_ptr] = + endpoint_snapshot[i]; + process_endpoint(endpoint, *objects_ptr); + } + }); + } + for (auto &w : workers) { + if (w.joinable()) w.join(); + } } } diff --git a/mooncake-store/src/storage_backend.cpp b/mooncake-store/src/storage_backend.cpp index 3612624a3f..fbf55d589c 100644 --- a/mooncake-store/src/storage_backend.cpp +++ b/mooncake-store/src/storage_backend.cpp @@ -9,11 +9,15 @@ #include #include +#include #include #include #include #include +#include #include +#include +#include #include #include @@ -1370,8 +1374,17 @@ tl::expected BucketStorageBackend::BatchOffload( << bucket->keys[i] << ", bucket_id=" << bucket_id; } } + // Set LRU timestamp to current time. With the upstream default of 0, + // newly written buckets appear as 'least recently used' and are the + // first candidates selected by SelectEvictionCandidate once the SSD + // fills up — fresh cache data is evicted before it is ever read, + // collapsing the hit rate from ~80% to ~27%. + int64_t now_ns = std::chrono::steady_clock::now() + .time_since_epoch() + .count(); + bucket->last_access_ns_.store(now_ns, std::memory_order_relaxed); buckets_.emplace(bucket_id, std::move(bucket)); - lru_index_.emplace(0LL, bucket_id); + lru_index_.emplace(now_ns, bucket_id); } return bucket_id; @@ -1396,6 +1409,8 @@ tl::expected BucketStorageBackend::BatchQuery( tl::expected BucketStorageBackend::BatchLoad( std::unordered_map& batch_object) { + auto batch_start = std::chrono::steady_clock::now(); + int64_t batch_total_bytes = 0; // Step 1: Build read plan by copying metadata under lock // BucketReadGuard increments inflight_reads_ to prevent deletion during IO. // When the guard goes out of scope, it decrements the counter. @@ -1467,81 +1482,331 @@ tl::expected BucketStorageBackend::BatchLoad( // which remain alive until this function returns (~line bucket_guards // destructor), keeping inflight_reads_ > 0 throughout the I/O phase. - // Step 2: Perform IO without holding any locks - for (auto& [bucket_id, read_plans] : bucket_read_plans) { + VLOG(1) << "action=batchload_start" + << " keys=" << batch_object.size() + << " buckets=" << bucket_read_plans.size(); + { + std::ostringstream oss; + oss << "action=batchload_plan"; + for (auto& [bid, plans] : bucket_read_plans) { + int64_t b_bytes = 0; + for (auto& p : plans) b_bytes += p.data_size; + batch_total_bytes += b_bytes; + oss << " b" << bid << "=" << plans.size() << "/" << b_bytes; + } + VLOG(1) << oss.str(); + } + + // Step 2: Perform IO without holding any locks. + // + // Buckets are independent (separate files, separate fds, disjoint key + // sets), so reads can safely overlap on a single NVMe device. Raising + // the in-flight pread count from 1 (serial) to N (parallel) lifts + // single-disk throughput from ~1.4 GB/s to ~4-5 GB/s — see commit + // e738cfd9 which applied the same pattern to the write path. + // + // Thread safety: + // - bucket_read_plans snapshot is taken once below; workers only read + // the snapshot (no concurrent unordered_map access). + // - Each bucket opens its own fd via OpenFile; StorageFile instances + // are not shared across worker threads. + // - batch_object.at(plan.key).ptr writes target disjoint keys (a key + // belongs to exactly one bucket), so different workers mutate + // distinct map values. This is safe on all mainstream libstdc++ + // implementations; the C++ standard does not strictly guarantee it. + // - VLOG is thread-safe via glog. + auto process_bucket = [&](int64_t bucket_id, + std::vector& read_plans) + -> std::optional { + auto bucket_io_start = std::chrono::steady_clock::now(); + int64_t bucket_bytes = 0; // Open file for this bucket (cheap syscall, no lock needed) auto filepath_res = GetBucketDataPath(bucket_id); if (!filepath_res) { LOG(ERROR) << "Failed to get bucket data path, bucket_id=" << bucket_id; - return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); + return ErrorCode::INTERNAL_ERROR; } auto file_res = OpenFile(filepath_res.value(), FileMode::Read); if (!file_res) { LOG(ERROR) << "Failed to open bucket file: " << filepath_res.value(); - return tl::make_unexpected(file_res.error()); + return file_res.error(); } auto& file = file_res.value(); - // Read each key's data - for (const auto& plan : read_plans) { - int64_t actual_offset = plan.offset + plan.key_size; - tl::expected read_res; + // ---- Merge optimization: sort by offset, merge adjacent reads ---- + // Group read_plans whose actual_offset values are within + // kMergeGapBytes of each other. Each group is read by a single + // pread into a temporary buffer, then scattered to dest slices + // via memcpy. This collapses N pread syscalls into 1 for groups + // of keys whose data sit close together in the bucket file. + constexpr int64_t kMergeGapBytes = 256 * 1024; // 256 KiB threshold + + // Working copy sorted by actual_offset (offset + key_size) + std::vector sorted_plans; + sorted_plans.reserve(read_plans.size()); + for (auto& p : read_plans) sorted_plans.push_back(&p); + std::sort(sorted_plans.begin(), sorted_plans.end(), + [](const ReadPlan* a, const ReadPlan* b) { + return (a->offset + a->key_size) < + (b->offset + b->key_size); + }); + + size_t merge_group_count = 0; + size_t pread_count_after_merge = 0; + + size_t idx = 0; + while (idx < sorted_plans.size()) { + size_t group_start = idx; + int64_t seg_offset = + sorted_plans[idx]->offset + sorted_plans[idx]->key_size; + int64_t seg_end = seg_offset + + static_cast(sorted_plans[idx]->dest_slice.size); + // Greedily extend the group while the next key is within + // kMergeGapBytes of the current segment end. + while (idx + 1 < sorted_plans.size()) { + int64_t next_offset = + sorted_plans[idx + 1]->offset + + sorted_plans[idx + 1]->key_size; + if (next_offset - seg_end > kMergeGapBytes) break; + seg_end = next_offset + + static_cast(sorted_plans[idx + 1]->dest_slice.size); + ++idx; + } + size_t group_end = idx; + size_t group_size = group_end - group_start + 1; + ++pread_count_after_merge; + if (group_size > 1) ++merge_group_count; + + if (group_size == 1) { + // Single key: original single-pread path + const auto& plan = *sorted_plans[group_start]; + int64_t actual_offset = plan.offset + plan.key_size; + tl::expected read_res; #ifdef USE_URING - // Try to use read_aligned for O_DIRECT I/O if file is UringFile - UringFile* uring_file = dynamic_cast(file.get()); - if (uring_file != nullptr) { - // Calculate aligned read range + UringFile* uring_file = dynamic_cast(file.get()); + if (uring_file != nullptr) { + int64_t aligned_offset = + align_down(actual_offset, kDirectIOAlignment); + int64_t data_end = actual_offset + + static_cast(plan.dest_slice.size); + int64_t aligned_end = static_cast(align_up( + static_cast(data_end), kDirectIOAlignment)); + size_t aligned_size = + static_cast(aligned_end - aligned_offset); + int64_t offset_in_buffer = actual_offset - aligned_offset; + read_res = uring_file->read_aligned( + plan.dest_slice.ptr, aligned_size, aligned_offset); + if (read_res) { + if (read_res.value() < plan.dest_slice.size) { + LOG(ERROR) << "uring read short of data for key: " + << plan.key + << ", bucket_id=" << plan.bucket_id + << ", expected>=" + << plan.dest_slice.size + << ", got=" << read_res.value(); + return ErrorCode::FILE_READ_FAIL; + } + batch_object.at(plan.key).ptr = + static_cast(plan.dest_slice.ptr) + + offset_in_buffer; + read_res = plan.dest_slice.size; + } + } else +#endif + { + iovec iov{plan.dest_slice.ptr, plan.dest_slice.size}; + read_res = file->vector_read(&iov, 1, actual_offset); + } + + if (!read_res) { + LOG(ERROR) << "vector_read failed for key: " << plan.key + << ", bucket_id=" << plan.bucket_id + << ", error: " << read_res.error(); + return read_res.error(); + } + if (read_res.value() != plan.dest_slice.size) { + LOG(ERROR) << "Read size mismatch for key: " << plan.key + << ", expected: " << plan.dest_slice.size + << ", got: " << read_res.value(); + return ErrorCode::FILE_READ_FAIL; + } + bucket_bytes += plan.data_size; + } else { + // Merged group: one big pread into an aligned temp buffer, + // then scatter to each key's dest_slice. + // + // O_DIRECT (UringFile) requires offset, length, and buffer + // address to all be 4096-aligned; seg_offset/seg_end are + // derived from key boundaries and are generally *not* + // aligned. Align the read range outward (matching the + // single-key path at the align_down/align_up calls above) + // and track where in the aligned buffer the actual data + // starts. Under buffered I/O (PosixFile) the extra bytes + // are simply ignored, so this path is safe for both backends. int64_t aligned_offset = - align_down(actual_offset, kDirectIOAlignment); - int64_t data_end = - actual_offset + static_cast(plan.dest_slice.size); + align_down(seg_offset, kDirectIOAlignment); int64_t aligned_end = static_cast(align_up( - static_cast(data_end), kDirectIOAlignment)); + static_cast(seg_end), kDirectIOAlignment)); size_t aligned_size = static_cast(aligned_end - aligned_offset); - int64_t offset_in_buffer = actual_offset - aligned_offset; - - // Zero-copy path: read directly into the slice buffer. - // dest_slice.ptr is 4096-aligned and oversized (from - // AllocateBatch) to accommodate the full aligned read range. - read_res = uring_file->read_aligned( - plan.dest_slice.ptr, aligned_size, aligned_offset); - - if (read_res) { - // Adjust ptr to point to actual data start (no memcpy) - batch_object.at(plan.key).ptr = - static_cast(plan.dest_slice.ptr) + - offset_in_buffer; - read_res = plan.dest_slice.size; + + void* raw_buf = nullptr; + int ret = + posix_memalign(&raw_buf, kDirectIOAlignment, aligned_size); + if (ret != 0) { + LOG(ERROR) << "merged posix_memalign failed for bucket_id=" + << bucket_id << ", aligned_size=" << aligned_size + << ": " << strerror(ret); + return ErrorCode::FILE_WRITE_FAIL; + } + std::unique_ptr merge_buf( + raw_buf, [](void* p) { free(p); }); + + iovec iov{merge_buf.get(), aligned_size}; + auto read_res = file->vector_read(&iov, 1, aligned_offset); + if (!read_res) { + LOG(ERROR) << "merged vector_read failed for bucket_id=" + << bucket_id << ", aligned_offset=" + << aligned_offset + << ", aligned_size=" << aligned_size + << ", error: " << read_res.error(); + return read_res.error(); + } + // The pread may short-read at EOF if the aligned read range + // overshoots the file end by up to (kDirectIOAlignment - 1) + // bytes. The actual data range we care about is + // [seg_offset, seg_end), which maps to + // [seg_offset - aligned_offset, seg_end - aligned_offset) + // inside the buffer; got >= (seg_end - aligned_offset) is + // therefore sufficient. Each key's memcpy stays in bounds + // because key_offset + dest_slice.size <= seg_end. + size_t min_needed = + static_cast(seg_end - aligned_offset); + if (read_res.value() < min_needed) { + LOG(ERROR) << "merged read short of data, bucket_id=" + << bucket_id << ", expected>=" << min_needed + << ", got=" << read_res.value(); + return ErrorCode::FILE_READ_FAIL; + } + // Scatter: for each key, memcpy from merge_buf to dest_slice. + // buf_offset accounts for the alignment padding before + // seg_offset (zero under buffered I/O when seg_offset happens + // to be aligned; positive under O_DIRECT otherwise). + char* merge_base = static_cast(merge_buf.get()); + for (size_t k = group_start; k <= group_end; ++k) { + const auto& plan = *sorted_plans[k]; + int64_t key_offset = plan.offset + plan.key_size; + int64_t buf_offset = key_offset - aligned_offset; + std::memcpy(plan.dest_slice.ptr, + merge_base + buf_offset, + plan.dest_slice.size); + bucket_bytes += plan.data_size; } - } else -#endif - { - // Fallback to vector_read for non-UringFile - iovec iov{plan.dest_slice.ptr, plan.dest_slice.size}; - read_res = file->vector_read(&iov, 1, actual_offset); } + ++idx; + } + + VLOG(1) << "bucket_id=" << bucket_id + << " preads_before_merge=" << read_plans.size() + << " preads_after_merge=" << pread_count_after_merge + << " merge_groups=" << merge_group_count; + auto bucket_elapsed_us = + std::chrono::duration_cast( + std::chrono::steady_clock::now() - bucket_io_start) + .count(); + VLOG(1) << "action=batchload_bucket_done" + << " bucket_id=" << bucket_id + << " keys=" << read_plans.size() + << " bytes=" << bucket_bytes + << " elapsed_us=" << bucket_elapsed_us + << " preads=" << read_plans.size() + << " per_key_us=" << (read_plans.empty() ? 0 : bucket_elapsed_us / static_cast(read_plans.size())); + return std::nullopt; + }; - if (!read_res) { - LOG(ERROR) << "vector_read failed for key: " << plan.key - << ", bucket_id=" << plan.bucket_id - << ", error: " << read_res.error(); - return tl::make_unexpected(read_res.error()); + // Snapshot bucket plans into a vector so workers can index concurrently + // without touching the unordered_map (its iterators are not standard- + // guaranteed safe under concurrent read). Pointers remain valid because + // bucket_read_plans is a local variable that outlives the workers. + std::vector*>> bucket_snapshot; + bucket_snapshot.reserve(bucket_read_plans.size()); + for (auto& [bid, plans] : bucket_read_plans) { + bucket_snapshot.emplace_back(bid, &plans); + } + const size_t n = bucket_snapshot.size(); + + const uint32_t read_threads = + std::max(1, file_storage_config_.offload_read_threads); + const bool use_parallel = (read_threads > 1 && n > 1); + + VLOG(1) << "action=batchload_dispatch buckets=" << n + << " threads=" << read_threads + << " parallel=" << (use_parallel ? "true" : "false"); + + if (!use_parallel) { + // Serial path — preserves prior behavior, no thread overhead. + for (auto& [bid, plans_ptr] : bucket_snapshot) { + auto err = process_bucket(bid, *plans_ptr); + if (err.has_value()) { + return tl::make_unexpected(err.value()); } - - if (read_res.value() != plan.dest_slice.size) { - LOG(ERROR) << "Read size mismatch for key: " << plan.key - << ", expected: " << plan.dest_slice.size - << ", got: " << read_res.value(); - return tl::make_unexpected(ErrorCode::FILE_READ_FAIL); + } + } else { + // Parallel path — dispatch independent buckets concurrently. + // Structure mirrors WriteBucket parallelization (commit e738cfd9): + // workers pull bucket indices from a shared atomic counter, stop + // early once one bucket returns a fatal error, and the caller + // joins all workers before returning. + std::atomic next_index{0}; + std::atomic first_error_set{false}; + std::optional first_error; + std::mutex result_mutex; + + auto worker = [&]() { + for (;;) { + if (first_error_set.load(std::memory_order_acquire)) return; + size_t i = next_index.fetch_add(1, std::memory_order_acq_rel); + if (i >= n) return; + auto& [bid, plans_ptr] = bucket_snapshot[i]; + auto err = process_bucket(bid, *plans_ptr); + if (err.has_value()) { + std::lock_guard lk(result_mutex); + if (!first_error_set.exchange(true)) { + first_error = err.value(); + } + return; + } } + }; + + std::vector workers; + workers.reserve(read_threads); + for (uint32_t i = 0; i < read_threads; ++i) { + workers.emplace_back(worker); + } + for (auto& w : workers) { + if (w.joinable()) w.join(); + } + if (first_error.has_value()) { + return tl::make_unexpected(first_error.value()); } } + auto batch_elapsed_us = + std::chrono::duration_cast( + std::chrono::steady_clock::now() - batch_start) + .count(); + VLOG(1) << "action=batchload_done" + << " keys=" << batch_object.size() + << " buckets=" << bucket_read_plans.size() + << " total_bytes=" << batch_total_bytes + << " elapsed_us=" << batch_elapsed_us + << " throughput_mbs=" << (batch_elapsed_us > 0 ? (double)batch_total_bytes / 1024.0 / 1024.0 / ((double)batch_elapsed_us / 1000000.0) : 0.0); + // bucket_guards go out of scope here, decrementing inflight_reads_ return {}; } @@ -1914,13 +2179,18 @@ tl::expected BucketStorageBackend::GroupOffloadingKeysByBucket( for (int64_t i = static_cast(bucket_keys.size()); i < bucket_backend_config_.bucket_keys_limit; ++i) { if (it == offloading_objects.cend()) { - for (const auto& bucket_object : bucket_objects) { - ungrouped_offloading_objects.emplace(bucket_object.first, - bucket_object.second); + if (!bucket_keys.empty()) { + auto bucket_keys_count = + static_cast(bucket_keys.size()); + residue_count -= bucket_keys_count; + buckets_keys.push_back(std::move(bucket_keys)); + VLOG(1) << "[OFFLOAD-PARTIAL] flushed partial bucket at " + "input end: keys=" + << bucket_keys_count + << " data_size=" << bucket_data_size + << " total=" << total_count + << " residue=" << residue_count; } - VLOG(1) << "Add offloading objects to ungrouped pool. " - << "Total ungrouped count: " - << ungrouped_offloading_objects.size(); return {}; } @@ -2031,33 +2301,25 @@ tl::expected BucketStorageBackend::WriteBucket( size_t total_size = static_cast(bucket_metadata->data_size); size_t aligned_size = align_up(total_size, kDirectIOAlignment); - // Allocate aligned buffer if needed + // Allocate a per-call aligned buffer. The shared member + // aligned_io_buffer_ is NOT safe for concurrent WriteBucket calls + // (parallel WriteBucket dispatches multiple buckets simultaneously + // via worker threads); using it here would cause data corruption. void* write_buffer = nullptr; - std::unique_ptr temp_buffer{nullptr, - [](void*) {}}; + std::unique_ptr temp_buffer{ + nullptr, [](void*) {}}; - if (aligned_size <= kAlignedBufferSize && aligned_io_buffer_) { - // Use the pre-allocated buffer - write_buffer = aligned_io_buffer_.get(); - } else { - // Allocate a temporary larger buffer - void* buf = nullptr; - int ret = posix_memalign(&buf, kDirectIOAlignment, aligned_size); - if (ret != 0) { - LOG(ERROR) - << "Failed to allocate aligned buffer for WriteBucket: " - << strerror(ret); - return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); - } - temp_buffer.reset(buf); - temp_buffer = std::unique_ptr( - buf, [](void* p) { free(p); }); - write_buffer = buf; - LOG(WARNING) << "WriteBucket: bucket_id=" << bucket_id - << " requires " << aligned_size - << " bytes, exceeds buffer size " << kAlignedBufferSize - << ", using temporary allocation"; + void* buf = nullptr; + int ret = posix_memalign(&buf, kDirectIOAlignment, aligned_size); + if (ret != 0) { + LOG(ERROR) + << "Failed to allocate aligned buffer for WriteBucket: " + << strerror(ret); + return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); } + temp_buffer = std::unique_ptr( + buf, [](void* p) { free(p); }); + write_buffer = buf; // Aggregate all iovs data into the aligned buffer char* dst = static_cast(write_buffer); @@ -2118,6 +2380,35 @@ tl::expected BucketStorageBackend::WriteBucket( return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); } + // Flush data to stable storage. O_DIRECT reads (UringFile) bypass + // the page cache; without fdatasync, a subsequent O_DIRECT read + // may race with in-progress writeback and see partial data at + // 4096 boundaries — matching the "uring read short of data" pattern + // observed in stage-4 SSD-offload tests. + auto sync_result = file->datasync(); + if (!sync_result) { + LOG(ERROR) << "datasync failed for bucket: " << bucket_id + << ", error: " << sync_result.error(); + return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); + } + + // Verify file size matches the metadata record. If the file is + // shorter than expected, the write was silently truncated by the + // filesystem (e.g. ENOSPC during delayed allocation). Catching it + // here prevents the "uring read short of data" error on the read + // path. + struct stat st; + if (::fstat(file->fd(), &st) == 0) { + if (static_cast(st.st_size) != + bucket_metadata->data_size) { + LOG(ERROR) << "File size mismatch after write+sync: " + << bucket_data_path + << ", expected=" << bucket_metadata->data_size + << ", actual=" << st.st_size; + return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); + } + } + // Invalidate cache for this file since content changed { MutexLocker cache_locker(&file_cache_mutex_); diff --git a/mooncake-store/src/uring_file.cpp b/mooncake-store/src/uring_file.cpp index 870eeba2f0..5ededa9b95 100644 --- a/mooncake-store/src/uring_file.cpp +++ b/mooncake-store/src/uring_file.cpp @@ -4,6 +4,7 @@ #include #include #include +#include #include #include #include @@ -242,6 +243,16 @@ class SharedUringRing { return std::max(s, MIN_CHUNK); } + static size_t max_rw_count() { + static const size_t value = [] { + long page_size = sysconf(_SC_PAGESIZE); + if (page_size <= 0) page_size = 4096; + return static_cast(INT_MAX) & + ~(static_cast(page_size) - 1); + }(); + return value; + } + // Drain exactly @expected CQEs and accumulate bytes. tl::expected collect(int expected) { int ret = io_uring_submit_and_wait(&ring_, expected); @@ -288,7 +299,7 @@ class SharedUringRing { static_cast((remaining + cs - 1) / cs), QUEUE_DEPTH); for (unsigned i = 0; i < n; ++i) { - size_t chunk = std::min(cs, remaining); + size_t chunk = std::min({cs, remaining, max_rw_count()}); struct io_uring_sqe* sqe = io_uring_get_sqe(&ring_); if (!sqe) { @@ -332,6 +343,7 @@ class SharedUringRing { off_t off) { const ErrorCode err_code = is_write ? ErrorCode::FILE_WRITE_FAIL : ErrorCode::FILE_READ_FAIL; + const size_t max_io = max_rw_count(); size_t total = 0; off_t cur = off; @@ -339,7 +351,23 @@ class SharedUringRing { int idx = 0; while (remaining > 0) { - int batch = std::min(remaining, static_cast(QUEUE_DEPTH)); + if (iovs[idx].iov_len > max_io) { + auto res = submit_rw(is_write, fd, iovs[idx].iov_base, + iovs[idx].iov_len, cur, + /*use_fixed_buf=*/false); + if (!res) return res; + total += res.value(); + cur += static_cast(iovs[idx].iov_len); + ++idx; + --remaining; + continue; + } + + int batch = 0; + while (batch < remaining && batch < static_cast(QUEUE_DEPTH) && + iovs[idx + batch].iov_len <= max_io) { + ++batch; + } for (int i = 0; i < batch; ++i) { struct io_uring_sqe* sqe = io_uring_get_sqe(&ring_); diff --git a/mooncake-wheel/mooncake/mooncake_store_service.py b/mooncake-wheel/mooncake/mooncake_store_service.py index a91584070d..601ed430b2 100644 --- a/mooncake-wheel/mooncake/mooncake_store_service.py +++ b/mooncake-wheel/mooncake/mooncake_store_service.py @@ -170,7 +170,10 @@ async def start_http_service(self, port: int = 8080): web.get('/api/get/{key}', _timed_handler("GET", self.handle_get)), web.get('/api/exist/{key}', _timed_handler("EXIST", self.handle_exist)), web.delete('/api/remove/{key}', _timed_handler("REMOVE", self.handle_remove)), - web.delete('/api/remove_all', _timed_handler("REMOVE_ALL", self.handle_remove_all)) + web.delete('/api/remove_all', _timed_handler("REMOVE_ALL", self.handle_remove_all)), + web.get('/metrics', _timed_handler("METRICS", self.handle_metrics)), + web.get('/metrics/summary', _timed_handler("METRICS_SUMMARY", self.handle_metrics_summary)), + web.get('/health', _timed_handler("HEALTH", self.handle_health)), ]) runner = web.AppRunner(app) await runner.setup() @@ -619,6 +622,58 @@ async def handle_remove_all(self, request): content_type='application/json' ) + async def handle_health(self, request): + try: + code = self.store.health_check() + status_map = {0: "healthy", 1: "not_initialized", 2: "master_unreachable"} + status_str = status_map.get(code, "unknown") + http_status = 200 if code == 0 else 503 + # logging.info("health_check: code=%d status=%s remote=%s", code, status_str, request.remote) + return web.Response( + status=http_status, + text=json.dumps({"status": status_str, "code": code}), + content_type="application/json" + ) + except Exception as e: + logging.error("HEALTH error: %s", e) + return web.Response(status=500, text=str(e)) + + async def handle_metrics(self, request): + try: + text = self.store.get_metrics() + if text is None: + return web.Response(status=503, text="metrics not available") + return web.Response( + status=200, text=text, + content_type="text/plain", charset="utf-8" + ) + except AttributeError: + return web.Response( + status=501, + text="get_metrics() not available; rebuild store.so with metric bindings" + ) + except Exception as e: + logging.error("METRICS error: %s", e) + return web.Response(status=500, text=str(e)) + + async def handle_metrics_summary(self, request): + try: + text = self.store.get_metrics_summary() + if text is None: + return web.Response(status=503, text="metrics not available") + return web.Response( + status=200, text=text, + content_type="text/plain", charset="utf-8" + ) + except AttributeError: + return web.Response( + status=501, + text="get_metrics_summary() not available; rebuild store.so with metric bindings" + ) + except Exception as e: + logging.error("METRICS_SUMMARY error: %s", e) + return web.Response(status=500, text=str(e)) + async def stop(self): if self.store: self.store.close()