diff --git a/docs/source/deployment/ssd/ssd-offload.md b/docs/source/deployment/ssd/ssd-offload.md index 34d2dd28a9..cd10b3c618 100644 --- a/docs/source/deployment/ssd/ssd-offload.md +++ b/docs/source/deployment/ssd/ssd-offload.md @@ -185,6 +185,32 @@ Groups multiple objects into bucket files. Reduces filesystem overhead, supports Best for: general-purpose use, large-scale deployments. +#### Explicit-Delete GC (tombstone compaction) + +To enable SSD space reclamation via `Remove`/`BatchRemove` (without LRU +eviction deleting live keys), set: + +| Environment Variable | Required Value | Description | +|---|---|---| +| `MOONCAKE_OFFLOAD_BUCKET_EVICTION_POLICY` | `lru` | Keeps `last_access_ns_` updated for GC cold-bucket selection | +| `MOONCAKE_OFFLOAD_DISABLE_SSD_EVICTION` | `true` | Makes `PrepareEviction` a no-op so no bucket is ever evicted by LRU | + +GC tuning (optional): + +| Environment Variable | Default | Description | +|---|---|---| +| `MOONCAKE_OFFLOAD_BUCKET_GC_ENABLE` | `true` | Enable background tombstone compaction | +| `MOONCAKE_OFFLOAD_BUCKET_GC_INTERVAL_MS` | `1000` | GC scan interval | +| `MOONCAKE_OFFLOAD_BUCKET_GC_DELETED_RATIO` | `0.25` | Compact bucket when deleted bytes / data size >= this | +| `MOONCAKE_OFFLOAD_BUCKET_GC_HIGH_WATERMARK_RATIO` | `0.90` | Force compaction of any tombstone bucket when total size / max >= this | +| `MOONCAKE_OFFLOAD_BUCKET_GC_MAX_BUCKETS_PER_ROUND` | `1` | Max old buckets collected per GC round for cross-bucket merge | +| `MOONCAKE_OFFLOAD_BUCKET_GC_MERGE_ENABLE` | `true` | Enable cross-bucket merge: collect live keys from multiple tombstone buckets into one new bucket. When false, each bucket is compacted independently | + +**Important:** Only keys removed via `Remove`/`BatchRemove` are reclaimed. +`RemoveByRegex`/`RemoveAll` do not trigger this GC. When no tombstone space +is reclaimable and SSD is full, `BatchOffload` returns an error instead of +deleting live keys. + ### `file_per_key_storage_backend` Stores each object in an individual file. Simple and easy to inspect, but generates many small files at scale. diff --git a/mooncake-integration/store/mooncake_perf_points.def b/mooncake-integration/store/mooncake_perf_points.def index e1b45a7e78..eaf6dd992d 100644 --- a/mooncake-integration/store/mooncake_perf_points.def +++ b/mooncake-integration/store/mooncake_perf_points.def @@ -179,3 +179,38 @@ PERF_KEY_DEF(MASTER_BG_DISCARD_EXPIRED, "master_service.cpp::EvictionThreadFun PERF_KEY_DEF(MASTER_BG_SNAPSHOT_PERSIST, "master_service.cpp::SnapshotThreadFunc", "SnapshotPersist") PERF_KEY_DEF(MASTER_BG_CLIENT_MONITOR, "master_service.cpp::ClientMonitorFunc", "ClientMonitorScan") PERF_KEY_DEF(MASTER_BG_CLIENT_UNMOUNT, "master_service.cpp::ClientMonitorFunc", "ExpiredClientUnmount") + +// ============================================================ +// === vLLM MooncakeStoreConnector 路径打点（vllm 0.26.1rc0）=== +// === 覆盖 S2-S9 入口 + T1-T8 下沉 + Client 服务层 === +// ============================================================ + +// === store_py.cpp Python 绑定层（vllm 直接调用入口，S2-S9）=== +PERF_KEY_DEF(STORE_PY_SETUP, "store_py.cpp::setup", "Setup") +PERF_KEY_DEF(STORE_PY_REGISTER_BUFFER, "store_py.cpp::register_buffer", "RegisterBuffer") +PERF_KEY_DEF(STORE_PY_BATCH_PUT_MULTI, "store_py.cpp::batch_put_from_multi_buffers","BatchPutMultiBuf") +PERF_KEY_DEF(STORE_PY_BATCH_GET_INTO_MULTI, "store_py.cpp::batch_get_into_multi_buffers","BatchGetIntoMultiBuf") +PERF_KEY_DEF(STORE_PY_BATCH_IS_EXIST, "store_py.cpp::batch_is_exist", "BatchIsExist") +PERF_KEY_DEF(STORE_PY_BATCH_GET_REPLICA_DESC, "store_py.cpp::batch_get_replica_desc", "BatchGetReplicaDesc") +PERF_KEY_DEF(STORE_PY_REMOVE_ALL, "store_py.cpp::remove_all", "RemoveAll") +PERF_KEY_DEF(STORE_PY_CLOSE, "store_py.cpp::close", "Close") + +// === RealClient 核心逻辑层（下沉路径，T1-T8）=== +PERF_KEY_DEF(RC_SETUP_REAL, "real_client.cpp::setup_real", "SetupReal") +PERF_KEY_DEF(RC_SETUP_INTERNAL, "real_client.cpp::setup_internal", "SetupInternal") +PERF_KEY_DEF(RC_TEARDOWN_ALL, "real_client.cpp::tearDownAll", "TeardownAll") +PERF_KEY_DEF(RC_TEARDOWN_ALL_INTERNAL, "real_client.cpp::tearDownAll_internal", "TeardownAllInternal") +PERF_KEY_DEF(RC_REMOVE_ALL, "real_client.cpp::removeAll", "RemoveAll") +PERF_KEY_DEF(RC_REMOVE_ALL_INTERNAL, "real_client.cpp::removeAll_internal", "RemoveAllInternal") +PERF_KEY_DEF(RC_BATCH_IS_EXIST, "real_client.cpp::batchIsExist", "BatchIsExist") +PERF_KEY_DEF(RC_BATCH_IS_EXIST_INTERNAL, "real_client.cpp::batchIsExist_internal", "BatchIsExistInternal") +PERF_KEY_DEF(RC_REGISTER_BUFFER, "real_client.cpp::register_buffer", "RegisterBuffer") +PERF_KEY_DEF(RC_REGISTER_BUFFER_INTERNAL, "real_client.cpp::register_buffer_internal","RegisterBufferInternal") +PERF_KEY_DEF(RC_BATCH_PUT_MULTI, "real_client.cpp::batch_put_from_multi_buffers", "BatchPutMultiBuf") +PERF_KEY_DEF(RC_BATCH_PUT_MULTI_INTERNAL, "real_client.cpp::batch_put_from_multi_buffers_internal","BatchPutMultiBufInternal") +PERF_KEY_DEF(RC_BATCH_GET_INTO_MULTI, "real_client.cpp::batch_get_into_multi_buffers", "BatchGetIntoMultiBuf") +PERF_KEY_DEF(RC_BATCH_GET_INTO_MULTI_INTERNAL, "real_client.cpp::batch_get_into_multi_buffers_internal","BatchGetIntoMultiBufInternal") +PERF_KEY_DEF(RC_BATCH_GET_REPLICA_DESC, "real_client.cpp::batch_get_replica_desc", "BatchGetReplicaDesc") + +// === Client 服务层（vllm 路径下沉，部分已有打点）=== +PERF_KEY_DEF(CLIENT_BATCH_QUERY, "client_service.cpp::BatchQuery", "BatchQuery") diff --git a/mooncake-integration/store/store_py.cpp b/mooncake-integration/store/store_py.cpp index 2970bbfd58..21b8e53286 100644 --- a/mooncake-integration/store/store_py.cpp +++ b/mooncake-integration/store/store_py.cpp @@ -16,6 +16,7 @@ #include "memory_alloc.h" #include "ssd_register_client.h" #include "device/accelerator_registry.h" +#include "mooncake_logging.h" // MC_LOG #include // for atexit #include @@ -2084,6 +2085,9 @@ PYBIND11_MODULE(store, m) { const std::string &tenant_id = "default", bool enable_client_http_server = false, int client_http_port = DEFAULT_CLIENT_HTTP_PORT) { + SpDiag::PerfPoint pt(PerfKey::STORE_PY_SETUP, + SpDiag::PerfLevel::KEY_MODULE); + pt.Start(); auto real_client = self.init_real_client(); std::shared_ptr transfer_engine = nullptr; @@ -2091,12 +2095,15 @@ PYBIND11_MODULE(store, m) { transfer_engine = engine.cast>(); } - return real_client->setup_real( + auto ret = real_client->setup_real( local_hostname, metadata_server, global_segment_size, local_buffer_size, protocol, rdma_devices, master_server_addr, transfer_engine, "", enable_ssd_offload, ssd_offload_path, tenant_id, enable_client_http_server, client_http_port); + pt.End(ret == 0 ? 0 : -1); + // MC_LOG 在下沉层 setup_real 输出（Q1b） + return ret; }, py::arg("local_hostname"), py::arg("metadata_server"), py::arg("global_segment_size"), py::arg("local_buffer_size"), @@ -2109,6 +2116,9 @@ PYBIND11_MODULE(store, m) { .def( "setup", [](MooncakeStorePyWrapper &self, const py::dict &config_dict) { + SpDiag::PerfPoint pt(PerfKey::STORE_PY_SETUP, + SpDiag::PerfLevel::KEY_MODULE); + pt.Start(); auto real_client = self.init_real_client(); // Convert py::dict to ConfigDict (all values as strings) @@ -2120,8 +2130,11 @@ PYBIND11_MODULE(store, m) { } auto result = real_client->setup_internal(config); - return result.has_value() ? 0 - : static_cast(result.error()); + int ret = result.has_value() ? 0 + : static_cast(result.error()); + pt.End(ret == 0 ? 0 : -1); + // MC_LOG 在下沉层 setup_real 输出（Q1b） + return ret; }, py::arg("config"), "Setup the store with a configuration dictionary.\n" @@ -2230,8 +2243,20 @@ PYBIND11_MODULE(store, m) { .def( "remove_all", [](MooncakeStorePyWrapper &self, bool force) { + SpDiag::PerfPoint pt(PerfKey::STORE_PY_REMOVE_ALL, + SpDiag::PerfLevel::KEY_MODULE); + pt.Start(); + auto t0 = std::chrono::steady_clock::now(); py::gil_scoped_release release; - return self.store_->removeAll(force); + auto ret = self.store_->removeAll(force); + auto t1 = std::chrono::steady_clock::now(); + auto elapsed_us = std::chrono::duration_cast( + t1 - t0).count(); + pt.End(ret == 0 ? 0 : -1); + // 下沉 removeAll 不加 MC_LOG，入口层输出汇总（Q1b） + MC_LOG(INFO) << "[remove_all] elapsed_us=" << elapsed_us + << " success=" << (ret == 0 ? 1 : 0); + return ret; }, py::arg("force") = false, "Remove all objects from the store. If force=True, skip lease " @@ -2255,17 +2280,43 @@ PYBIND11_MODULE(store, m) { "batch_is_exist", [](MooncakeStorePyWrapper &self, const std::vector &keys) { + SpDiag::PerfPoint pt(PerfKey::STORE_PY_BATCH_IS_EXIST, + SpDiag::PerfLevel::KEY_MODULE); + pt.Start(); py::gil_scoped_release release; - return self.store_->batchIsExist(keys); + auto ret = self.store_->batchIsExist(keys); + pt.End(0); + // MC_LOG 在下沉层 batchIsExist 输出 per-key（Q1b） + return ret; }, py::arg("keys"), "Check if multiple objects exist. Returns list of results: 1 if " "exists, 0 if not exists, -1 if error") .def("close", [](MooncakeStorePyWrapper &self) { - if (!self.store_) return 0; + SpDiag::PerfPoint pt(PerfKey::STORE_PY_CLOSE, + SpDiag::PerfLevel::KEY_MODULE); + pt.Start(); + auto t0 = std::chrono::steady_clock::now(); + if (!self.store_) { + auto t1 = std::chrono::steady_clock::now(); + auto elapsed_us = std::chrono::duration_cast( + t1 - t0).count(); + pt.End(0); + // 无下沉 MC_LOG，入口层输出汇总（Q1b） + MC_LOG(INFO) << "[close] elapsed_us=" << elapsed_us + << " success=1"; + return 0; + } int rc = self.store_->tearDownAll(); self.store_.reset(); + auto t1 = std::chrono::steady_clock::now(); + auto elapsed_us = std::chrono::duration_cast( + t1 - t0).count(); + pt.End(rc == 0 ? 0 : -1); + // 下沉 tearDownAll 不加 MC_LOG，入口层输出汇总（Q1b） + MC_LOG(INFO) << "[close] elapsed_us=" << elapsed_us + << " success=" << (rc == 0 ? 1 : 0); return rc; }) .def("health_check", &MooncakeStorePyWrapper::health_check, @@ -2649,10 +2700,16 @@ PYBIND11_MODULE(store, m) { "register_buffer", [](MooncakeStorePyWrapper &self, uintptr_t buffer_ptr, size_t size) { + SpDiag::PerfPoint pt(PerfKey::STORE_PY_REGISTER_BUFFER, + SpDiag::PerfLevel::KEY_MODULE); + pt.Start(); // Register memory buffer for RDMA operations void *buffer = reinterpret_cast(buffer_ptr); py::gil_scoped_release release; - return self.store_->register_buffer(buffer, size); + auto ret = self.store_->register_buffer(buffer, size); + pt.End(ret == 0 ? 0 : -1); + // MC_LOG 在下沉层 register_buffer 输出汇总（Q1b） + return ret; }, py::arg("buffer_ptr"), py::arg("size"), "Register a memory buffer for direct access operations") @@ -2897,13 +2954,20 @@ PYBIND11_MODULE(store, m) { const std::vector> &all_buffer_ptrs, const std::vector> &all_sizes, const ReplicateConfig &config = ReplicateConfig{}) { + SpDiag::PerfPoint pt(PerfKey::STORE_PY_BATCH_PUT_MULTI, + SpDiag::PerfLevel::KEY_MODULE); + pt.Start(); if (!self.is_client_initialized()) { LOG(ERROR) << "Client is not initialized"; + pt.End(-1); return std::vector{}; } py::gil_scoped_release release; - return self.store_->batch_put_from_multi_buffers( + auto ret = self.store_->batch_put_from_multi_buffers( keys, CastAddrs2Ptrs(all_buffer_ptrs), all_sizes, config); + pt.End(ret.empty() ? -1 : 0); + // MC_LOG 在下沉层 *_internal 输出 汇总+per-key（Q1b） + return ret; }, py::arg("keys"), py::arg("all_buffer_ptrs"), py::arg("all_sizes"), py::arg("config") = ReplicateConfig{}, @@ -2917,10 +2981,16 @@ PYBIND11_MODULE(store, m) { const std::vector> &all_buffer_ptrs, const std::vector> &all_sizes, bool prefer_alloc_in_same_node = false) { + SpDiag::PerfPoint pt(PerfKey::STORE_PY_BATCH_GET_INTO_MULTI, + SpDiag::PerfLevel::KEY_MODULE); + pt.Start(); py::gil_scoped_release release; - return self.store_->batch_get_into_multi_buffers( + auto ret = self.store_->batch_get_into_multi_buffers( keys, CastAddrs2Ptrs(all_buffer_ptrs), all_sizes, prefer_alloc_in_same_node); + pt.End(ret.empty() ? -1 : 0); + // MC_LOG 在下沉层 *_internal 输出 汇总+per-key（Q1b） + return ret; }, py::arg("keys"), py::arg("all_buffer_ptrs"), py::arg("all_sizes"), py::arg("prefer_alloc_in_same_node") = false, @@ -2938,8 +3008,14 @@ PYBIND11_MODULE(store, m) { "batch_get_replica_desc", [](MooncakeStorePyWrapper &self, const std::vector &keys) { + SpDiag::PerfPoint pt(PerfKey::STORE_PY_BATCH_GET_REPLICA_DESC, + SpDiag::PerfLevel::KEY_MODULE); + pt.Start(); py::gil_scoped_release release; - return self.store_->batch_get_replica_desc(keys); + auto ret = self.store_->batch_get_replica_desc(keys); + pt.End(0); + // MC_LOG 在下沉层 batch_get_replica_desc 输出 per-key（Q1b） + return ret; }, py::arg("keys")) .def( diff --git a/mooncake-store/benchmarks/CMakeLists.txt b/mooncake-store/benchmarks/CMakeLists.txt index 003f0b8724..817230b940 100644 --- a/mooncake-store/benchmarks/CMakeLists.txt +++ b/mooncake-store/benchmarks/CMakeLists.txt @@ -38,6 +38,16 @@ target_link_libraries( stress_cluster_bench PRIVATE mooncake_store transfer_engine asio_shared gflags::gflags glog::glog pthread) +# Benchmark for vLLM Store Connector path +# Triggers: batch_put_from_multi_buffers / batchIsExist / +# batch_get_into_multi_buffers (and setup/register_buffer/tearDownAll). +# Used to verify SpDiag perf points and MC_LOG logging. +add_executable(store_connector_bench store_connector_bench.cpp) +target_link_libraries( + store_connector_bench PRIVATE mooncake_store transfer_engine asio_shared + gflags::gflags glog::glog pthread) + + # Benchmark for RealClient::get_into_ranges with configurable value size, # fragments per key and keys per query. add_executable(stress_cluster_ranges_bench stress_cluster_ranges_bench.cpp) diff --git a/mooncake-store/benchmarks/store_connector_bench.cpp b/mooncake-store/benchmarks/store_connector_bench.cpp new file mode 100644 index 0000000000..31a1b88b14 --- /dev/null +++ b/mooncake-store/benchmarks/store_connector_bench.cpp @@ -0,0 +1,495 @@ +// Store Connector Benchmark: 专测 vLLM 调用路径（batch_put_from_multi_buffers +// / batchIsExist / batch_get_into_multi_buffers），触发新增的 SpDiag 打点。 +// 参考方案文档 vllm_spdiag_logging_plan.md 第 6 章。 + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include "gflags/gflags.h" +#include "glog/logging.h" +#include "mooncake_logging.h" +#include "real_client.h" + +namespace { +constexpr size_t KB = 1024; +constexpr size_t MB = 1024 * KB; +constexpr size_t GB = 1024 * MB; + +using Clock = std::chrono::steady_clock; +using Nanos = std::chrono::nanoseconds; + +inline int64_t ElapsedNanos(Clock::time_point t0, Clock::time_point t1) { + return std::chrono::duration_cast(t1 - t0).count(); +} +inline double NanosToUs(int64_t ns) { return static_cast(ns) / 1000.0; } +inline double NanosToSec(int64_t ns) { + return static_cast(ns) / 1e9; +} + +static std::string FormatBytes(size_t bytes) { + if (bytes == 0) return "0 B"; + const char* units[] = {"B", "KB", "MB", "GB", "TB"}; + int i = static_cast(std::floor(std::log2(bytes) / 10)); + if (i > 4) i = 4; + double val = static_cast(bytes) / std::pow(1024, i); + std::ostringstream oss; + oss << std::fixed << std::setprecision(2) << val << " " << units[i]; + return oss.str(); +} +} // namespace + +DEFINE_string(local_hostname, "localhost", "Local hostname"); +DEFINE_string( + metadata_server, "http://127.0.0.1:8080/metadata", + "Metadata server URL (e.g. http://127.0.0.1:8080/metadata or etcd://...)"); +DEFINE_string(master_server, "127.0.0.1:50051", "Master server address"); +DEFINE_string(protocol, "tcp", "Transport protocol: tcp, rdma, ub"); +DEFINE_string(device_name, "", "RDMA/UB device name (comma-separated)"); +DEFINE_uint64(global_segment_size, 4 * GB, "Global segment size in bytes"); +DEFINE_uint64(local_buffer_size, 512 * MB, "Local buffer size in bytes"); +DEFINE_bool(enable_ssd_offload, false, "Enable SSD offload on this client"); +DEFINE_string(ssd_offload_path, "", "SSD offload directory path"); + +DEFINE_string(scenario, "all", + "Benchmark scenario: write, is_exist, get, all"); +DEFINE_uint64(num_requests, 100, + "Number of requests (each has num_layers keys)"); +DEFINE_uint64(num_layers, 32, + "Layers per request (simulate vLLM transformer)"); +DEFINE_uint64(layer_size, 1 * MB, "Size of each layer buffer in bytes"); +DEFINE_uint64(num_threads, 1, "Number of concurrent threads"); +DEFINE_uint64(warmup_requests, 5, "Warmup requests (not counted)"); + +enum class Phase { WRITE, IS_EXIST, GET }; + +static std::string PhaseName(Phase p) { + switch (p) { + case Phase::WRITE: + return "WRITE [batch_put_from_multi_buffers]"; + case Phase::IS_EXIST: + return "IS_EXIST [batchIsExist]"; + case Phase::GET: + return "GET [batch_get_into_multi_buffers]"; + } + return "UNKNOWN"; +} + +struct ThreadResult { + std::vector latencies_ns; + size_t total_bytes = 0; + size_t total_keys = 0; + size_t total_queries = 0; + size_t failed_ops = 0; +}; + +class BenchmarkStats { + public: + void InitThreads(size_t n) { thread_results_.resize(n); } + ThreadResult& GetThreadResult(size_t tid) { return thread_results_[tid]; } + void StartTimer() { start_ = Clock::now(); } + void StopTimer() { end_ = Clock::now(); } + double WallSeconds() const { return NanosToSec(ElapsedNanos(start_, end_)); } + + void Finalize() { + merged_latencies_ns_.clear(); + total_bytes_ = total_keys_ = total_queries_ = total_failed_ = 0; + for (auto& tr : thread_results_) { + merged_latencies_ns_.insert(merged_latencies_ns_.end(), + tr.latencies_ns.begin(), + tr.latencies_ns.end()); + total_bytes_ += tr.total_bytes; + total_keys_ += tr.total_keys; + total_queries_ += tr.total_queries; + total_failed_ += tr.failed_ops; + } + std::sort(merged_latencies_ns_.begin(), merged_latencies_ns_.end()); + } + + double PercentileUs(double p) const { + if (merged_latencies_ns_.empty()) return 0.0; + double rank = (p / 100.0) * (merged_latencies_ns_.size() - 1); + size_t lo = static_cast(rank); + size_t hi = std::min(lo + 1, merged_latencies_ns_.size() - 1); + double frac = rank - lo; + int64_t ns_val = static_cast( + merged_latencies_ns_[lo] * (1.0 - frac) + + merged_latencies_ns_[hi] * frac); + return NanosToUs(ns_val); + } + + double MeanLatencyUs() const { + if (merged_latencies_ns_.empty()) return 0.0; + double sum = static_cast(std::accumulate( + merged_latencies_ns_.begin(), merged_latencies_ns_.end(), + int64_t(0))); + return NanosToUs(sum / static_cast(merged_latencies_ns_.size())); + } + + double ThroughputMBps() const { + double wall = WallSeconds(); + return (wall > 0) ? (static_cast(total_bytes_) / MB) / wall : 0; + } + + double KeysPerSec() const { + double wall = WallSeconds(); + return (wall > 0) ? static_cast(total_keys_) / wall : 0; + } + + void Print(const std::string& title, bool show_bandwidth) const { + std::cout << "\n========================================" + "========================================\n"; + std::cout << " " << title << "\n"; + std::cout << "========================================" + "========================================\n"; + std::cout << std::fixed << std::setprecision(2); + std::cout << " Wall time: " << WallSeconds() << " s\n"; + std::cout << " Total queries: " << total_queries_ + << " (failed: " << total_failed_ << ")\n"; + std::cout << " Total keys: " << total_keys_ << "\n"; + if (show_bandwidth) { + std::cout << " Total data: " << FormatBytes(total_bytes_) + << "\n"; + std::cout << " Throughput: " << ThroughputMBps() + << " MB/s"; + if (ThroughputMBps() > 1024) + std::cout << " (" << ThroughputMBps() / 1024 << " GB/s)"; + std::cout << "\n"; + } + std::cout << " Keys/sec: " << KeysPerSec() << "\n"; + + if (!merged_latencies_ns_.empty()) { + size_t n = merged_latencies_ns_.size(); + std::cout << "\n Latency (us) [n=" << n << ", per-query]\n"; + std::cout << " Min: " << std::setw(12) + << NanosToUs(merged_latencies_ns_.front()) << "\n"; + std::cout << " Avg: " << std::setw(12) << MeanLatencyUs() + << "\n"; + std::cout << " P50: " << std::setw(12) << PercentileUs(50) + << "\n"; + std::cout << " P90: " << std::setw(12) << PercentileUs(90) + << "\n"; + std::cout << " P99: " << std::setw(12) << PercentileUs(99) + << "\n"; + std::cout << " Max: " << std::setw(12) + << NanosToUs(merged_latencies_ns_.back()) << "\n"; + } + std::cout << "========================================" + "========================================\n\n"; + } + + private: + std::vector thread_results_; + std::vector merged_latencies_ns_; + size_t total_bytes_ = 0, total_keys_ = 0, total_queries_ = 0, + total_failed_ = 0; + Clock::time_point start_, end_; +}; + +class StoreConnectorBench { + public: + StoreConnectorBench() : client_(mooncake::RealClient::create()) {} + + ~StoreConnectorBench() { + for (auto& tb : thread_buffers_) { + if (tb.ptr) { + try { + client_->unregister_buffer(tb.ptr); + } catch (...) { + LOG(WARNING) + << "Failed to unregister thread buffer, ignoring"; + } + numa_free(tb.ptr, tb.size); + } + } + if (main_buffer_) { + try { + client_->unregister_buffer(main_buffer_); + } catch (...) { + LOG(WARNING) << "Failed to unregister main buffer, ignoring"; + } + numa_free(main_buffer_, main_buffer_size_); + } + } + + int Setup() { + int ret = client_->setup_real( + FLAGS_local_hostname, FLAGS_metadata_server, + FLAGS_global_segment_size, FLAGS_local_buffer_size, + FLAGS_protocol, FLAGS_device_name, FLAGS_master_server, nullptr, + "", FLAGS_enable_ssd_offload, FLAGS_ssd_offload_path); + if (ret != 0) { + LOG(ERROR) << "RealClient setup_real failed, ret=" << ret; + return ret; + } + LOG(INFO) << "RealClient setup succeeded"; + + // 主 buffer 用于 PrepareData 阶段 + main_buffer_size_ = FLAGS_num_layers * FLAGS_layer_size; + main_buffer_ = reinterpret_cast( + numa_alloc_local(main_buffer_size_)); + if (!main_buffer_) { + LOG(ERROR) << "Failed to allocate main buffer"; + return -1; + } + memset(main_buffer_, 0xAB, main_buffer_size_); // 填充测试数据 + ret = client_->register_buffer(main_buffer_, main_buffer_size_); + if (ret != 0) { + LOG(ERROR) << "register_buffer failed for main buffer"; + return ret; + } + + return AllocateThreadBuffers(FLAGS_num_threads); + } + + // is_exist/get 模式前先写入数据（不计入统计） + int PrepareData() { + LOG(INFO) << "Preparing data: writing " << FLAGS_num_requests + << " requests..."; + mooncake::ReplicateConfig config; + config.replica_num = 1; + + for (size_t r = 0; r < FLAGS_num_requests; ++r) { + auto keys = MakeRequestKeys(r); + auto all_buffers = MakeBufferList(main_buffer_); + auto all_sizes = MakeSizeList(); + auto ret = client_->batch_put_from_multi_buffers( + keys, all_buffers, all_sizes, config); + for (int v : ret) + if (v != 0) { + LOG(ERROR) << "PrepareData failed at request " << r; + return -1; + } + if ((r + 1) % 20 == 0) + LOG(INFO) << " Prepared " << (r + 1) << "/" + << FLAGS_num_requests; + } + LOG(INFO) << "Data preparation complete"; + return 0; + } + + int RunPhase(Phase phase, bool is_warmup) { + BenchmarkStats stats; + stats.InitThreads(FLAGS_num_threads); + stats.StartTimer(); + + std::latch start_latch(static_cast(FLAGS_num_threads)); + std::latch done_latch(static_cast(FLAGS_num_threads)); + + size_t total = FLAGS_num_requests; + std::vector threads; + for (size_t t = 0; t < FLAGS_num_threads; ++t) { + size_t my = total / FLAGS_num_threads + + (t < total % FLAGS_num_threads ? 1 : 0); + size_t offset = t * (total / FLAGS_num_threads) + + std::min(t, total % FLAGS_num_threads); + threads.emplace_back([&, t, my, offset]() { + PhaseWorker(t, my, offset, phase, stats, start_latch, + done_latch); + }); + } + done_latch.wait(); + stats.StopTimer(); + for (auto& th : threads) th.join(); + stats.Finalize(); + + if (!is_warmup) { + bool show_bw = (phase != Phase::IS_EXIST); + stats.Print("BENCHMARK " + PhaseName(phase), show_bw); + } + return 0; + } + + int Run() { + // Warmup（所有模式都 warmup get，确保连接建立） + if (FLAGS_warmup_requests > 0 && FLAGS_scenario != "write") { + LOG(INFO) << "Warmup: " << FLAGS_warmup_requests << " requests"; + size_t saved = FLAGS_num_requests; + FLAGS_num_requests = FLAGS_warmup_requests; + RunPhase(Phase::GET, true); + FLAGS_num_requests = saved; + } + + if (FLAGS_scenario == "all") { + RunPhase(Phase::WRITE, false); + RunPhase(Phase::IS_EXIST, false); + RunPhase(Phase::GET, false); + } else if (FLAGS_scenario == "write") { + RunPhase(Phase::WRITE, false); + } else if (FLAGS_scenario == "is_exist") { + if (PrepareData() != 0) return -1; + RunPhase(Phase::IS_EXIST, false); + } else if (FLAGS_scenario == "get") { + if (PrepareData() != 0) return -1; + RunPhase(Phase::GET, false); + } else { + LOG(ERROR) << "Unknown scenario: " << FLAGS_scenario; + return -1; + } + return 0; + } + + private: + static std::vector MakeRequestKeys(size_t req_id) { + std::vector keys; + keys.reserve(FLAGS_num_layers); + for (size_t l = 0; l < FLAGS_num_layers; ++l) + keys.push_back("layer." + std::to_string(l) + ".req_" + + std::to_string(req_id)); + return keys; + } + + // 每 key 对应 1 个 buffer（vLLM 场景），从大 buffer 切片 + std::vector> MakeBufferList(char* base) { + std::vector> all_buffers(FLAGS_num_layers); + for (size_t l = 0; l < FLAGS_num_layers; ++l) + all_buffers[l] = {base + l * FLAGS_layer_size}; + return all_buffers; + } + + std::vector> MakeSizeList() { + return std::vector>( + FLAGS_num_layers, {static_cast(FLAGS_layer_size)}); + } + + void PhaseWorker(size_t tid, size_t my_requests, size_t offset, + Phase phase, BenchmarkStats& stats, + std::latch& start_latch, std::latch& done_latch) { + ThreadResult& result = stats.GetThreadResult(tid); + result.latencies_ns.reserve(my_requests); + char* my_buf = thread_buffers_[tid].ptr; + + mooncake::ReplicateConfig config; + config.replica_num = 1; + + start_latch.arrive_and_wait(); + + size_t bytes_per_req = FLAGS_num_layers * FLAGS_layer_size; + + for (size_t i = 0; i < my_requests; ++i) { + size_t req_id = offset + i; + auto keys = MakeRequestKeys(req_id); + + auto t0 = Clock::now(); + std::vector ret; + + if (phase == Phase::WRITE) { + auto bufs = MakeBufferList(my_buf); + auto sizes = MakeSizeList(); + ret = client_->batch_put_from_multi_buffers(keys, bufs, sizes, + config); + } else if (phase == Phase::IS_EXIST) { + ret = client_->batchIsExist(keys); + } else { // GET + auto bufs = MakeBufferList(my_buf); + auto sizes = MakeSizeList(); + ret = client_->batch_get_into_multi_buffers(keys, bufs, sizes, + false); + } + auto t1 = Clock::now(); + result.latencies_ns.push_back(ElapsedNanos(t0, t1)); + + bool ok = true; + for (int v : ret) { + if (phase == Phase::IS_EXIST) { + // batchIsExist: 1=存在(成功), 0=不存在(失败) + if (v != 1) ok = false; + } else if (phase == Phase::GET) { + // batch_get: >0=字节数(成功), <0=错误码(失败) + if (v <= 0) ok = false; + } else { + // WRITE: 0=成功, 非0=失败 + if (v != 0) ok = false; + } + } + if (ok) { + result.total_keys += FLAGS_num_layers; + if (phase != Phase::IS_EXIST) + result.total_bytes += bytes_per_req; + } else { + result.failed_ops++; + } + result.total_queries++; + } + + done_latch.arrive_and_wait(); + } + + int AllocateThreadBuffers(size_t num_threads) { + thread_buffers_.resize(num_threads); + size_t per_buf = FLAGS_num_layers * FLAGS_layer_size; + for (size_t t = 0; t < num_threads; ++t) { + thread_buffers_[t].size = per_buf; + thread_buffers_[t].ptr = + reinterpret_cast(numa_alloc_local(per_buf)); + if (!thread_buffers_[t].ptr) { + LOG(ERROR) << "Failed to allocate buffer for thread " << t; + return -1; + } + memset(thread_buffers_[t].ptr, 0, per_buf); + int ret = + client_->register_buffer(thread_buffers_[t].ptr, per_buf); + if (ret != 0) { + LOG(ERROR) << "register_buffer failed for thread " << t; + return ret; + } + } + LOG(INFO) << "Allocated " << num_threads << " thread buffers, each " + << FormatBytes(per_buf); + return 0; + } + + std::shared_ptr client_; + char* main_buffer_ = nullptr; + size_t main_buffer_size_ = 0; + struct ThreadBuf { + char* ptr = nullptr; + size_t size = 0; + }; + std::vector thread_buffers_; +}; + +int main(int argc, char* argv[]) { + if (!google::IsGoogleLoggingInitialized()) { + google::InitGoogleLogging(argv[0]); + } + gflags::ParseCommandLineFlags(&argc, &argv, true); + + if (std::getenv("MC_LOG_DIR") == nullptr) { + FLAGS_logtostderr = true; + } + mooncake::logging::ApplyMooncakeLogEnableToGlog(); + + LOG(INFO) << "Mooncake Store Connector Benchmark (vLLM path)"; + LOG(INFO) << " Scenario: " << FLAGS_scenario; + LOG(INFO) << " Protocol: " << FLAGS_protocol; + LOG(INFO) << " Requests: " << FLAGS_num_requests; + LOG(INFO) << " Layers/req: " << FLAGS_num_layers; + LOG(INFO) << " Layer size: " << FormatBytes(FLAGS_layer_size); + LOG(INFO) << " Threads: " << FLAGS_num_threads; + size_t total_data = + FLAGS_num_requests * FLAGS_num_layers * FLAGS_layer_size; + LOG(INFO) << " Total data: " << FormatBytes(total_data); + + StoreConnectorBench bench; + int ret = bench.Setup(); + if (ret != 0) { + LOG(ERROR) << "Setup failed"; + return ret; + } + return bench.Run(); +} diff --git a/mooncake-store/benchmarks/stress_cluster_bench.cpp b/mooncake-store/benchmarks/stress_cluster_bench.cpp index 5da134b429..fcd4e445ce 100644 --- a/mooncake-store/benchmarks/stress_cluster_bench.cpp +++ b/mooncake-store/benchmarks/stress_cluster_bench.cpp @@ -190,7 +190,7 @@ DEFINE_string(ssd_offload_path, "", "SSD offload directory path"); DEFINE_string(scenario, "local_memory", "Benchmark scenario: local_memory, remote_memory, local_disk, " - "remote_disk, segment_write, segment_read"); + "remote_disk, segment_write, segment_read, remove, batch_remove"); DEFINE_string(role, "writer", "Node role: writer (prefill data) or reader (benchmark reads)"); DEFINE_uint64(value_size, 4 * MB, "Size of each value in bytes"); @@ -719,6 +719,21 @@ class StressBenchmark { return "seg_" + sanitized + "_key_" + std::to_string(idx); } + // Key generator for remove/batch_remove scenarios. Uses a distinct + // "rmv_" prefix so remove keys never collide with segment_write keys + // ("seg_"). This allows remove and write benchmarks to run independently + // without cross-contamination. + static std::string MakeRemoveKey(const std::string& segment, size_t idx) { + static const char* kSpecialChars = ".:-/\\[]{}()@#$%^&*+=|<>,;!?`'\"~"; + std::string sanitized = segment; + for (char& c : sanitized) { + if (std::strchr(kSpecialChars, c) != nullptr || std::isspace(c)) { + c = '_'; + } + } + return "rmv_" + sanitized + "_key_" + std::to_string(idx); + } + int RunSegmentWrite() { auto segments = DiscoverSegmentsIfNeeded( "--segments not specified, auto-discovering"); @@ -1253,6 +1268,113 @@ class StressBenchmark { return 0; } + int RunSegmentRemove(bool use_batch) { + auto segments = DiscoverSegmentsIfNeeded( + "--segments not specified, auto-discovering"); + if (segments.empty()) { + return -1; + } + LOG(INFO) << "Discovered " << segments.size() + << " segments from master"; + + size_t remove_segment_nums = FLAGS_read_segment_nums; + if (remove_segment_nums == 0 || + remove_segment_nums > segments.size()) { + remove_segment_nums = segments.size(); + } + std::vector remove_segments( + segments.begin(), segments.begin() + remove_segment_nums); + + LOG(INFO) << "=== SEGMENT REMOVE MODE ===" + << (use_batch ? " (batch)" : " (single key)"); + LOG(INFO) << "Removing from " << remove_segment_nums << " segments (" + << remove_segment_nums << " nodes)"; + for (size_t s = 0; s < remove_segments.size(); ++s) { + LOG(INFO) << " Segment [" << s << "]: " << remove_segments[s]; + } + LOG(INFO) << "Keys per segment: " << FLAGS_num_keys; + LOG(INFO) << "Batch size: " << FLAGS_batch_size; + + if (FLAGS_duration > 0) { + LOG(WARNING) << "--duration is ignored for remove scenarios: " + << "removal is not idempotent, a single pass is used"; + } + + // Phase 1: prefill each segment. Key layout and preferred_segments + // pinning mirror RunSegmentWrite exactly. + std::vector configs(remove_segments.size()); + for (size_t s = 0; s < remove_segments.size(); ++s) { + configs[s].replica_num = FLAGS_replica_num; + configs[s].with_hard_pin = FLAGS_hard_pin; + configs[s].preferred_segments = {remove_segments[s]}; + } + + LOG(INFO) << "Phase 1: Prefilling " << FLAGS_num_keys + << " keys to " << remove_segment_nums + << " segments (interleaved), each " + << FLAGS_value_size / MB << " MB"; + for (size_t i = 0; i < FLAGS_num_keys; ++i) { + for (size_t s = 0; s < remove_segments.size(); ++s) { + const auto& segment = remove_segments[s]; + std::string key = MakeRemoveKey(segment, i); + FillBuffer(i); + int ret = client_->put_from(key, buffer_, FLAGS_value_size, + configs[s]); + if (ret != 0) { + LOG(ERROR) << "put_from failed for key=" << key + << " segment=" << segment << " ret=" << ret; + return ret; + } + } + if ((i + 1) % 10 == 0 || i == FLAGS_num_keys - 1) { + LOG(INFO) << " Prefilled " << (i + 1) << "/" << FLAGS_num_keys + << " keys to all " << remove_segment_nums + << " segments"; + } + } + LOG(INFO) << "Prefill phase complete"; + + // Phase 2: assemble the full key list in the same order as + // RunSegmentRead (outer key index, inner segment). + std::vector all_keys; + for (size_t i = 0; i < FLAGS_num_keys; ++i) { + for (size_t s = 0; s < remove_segments.size(); ++s) { + all_keys.push_back(MakeRemoveKey(remove_segments[s], i)); + } + } + LOG(INFO) << "Total keys to remove: " << all_keys.size(); + + // Phase 3: concurrent remove (single pass, mirrors segment_read). + LOG(INFO) << "Phase 3: Concurrent " << (use_batch ? "batch " : "") + << "remove with " << FLAGS_num_threads << " threads"; + + BenchmarkStats stats; + stats.InitThreads(FLAGS_num_threads, + all_keys.size() / FLAGS_num_threads); + stats.StartTimer(); + + std::latch start_latch(static_cast(FLAGS_num_threads)); + std::latch done_latch(static_cast(FLAGS_num_threads)); + auto threads = LaunchRemoveWorkers( + FLAGS_num_threads, all_keys.size(), stats, start_latch, done_latch, + use_batch, [&all_keys](size_t idx) { + return all_keys[idx % all_keys.size()]; + }); + + done_latch.wait(); + stats.StopTimer(); + + for (auto& th : threads) { + th.join(); + } + + stats.Finalize(); + stats.Print(use_batch ? "SEGMENT BATCH REMOVE BENCHMARK" + : "SEGMENT REMOVE BENCHMARK"); + + return 0; + } + int RunListSegments() { LOG(INFO) << "Discovering segments from master at " << FLAGS_master_server << ":" << FLAGS_master_admin_port; @@ -1300,6 +1422,10 @@ class StressBenchmark { return RunSegmentRead(); } else if (FLAGS_scenario == "list_segments") { return RunListSegments(); + } else if (FLAGS_scenario == "remove") { + return RunSegmentRemove(false); + } else if (FLAGS_scenario == "batch_remove") { + return RunSegmentRemove(true); } else if (FLAGS_scenario == "remote_memory" || FLAGS_scenario == "remote_disk") { if (FLAGS_role == "writer") { @@ -1473,6 +1599,108 @@ class StressBenchmark { return threads; } + void RemoveWorker(size_t tid, size_t my_keys, size_t key_offset, + BenchmarkStats& stats, std::latch& start_latch, + std::latch& done_latch, bool use_batch, + const std::function& key_func) { + bindToSocket(tid % NR_SOCKETS); + + ThreadResult& result = stats.GetThreadResult(tid); + result.latencies_ns.reserve(my_keys); + + start_latch.arrive_and_wait(); + + size_t keys = 0; + size_t queries = 0; + size_t failed = 0; + size_t bytes = 0; + + if (!use_batch) { + for (size_t i = 0; i < my_keys; ++i) { + size_t key_idx = key_offset + i; + std::string key = key_func(key_idx); + + auto t0 = Clock::now(); + int ret = client_->remove(key, /*force=*/true); + auto t1 = Clock::now(); + + int64_t lat_ns = ElapsedNanos(t0, t1); + result.latencies_ns.push_back(lat_ns); + + if (ret != 0) { + ++failed; + LOG_EVERY_N(ERROR, 100) + << "remove failed key=" << key << " ret=" << ret; + } else { + // Account removed payload as transferred bytes so throughput + // stats stay comparable with the read benchmark. + bytes += FLAGS_value_size; + } + ++keys; + ++queries; + } + } else { + size_t per_key_buf = FLAGS_value_size; + size_t i = 0; + while (i < my_keys) { + std::vector key_list; + size_t batch_end = std::min(i + FLAGS_batch_size, my_keys); + key_list.reserve(batch_end - i); + + for (size_t j = i; j < batch_end; ++j) { + size_t key_idx = key_offset + j; + key_list.push_back(key_func(key_idx)); + } + + auto t0 = Clock::now(); + auto results = client_->batchRemove(key_list, /*force=*/true); + auto t1 = Clock::now(); + + int64_t lat_ns = ElapsedNanos(t0, t1); + result.latencies_ns.push_back(lat_ns); + + for (size_t k = 0; k < results.size(); ++k) { + if (results[k] != 0) { + ++failed; + } else { + bytes += per_key_buf; + } + ++keys; + } + ++queries; + + i = batch_end; + } + } + + result.total_bytes = bytes; + result.total_keys = keys; + result.total_queries = queries; + result.failed_ops = failed; + + done_latch.arrive_and_wait(); + } + + std::vector LaunchRemoveWorkers( + size_t num_threads, size_t total_keys, BenchmarkStats& stats, + std::latch& start_latch, std::latch& done_latch, bool use_batch, + const std::function& key_func) { + std::vector threads; + size_t keys_per_thread = total_keys / num_threads; + size_t remainder = total_keys % num_threads; + + for (size_t t = 0; t < num_threads; ++t) { + size_t my_keys = keys_per_thread + (t < remainder ? 1 : 0); + size_t key_offset = t * keys_per_thread + std::min(t, remainder); + + threads.emplace_back([&, t, my_keys, key_offset, use_batch]() { + RemoveWorker(t, my_keys, key_offset, stats, start_latch, + done_latch, use_batch, key_func); + }); + } + return threads; + } + std::vector DiscoverSegmentsIfNeeded( const std::string& context) { auto segments = ParseSegments(); diff --git a/mooncake-store/include/client_service.h b/mooncake-store/include/client_service.h index f4ab18fccc..631e4328bd 100644 --- a/mooncake-store/include/client_service.h +++ b/mooncake-store/include/client_service.h @@ -454,6 +454,18 @@ class Client { virtual tl::expected PromotionObjectHeartbeat( std::vector& promotion_objects); + /** + * @brief Drain the removed_keys queue from master. Returns {tenant_id, + * key} pairs that were removed via Remove/BatchRemove and had LOCAL_DISK + * replicas on this client. The caller should MarkRemoved each key to + * trigger SSD tombstone + GC compaction. + */ + [[nodiscard]] tl::expected, ErrorCode> + RemoveObjectHeartbeat(const UUID& client_id); + + tl::expected AckRemoveObjectHeartbeat( + const UUID& client_id, const std::vector& tasks); + /** * @brief Stage a PROCESSING MEMORY replica for an existing key during * L2->L1 promotion. Returns the new replica's descriptor that the caller diff --git a/mooncake-store/include/file_storage.h b/mooncake-store/include/file_storage.h index d76d50145b..2439a2a082 100644 --- a/mooncake-store/include/file_storage.h +++ b/mooncake-store/include/file_storage.h @@ -51,6 +51,13 @@ class FileStorage { */ bool ReleaseBuffer(uint64_t batch_id); + // Forward explicit-delete tombstone to the storage backend. + // For BucketStorageBackend: marks tombstone + enables GC. + // For other backends: no-op (default in StorageBackendInterface). + tl::expected MarkRemoved(const std::string& key); + tl::expected BatchMarkRemoved( + const std::vector& keys); + private: friend class FileStorageTest; friend class FileStoragePromotionTest; diff --git a/mooncake-store/include/master_client.h b/mooncake-store/include/master_client.h index 8210b7d882..be6e1b9c24 100644 --- a/mooncake-store/include/master_client.h +++ b/mooncake-store/include/master_client.h @@ -493,6 +493,12 @@ class MasterClient { [[nodiscard]] tl::expected, ErrorCode> PromotionObjectHeartbeat(const UUID& client_id); + /** Fetch pending remove tasks without removing them from the queue. */ + [[nodiscard]] tl::expected, ErrorCode> + RemoveObjectHeartbeat(const UUID& client_id); + tl::expected AckRemoveObjectHeartbeat( + const UUID& client_id, const std::vector& tasks); + /** * @brief Stage a PROCESSING MEMORY replica for an existing key during * promotion. Returns the new replica's descriptor that the caller writes diff --git a/mooncake-store/include/master_service.h b/mooncake-store/include/master_service.h index 9b758ea7ff..f1f14e1990 100644 --- a/mooncake-store/include/master_service.h +++ b/mooncake-store/include/master_service.h @@ -727,6 +727,13 @@ class MasterService { auto PromotionObjectHeartbeat(const UUID& client_id) -> tl::expected, ErrorCode>; + /** Fetch pending remove tasks without removing them from the queue. */ + auto RemoveObjectHeartbeat(const UUID& client_id) + -> tl::expected, ErrorCode>; + auto AckRemoveObjectHeartbeat( + const UUID& client_id, const std::vector& tasks) + -> tl::expected; + /** * @brief Stage a PROCESSING MEMORY replica for an existing key. Allocates * DRAM via the existing AllocationStrategy, optionally biased toward the @@ -1533,6 +1540,8 @@ class MasterService { void FinalizeRemovedReplicasAfterDurable( const OpLogEntry& durable_entry, const std::vector& replica_ids, QuotaEraseMode quota_mode); + void EnqueueRemoveTasks(const std::vector& holder_ids, + const RemoveTaskItem& task); void FinalizeMetadataEraseAfterDurable(const OpLogEntry& durable_entry, QuotaEraseMode quota_mode); void FinalizeExpiredProcessingReplicasAfterDurable( diff --git a/mooncake-store/include/rpc_service.h b/mooncake-store/include/rpc_service.h index dc38a091e2..7abf3a34e1 100644 --- a/mooncake-store/include/rpc_service.h +++ b/mooncake-store/include/rpc_service.h @@ -219,6 +219,10 @@ class WrappedMasterService { tl::expected PollRemoveAll(const UUID& client_id); + tl::expected, ErrorCode> RemoveObjectHeartbeat( + const UUID& client_id); + tl::expected AckRemoveObjectHeartbeat( + const UUID& client_id, const std::vector& tasks); tl::expected ReportSsdCapacity( const UUID& client_id, int64_t ssd_total_capacity_bytes); diff --git a/mooncake-store/include/segment.h b/mooncake-store/include/segment.h index 75b9b0c0f4..0fc34a841a 100644 --- a/mooncake-store/include/segment.h +++ b/mooncake-store/include/segment.h @@ -100,6 +100,12 @@ struct LocalDiskSegment { // offloading_objects (offloading_mutex_). std::unordered_map GUARDED_BY( offloading_mutex_) promotion_objects; + // Keys removed via Remove/BatchRemove that had LOCAL_DISK replicas on + // this client. Populated by master's Remove when the key has a + // LOCAL_DISK replica. Drained by RemoveObjectHeartbeat RPC. Same locking as + // offloading_objects (offloading_mutex_). + std::vector GUARDED_BY( + offloading_mutex_) removed_keys; // Set by master's RemoveAll. When the client sees this flag via // PollRemoveAll, it calls FileStorage::RemoveAll() to physically // delete all SSD files. Same locking as offloading_objects diff --git a/mooncake-store/include/storage_backend.h b/mooncake-store/include/storage_backend.h index d9863896d0..b417c3c950 100644 --- a/mooncake-store/include/storage_backend.h +++ b/mooncake-store/include/storage_backend.h @@ -4,6 +4,7 @@ #include #include +#include #include #include #include @@ -12,6 +13,7 @@ #include #include #include +#include #include #include @@ -40,6 +42,10 @@ struct BucketMetadata { std::vector keys; std::vector metadatas; + // Persisted tombstones. Init filters these keys before rebuilding the + // in-memory object index. + std::vector tombstones; + // Runtime-only fields (not serialized) for safe deletion support // Tracks number of in-flight reads to enable safe bucket deletion mutable std::atomic inflight_reads_{0}; @@ -47,6 +53,17 @@ struct BucketMetadata { // Updated on every read with relaxed ordering (approximate is sufficient). mutable std::atomic last_access_ns_{0}; + // Runtime-only (not serialized): bytes marked removed via MarkRemoved. + // Drives GC candidate selection (compact when deleted_bytes_ > 0). + mutable std::atomic deleted_bytes_{0}; + // Runtime-only (not serialized): true while a GC compaction is in flight + // for this bucket, preventing re-entrant compaction. + mutable std::atomic compacting_{false}; + // Runtime-only version of the bucket's deletion state. Compaction captures + // this value with its read snapshot and validates it before publishing a + // new bucket, so a concurrent deletion cannot publish stale data. + uint64_t generation_{0}; + // Default constructor BucketMetadata() = default; @@ -56,8 +73,12 @@ struct BucketMetadata { data_size(other.data_size), keys(other.keys), metadatas(other.metadatas), + tombstones(other.tombstones), inflight_reads_(0), - last_access_ns_(0) {} + last_access_ns_(0), + deleted_bytes_(0), + compacting_(false), + generation_(0) {} // Move constructor BucketMetadata(BucketMetadata&& other) noexcept @@ -65,8 +86,12 @@ struct BucketMetadata { data_size(other.data_size), keys(std::move(other.keys)), metadatas(std::move(other.metadatas)), + tombstones(std::move(other.tombstones)), inflight_reads_(0), - last_access_ns_(0) {} + last_access_ns_(0), + deleted_bytes_(0), + compacting_(false), + generation_(0) {} // Copy assignment BucketMetadata& operator=(const BucketMetadata& other) { @@ -75,6 +100,7 @@ struct BucketMetadata { data_size = other.data_size; keys = other.keys; metadatas = other.metadatas; + tombstones = other.tombstones; // Don't copy runtime state } return *this; @@ -87,12 +113,13 @@ struct BucketMetadata { data_size = other.data_size; keys = std::move(other.keys); metadatas = std::move(other.metadatas); + tombstones = std::move(other.tombstones); // Don't move runtime state } return *this; } }; -YLT_REFL(BucketMetadata, data_size, keys, metadatas); +YLT_REFL(BucketMetadata, data_size, keys, metadatas, tombstones); /** * @brief RAII guard for tracking in-flight bucket reads. @@ -204,6 +231,22 @@ struct BucketBackendConfig { // eviction_policy. Set via // MOONCAKE_OFFLOAD_DISABLE_SSD_EVICTION. + // --- Explicit-delete-only GC config --- + // Enable background tombstone compaction GC. + bool gc_enable = true; + // GC scan interval in milliseconds. + int64_t gc_interval_ms = 1000; + // Compact a bucket when deleted bytes / bucket data size >= this ratio. + double gc_deleted_ratio = 0.25; + // Trigger GC when total_size / max_total_size >= this ratio. + double gc_high_watermark_ratio = 0.90; + // Max old buckets collected per GC round for cross-bucket merge. + int64_t gc_max_buckets_per_round = 1; + // Enable cross-bucket merge compaction (collect live keys from multiple + // tombstone buckets into one new bucket). When false, each bucket is + // compacted independently (no merge). + bool gc_merge_enable = true; + bool Validate() const; static BucketBackendConfig FromEnvironment(); @@ -412,6 +455,20 @@ class StorageBackendInterface { // Default: no-op (no test failures injected) } + // Mark a key as removed (tombstone) for explicit-delete-only GC. + // Default no-op: only BucketStorageBackend implements tombstone + GC. + // File-per-key and other backends inherit the no-op (do not delete files). + // Safe to call for keys not present in local storage (idempotent). + virtual tl::expected MarkRemoved( + const std::string& /* key */) { + return {}; + } + + // Batch variant: mark multiple keys as removed in one lock acquisition. + virtual tl::expected BatchMarkRemoved( + const std::vector& /* keys */) { + return {}; + } // Remove all persisted objects from disk. Called during RemoveAll to // clean up physical SSD files alongside master metadata deletion. virtual void RemoveAll() {} @@ -1005,11 +1062,53 @@ class BucketStorageBackend : public StorageBackendInterface { */ tl::expected DeleteBucket(int64_t bucket_id); + // Explicit-delete-only GC: mark a key as tombstone (no disk IO). + // Removes key from object_bucket_map_ (immediately invisible to + // BatchLoad/IsExist) and bumps bucket deleted_bytes_. + // Idempotent: no-op if key not in local storage. + tl::expected MarkRemoved( + const std::string& key) override; + tl::expected BatchMarkRemoved( + const std::vector& keys) override; + + // Compact a single bucket: copy-on-write live keys to a new bucket, + // atomically swap mappings, delete old bucket file after reads drain. + // Returns true on success (or no-op), false on transient failure + // (will retry next round). Public to allow explicit compaction and + // testing (analogous to DeleteBucket). + bool CompactBucket(int64_t bucket_id); + + // Compact multiple buckets into one new bucket (cross-bucket merge). + // Collects live keys from all given old buckets, groups them by + // bucket_keys_limit/bucket_size_limit, and writes ONE new bucket per + // round (the first group that fills up). If the first group doesn't + // fill a full bucket and there's no space pressure, the merge is + // deferred to the next round. Old buckets whose live keys are all + // migrated are deleted. + // Returns true on success (or deferred), false on transient failure. + bool CompactBuckets(const std::vector& bucket_ids, + bool space_pressure = false); tl::expected, ErrorCode> EvictAboveDiskWatermark( double high_watermark_ratio, double low_watermark_ratio, EvictionHandler eviction_handler = nullptr) override; private: + // --- Background GC --- + // Background GC thread entry point. + void GCThreadFunc(); + + // Wait for in-flight reads on a bucket to drain (up to 10s). + void WaitForInflightReads(std::shared_ptr bucket); + + // Delete .bucket and .meta files for a bucket_id, ignore missing. + void DeleteBucketFiles(int64_t bucket_id); + + // GC thread lifecycle members + std::atomic gc_running_{false}; + std::thread gc_thread_; + std::mutex gc_mutex_; + std::condition_variable gc_cv_; + tl::expected, ErrorCode> BuildBucket( int64_t bucket_id, const std::unordered_map>& batch_object, diff --git a/mooncake-store/include/types.h b/mooncake-store/include/types.h index b7c983786d..eb53a52cab 100644 --- a/mooncake-store/include/types.h +++ b/mooncake-store/include/types.h @@ -258,6 +258,16 @@ struct PromotionTaskItem { }; YLT_REFL(PromotionTaskItem, tenant_id, key, size); +struct RemoveTaskItem { + std::string tenant_id; + std::string key; + + bool operator==(const RemoveTaskItem& other) const { + return tenant_id == other.tenant_id && key == other.key; + } +}; +YLT_REFL(RemoveTaskItem, tenant_id, key); + // Store client configuration validation limits static constexpr size_t MIN_SEGMENT_SIZE = 1024; // 1KB static constexpr size_t MAX_SEGMENT_SIZE = 1024ULL * 1024 * 1024 * 1024; // 1TB diff --git a/mooncake-store/src/client_service.cpp b/mooncake-store/src/client_service.cpp index 501e5567fc..0dca231eca 100644 --- a/mooncake-store/src/client_service.cpp +++ b/mooncake-store/src/client_service.cpp @@ -1116,6 +1116,9 @@ std::vector> Client::BatchQuery( std::vector> Client::BatchQuery( const std::vector& object_keys, const std::string& tenant_id) { + SpDiag::PerfPoint pt(PerfKey::CLIENT_BATCH_QUERY, + SpDiag::PerfLevel::MODULE); + pt.Start(); std::chrono::steady_clock::time_point start_time = std::chrono::steady_clock::now(); auto response = master_client_.BatchGetReplicaList(object_keys, tenant_id); @@ -1130,6 +1133,7 @@ std::vector> Client::BatchQuery( for (size_t i = 0; i < object_keys.size(); ++i) { results.emplace_back(tl::unexpected(ErrorCode::RPC_FAIL)); } + pt.End(-1); return results; } std::vector> results; @@ -1145,6 +1149,7 @@ std::vector> Client::BatchQuery( results.emplace_back(tl::unexpected(response[i].error())); } } + pt.End(0); return results; } @@ -3505,6 +3510,16 @@ tl::expected Client::PromotionObjectHeartbeat( return {}; } +tl::expected, ErrorCode> +Client::RemoveObjectHeartbeat(const UUID& client_id) { + return master_client_.RemoveObjectHeartbeat(client_id); +} + +tl::expected Client::AckRemoveObjectHeartbeat( + const UUID& client_id, const std::vector& tasks) { + return master_client_.AckRemoveObjectHeartbeat(client_id, tasks); +} + tl::expected Client::PromotionAllocStart( const std::string& key, uint64_t size, diff --git a/mooncake-store/src/file_storage.cpp b/mooncake-store/src/file_storage.cpp index 02f22968c4..c637b096c6 100644 --- a/mooncake-store/src/file_storage.cpp +++ b/mooncake-store/src/file_storage.cpp @@ -543,7 +543,6 @@ tl::expected FileStorage::OffloadObjects( } buckets_keys.emplace_back(std::move(keys)); } - auto complete_handler = [this, &task_by_storage_key]( const std::vector& keys, @@ -816,6 +815,16 @@ tl::expected FileStorage::IsEnableOffloading() { return enable_offloading; } +tl::expected FileStorage::MarkRemoved( + const std::string& key) { + return storage_backend_->MarkRemoved(key); +} + +tl::expected FileStorage::BatchMarkRemoved( + const std::vector& keys) { + return storage_backend_->BatchMarkRemoved(keys); +} + tl::expected FileStorage::Heartbeat() { if (client_ == nullptr) { LOG(ERROR) << "client is nullptr"; @@ -842,6 +851,41 @@ tl::expected FileStorage::Heartbeat() { }); } + // === STEP 0: Drain removed keys from master === + // Master pushes {tenant_id, key} pairs to this client's removed_keys + // queue when a Remove/BatchRemove deletes a key that had a LOCAL_DISK + // replica here. We mark each as a tombstone so GC can reclaim SSD space. + { + auto remove_result = + client_->RemoveObjectHeartbeat(client_->getClientId()); + if (remove_result) { + bool all_marked = true; + for (const auto& item : remove_result.value()) { + auto storage_key = + TenantId(item.tenant_id).MakeScopedKey(item.key); + auto mark_result = storage_backend_->MarkRemoved(storage_key); + if (!mark_result) { + all_marked = false; + LOG(ERROR) << "Failed to persist remove tombstone: " + << mark_result.error(); + break; + } + } + if (all_marked && !remove_result.value().empty()) { + auto ack_result = client_->AckRemoveObjectHeartbeat( + client_->getClientId(), remove_result.value()); + if (!ack_result) { + LOG(ERROR) << "Failed to ACK remove tasks: " + << ack_result.error(); + } + VLOG(1) << "RemoveObjectHeartbeat processed " + << remove_result.value().size() + << " removed key(s) from master"; + } + } + // Errors are non-fatal: removed keys will be retried next heartbeat. + } + std::vector offloading_objects; // Objects selected for offloading diff --git a/mooncake-store/src/master_client.cpp b/mooncake-store/src/master_client.cpp index 7cc1b28c28..5a266a696f 100644 --- a/mooncake-store/src/master_client.cpp +++ b/mooncake-store/src/master_client.cpp @@ -249,6 +249,15 @@ template <> struct RpcNameTraits<&WrappedMasterService::PromotionObjectHeartbeat> { static constexpr const char* value = "PromotionObjectHeartbeat"; }; +template <> +struct RpcNameTraits<&WrappedMasterService::RemoveObjectHeartbeat> { + static constexpr const char* value = "RemoveObjectHeartbeat"; +}; + +template <> +struct RpcNameTraits<&WrappedMasterService::AckRemoveObjectHeartbeat> { + static constexpr const char* value = "AckRemoveObjectHeartbeat"; +}; template <> struct RpcNameTraits<&WrappedMasterService::PromotionAllocStart> { @@ -1196,6 +1205,23 @@ MasterClient::PromotionObjectHeartbeat(const UUID& client_id) { std::vector>(client_id); } +tl::expected, ErrorCode> +MasterClient::RemoveObjectHeartbeat(const UUID& client_id) { + ScopedVLogTimer timer(1, "MasterClient::RemoveObjectHeartbeat"); + timer.LogRequest("client_id=", client_id.first, ":", client_id.second); + return invoke_rpc<&WrappedMasterService::RemoveObjectHeartbeat, + std::vector>(client_id); +} + +tl::expected MasterClient::AckRemoveObjectHeartbeat( + const UUID& client_id, const std::vector& tasks) { + ScopedVLogTimer timer(1, "MasterClient::AckRemoveObjectHeartbeat"); + timer.LogRequest("client_id=", client_id.first, ":", client_id.second, + " tasks=", tasks.size()); + return invoke_rpc<&WrappedMasterService::AckRemoveObjectHeartbeat, void>( + client_id, tasks); +} + tl::expected MasterClient::PromotionAllocStart( const UUID& client_id, const std::string& key, uint64_t size, diff --git a/mooncake-store/src/master_service.cpp b/mooncake-store/src/master_service.cpp index d5660e8e74..2aef820aa5 100644 --- a/mooncake-store/src/master_service.cpp +++ b/mooncake-store/src/master_service.cpp @@ -1,6 +1,5 @@ #include "master_service.h" -#include #include #include #include @@ -1764,6 +1763,14 @@ void MasterService::FinalizeRemovedReplicasAfterDurable( const bool erased_local_disk = std::any_of( erased_replicas.begin(), erased_replicas.end(), [](const Replica& replica) { return replica.is_local_disk_replica(); }); + std::vector local_disk_holders; + for (const auto& replica : erased_replicas) { + if (!replica.is_local_disk_replica()) continue; + auto client_id = replica.get_local_disk_client_id(); + if (client_id.has_value()) { + local_disk_holders.push_back(client_id.value()); + } + } ReleaseLocalDiskUsage(erased_replicas); if (erased_local_disk) { shard.OnDiskReplicaRemoved(erased_local_disk, metadata); @@ -1774,6 +1781,11 @@ void MasterService::FinalizeRemovedReplicasAfterDurable( shard->tenants.erase(tenant_it); } } + if (erased_local_disk) { + EnqueueRemoveTasks( + local_disk_holders, + RemoveTaskItem{tenant_id.value(), durable_entry.object_key}); + } } void MasterService::FinalizeMetadataEraseAfterDurable( @@ -1810,6 +1822,7 @@ void MasterService::FinalizeExpiredProcessingReplicasAfterDurable( } auto& metadata = accessor.Get(); + auto replicas = PopReplicasWithCacheTotalAccounting( metadata, &Replica::fn_is_processing); if (!replicas.empty()) { @@ -5599,6 +5612,17 @@ auto MasterService::Remove(const std::string& key, const TenantId& tenant_id, } auto& metadata = accessor.Get(); + std::vector local_disk_holders; + metadata.VisitReplicas( + [](const Replica& replica) { + return replica.is_local_disk_replica(); + }, + [&local_disk_holders](Replica& replica) { + auto client_id = replica.get_local_disk_client_id(); + if (client_id.has_value()) { + local_disk_holders.push_back(client_id.value()); + } + }); if (!force && !metadata.IsLeaseExpired()) { VLOG(1) << "key=" << key << ", error=object_has_lease"; @@ -5636,10 +5660,15 @@ auto MasterService::Remove(const std::string& key, const TenantId& tenant_id, auto persist_result = AppendReservedOpLogWithDurableFinalize( std::move(reservation.value()), OpType::REMOVE, object_id.tenant_id.value(), key, {}, - [this, removed_ids = std::move(removed_ids)]( + [this, removed_ids = std::move(removed_ids), + local_disk_holders, + tenant_id_for_task = object_id.tenant_id.value(), key]( const OpLogEntry& durable_entry) { FinalizeRemovedReplicasAfterDurable( durable_entry, removed_ids, QuotaEraseMode::kFull); + EnqueueRemoveTasks( + local_disk_holders, + RemoveTaskItem{tenant_id_for_task, key}); }); if (!persist_result) { return tl::make_unexpected(persist_result.error()); @@ -5648,7 +5677,14 @@ auto MasterService::Remove(const std::string& key, const TenantId& tenant_id, } } PublishKvRemoved(key, metadata, object_id.tenant_id); + + // Before erasing metadata, collect LOCAL_DISK replica holders so we + // can notify them to reclaim SSD space via RemoveObjectHeartbeat. accessor.Erase(); + + // Push removed key to each LOCAL_DISK holder's removed_keys queue. + EnqueueRemoveTasks(local_disk_holders, RemoveTaskItem{tenant_id.value(), key}); + return {}; } @@ -6045,6 +6081,18 @@ auto MasterService::BatchRemove(const std::vector& keys, auto& metadata = it->second; + std::vector batch_local_disk_holders; + metadata.VisitReplicas( + [](const Replica& replica) { + return replica.is_local_disk_replica(); + }, + [&batch_local_disk_holders](Replica& replica) { + auto cid = replica.get_local_disk_client_id(); + if (cid.has_value()) { + batch_local_disk_holders.push_back(cid.value()); + } + }); + if (!force && !metadata.IsLeaseExpired(now)) { VLOG(1) << "key=" << key << ", error=object_has_lease"; results[original_idx] = @@ -6087,11 +6135,16 @@ auto MasterService::BatchRemove(const std::vector& keys, AppendReservedOpLogWithDurableFinalize( std::move(reservation.value()), OpType::REMOVE, normalized_tenant.value(), key, {}, - [this, removed_ids = std::move(removed_ids)]( + [this, removed_ids = std::move(removed_ids), + batch_local_disk_holders, + tenant_id = normalized_tenant.value(), key]( const OpLogEntry& durable_entry) { FinalizeRemovedReplicasAfterDurable( durable_entry, removed_ids, QuotaEraseMode::kFull); + EnqueueRemoveTasks( + batch_local_disk_holders, + RemoveTaskItem{tenant_id, key}); }); if (!persist_result) { results[original_idx] = @@ -6102,11 +6155,20 @@ auto MasterService::BatchRemove(const std::vector& keys, continue; } } + + // Collect LOCAL_DISK replica holders before erasing, so we + // can notify them to reclaim SSD space via RemoveObjectHeartbeat. EraseMetadata(tenant_state, it, normalized_tenant, QuotaEraseMode::kFull, &shard); if (tenant_state.Empty()) { shard->tenants.erase(tenant_it); } + + // Push removed key to each LOCAL_DISK holder's removed_keys queue. + EnqueueRemoveTasks( + batch_local_disk_holders, + RemoveTaskItem{normalized_tenant.value(), key}); + results[original_idx] = {}; // Success } } @@ -6331,6 +6393,64 @@ auto MasterService::PollRemoveAll(const UUID& client_id) return result; } +auto MasterService::RemoveObjectHeartbeat(const UUID& client_id) + -> tl::expected, ErrorCode> { + std::shared_lock shared_lock(snapshot_mutex_); + ScopedLocalDiskSegmentAccess local_disk_segment_access = + segment_manager_.getLocalDiskSegmentAccess(); + auto& client_local_disk_segment = + local_disk_segment_access.getClientLocalDiskSegment(); + auto local_disk_segment_it = client_local_disk_segment.find(client_id); + if (local_disk_segment_it == client_local_disk_segment.end()) { + return tl::make_unexpected(ErrorCode::SEGMENT_NOT_FOUND); + } + { + MutexLocker locker(&local_disk_segment_it->second->offloading_mutex_); + return local_disk_segment_it->second->removed_keys; + } +} + +void MasterService::EnqueueRemoveTasks( + const std::vector& holder_ids, const RemoveTaskItem& task) { + if (holder_ids.empty()) return; + ScopedLocalDiskSegmentAccess access = + segment_manager_.getLocalDiskSegmentAccess(); + auto& segments = access.getClientLocalDiskSegment(); + for (const auto& holder_id : holder_ids) { + auto it = segments.find(holder_id); + if (it == segments.end()) continue; + MutexLocker locker(&it->second->offloading_mutex_); + if (std::find(it->second->removed_keys.begin(), + it->second->removed_keys.end(), task) == + it->second->removed_keys.end()) { + it->second->removed_keys.push_back(task); + } + } +} + +auto MasterService::AckRemoveObjectHeartbeat( + const UUID& client_id, const std::vector& tasks) + -> tl::expected { + std::shared_lock shared_lock(snapshot_mutex_); + ScopedLocalDiskSegmentAccess access = + segment_manager_.getLocalDiskSegmentAccess(); + auto& segments = access.getClientLocalDiskSegment(); + auto it = segments.find(client_id); + if (it == segments.end()) { + return tl::make_unexpected(ErrorCode::SEGMENT_NOT_FOUND); + } + MutexLocker locker(&it->second->offloading_mutex_); + auto& pending = it->second->removed_keys; + pending.erase(std::remove_if(pending.begin(), pending.end(), + [&tasks](const RemoveTaskItem& task) { + return std::find(tasks.begin(), + tasks.end(), task) != + tasks.end(); + }), + pending.end()); + return {}; +} + auto MasterService::ReportSsdCapacity(const UUID& client_id, int64_t ssd_total_capacity_bytes) -> tl::expected { diff --git a/mooncake-store/src/real_client.cpp b/mooncake-store/src/real_client.cpp index c7eb5df178..e26b6ea96a 100644 --- a/mooncake-store/src/real_client.cpp +++ b/mooncake-store/src/real_client.cpp @@ -1205,11 +1205,26 @@ int RealClient::setup_real( const std::string &ipc_socket_path, bool enable_ssd_offload, const std::string &ssd_offload_path, const std::string &tenant_id, bool enable_client_http_server, int client_http_port) { - return to_py_ret(setup_internal( + SpDiag::PerfPoint pt(PerfKey::RC_SETUP_REAL, + SpDiag::PerfLevel::KEY_MODULE); + pt.Start(); + auto t0 = std::chrono::steady_clock::now(); + auto ret = to_py_ret(setup_internal( local_hostname, metadata_server, global_segment_size, local_buffer_size, protocol, rdma_devices, master_server_addr, transfer_engine, ipc_socket_path, 50052, enable_ssd_offload, true, ssd_offload_path, tenant_id, Environ::Get().GetOffloadRpcThreadNum(8), enable_client_http_server, client_http_port)); + auto t1 = std::chrono::steady_clock::now(); + auto elapsed_us = std::chrono::duration_cast( + t1 - t0).count(); + pt.End(ret == 0 ? 0 : -1); + // S2 setup 下沉层 MC_LOG 汇总（Q1b：入口层只 PerfPoint） + MC_LOG(INFO) << "[setup] elapsed_us=" << elapsed_us + << " success=" << (ret == 0 ? 1 : 0) + << " hostname=" << local_hostname + << " metadata_server=" << metadata_server + << " protocol=" << protocol; + return ret; } namespace { @@ -2397,6 +2412,7 @@ tl::expected RealClient::remove_internal( if (!remove_result) { return tl::unexpected(remove_result.error()); } + // SSD tombstone is handled by the storage node via RemoveObjectHeartbeat. return {}; } @@ -2436,7 +2452,9 @@ std::vector> RealClient::batchRemove_internal( return std::vector>( keys.size(), tl::unexpected(ErrorCode::INVALID_PARAMS)); } - return client_->BatchRemove(keys, force); + auto results = client_->BatchRemove(keys, force); + // SSD tombstone is handled by the storage node via RemoveObjectHeartbeat. + return results; } std::vector RealClient::batchRemove(const std::vector &keys, @@ -2476,7 +2494,14 @@ int RealClient::isExist(const std::string &key) { std::vector RealClient::batchIsExist( const std::vector &keys) { + SpDiag::PerfPoint pt(PerfKey::RC_BATCH_IS_EXIST, + SpDiag::PerfLevel::KEY_MODULE); + pt.Start(); + auto t0 = std::chrono::steady_clock::now(); auto internal_results = batchIsExist_internal(keys); + auto t1 = std::chrono::steady_clock::now(); + auto elapsed_us = std::chrono::duration_cast( + t1 - t0).count(); std::vector results; results.reserve(internal_results.size()); @@ -2487,7 +2512,17 @@ std::vector RealClient::batchIsExist( results.push_back(toInt(result.error())); } } - + pt.End(0); + // S6 batch_is_exist 下沉层 MC_LOG 汇总+per-key（Q1b：入口层只 PerfPoint） + // 字段：success+key（无 size/replica/endpoint，缺失不输出） + std::ostringstream oss; + oss << "[batch_is_exist] elapsed_us=" << elapsed_us + << " num_keys=" << keys.size(); + for (size_t i = 0; i < keys.size(); ++i) { + int success = (i < results.size()) ? results[i] : 0; + oss << "\n success=" << success << " key=" << keys[i]; + } + MC_LOG(INFO) << oss.str(); return results; } @@ -3768,7 +3803,20 @@ tl::expected RealClient::register_buffer_internal( } int RealClient::register_buffer(void *buffer, size_t size) { - return to_py_ret(register_buffer_internal(buffer, size)); + SpDiag::PerfPoint pt(PerfKey::RC_REGISTER_BUFFER, + SpDiag::PerfLevel::KEY_MODULE); + pt.Start(); + auto t0 = std::chrono::steady_clock::now(); + auto ret = to_py_ret(register_buffer_internal(buffer, size)); + auto t1 = std::chrono::steady_clock::now(); + auto elapsed_us = std::chrono::duration_cast( + t1 - t0).count(); + pt.End(ret == 0 ? 0 : -1); + // S3 register_buffer 下沉层 MC_LOG 汇总（Q1b：入口层只 PerfPoint） + MC_LOG(INFO) << "[register_buffer] elapsed_us=" << elapsed_us + << " success=" << (ret == 0 ? 1 : 0) + << " base_addr=" << buffer << " size=" << size; + return ret; } tl::expected RealClient::unregister_buffer_internal( @@ -5854,19 +5902,26 @@ RealClient::batch_get_into_internal(const std::vector &keys, std::vector> RealClient::batchIsExist_internal( const std::vector &keys) { + SpDiag::PerfPoint pt(PerfKey::RC_BATCH_IS_EXIST_INTERNAL, + SpDiag::PerfLevel::MODULE); + pt.Start(); if (!client_) { LOG(ERROR) << "Client is not initialized"; + pt.End(-1); return std::vector>( keys.size(), tl::unexpected(ErrorCode::INVALID_PARAMS)); } if (keys.empty()) { LOG(WARNING) << "Empty keys vector provided to batchIsExist_internal"; + pt.End(0); return std::vector>(); } // Call client BatchIsExist and return the vector directly - return client_->BatchIsExist(keys); + auto ret = client_->BatchIsExist(keys); + pt.End(0); + return ret; } int RealClient::put_from_with_metadata(const std::string &key, void *buffer, @@ -5913,6 +5968,9 @@ std::vector RealClient::batch_put_from_multi_buffers( const std::vector> &all_buffers, const std::vector> &sizes, const ReplicateConfig &config) { + SpDiag::PerfPoint pt(PerfKey::RC_BATCH_PUT_MULTI, + SpDiag::PerfLevel::KEY_MODULE); + pt.Start(); mooncake::logging::ScopedTraceId trace(mooncake::logging::NewTraceId()); auto internal_results = execute_timed_operation>>( @@ -5938,7 +5996,8 @@ std::vector RealClient::batch_put_from_multi_buffers( for (const auto &result : internal_results) { results.push_back(to_py_ret(result)); } - + pt.End(results.empty() ? -1 : 0); + // MC_LOG 在 *_internal 输出 汇总+per-key（Q1b） return results; } @@ -5948,8 +6007,13 @@ RealClient::batch_put_from_multi_buffers_internal( const std::vector> &all_buffers, const std::vector> &all_sizes, const ReplicateConfig &config) { + SpDiag::PerfPoint pt(PerfKey::RC_BATCH_PUT_MULTI_INTERNAL, + SpDiag::PerfLevel::MODULE); + pt.Start(); + auto t0 = std::chrono::steady_clock::now(); if (!client_) { LOG(ERROR) << "Client is not initialized"; + pt.End(-1); return std::vector>( keys.size(), tl::unexpected(ErrorCode::INVALID_PARAMS)); } @@ -5957,6 +6021,7 @@ RealClient::batch_put_from_multi_buffers_internal( if ((keys.size() != all_buffers.size()) || (all_buffers.size() != all_sizes.size())) { LOG(ERROR) << "Mismatched sizes for keys, buffers, and sizes"; + pt.End(-1); return std::vector>( keys.size(), tl::unexpected(ErrorCode::INVALID_PARAMS)); } @@ -5967,6 +6032,7 @@ RealClient::batch_put_from_multi_buffers_internal( const auto &sizes = all_sizes[i]; if (buffers.size() != sizes.size()) { LOG(ERROR) << "Mismatched buffers and sizes of key:" << keys[i]; + pt.End(-1); return std::vector>( keys.size(), tl::unexpected(ErrorCode::INVALID_PARAMS)); } @@ -5976,7 +6042,32 @@ RealClient::batch_put_from_multi_buffers_internal( } } // Call client BatchPut and return the vector directly - return client_->BatchPut(keys, batched_slices, config); + auto result = client_->BatchPut(keys, batched_slices, config); + auto t1 = std::chrono::steady_clock::now(); + auto total_us = std::chrono::duration_cast( + t1 - t0).count(); + pt.End(result.empty() ? -1 : 0); + // S4 batch_put_from_multi_buffers 下沉层 MC_LOG 汇总+per-key（Q1b） + // 字段：success+key+size（无 replica/endpoint，缺失不输出） + size_t total_bytes = 0; + for (const auto &sizes : all_sizes) + for (auto s : sizes) total_bytes += s; + std::ostringstream oss; + oss << "[batch_put_from_multi_buffers] elapsed_us=" << total_us + << " num_keys=" << keys.size() + << " total_bytes=" << total_bytes; + for (size_t i = 0; i < keys.size(); ++i) { + int success = (i < result.size() && result[i].has_value()) ? 1 : 0; + oss << "\n success=" << success << " key=" << keys[i]; + if (i < all_sizes.size()) { + size_t key_size = 0; + for (auto s : all_sizes[i]) key_size += s; + oss << " size=" << key_size; + } + // 无 replica/endpoint，不输出 + } + MC_LOG(INFO) << oss.str(); + return result; } std::vector RealClient::batch_get_into_multi_buffers( @@ -5984,6 +6075,9 @@ std::vector RealClient::batch_get_into_multi_buffers( const std::vector> &all_buffers, const std::vector> &all_sizes, bool prefer_alloc_in_same_node) { + SpDiag::PerfPoint pt(PerfKey::RC_BATCH_GET_INTO_MULTI, + SpDiag::PerfLevel::KEY_MODULE); + pt.Start(); auto internal_results = execute_timed_operation>>( [&]() { @@ -6008,6 +6102,8 @@ std::vector RealClient::batch_get_into_multi_buffers( for (const auto &result : internal_results) { results.push_back(to_py_ret(result)); } + pt.End(results.empty() ? -1 : 0); + // MC_LOG 在 *_internal 输出 汇总+per-key（Q1b） return results; } @@ -6017,9 +6113,14 @@ RealClient::batch_get_into_multi_buffers_internal( const std::vector> &all_buffers, const std::vector> &all_sizes, bool prefer_alloc_in_same_node) { + SpDiag::PerfPoint pt(PerfKey::RC_BATCH_GET_INTO_MULTI_INTERNAL, + SpDiag::PerfLevel::MODULE); + pt.Start(); + auto t0 = std::chrono::steady_clock::now(); // Validate preconditions if (!client_) { LOG(ERROR) << "Client is not initialized"; + pt.End(-1); return std::vector>( keys.size(), tl::unexpected(ErrorCode::INVALID_PARAMS)); } @@ -6028,6 +6129,7 @@ RealClient::batch_get_into_multi_buffers_internal( LOG(ERROR) << "Input vector sizes mismatch: keys=" << keys.size() << ", buffers=" << all_buffers.size() << ", sizes=" << all_sizes.size(); + pt.End(-1); return std::vector>( keys.size(), tl::unexpected(ErrorCode::INVALID_PARAMS)); } @@ -6035,7 +6137,11 @@ RealClient::batch_get_into_multi_buffers_internal( const size_t num_keys = keys.size(); std::vector> results; results.reserve(num_keys); + // Per-key replica/endpoint tracking for MC_LOG（Q1b：下沉层 per-key 全字段） + std::vector per_key_replica_type(num_keys); + std::vector per_key_endpoint(num_keys); if (num_keys == 0) { + pt.End(0); return results; } // Query metadata for all keys @@ -6094,6 +6200,21 @@ RealClient::batch_get_into_multi_buffers_internal( continue; } const auto replica = *best_replica; + // Capture replica info for MC_LOG per-key output + if (replica.is_memory_replica()) { + per_key_replica_type[i] = "memory"; + per_key_endpoint[i] = std::string(replica.get_memory_descriptor() + .buffer_descriptor + .transport_endpoint_); + } else if (replica.is_local_disk_replica()) { + per_key_replica_type[i] = "local_disk"; + per_key_endpoint[i] = + std::string(replica.get_local_disk_descriptor() + .transport_endpoint); + } else if (replica.is_disk_replica()) { + per_key_replica_type[i] = "disk"; + // DISK 副本无 transport_endpoint，缺失不输出 + } uint64_t total_size = calculate_total_size(replica); const auto &sizes = all_sizes[i]; uint64_t dst_total_size = 0; @@ -6149,6 +6270,7 @@ RealClient::batch_get_into_multi_buffers_internal( } // Early return if no valid operations if (valid_operations.empty() && valid_local_disk_ops.empty()) { + pt.End(0); return results; } @@ -6366,6 +6488,32 @@ RealClient::batch_get_into_multi_buffers_internal( } } + auto t1 = std::chrono::steady_clock::now(); + auto total_us = std::chrono::duration_cast( + t1 - t0).count(); + pt.End(results.empty() ? -1 : 0); + // S5 batch_get_into_multi_buffers 下沉层 MC_LOG 汇总+per-key（Q1b） + // 字段：success+key+size+replica+endpoint（缺失字段不输出） + std::ostringstream oss; + oss << "[batch_get_into_multi_buffers] elapsed_us=" << total_us + << " num_keys=" << keys.size(); + for (size_t i = 0; i < keys.size(); ++i) { + int success = + (i < results.size() && results[i].has_value()) ? 1 : 0; + oss << "\n success=" << success << " key=" << keys[i]; + if (i < all_sizes.size()) { + size_t key_size = 0; + for (auto s : all_sizes[i]) key_size += s; + oss << " size=" << key_size; + } + if (i < per_key_replica_type.size() && !per_key_replica_type[i].empty()) { + oss << " replica=" << per_key_replica_type[i]; + } + if (i < per_key_endpoint.size() && !per_key_endpoint[i].empty()) { + oss << " endpoint=" << per_key_endpoint[i]; + } + } + MC_LOG(INFO) << oss.str(); return results; } diff --git a/mooncake-store/src/rpc_service.cpp b/mooncake-store/src/rpc_service.cpp index e3b7891086..57abd06492 100644 --- a/mooncake-store/src/rpc_service.cpp +++ b/mooncake-store/src/rpc_service.cpp @@ -1728,6 +1728,21 @@ WrappedMasterService::PromotionObjectHeartbeat(const UUID& client_id) { return master_service_.PromotionObjectHeartbeat(client_id); } +tl::expected, ErrorCode> +WrappedMasterService::RemoveObjectHeartbeat(const UUID& client_id) { + ScopedVLogTimer timer(1, "RemoveObjectHeartbeat"); + timer.LogRequest("action=remove_heartbeat"); + return master_service_.RemoveObjectHeartbeat(client_id); +} + +tl::expected +WrappedMasterService::AckRemoveObjectHeartbeat( + const UUID& client_id, const std::vector& tasks) { + ScopedVLogTimer timer(1, "AckRemoveObjectHeartbeat"); + timer.LogRequest("action=ack_remove_heartbeat"); + return master_service_.AckRemoveObjectHeartbeat(client_id, tasks); +} + tl::expected WrappedMasterService::PromotionAllocStart( const UUID& client_id, const std::string& key, const std::string& tenant_id, @@ -1922,6 +1937,12 @@ void RegisterRpcService( server.register_handler< &mooncake::WrappedMasterService::PromotionObjectHeartbeat>( &wrapped_master_service); + server.register_handler< + &mooncake::WrappedMasterService::RemoveObjectHeartbeat>( + &wrapped_master_service); + server.register_handler< + &mooncake::WrappedMasterService::AckRemoveObjectHeartbeat>( + &wrapped_master_service); server .register_handler<&mooncake::WrappedMasterService::PromotionAllocStart>( &wrapped_master_service); diff --git a/mooncake-store/src/storage_backend.cpp b/mooncake-store/src/storage_backend.cpp index 0084902d39..ab2723b897 100644 --- a/mooncake-store/src/storage_backend.cpp +++ b/mooncake-store/src/storage_backend.cpp @@ -14,11 +14,13 @@ #include #include #include +#include #include #include #include #include #include +#include #include #include @@ -133,6 +135,40 @@ BucketBackendConfig BucketBackendConfig::FromEnvironment() { config.disable_ssd_eviction = GetEnvOr("MOONCAKE_OFFLOAD_DISABLE_SSD_EVICTION", false); + config.gc_enable = GetEnvOr("MOONCAKE_OFFLOAD_BUCKET_GC_ENABLE", + config.gc_enable); + config.gc_interval_ms = + GetEnvOr("MOONCAKE_OFFLOAD_BUCKET_GC_INTERVAL_MS", + config.gc_interval_ms); + // Parse doubles manually: GetEnvOr uses std::stoll which cannot parse + // fractional values like "0.25". + { + const char* ratio_env = + std::getenv("MOONCAKE_OFFLOAD_BUCKET_GC_DELETED_RATIO"); + if (ratio_env && !std::string(ratio_env).empty()) { + try { + config.gc_deleted_ratio = std::stod(std::string(ratio_env)); + } catch (...) { + // keep default + } + } + const char* wm_env = std::getenv( + "MOONCAKE_OFFLOAD_BUCKET_GC_HIGH_WATERMARK_RATIO"); + if (wm_env && !std::string(wm_env).empty()) { + try { + config.gc_high_watermark_ratio = + std::stod(std::string(wm_env)); + } catch (...) { + // keep default + } + } + } + config.gc_max_buckets_per_round = GetEnvOr( + "MOONCAKE_OFFLOAD_BUCKET_GC_MAX_BUCKETS_PER_ROUND", + config.gc_max_buckets_per_round); + config.gc_merge_enable = GetEnvOr( + "MOONCAKE_OFFLOAD_BUCKET_GC_MERGE_ENABLE", config.gc_merge_enable); + return config; } @@ -1739,6 +1775,12 @@ BucketStorageBackend::BucketStorageBackend( } BucketStorageBackend::~BucketStorageBackend() { + // Stop background GC thread first, before clearing any state it touches. + if (gc_running_.load(std::memory_order_acquire)) { + gc_running_.store(false, std::memory_order_release); + gc_cv_.notify_all(); + if (gc_thread_.joinable()) gc_thread_.join(); + } // Clear file cache to release UringFile instances before destruction // This ensures orderly cleanup of io_uring resources ClearFileCache(); @@ -2068,8 +2110,8 @@ tl::expected BucketStorageBackend::BatchLoad( // Calculate aligned read range int64_t aligned_offset = align_down(actual_offset, kDirectIOAlignment); - int64_t data_end = - actual_offset + static_cast(plan.dest_slice.size); + 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 = @@ -2095,15 +2137,14 @@ tl::expected BucketStorageBackend::BatchLoad( } } else #endif - { - // Fallback to vector_read for non-UringFile - iovec iov{plan.dest_slice.ptr, plan.dest_slice.size}; - SpDiag::PerfPoint pt_posix(PerfKey::GET_SSD_OWNER_LOAD_POSIX, - SpDiag::PerfLevel::MODULE); - pt_posix.Start(); - read_res = file->vector_read(&iov, 1, actual_offset); - pt_posix.End(read_res ? 0 : -1); - } + { + // Fallback to per-key vector_read for non-UringFile (PosixFile). + iovec iov{plan.dest_slice.ptr, plan.dest_slice.size}; + SpDiag::PerfPoint pt_posix(PerfKey::GET_SSD_OWNER_LOAD_POSIX, + SpDiag::PerfLevel::MODULE); + pt_posix.Start(); + read_res = file->vector_read(&iov, 1, actual_offset); + pt_posix.End(read_res ? 0 : -1); if (stats) { const auto read_us = std::chrono::duration_cast( @@ -2145,6 +2186,7 @@ tl::expected BucketStorageBackend::BatchLoad( return tl::make_unexpected(ErrorCode::FILE_READ_FAIL); } } + } } // bucket_guards go out of scope here, decrementing inflight_reads_ @@ -2292,13 +2334,23 @@ tl::expected BucketStorageBackend::Init() { total_size_ += metadata_it->second->data_size + metadata_it->second->meta_size; for (size_t i = 0; i < metadata_it->second->keys.size(); i++) { + const auto& key = metadata_it->second->keys[i]; + if (std::find(metadata_it->second->tombstones.begin(), + metadata_it->second->tombstones.end(), key) != + metadata_it->second->tombstones.end()) { + metadata_it->second->deleted_bytes_.fetch_add( + metadata_it->second->metadatas[i].key_size + + metadata_it->second->metadatas[i].data_size, + std::memory_order_relaxed); + continue; + } object_bucket_map_.emplace( - metadata_it->second->keys[i], - StorageObjectMetadata{ - metadata_it->first, - metadata_it->second->metadatas[i].offset, - metadata_it->second->metadatas[i].key_size, - metadata_it->second->metadatas[i].data_size, ""}); + key, StorageObjectMetadata{ + metadata_it->first, + metadata_it->second->metadatas[i].offset, + metadata_it->second->metadatas[i].key_size, + metadata_it->second->metadatas[i].data_size, + ""}); } } } @@ -2399,6 +2451,12 @@ tl::expected BucketStorageBackend::Init() { return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); } + // Start background GC thread if enabled. + if (bucket_backend_config_.gc_enable) { + gc_running_.store(true, std::memory_order_release); + gc_thread_ = std::thread(&BucketStorageBackend::GCThreadFunc, this); + } + return {}; } @@ -3432,6 +3490,531 @@ void BucketStorageBackend::RemoveAll() { LOG(INFO) << "RemoveAll: removed " << bucket_ids.size() << " bucket(s)"; } +// --- Explicit-delete-only GC --- +tl::expected BucketStorageBackend::MarkRemoved( + const std::string& key) { + SharedMutexLocker lock(&mutex_); + auto it = object_bucket_map_.find(key); + if (it == object_bucket_map_.end()) { + return {}; + } + int64_t bucket_id = it->second.bucket_id; + int64_t freed = it->second.data_size + it->second.key_size; + auto bucket_it = buckets_.find(bucket_id); + if (bucket_it == buckets_.end()) { + return tl::make_unexpected(ErrorCode::BUCKET_NOT_FOUND); + } + auto bucket = bucket_it->second; + const auto object_metadata = it->second; + object_bucket_map_.erase(it); + bucket->tombstones.push_back(key); + auto persist_result = StoreBucketMetadata(bucket_id, bucket); + if (!persist_result) { + object_bucket_map_.emplace(key, object_metadata); + bucket->tombstones.pop_back(); + return tl::make_unexpected(persist_result.error()); + } + bucket->deleted_bytes_.fetch_add(freed, std::memory_order_relaxed); + ++bucket->generation_; + return {}; +} + +tl::expected BucketStorageBackend::BatchMarkRemoved( + const std::vector& keys) { + for (const auto& key : keys) { + auto result = MarkRemoved(key); + if (!result) { + return result; + } + } + return {}; +} + +bool BucketStorageBackend::CompactBucket(int64_t bucket_id) { + return CompactBuckets({bucket_id}, true); +} + +bool BucketStorageBackend::CompactBuckets( + const std::vector& bucket_ids, bool space_pressure) { + if (bucket_ids.empty()) return true; + + // Step 1: lock-snapshot live keys from ALL old buckets + mark compacting_ + struct LiveKeyInfo { + std::string key; + int64_t old_bucket_id; + BucketObjectMetadata meta; + }; + std::vector live_keys_info; + // old_bucket_id -> {shared_ptr, last_access_ns} + std::unordered_map, int64_t>> + old_buckets; + std::unordered_map old_bucket_generations; + { + SharedMutexLocker lock(&mutex_); + for (int64_t bid : bucket_ids) { + auto it = buckets_.find(bid); + if (it == buckets_.end()) continue; + auto& bucket = it->second; + if (bucket->compacting_.load(std::memory_order_relaxed)) continue; + bucket->compacting_.store(true, std::memory_order_relaxed); + int64_t ts = bucket->last_access_ns_.load( + std::memory_order_relaxed); + old_buckets[bid] = {bucket, ts}; + old_bucket_generations[bid] = bucket->generation_; + + for (size_t i = 0; i < bucket->keys.size(); ++i) { + const auto& key = bucket->keys[i]; + auto map_it = object_bucket_map_.find(key); + if (map_it != object_bucket_map_.end() && + map_it->second.bucket_id == bid) { + live_keys_info.push_back( + {key, bid, bucket->metadatas[i]}); + } + } + } + } + + if (old_buckets.empty()) return true; + + auto reset_compacting = [&]() { + for (auto& [bid, pr] : old_buckets) { + pr.first->compacting_.store(false, std::memory_order_relaxed); + } + }; + + // Identify buckets with zero live keys — delete them immediately + // without waiting for merge. Buckets with live keys participate in + // the merge. + std::set buckets_with_live; + for (const auto& info : live_keys_info) { + buckets_with_live.insert(info.old_bucket_id); + } + + std::vector empty_buckets; + for (auto& [bid, pr] : old_buckets) { + if (buckets_with_live.find(bid) == buckets_with_live.end()) { + empty_buckets.push_back(bid); + } + } + + if (!empty_buckets.empty()) { + std::vector> to_drain; + { + SharedMutexLocker lock(&mutex_); + for (int64_t bid : empty_buckets) { + auto it = buckets_.find(bid); + if (it != buckets_.end()) { + total_size_ -= + it->second->data_size + it->second->meta_size; + buckets_.erase(it); + int64_t ts = old_buckets[bid].second; + lru_index_.erase({ts, bid}); + to_drain.push_back(old_buckets[bid].first); + } + } + } + for (auto& bucket : to_drain) { + WaitForInflightReads(bucket); + } + for (int64_t bid : empty_buckets) { + DeleteBucketFiles(bid); + } + // Remove deleted buckets from old_buckets so they don't interfere + // with the merge logic below. + for (int64_t bid : empty_buckets) { + old_buckets.erase(bid); + } + } + + // If no live keys remain (all buckets were empty), we're done. + if (live_keys_info.empty()) { + return true; + } + + // Step 2: group live keys by bucket_keys_limit / bucket_size_limit + // using metadata only (no file IO). Only the FIRST group that fills up + // will be written as a new bucket; remaining keys are deferred to the + // next round. + struct GroupedKey { + std::string key; + int64_t old_bucket_id; + BucketObjectMetadata meta; + int64_t total_size; // key_size + data_size + }; + std::vector first_group_keys; + int64_t group_count = 0; + int64_t group_size = 0; + + for (const auto& info : live_keys_info) { + int64_t key_total = + info.meta.key_size + info.meta.data_size; + if (group_count >= bucket_backend_config_.bucket_keys_limit || + (group_count > 0 && group_size + key_total > + bucket_backend_config_.bucket_size_limit)) { + break; // first group is full + } + first_group_keys.push_back( + {info.key, info.old_bucket_id, info.meta, key_total}); + group_size += key_total; + ++group_count; + } + + // Check if first group is full enough to write. + bool group_full = + (group_count >= bucket_backend_config_.bucket_keys_limit) || + (group_size >= bucket_backend_config_.bucket_size_limit); + if (!group_full && !space_pressure) { + // Not enough live keys to fill a bucket; defer to next round. + LOG(INFO) << "[GC] CompactBuckets deferred: group_count=" + << group_count + << " group_size=" << group_size + << " bucket_keys_limit=" + << bucket_backend_config_.bucket_keys_limit + << " bucket_size_limit=" + << bucket_backend_config_.bucket_size_limit + << " total_live_keys=" << live_keys_info.size(); + reset_compacting(); + return true; + } + + // Step 3: read ONLY the first group's live key data from old buckets + // (lock-free IO). Open each old bucket file once, read only the keys + // that are in the first group. + std::unordered_map live_data_buffers; + std::unordered_map> first_group; + + // Group first_group_keys by old_bucket_id to open each file once. + std::unordered_map> keys_by_bucket; + for (const auto& gk : first_group_keys) { + keys_by_bucket[gk.old_bucket_id].push_back(&gk); + } + + for (auto& [bid, keys_ptr] : keys_by_bucket) { + auto& old_bucket = old_buckets[bid].first; + bool read_ok = true; + { + BucketReadGuard guard(old_bucket); + auto data_path = GetBucketDataPath(bid); + if (!data_path) { + read_ok = false; + } else { + auto file_result = + OpenFile(data_path.value(), FileMode::Read); + if (!file_result) { + read_ok = false; + } else { + auto& file = file_result.value(); + for (const auto* gk : keys_ptr) { + std::string data; + data.resize(gk->meta.data_size); + int64_t actual_offset = + gk->meta.offset + gk->meta.key_size; + iovec iov{data.data(), + static_cast(gk->meta.data_size)}; + auto read_res = + file->vector_read(&iov, 1, actual_offset); + if (!read_res || + read_res.value() != + static_cast(gk->meta.data_size)) { + LOG(ERROR) + << "CompactBuckets: read failed for key: " + << gk->key << ", bucket_id=" << bid; + read_ok = false; + break; + } + live_data_buffers[gk->key] = std::move(data); + first_group.emplace( + gk->key, + std::vector{Slice{ + live_data_buffers[gk->key].data(), + live_data_buffers[gk->key].size()}}); + } + } + } + } // guard released + + if (!read_ok) { + reset_compacting(); + return false; + } + } + + // Step 4: write the first group as a new bucket. + int64_t new_bucket_id = bucket_id_generator_->NextId(); + std::vector iovs; + std::vector new_metas; + auto build_result = + BuildBucket(new_bucket_id, first_group, iovs, new_metas); + if (!build_result) { + LOG(ERROR) << "CompactBuckets: BuildBucket failed"; + reset_compacting(); + return false; + } + auto write_result = + WriteBucket(new_bucket_id, build_result.value(), iovs); + if (!write_result) { + LOG(ERROR) << "CompactBuckets: WriteBucket failed"; + reset_compacting(); + return false; + } + + // Step 5: atomic swap under lock with re-validation. + auto& new_bucket = build_result.value(); + // Determine which old buckets have ALL their live keys migrated. + // A key is "migrated" if it's in the first_group AND still maps to an + // old bucket being compacted. Old buckets with no remaining live keys + // (all migrated or removed) can be deleted. + std::set old_bucket_ids_set(bucket_ids.begin(), + bucket_ids.end()); + // Use the coldest last_access_ns among old buckets for the new bucket. + int64_t new_last_access_ns = std::numeric_limits::max(); + for (auto& [bid, pr] : old_buckets) { + new_last_access_ns = std::min(new_last_access_ns, pr.second); + } + if (new_last_access_ns == std::numeric_limits::max()) { + new_last_access_ns = 0; + } + + std::vector buckets_to_delete; + { + SharedMutexLocker lock(&mutex_); + for (const auto& [bid, pr] : old_buckets) { + auto current_it = buckets_.find(bid); + auto generation_it = old_bucket_generations.find(bid); + if (current_it == buckets_.end() || + generation_it == old_bucket_generations.end() || + current_it->second->generation_ != generation_it->second) { + // A delete changed the source bucket after the snapshot. The + // data file was built from a stale view, so do not publish it + // or replace any object mappings with it. + lock.unlock(); + CleanupOrphanedBucket(new_bucket_id); + reset_compacting(); + LOG(INFO) << "CompactBuckets discarded stale snapshot for " + << "bucket_id=" << bid; + return true; + } + } + // Re-validate and remap each key in the new bucket. + for (size_t i = 0; i < new_bucket->keys.size(); ++i) { + const auto& key = new_bucket->keys[i]; + auto map_it = object_bucket_map_.find(key); + if (map_it != object_bucket_map_.end() && + old_bucket_ids_set.count(map_it->second.bucket_id)) { + // Still live at an old bucket -> remap to new bucket. + map_it->second = new_metas[i]; + } + // else: key was removed or remapped -> skip. + } + + // Insert new bucket. + total_size_ += new_bucket->data_size + new_bucket->meta_size; + new_bucket->last_access_ns_.store( + new_last_access_ns, std::memory_order_relaxed); + buckets_.emplace(new_bucket_id, new_bucket); + lru_index_.emplace(new_last_access_ns, new_bucket_id); + + // Determine which old buckets can be deleted: those whose live keys + // are all now absent from object_bucket_map_ or remapped to the new + // bucket (i.e., no key still points at the old bucket). + for (auto& [bid, pr] : old_buckets) { + auto& old_bucket = pr.first; + bool has_remaining_live = false; + for (const auto& key : old_bucket->keys) { + auto map_it = object_bucket_map_.find(key); + if (map_it != object_bucket_map_.end() && + map_it->second.bucket_id == bid) { + // This key is still live at the old bucket — it was not + // in the first group (deferred to next round). + has_remaining_live = true; + break; + } + } + if (!has_remaining_live) { + buckets_to_delete.push_back(bid); + } else { + // Reset compacting_ so this bucket can be compacted later. + old_bucket->compacting_.store(false, + std::memory_order_relaxed); + } + } + + // Remove old buckets that are fully migrated. + for (int64_t bid : buckets_to_delete) { + auto it = buckets_.find(bid); + if (it != buckets_.end()) { + total_size_ -= + it->second->data_size + it->second->meta_size; + buckets_.erase(it); + int64_t ts = old_buckets[bid].second; + lru_index_.erase({ts, bid}); + } + } + } + + // Step 6: wait for inflight reads + delete old bucket files. + for (int64_t bid : buckets_to_delete) { + WaitForInflightReads(old_buckets[bid].first); + DeleteBucketFiles(bid); + } + + return true; +} + +void BucketStorageBackend::WaitForInflightReads( + std::shared_ptr bucket) { + constexpr int kMaxSpinIterations = 1000; + constexpr auto kMaxWaitTime = std::chrono::seconds(10); + int spin_count = 0; + auto wait_start = std::chrono::steady_clock::now(); + while (bucket->inflight_reads_.load(std::memory_order_acquire) > 0) { + if (++spin_count > kMaxSpinIterations) { + std::this_thread::yield(); + spin_count = 0; + if (std::chrono::steady_clock::now() - wait_start > + kMaxWaitTime) { + LOG(ERROR) << "CompactBucket: timed out waiting for " + "in-flight reads, inflight_reads=" + << bucket->inflight_reads_.load( + std::memory_order_relaxed); + break; + } + } else { + PAUSE(); + } + } +} + +void BucketStorageBackend::DeleteBucketFiles(int64_t bucket_id) { + namespace fs = std::filesystem; + std::error_code ec; + auto data_path = GetBucketDataPath(bucket_id); + if (data_path) { + { + MutexLocker cache_locker(&file_cache_mutex_); + file_cache_.erase(data_path.value()); + } + fs::remove(data_path.value(), ec); + if (ec && ec != std::errc::no_such_file_or_directory) { + LOG(ERROR) << "CompactBucket: failed to remove data file: " + << data_path.value() << ", error: " << ec.message(); + } + } + auto meta_path = GetBucketMetadataPath(bucket_id); + if (meta_path) { + ec.clear(); + fs::remove(meta_path.value(), ec); + if (ec && ec != std::errc::no_such_file_or_directory) { + LOG(ERROR) << "CompactBucket: failed to remove meta file: " + << meta_path.value() << ", error: " << ec.message(); + } + } +} + +// GCThreadFunc: background tombstone compaction loop. +void BucketStorageBackend::GCThreadFunc() { + LOG(INFO) << "[GC] background compaction thread started"; + while (gc_running_.load(std::memory_order_acquire)) { + // Sleep for gc_interval_ms or until woken for shutdown. + { + std::unique_lock lock(gc_mutex_); + gc_cv_.wait_for( + lock, + std::chrono::milliseconds( + bucket_backend_config_.gc_interval_ms), + [this]() { + return !gc_running_.load(std::memory_order_relaxed); + }); + } + if (!gc_running_.load(std::memory_order_acquire)) break; + + if (!bucket_backend_config_.gc_enable) continue; + + // Check space pressure under shared lock (total_size_ is + // GUARDED_BY(mutex_)). + bool space_pressure = false; + { + SharedMutexLocker lock(&mutex_, shared_lock); + if (bucket_backend_config_.max_total_size > 0) { + double used_ratio = + static_cast(total_size_) / + static_cast( + bucket_backend_config_.max_total_size); + space_pressure = used_ratio >= + bucket_backend_config_ + .gc_high_watermark_ratio; + } + } + + // Collect GC candidate buckets (up to gc_max_buckets_per_round). + std::vector candidates; + { + SharedMutexLocker lock(&mutex_); + int64_t count = 0; + for (auto it = buckets_.begin(); + it != buckets_.end() && + count < bucket_backend_config_.gc_max_buckets_per_round; + ++it) { + int64_t deleted = + it->second->deleted_bytes_.load( + std::memory_order_relaxed); + if (deleted <= 0) continue; + if (it->second->compacting_.load( + std::memory_order_relaxed)) + continue; + int64_t data_size = it->second->data_size; + double ratio = + (data_size > 0) + ? static_cast(deleted) / + static_cast(data_size) + : 0.0; + if (!space_pressure && + ratio < bucket_backend_config_.gc_deleted_ratio) { + continue; + } + candidates.push_back(it->first); + ++count; + } + } + + if (!candidates.empty()) { + if (bucket_backend_config_.gc_merge_enable && + candidates.size() > 1) { + // Cross-bucket merge: collect live keys from multiple + // tombstone buckets into one new bucket. + if (CompactBuckets(candidates, space_pressure)) { + LOG(INFO) << "[GC] merged " << candidates.size() + << " bucket(s)"; + } else { + LOG(WARNING) << "[GC] CompactBuckets failed for " + << candidates.size() + << " bucket(s), will retry next round"; + } + } else { + // Single-bucket compaction (one at a time). + int64_t compacted = 0; + for (int64_t bid : candidates) { + if (CompactBuckets({bid}, space_pressure)) { + ++compacted; + } else { + LOG(WARNING) + << "[GC] CompactBuckets failed for bucket " + << bid << ", will retry next round"; + break; + } + } + if (compacted > 0) { + LOG(INFO) << "[GC] compacted " << compacted + << " bucket(s)"; + } + } + } + } + LOG(INFO) << "[GC] background compaction thread stopped"; +} + tl::expected BucketStorageBackend::StoreBucketMetadata( int64_t id, std::shared_ptr metadata) { auto meta_path_res = GetBucketMetadataPath(id); diff --git a/mooncake-store/tests/e2e/CMakeLists.txt b/mooncake-store/tests/e2e/CMakeLists.txt index 7f477bd227..2606c11cfe 100644 --- a/mooncake-store/tests/e2e/CMakeLists.txt +++ b/mooncake-store/tests/e2e/CMakeLists.txt @@ -92,6 +92,21 @@ target_link_libraries(storage_backend_e2e_test PUBLIC add_test(NAME storage_backend_e2e_test COMMAND storage_backend_e2e_test) +add_executable(gc_e2e_test gc_e2e_test.cpp) +target_include_directories(gc_e2e_test PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/.. +) +target_link_libraries(gc_e2e_test PUBLIC + mooncake_store + transfer_engine + cachelib_memory_allocator + glog + gtest + pthread + ${ETCD_WRAPPER_LIB} +) +add_test(NAME gc_e2e_test COMMAND gc_e2e_test) + add_executable(oplog_batch_e2e_test oplog_batch_e2e_test.cpp process_handler.cpp) target_link_libraries(oplog_batch_e2e_test PUBLIC glog gtest pthread) add_test(NAME oplog_batch_e2e_test COMMAND oplog_batch_e2e_test) diff --git a/mooncake-store/tests/e2e/gc_e2e_test.cpp b/mooncake-store/tests/e2e/gc_e2e_test.cpp new file mode 100644 index 0000000000..b39f03269b --- /dev/null +++ b/mooncake-store/tests/e2e/gc_e2e_test.cpp @@ -0,0 +1,535 @@ +// gc_e2e_test.cpp +// End-to-end integration tests for the explicit-delete-only SSD GC. +// +// Verifies the full pipeline that unit tests cannot cover: +// RealClient::remove -> master metadata erase +// -> FileStorage::MarkRemoved (tombstone) +// -> BucketStorageBackend GC compaction +// -> SSD bucket file reclamation +// +// Unlike storage_backend_e2e_test (which uses the Client base class + +// file-per-key backend), this suite uses RealClient with +// enable_ssd_offload=true so the BucketStorageBackend + FileStorage +// offload path is exercised, and remove goes through +// RealClient::remove_internal -> MarkRemoved. + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "client_buffer.h" +#include "real_client.h" +#include "test_server_helpers.h" +#include "types.h" + +DEFINE_string(protocol, "tcp", "Transfer protocol: rdma|tcp"); +DEFINE_string(device_name, "", "Device name to use, valid if protocol=rdma"); + +namespace mooncake { +namespace testing { + +namespace fs = std::filesystem; + +static constexpr size_t kMB = 1024ULL * 1024; + +// Count regular files with a given suffix in dir. +static int CountFilesWithSuffix(const fs::path& dir, + const std::string& suffix) { + int count = 0; + std::error_code ec; + for (auto& entry : fs::directory_iterator(dir, ec)) { + if (entry.is_regular_file()) { + auto name = entry.path().filename().string(); + if (name.size() >= suffix.size() && + name.compare(name.size() - suffix.size(), suffix.size(), + suffix) == 0) { + ++count; + } + } + } + return count; +} + +// List all .bucket file names (without directory) in dir. +static std::vector ListBucketFiles(const fs::path& dir) { + std::vector names; + std::error_code ec; + for (auto& entry : fs::directory_iterator(dir, ec)) { + if (entry.is_regular_file()) { + auto name = entry.path().filename().string(); + if (name.size() >= 6 && + name.compare(name.size() - 6, 6, ".bucket") == 0) { + names.push_back(name); + } + } + } + return names; +} + +// Read a key via RealClient::get_buffer into a std::string. Returns +// std::nullopt on failure. +static std::optional ReadKey( + const std::shared_ptr& client, const std::string& key) { + auto buf = client->get_buffer(key); + if (!buf) return std::nullopt; + return std::string(static_cast(buf->ptr()), buf->size()); +} + +class GCE2ETest : public ::testing::Test { + protected: + static void SetUpTestSuite() { + google::InitGoogleLogging("GCE2ETest"); + FLAGS_logtostderr = 1; + } + + static void TearDownTestSuite() { google::ShutdownGoogleLogging(); } + + void SetUp() override { + if (getenv("PROTOCOL")) FLAGS_protocol = getenv("PROTOCOL"); + if (getenv("DEVICE_NAME")) FLAGS_device_name = getenv("DEVICE_NAME"); + + tmp_dir_ = fs::temp_directory_path() / + ("mc_gc_e2e_" + std::to_string(::getpid())); + fs::create_directories(tmp_dir_); + + // Save and set the GC-required bucket backend env vars. + // eviction_policy=LRU keeps last_access_ns_ updated for GC candidate + // coldness; disable_ssd_eviction=true makes PrepareEviction a no-op + // so no live bucket is ever evicted. + saved_policy_ = GetEnvOpt("MOONCAKE_OFFLOAD_BUCKET_EVICTION_POLICY"); + setenv("MOONCAKE_OFFLOAD_BUCKET_EVICTION_POLICY", "lru", 1); + saved_disable_ = GetEnvOpt("MOONCAKE_OFFLOAD_DISABLE_SSD_EVICTION"); + setenv("MOONCAKE_OFFLOAD_DISABLE_SSD_EVICTION", "true", 1); + // Tighten GC so compaction runs after remove. Set interval long + // enough that GC doesn't fire during offload settlement (which + // could compact a bucket before remove creates a tombstone). + saved_gc_interval_ = GetEnvOpt("MOONCAKE_OFFLOAD_BUCKET_GC_INTERVAL_MS"); + setenv("MOONCAKE_OFFLOAD_BUCKET_GC_INTERVAL_MS", "15000", 1); + saved_gc_ratio_ = GetEnvOpt("MOONCAKE_OFFLOAD_BUCKET_GC_DELETED_RATIO"); + setenv("MOONCAKE_OFFLOAD_BUCKET_GC_DELETED_RATIO", "0.01", 1); + // Set bucket_keys_limit=1 so each offloaded key fills a bucket + // immediately. This ensures .bucket files are written on the first + // heartbeat after put. With limit=2, keys may sit in the ungrouped + // pool and no .bucket file is written. + saved_bucket_keys_limit_ = + GetEnvOpt("MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT"); + setenv("MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT", "1", 1); + } + + void TearDown() override { + if (real_client_) real_client_->tearDownAll(); + master_.Stop(); + easylog::set_min_severity(easylog::Severity::WARN); + + // Restore env. + RestoreEnv("MOONCAKE_OFFLOAD_BUCKET_EVICTION_POLICY", saved_policy_); + RestoreEnv("MOONCAKE_OFFLOAD_DISABLE_SSD_EVICTION", saved_disable_); + RestoreEnv("MOONCAKE_OFFLOAD_BUCKET_GC_INTERVAL_MS", saved_gc_interval_); + RestoreEnv("MOONCAKE_OFFLOAD_BUCKET_GC_DELETED_RATIO", saved_gc_ratio_); + RestoreEnv("MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT", + saved_bucket_keys_limit_); + + std::error_code ec; + fs::remove_all(tmp_dir_, ec); + } + + static void RestoreEnv(const char* name, + const std::optional& saved) { + if (saved.has_value()) { + setenv(name, saved->c_str(), 1); + } else { + unsetenv(name); + } + } + + // Safely capture an env var as optional (getenv may return nullptr). + static std::optional GetEnvOpt(const char* name) { + const char* val = getenv(name); + if (val) return std::string(val); + return std::nullopt; + } + + bool StartMasterWithOffload() { + // Match production config: enable_offload=true, no root_fs_dir + // (master doesn't do disk caching; offload tasks are pushed to the + // client's FileStorage via heartbeat). Set a long lease TTL so + // objects aren't evicted before offload completes. + auto config = InProcMasterConfigBuilder() + .set_enable_offload(true) + .set_default_kv_lease_ttl(300000) + .build(); + return master_.Start(config); + } + + bool StartRealClient() { + real_client_ = RealClient::create(); + if (!real_client_) return false; + const std::string rdma_devices = + (FLAGS_protocol == "rdma") ? FLAGS_device_name : ""; + std::string ssd_path = tmp_dir_.string() + "/ssd_offload"; + fs::create_directories(ssd_path); + // Set MOONCAKE_OFFLOAD_FILE_STORAGE_PATH env var (same as production) + // so FileStorageConfig::FromEnvironment picks it up. This matches the + // production deployment pattern where the env var is set before + // launching the client. + setenv("MOONCAKE_OFFLOAD_FILE_STORAGE_PATH", ssd_path.c_str(), 1); + setenv("MOONCAKE_OFFLOAD_STORAGE_BACKEND_DESCRIPTOR", + "bucket_storage_backend", 1); + // enable_ssd_offload=true creates FileStorage + BucketStorageBackend. + int ret = real_client_->setup_real( + "localhost:17890", "P2PHANDSHAKE", + /*global_segment_size=*/512 * kMB, + /*local_buffer_size=*/256 * kMB, FLAGS_protocol, rdma_devices, + master_.master_address(), nullptr, + /*ipc_socket_path=*/"", + /*enable_ssd_offload=*/true, + /*ssd_offload_path=*/ssd_path, + /*tenant_id=*/"default"); + return ret == 0; + } + + // Put a key via RealClient and wait until it has been offloaded to the + // BucketStorageBackend (a .bucket file appears on SSD). Returns false on + // timeout. Waiting on memory reads is insufficient — offload is async + // (PutEnd queues, heartbeat drains) and MarkRemoved is a no-op until the + // key lands in object_bucket_map_. + bool PutAndWaitOffloaded(const std::string& key, + const std::string& value, + const fs::path& ssd_dir) { + std::span span(value.data(), value.size()); + ReplicateConfig config; + config.replica_num = 1; + if (real_client_->put(key, span, config) != 0) { + return false; + } + // Wait for offload: a .bucket file must appear in ssd_dir, AND the + // key must be readable via get_buffer (confirms data integrity). + // Heartbeat interval is 10s, so wait up to 40s. + for (int i = 0; i < 400; ++i) { // up to 40s + int buckets = CountFilesWithSuffix(ssd_dir, ".bucket"); + if (buckets > 0) { + auto got = ReadKey(real_client_, key); + if (got.has_value() && got.value() == value) return true; + } + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + } + return false; + } + + // After all keys are put and individually confirmed offloaded, wait an + // extra heartbeat cycle to ensure ALL keys have been drained from the + // offloading queue into object_bucket_map_. Without this, a key put + // after the first heartbeat may not yet be in object_bucket_map_ when + // MarkRemoved is called, making the tombstone a no-op. + void WaitForAllOffloadsSettled() { + // Wait less than gc_interval_ms (15s) so GC doesn't fire during + // settlement. Two heartbeat cycles (10s each) would be ideal, but + // 12s is enough for the 2nd heartbeat to drain remaining tasks + // while staying under the GC interval. + std::this_thread::sleep_for(std::chrono::seconds(12)); + } + + // Put multiple keys, then wait until ALL are offloaded (a .bucket file + // appears and each key is readable). Keys put before the next heartbeat + // are grouped into the same bucket (up to bucket_keys_limit). + bool PutBatchAndWaitOffloaded( + const std::vector>& kvs, + const fs::path& ssd_dir) { + ReplicateConfig config; + config.replica_num = 1; + for (const auto& [key, value] : kvs) { + std::span span(value.data(), value.size()); + if (real_client_->put(key, span, config) != 0) return false; + } + // Wait for offload of all keys: a .bucket file MUST appear (offload + // completed) AND each key must be readable via get_buffer. + // Heartbeat interval is 10s; with bucket_keys_limit=2, 2 keys fill + // a bucket on the first heartbeat after put. Wait up to 40s. + for (int i = 0; i < 400; ++i) { // up to 40s + if (CountFilesWithSuffix(ssd_dir, ".bucket") > 0) { + bool all_ok = true; + for (const auto& [key, value] : kvs) { + auto got = ReadKey(real_client_, key); + if (!got.has_value() || got.value() != value) { + all_ok = false; + break; + } + } + if (all_ok) return true; + } + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + } + return false; + } + + // Wait for GC compaction: detect by old bucket file(s) disappearing and + // new one(s) appearing. Returns the set of bucket file names before and + // after for caller verification. Polls up to 30s. + bool WaitForCompaction(const fs::path& ssd_dir, + const std::vector& buckets_before, + std::vector& buckets_after) { + for (int i = 0; i < 150; ++i) { // up to 30s + buckets_after = ListBucketFiles(ssd_dir); + // Compaction: at least one old bucket file gone, or set changed. + if (buckets_after != buckets_before) { + return true; + } + std::this_thread::sleep_for(std::chrono::milliseconds(200)); + } + buckets_after = ListBucketFiles(ssd_dir); + return false; + } + + fs::path tmp_dir_; + InProcMaster master_; + std::shared_ptr real_client_; + std::optional saved_policy_; + std::optional saved_disable_; + std::optional saved_gc_interval_; + std::optional saved_gc_ratio_; + std::optional saved_bucket_keys_limit_; +}; + +// ------------------------------------------------------------------- +// Test 1: RemoveReclaimsSSDSpace +// +// Put 2 keys (each in its own bucket, bucket_keys_limit=1), remove k1. +// GC compaction must delete k1's (now-empty) bucket file. Validates: +// - k1's .bucket file is eventually deleted (space reclaimed) +// - k2 remains readable with correct data throughout +// - k1 stays gone +// ------------------------------------------------------------------- +TEST_F(GCE2ETest, RemoveReclaimsSSDSpace) { + ASSERT_TRUE(StartMasterWithOffload()); + ASSERT_TRUE(StartRealClient()); + + const std::string k1 = "gc_e2e_k1"; + const std::string k2 = "gc_e2e_k2"; + const std::string v1(4 * kMB, 'A'); + const std::string v2(4 * kMB, 'B'); + + fs::path ssd_dir = tmp_dir_ / "ssd_offload"; + ASSERT_TRUE(PutAndWaitOffloaded(k1, v1, ssd_dir)) + << "k1 offload timed out"; + ASSERT_TRUE(PutAndWaitOffloaded(k2, v2, ssd_dir)) + << "k2 offload timed out"; + + int buckets_before = CountFilesWithSuffix(ssd_dir, ".bucket"); + ASSERT_GT(buckets_before, 0) << "No bucket files after offload"; + + // Wait for all offload tasks to settle into object_bucket_map_. + WaitForAllOffloadsSettled(); + + // Snapshot bucket file names AFTER settle (all buckets written) and + // BEFORE remove. This is the baseline for detecting GC compaction. + auto bucket_files_before = ListBucketFiles(ssd_dir); + int buckets_after_settle = CountFilesWithSuffix(ssd_dir, ".bucket"); + + // Remove k1. GC should compact (delete k1's empty bucket file). + ASSERT_EQ(real_client_->remove(k1, /*force=*/true), 0); + + // Wait for GC: bucket file set should change (old file deleted, and/or + // new file written if compaction rewrote surviving keys). + bool reclaimed = false; + for (int i = 0; i < 150; ++i) { // up to 30s + // k2 must stay readable with correct data throughout GC. + auto got2 = ReadKey(real_client_, k2); + ASSERT_TRUE(got2.has_value()) + << "Surviving key k2 became unreadable during GC"; + ASSERT_EQ(got2.value(), v2) + << "Surviving key k2 data corrupted during GC"; + // k1 must stay gone. + auto got1 = ReadKey(real_client_, k1); + ASSERT_FALSE(got1.has_value()) + << "Removed key k1 became readable again"; + + // Detect: any file in before-set gone, OR count decreased. + auto bucket_files_now = ListBucketFiles(ssd_dir); + for (const auto& old_name : bucket_files_before) { + if (std::find(bucket_files_now.begin(), + bucket_files_now.end(), + old_name) == bucket_files_now.end()) { + reclaimed = true; + break; + } + } + int buckets_now = CountFilesWithSuffix(ssd_dir, ".bucket"); + if (!reclaimed && buckets_now < buckets_after_settle) { + reclaimed = true; + } + if (reclaimed) break; + std::this_thread::sleep_for(std::chrono::milliseconds(200)); + } + + // Final checks. + auto got = ReadKey(real_client_, k2); + ASSERT_TRUE(got.has_value()); + EXPECT_EQ(got.value(), v2) << "Surviving key data corrupted after GC"; + + EXPECT_TRUE(reclaimed) << "GC did not reduce bucket file count within 30s"; +} + +// ------------------------------------------------------------------- +// Test 2: RemoveMiddleKeyPreservesSurvivors +// +// Put 2 keys, remove 1, wait for GC. The surviving key must remain +// readable with correct data. This is the core "don't lose un-removed +// keys" invariant. +// ------------------------------------------------------------------- +TEST_F(GCE2ETest, RemoveMiddleKeyPreservesSurvivors) { + ASSERT_TRUE(StartMasterWithOffload()); + ASSERT_TRUE(StartRealClient()); + + const std::string k1 = "gc_mid_k1"; + const std::string k2 = "gc_mid_k2"; + const std::string v1(4 * kMB, 'X'); + const std::string v2(4 * kMB, 'Y'); + + fs::path ssd_dir = tmp_dir_ / "ssd_offload"; + ASSERT_TRUE(PutAndWaitOffloaded(k1, v1, ssd_dir)) + << "k1 offload timed out"; + ASSERT_TRUE(PutAndWaitOffloaded(k2, v2, ssd_dir)) + << "k2 offload timed out"; + + int buckets_before = CountFilesWithSuffix(ssd_dir, ".bucket"); + ASSERT_GT(buckets_before, 0); + + // Wait for all offload tasks to settle into object_bucket_map_. + WaitForAllOffloadsSettled(); + + // Snapshot bucket file names AFTER settle (all buckets written) and + // BEFORE remove. This is the baseline for detecting GC compaction. + auto bucket_files_before = ListBucketFiles(ssd_dir); + // Re-count after settle — more buckets may have appeared. + int buckets_after_settle = CountFilesWithSuffix(ssd_dir, ".bucket"); + + // Remove k2 (force=true to bypass lease). + ASSERT_EQ(real_client_->remove(k2, /*force=*/true), 0); + + // Wait for GC: bucket file set should change. + bool compacted = false; + for (int i = 0; i < 150; ++i) { + auto got1 = ReadKey(real_client_, k1); + if (got1.has_value() && got1.value() == v1) { + auto bucket_files_now = ListBucketFiles(ssd_dir); + // Detect: any file in before-set gone, OR count decreased. + for (const auto& old_name : bucket_files_before) { + if (std::find(bucket_files_now.begin(), + bucket_files_now.end(), + old_name) == bucket_files_now.end()) { + compacted = true; + break; + } + } + int buckets_now = CountFilesWithSuffix(ssd_dir, ".bucket"); + if (!compacted && buckets_now < buckets_after_settle) { + compacted = true; + } + if (compacted) break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(200)); + } + + // k1 must survive with correct data. + auto got1 = ReadKey(real_client_, k1); + ASSERT_TRUE(got1.has_value()); + EXPECT_EQ(got1.value(), v1) << "k1 data corrupted after GC"; + + // k2 must remain gone. + auto got2 = ReadKey(real_client_, k2); + EXPECT_FALSE(got2.has_value()) + << "Removed key k2 should not be readable"; + + EXPECT_TRUE(compacted) << "GC compaction not detected within 30s"; +} + +// ------------------------------------------------------------------- +// Test 3: BatchRemoveMixedExistingAndAbsent +// +// BatchRemove with a non-existent key mixed in: the existing key should +// be tombstoned + GC'd, the non-existent one ignored. +// ------------------------------------------------------------------- +TEST_F(GCE2ETest, BatchRemoveMixedExistingAndAbsent) { + ASSERT_TRUE(StartMasterWithOffload()); + ASSERT_TRUE(StartRealClient()); + + const std::string k1 = "gc_batch_k1"; + const std::string k2 = "gc_batch_k2"; + const std::string v1(4 * kMB, 'Q'); + const std::string v2(4 * kMB, 'R'); + + fs::path ssd_dir = tmp_dir_ / "ssd_offload"; + ASSERT_TRUE(PutAndWaitOffloaded(k1, v1, ssd_dir)) + << "k1 offload timed out"; + ASSERT_TRUE(PutAndWaitOffloaded(k2, v2, ssd_dir)) + << "k2 offload timed out"; + + int buckets_before = CountFilesWithSuffix(ssd_dir, ".bucket"); + ASSERT_GT(buckets_before, 0); + + // Wait for all offload tasks to settle into object_bucket_map_. + WaitForAllOffloadsSettled(); + + // Snapshot bucket file names AFTER settle (all buckets written) and + // BEFORE remove. + auto bucket_files_before = ListBucketFiles(ssd_dir); + int buckets_after_settle = CountFilesWithSuffix(ssd_dir, ".bucket"); + + // Batch remove: k1 exists, k_absent does not. force=true bypasses lease. + std::vector keys{k1, "gc_batch_absent"}; + auto results = real_client_->batchRemove(keys, /*force=*/true); + ASSERT_EQ(results.size(), 2u); + + // Wait for GC: bucket file set should change. + bool compacted = false; + for (int i = 0; i < 150; ++i) { + auto got2 = ReadKey(real_client_, k2); + if (got2.has_value() && got2.value() == v2) { + auto bucket_files_now = ListBucketFiles(ssd_dir); + for (const auto& old_name : bucket_files_before) { + if (std::find(bucket_files_now.begin(), + bucket_files_now.end(), + old_name) == bucket_files_now.end()) { + compacted = true; + break; + } + } + int buckets_now = CountFilesWithSuffix(ssd_dir, ".bucket"); + if (!compacted && buckets_now < buckets_after_settle) { + compacted = true; + } + if (compacted) break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(200)); + } + + // k1 must be gone. + auto got1 = ReadKey(real_client_, k1); + EXPECT_FALSE(got1.has_value()) + << "Batch-removed key k1 should not be readable"; + + // k2 must survive. + auto got2 = ReadKey(real_client_, k2); + ASSERT_TRUE(got2.has_value()); + EXPECT_EQ(got2.value(), v2) << "k2 data corrupted after GC"; + + EXPECT_TRUE(compacted) << "GC compaction not detected within 30s"; +} + +} // namespace testing +} // namespace mooncake + +int main(int argc, char** argv) { + ::testing::InitGoogleTest(&argc, argv); + gflags::ParseCommandLineFlags(&argc, &argv, false); + return RUN_ALL_TESTS(); +} diff --git a/mooncake-store/tests/storage_backend_test.cpp b/mooncake-store/tests/storage_backend_test.cpp index 1f3e1ae246..77acec964d 100644 --- a/mooncake-store/tests/storage_backend_test.cpp +++ b/mooncake-store/tests/storage_backend_test.cpp @@ -2821,9 +2821,319 @@ TEST_F(StorageBackendTest, BucketStorageBackend_ConcurrentReadWriteDelete) { } //----------------------------------------------------------------------------- -// Tests for FileRecord key tracking and eviction return values +// Explicit-delete-only GC tests (tombstone + compaction) //----------------------------------------------------------------------------- +TEST_F(StorageBackendTest, BucketStorageBackend_MarkRemovedHidesKey) { + FileStorageConfig config; + config.storage_filepath = data_path; + BucketBackendConfig bucket_config; + BucketStorageBackend storage_backend(config, bucket_config); + ASSERT_TRUE(storage_backend.Init()); + + std::string k1 = "mark_k1"; + std::string k2 = "mark_k2"; + std::string v1 = "value1"; + std::string v2 = "value2"; + + std::unordered_map> batch; + auto buf1 = std::make_unique(v1.size()); + auto buf2 = std::make_unique(v2.size()); + std::memcpy(buf1.get(), v1.data(), v1.size()); + std::memcpy(buf2.get(), v2.data(), v2.size()); + batch.emplace(k1, std::vector{Slice{buf1.get(), v1.size()}}); + batch.emplace(k2, std::vector{Slice{buf2.get(), v2.size()}}); + + auto offload_result = storage_backend.BatchOffload( + batch, + [](const std::vector&, + std::vector&) { return ErrorCode::OK; }); + ASSERT_TRUE(offload_result.has_value()); + + EXPECT_TRUE(storage_backend.IsExist(k1).value()); + EXPECT_TRUE(storage_backend.IsExist(k2).value()); + + // Mark k1 removed + storage_backend.MarkRemoved(k1); + + // k1 invisible, k2 still visible + EXPECT_FALSE(storage_backend.IsExist(k1).value()); + EXPECT_TRUE(storage_backend.IsExist(k2).value()); + + // MarkRemoved is idempotent on absent key + storage_backend.MarkRemoved("nonexistent_key"); // no crash + storage_backend.MarkRemoved(k1); // already removed, idempotent +} + +TEST_F(StorageBackendTest, BucketStorageBackend_CompactReclaimsDeletedKeys) { + FileStorageConfig config; + config.storage_filepath = data_path; + BucketBackendConfig bucket_config; + bucket_config.eviction_policy = BucketEvictionPolicy::LRU; + bucket_config.disable_ssd_eviction = true; + BucketStorageBackend storage_backend(config, bucket_config); + ASSERT_TRUE(storage_backend.Init()); + + // Offload 3 keys into one bucket + std::string k1 = "compact_k1", k2 = "compact_k2", k3 = "compact_k3"; + std::string v1(1024, 'a'), v2(1024, 'b'), v3(1024, 'c'); + + std::unordered_map> batch; + auto buf1 = std::make_unique(v1.size()); + auto buf2 = std::make_unique(v2.size()); + auto buf3 = std::make_unique(v3.size()); + std::memcpy(buf1.get(), v1.data(), v1.size()); + std::memcpy(buf2.get(), v2.data(), v2.size()); + std::memcpy(buf3.get(), v3.data(), v3.size()); + batch.emplace(k1, std::vector{Slice{buf1.get(), v1.size()}}); + batch.emplace(k2, std::vector{Slice{buf2.get(), v2.size()}}); + batch.emplace(k3, std::vector{Slice{buf3.get(), v3.size()}}); + + auto offload_result = storage_backend.BatchOffload( + batch, + [](const std::vector&, + std::vector&) { return ErrorCode::OK; }); + ASSERT_TRUE(offload_result.has_value()); + int64_t old_bucket_id = offload_result.value(); + + // Mark k2 removed + storage_backend.MarkRemoved(k2); + + // Compact the bucket + ASSERT_TRUE(storage_backend.CompactBucket(old_bucket_id)); + + // k1, k3 still loadable with correct data; k2 gone + EXPECT_TRUE(storage_backend.IsExist(k1).value()); + EXPECT_TRUE(storage_backend.IsExist(k3).value()); + EXPECT_FALSE(storage_backend.IsExist(k2).value()); + + // Verify k1, k3 data integrity + auto alloc = SimpleAllocator(128 * 1024 * 1024); + void* b1 = alloc.allocate(v1.size()); + void* b3 = alloc.allocate(v3.size()); + std::unordered_map load_batch; + load_batch.emplace(k1, Slice{b1, v1.size()}); + load_batch.emplace(k3, Slice{b3, v3.size()}); + ASSERT_TRUE(storage_backend.BatchLoad(load_batch)); + EXPECT_EQ(std::string((char*)b1, v1.size()), v1); + EXPECT_EQ(std::string((char*)b3, v3.size()), v3); + + // Old bucket file should be deleted + std::string old_data_path = + data_path + "/" + std::to_string(old_bucket_id) + ".bucket"; + EXPECT_FALSE(fs::exists(old_data_path)) + << "Old bucket file should be deleted after compaction"; +} + +TEST_F(StorageBackendTest, BucketStorageBackend_MarkRemovedConcurrentLoad) { + FileStorageConfig config; + config.storage_filepath = data_path; + BucketBackendConfig bucket_config; + bucket_config.eviction_policy = BucketEvictionPolicy::LRU; + bucket_config.disable_ssd_eviction = true; + BucketStorageBackend storage_backend(config, bucket_config); + ASSERT_TRUE(storage_backend.Init()); + + std::string k1 = "conc_k1", k2 = "conc_k2"; + std::string v1(4096, 'a'), v2(4096, 'b'); + std::unordered_map> batch; + auto buf1 = std::make_unique(v1.size()); + auto buf2 = std::make_unique(v2.size()); + std::memcpy(buf1.get(), v1.data(), v1.size()); + std::memcpy(buf2.get(), v2.data(), v2.size()); + batch.emplace(k1, std::vector{Slice{buf1.get(), v1.size()}}); + batch.emplace(k2, std::vector{Slice{buf2.get(), v2.size()}}); + ASSERT_TRUE(storage_backend.BatchOffload( + batch, [](const std::vector&, + std::vector&) { + return ErrorCode::OK; + })); + + // Concurrent: load k1 in a thread while marking k2 removed. + auto alloc = SimpleAllocator(128 * 1024 * 1024); + void* b1 = alloc.allocate(v1.size()); + std::thread loader([&]() { + std::unordered_map load; + load.emplace(k1, Slice{b1, v1.size()}); + auto r = storage_backend.BatchLoad(load); + ASSERT_TRUE(r); + }); + + storage_backend.MarkRemoved(k2); + loader.join(); + + // k1 data intact, k2 gone. + EXPECT_EQ(std::string((char*)b1, v1.size()), v1); + EXPECT_FALSE(storage_backend.IsExist(k2).value()); + EXPECT_TRUE(storage_backend.IsExist(k1).value()); +} + +TEST_F(StorageBackendTest, + BucketStorageBackend_DisableEvictionNoopUnderPressure) { + FileStorageConfig config; + config.storage_filepath = data_path; + // Set global size limit smaller than a single bucket's reservation so + // IsEnableOffloading's quota check rejects the offload. + config.total_size_limit = 512; + BucketBackendConfig bucket_config; + bucket_config.eviction_policy = BucketEvictionPolicy::LRU; + bucket_config.disable_ssd_eviction = true; + // bucket_size_limit (default 256MB) >> total_size_limit (512), so the + // quota check in IsEnableOffloading rejects without eviction. + BucketStorageBackend storage_backend(config, bucket_config); + ASSERT_TRUE(storage_backend.Init()); + + // disable_ssd_eviction=true means PrepareEviction is a no-op, so under + // space pressure no bucket is deleted. IsEnableOffloading rejects via + // the quota check instead. + std::string k = "pressure_k"; + std::string v(2048, 'z'); + auto buf = std::make_unique(v.size()); + std::memcpy(buf.get(), v.data(), v.size()); + std::unordered_map> batch; + batch.emplace(k, std::vector{Slice{buf.get(), v.size()}}); + + auto result = storage_backend.BatchOffload( + batch, [](const std::vector&, + std::vector&) { + return ErrorCode::OK; + }); + // Offload should be rejected by quota (no eviction to reclaim space). + EXPECT_FALSE(result.has_value()); +} + +//----------------------------------------------------------------------------- +// Cross-bucket merge compaction tests +//----------------------------------------------------------------------------- + +TEST_F(StorageBackendTest, BucketStorageBackend_CrossBucketMergeCompaction) { + FileStorageConfig config; + config.storage_filepath = data_path; + BucketBackendConfig bucket_config; + bucket_config.eviction_policy = BucketEvictionPolicy::LRU; + bucket_config.disable_ssd_eviction = true; + // Set keys_limit=2 so 2 keys fill a bucket. We'll create 2 buckets with + // 2 keys each (4 keys total), remove 1 key from each (2 tombstones), + // then CompactBuckets should merge the 2 remaining live keys from each + // bucket into 1 new bucket. + bucket_config.bucket_keys_limit = 2; + BucketStorageBackend storage_backend(config, bucket_config); + ASSERT_TRUE(storage_backend.Init()); + + // Offload 4 keys into 2 buckets (bucket_keys_limit=2). + // Batch 1: k1, k2 -> bucket A + // Batch 2: k3, k4 -> bucket B + auto offload_batch = [&](const std::vector>& kvs) { + std::unordered_map> batch; + std::vector> bufs; + for (const auto& [k, v] : kvs) { + auto buf = std::make_unique(v.size()); + std::memcpy(buf.get(), v.data(), v.size()); + bufs.push_back(std::move(buf)); + batch.emplace(k, std::vector{ + Slice{bufs.back().get(), v.size()}}); + } + return storage_backend.BatchOffload( + batch, + [](const std::vector&, + std::vector&) { + return ErrorCode::OK; + }); + }; + + std::string k1 = "merge_k1", k2 = "merge_k2"; + std::string k3 = "merge_k3", k4 = "merge_k4"; + std::string v1(1024, 'a'), v2(1024, 'b'); + std::string v3(1024, 'c'), v4(1024, 'd'); + + auto offload_a = offload_batch({{k1, v1}, {k2, v2}}); + ASSERT_TRUE(offload_a.has_value()); + int64_t bucket_a = offload_a.value(); + auto offload_b = offload_batch({{k3, v3}, {k4, v4}}); + ASSERT_TRUE(offload_b.has_value()); + int64_t bucket_b = offload_b.value(); + + // Remove k2 from bucket A and k4 from bucket B (tombstones). + storage_backend.MarkRemoved(k2); + storage_backend.MarkRemoved(k4); + + // CompactBuckets should merge live keys k1, k3 into a new bucket. + // bucket_keys_limit=2, so k1+k3 fills exactly one bucket. + ASSERT_TRUE(storage_backend.CompactBuckets({bucket_a, bucket_b})); + + // k1 and k3 must still be readable with correct data. + EXPECT_TRUE(storage_backend.IsExist(k1).value()); + EXPECT_TRUE(storage_backend.IsExist(k3).value()); + // k2 and k4 must be gone. + EXPECT_FALSE(storage_backend.IsExist(k2).value()); + EXPECT_FALSE(storage_backend.IsExist(k4).value()); + + // Verify data integrity. + auto alloc = SimpleAllocator(128 * 1024 * 1024); + void* b1 = alloc.allocate(v1.size()); + void* b3 = alloc.allocate(v3.size()); + std::unordered_map load_batch; + load_batch.emplace(k1, Slice{b1, v1.size()}); + load_batch.emplace(k3, Slice{b3, v3.size()}); + ASSERT_TRUE(storage_backend.BatchLoad(load_batch)); + EXPECT_EQ(std::string((char*)b1, v1.size()), v1); + EXPECT_EQ(std::string((char*)b3, v3.size()), v3); + + // Old bucket files should be deleted (both A and B fully migrated). + std::string path_a = data_path + "/" + std::to_string(bucket_a) + ".bucket"; + std::string path_b = data_path + "/" + std::to_string(bucket_b) + ".bucket"; + EXPECT_FALSE(fs::exists(path_a)) + << "Old bucket A file should be deleted after merge"; + EXPECT_FALSE(fs::exists(path_b)) + << "Old bucket B file should be deleted after merge"; +} + +TEST_F(StorageBackendTest, BucketStorageBackend_MergeDeferredWhenNotFull) { + FileStorageConfig config; + config.storage_filepath = data_path; + BucketBackendConfig bucket_config; + bucket_config.eviction_policy = BucketEvictionPolicy::LRU; + bucket_config.disable_ssd_eviction = true; + // keys_limit=2: need 2 live keys to fill a bucket. + bucket_config.bucket_keys_limit = 2; + BucketStorageBackend storage_backend(config, bucket_config); + ASSERT_TRUE(storage_backend.Init()); + + // Offload 2 keys into 1 bucket. + std::unordered_map> batch; + auto buf1 = std::make_unique(1024); + auto buf2 = std::make_unique(1024); + std::memset(buf1.get(), 'x', 1024); + std::memset(buf2.get(), 'y', 1024); + batch.emplace("defer_k1", std::vector{Slice{buf1.get(), 1024}}); + batch.emplace("defer_k2", std::vector{Slice{buf2.get(), 1024}}); + auto offload_result = storage_backend.BatchOffload( + batch, + [](const std::vector&, + std::vector&) { return ErrorCode::OK; }); + ASSERT_TRUE(offload_result.has_value()); + int64_t bucket_id = offload_result.value(); + + // Remove k2 — only 1 live key (k1) remains. Not enough for keys_limit=2. + storage_backend.MarkRemoved("defer_k2"); + + // CompactBuckets without space_pressure should defer (return true, no + // compaction). The old bucket should still exist. + ASSERT_TRUE(storage_backend.CompactBuckets({bucket_id}, false)); + EXPECT_TRUE(storage_backend.IsExist("defer_k1").value()); + // Old bucket file should still exist (not compacted). + std::string path = data_path + "/" + std::to_string(bucket_id) + ".bucket"; + EXPECT_TRUE(fs::exists(path)) + << "Bucket should not be compacted when live keys don't fill a bucket"; + + // With space_pressure=true, compaction should proceed even if not full. + ASSERT_TRUE(storage_backend.CompactBuckets({bucket_id}, true)); + EXPECT_TRUE(storage_backend.IsExist("defer_k1").value()); + EXPECT_FALSE(fs::exists(path)) + << "Bucket should be compacted under space pressure"; +} + + TEST_F(StorageBackendTest, StoreObjectReturnsEvictedKeys) { std::string test_dir = data_path + "/evict_return_test"; std::filesystem::create_directories(test_dir); diff --git a/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h b/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h index 07d0cb5711..1f119e7679 100644 --- a/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h +++ b/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h @@ -195,15 +195,18 @@ class UbContext { // Polls one JFC and processes the slices internally: // * Aggregates each completion's jetty depth into jetty_depth_set. - // * Successful slices have markSuccess() called in place and are NOT - // returned (they may be recycled by the submitting thread the + // * Successful normal slices have markSuccess() called in place and are + // NOT returned (they may be recycled by the submitting thread the // moment markSuccess() runs). + // * Successful slices that need post-URMA staging work are returned in + // deferred_success_slices and are NOT marked successful yet. // * Failed slices are returned in failed_slices[0..num_failed-1] for // the caller to apply retry / markFailed. // Returns the total number of completions polled (>= 0), or a negative // error code. virtual int poll(int num_entries, Transport::Slice** failed_slices, int& num_failed, + std::vector& deferred_success_slices, std::unordered_map& jetty_depth_set, int jfc_index = 0) = 0; diff --git a/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_transport.h b/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_transport.h index d38c72083e..1ba12c0815 100644 --- a/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_transport.h +++ b/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_transport.h @@ -14,8 +14,15 @@ #ifndef UB_TRANSPORT_H #define UB_TRANSPORT_H +#include +#include +#include +#include +#include #include #include +#include +#include #include #include "topology.h" #include "transfer_metadata.h" @@ -24,6 +31,7 @@ namespace mooncake { class UbContext; class UbEndPoint; +class UrmaContext; class TransferMetadata; class UbWorkerpool; @@ -36,6 +44,7 @@ enum UB_ENDPOINT_TYPE { URMA_ENDPOINT = 0, OBMM_ENDPOINT = 1 }; class UbTransport : public Transport { friend class UbContext; friend class UbEndPoint; + friend class UrmaContext; friend class UbWorkerPool; public: @@ -82,6 +91,55 @@ class UbTransport : public Transport { private: int allocateLocalSegmentID(); + struct StagingLease { + void* host_ptr = nullptr; + size_t size = 0; + }; + + struct DeviceRegion { + uint64_t addr = 0; + size_t length = 0; + std::string location; + bool remote_accessible = true; + }; + + struct StagingState { + void* original_device_ptr = nullptr; + void* staging_ptr = nullptr; + size_t size = 0; + size_t lease_size = 0; + TransferRequest::OpCode opcode = TransferRequest::WRITE; + std::atomic completed_slices{0}; + uint64_t total_slices = 0; + std::atomic failed{false}; + std::mutex deferred_mutex; + std::vector deferred_success_slices; + }; + + bool stagingEnabled() const; + bool isDevicePointer(const void* ptr) const; + bool isLogicalDeviceRange(const void* ptr, size_t length) const; + int registerLogicalDeviceRegion(void* addr, size_t length, + const std::string& location, + bool remote_accessible); + int unregisterLogicalDeviceRegion(void* addr); + bool copyDeviceToHost(void* dst, const void* src, size_t size) const; + bool copyHostToDevice(void* dst, const void* src, size_t size) const; + Status acquireStaging(size_t size, StagingLease& lease); + void releaseStaging(const StagingLease& lease); + bool isStagedSlice(Slice* slice); + bool shouldDeferSuccess(Slice* slice); + void attachStaging(TransferTask* task, + std::shared_ptr state); + void attachStagingSlice(Slice* slice, + const std::shared_ptr& state); + void detachStagingSlice(Slice* slice); + std::shared_ptr stagingStateForSlice(Slice* slice); + void cleanupStagingForTask(TransferTask* task, + bool detach_all_slices = false); + void onStagedSliceSuccess(Slice* slice); + void onStagedSliceFinalFailure(Slice* slice); + public: int onSetupConnections(const HandShakeDesc& peer_desc, HandShakeDesc& local_desc); @@ -118,6 +176,21 @@ class UbTransport : public Transport { std::shared_ptr local_topology_; UB_ENDPOINT_TYPE endpoint_type_; bool runtime_initialized_ = false; + + mutable std::once_flag staging_config_once_; + mutable bool staging_enabled_ = false; + mutable std::mutex device_region_mutex_; + std::vector device_regions_; + std::mutex staging_pool_mutex_; + void* staging_pool_base_ = nullptr; + size_t staging_pool_size_ = 0; + size_t staging_pool_offset_ = 0; + std::deque> staging_free_list_; + std::mutex staging_state_mutex_; + std::unordered_map> + task_staging_map_; + std::unordered_map> + slice_staging_map_; }; } // namespace mooncake diff --git a/mooncake-transfer-engine/include/transport/kunpeng_transport/urma/urma_endpoint.h b/mooncake-transfer-engine/include/transport/kunpeng_transport/urma/urma_endpoint.h index 846ac640c9..59fb599600 100644 --- a/mooncake-transfer-engine/include/transport/kunpeng_transport/urma/urma_endpoint.h +++ b/mooncake-transfer-engine/include/transport/kunpeng_transport/urma/urma_endpoint.h @@ -62,6 +62,7 @@ class UrmaContext : public UbContext { int doProcessContextEvents() override; void* retrieveRemoteSeg(const std::string& value) override; int poll(int num_entries, Transport::Slice** failed_slices, int& num_failed, + std::vector& deferred_success_slices, std::unordered_map& jetty_depth_set, int jfc_index) override; volatile int* outstandingCount(int jfc_index) override; diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp index cda2b97cd6..61b8108882 100644 --- a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp @@ -442,18 +442,19 @@ void UbWorkerPool::performPostSend(int thread_id) { void UbWorkerPool::performPoll(int thread_id) { int processed_slice_count = 0; const static size_t kPollCount = 64; - // context_.poll() aggregates each completion's jetty_depth here and - // calls markSuccess() on successful slices in place. Successful slices - // are NOT returned from poll(), so this worker never dereferences them - // after they may have been recycled by the submitting thread. + // context_.poll() aggregates each completion's jetty_depth here. Normal + // successful slices are published in place; staged READ successes are + // returned in deferred_success_slices so H2D can finish before publishing. std::unordered_map jetty_depth_set; std::vector failed_slices; + std::vector deferred_success_slices; for (int jfc_index = thread_id; jfc_index < context_.jfcCount(); jfc_index += kTransferWorkerCount) { UbTransport::Slice* failed[kPollCount]; int num_failed = 0; int nr_poll = context_.poll(kPollCount, failed, num_failed, - jetty_depth_set, jfc_index); + deferred_success_slices, jetty_depth_set, + jfc_index); if (nr_poll < 0) { LOG(ERROR) << "Worker: Failed to poll jetty for complete"; continue; @@ -488,6 +489,10 @@ void UbWorkerPool::performPoll(int thread_id) { for (auto& entry : jetty_depth_set) __sync_fetch_and_sub(entry.first, entry.second); + for (auto& slice : deferred_success_slices) { + context_.engine().onStagedSliceSuccess(slice); + } + // Slices that hit max_retry: final markFailed() after all reads (and the // jetty depth returns above) are done. Failed slices were never published // by poll(), so they remained safe to deref up to this point. @@ -496,7 +501,11 @@ void UbWorkerPool::performPoll(int thread_id) { auto ptr = static_cast(slice->ub.endpoint); context_.deleteEndpointByPtr(ptr); } - slice->markFailed(); + if (context_.engine().isStagedSlice(slice)) { + context_.engine().onStagedSliceFinalFailure(slice); + } else { + slice->markFailed(); + } processed_slice_count_++; } @@ -633,4 +642,4 @@ void UbWorkerPool::monitorWorker() { int UbWorkerPool::doProcessContextEvents() { return context_.doProcessContextEvents(); } -} // namespace mooncake \ No newline at end of file +} // namespace mooncake diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp index 78a5c205dd..6016f828d1 100644 --- a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp @@ -13,10 +13,18 @@ // limitations under the License. #include +#include +#include +#include +#include #include +#include +#include #include "config.h" +#include "cuda_alike.h" #include "memory_location.h" #include +#include "ub_allocator.h" #include "transport/kunpeng_transport/ub_context.h" #include "transport/kunpeng_transport/ub_transport.h" #include "transport/kunpeng_transport/ub_endpoint.h" @@ -26,6 +34,37 @@ namespace mooncake { namespace { constexpr uint64_t kNumaAffinitySampleInterval = 10000; +constexpr size_t kDefaultUbStagingPoolSize = 1ull << 30; + +size_t alignUp(size_t value, size_t alignment) { + return (value + alignment - 1) / alignment * alignment; +} + +uint64_t countUbSlices(size_t length, size_t block_size, + size_t fragment_size) { + uint64_t count = 0; + for (uint64_t offset = 0; offset < length; offset += block_size) { + ++count; + if (length - offset <= block_size + fragment_size) break; + } + return count; +} + +size_t getUbStagingPoolSize() { + static const size_t pool_size = [] { + const char* env = std::getenv("MC_UB_STAGING_POOL_SIZE"); + if (!env) return kDefaultUbStagingPoolSize; + try { + size_t value = std::stoull(env); + if (value > 0) return value; + } catch (const std::exception& e) { + LOG(WARNING) << "Invalid MC_UB_STAGING_POOL_SIZE value: " << env + << ", error: " << e.what(); + } + return kDefaultUbStagingPoolSize; + }(); + return pool_size; +} } // namespace UbTransport::UbTransport(UB_ENDPOINT_TYPE endpoint_type) @@ -35,11 +74,323 @@ UbTransport::~UbTransport() { #ifdef CONFIG_USE_BATCH_DESC_SET batch_desc_set_.clear(); #endif + if (staging_pool_base_) { + unregisterLocalMemory(staging_pool_base_, true); + ub_free_memory(staging_pool_base_); + staging_pool_base_ = nullptr; + staging_pool_size_ = 0; + } metadata_->removeSegmentDesc(local_server_name_); batch_desc_set_.clear(); context_list_.clear(); } +bool UbTransport::stagingEnabled() const { + std::call_once(staging_config_once_, [this] { + const char* env = std::getenv("MC_UB_TRANSPORT_CPU_STAGING"); + staging_enabled_ = !env || std::string(env) != "0"; + LOG(INFO) << "UbTransport CPU staging " + << (staging_enabled_ ? "enabled" : "disabled"); + }); + return staging_enabled_; +} + +bool UbTransport::isDevicePointer(const void* ptr) const { + if (!ptr) return false; +#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \ + defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \ + defined(USE_COREX) + cudaPointerAttributes attributes; + auto status = cudaPointerGetAttributes(&attributes, ptr); + if (status != cudaSuccess) { + return false; + } + return attributes.type == cudaMemoryTypeDevice; +#else + return false; +#endif +} + +bool UbTransport::isLogicalDeviceRange(const void* ptr, size_t length) const { + if (!ptr || length == 0) return false; + uint64_t addr = reinterpret_cast(ptr); + std::lock_guard lock(device_region_mutex_); + for (const auto& region : device_regions_) { + if (addr < region.addr || length > region.length) continue; + if (addr - region.addr <= region.length - length) return true; + } + return false; +} + +int UbTransport::registerLogicalDeviceRegion(void* addr, size_t length, + const std::string& location, + bool remote_accessible) { + if (!addr || length == 0) return ERR_INVALID_ARGUMENT; + uint64_t start = reinterpret_cast(addr); + if (start + length < start) return ERR_INVALID_ARGUMENT; + uint64_t end = start + length; + std::lock_guard lock(device_region_mutex_); + for (const auto& region : device_regions_) { + uint64_t region_start = region.addr; + uint64_t region_end = region.addr + region.length; + if (start < region_end && region_start < end) { + LOG(ERROR) << "UbTransport: logical device region overlaps, addr=" + << addr << " length=" << length; + return ERR_ADDRESS_OVERLAPPED; + } + } + device_regions_.push_back( + DeviceRegion{start, length, location, remote_accessible}); + LOG(INFO) << "UbTransport: registered logical device region addr=" << addr + << " length=" << length << " location=" << location; + return 0; +} + +int UbTransport::unregisterLogicalDeviceRegion(void* addr) { + if (!addr) return ERR_INVALID_ARGUMENT; + uint64_t start = reinterpret_cast(addr); + std::lock_guard lock(device_region_mutex_); + for (auto it = device_regions_.begin(); it != device_regions_.end(); ++it) { + if (it->addr != start) continue; + LOG(INFO) << "UbTransport: unregistered logical device region addr=" + << addr << " length=" << it->length; + device_regions_.erase(it); + return 0; + } + return ERR_ADDRESS_NOT_REGISTERED; +} + +bool UbTransport::copyDeviceToHost(void* dst, const void* src, + size_t size) const { +#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \ + defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \ + defined(USE_COREX) + return cudaMemcpy(dst, src, size, cudaMemcpyDeviceToHost) == cudaSuccess; +#else + (void)dst; + (void)src; + (void)size; + return false; +#endif +} + +bool UbTransport::copyHostToDevice(void* dst, const void* src, + size_t size) const { +#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \ + defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \ + defined(USE_COREX) + return cudaMemcpy(dst, src, size, cudaMemcpyHostToDevice) == cudaSuccess; +#else + (void)dst; + (void)src; + (void)size; + return false; +#endif +} + +Status UbTransport::acquireStaging(size_t size, StagingLease& lease) { + if (size == 0) return Status::InvalidArgument("zero-sized UB staging"); + const size_t alignment = 4096; + const size_t aligned_size = alignUp(size, alignment); + + std::lock_guard lock(staging_pool_mutex_); + if (!staging_pool_base_) { + staging_pool_size_ = std::max(getUbStagingPoolSize(), aligned_size); + staging_pool_base_ = ub_allocate_memory(alignment, staging_pool_size_); + if (!staging_pool_base_) { + return Status::Memory("UbTransport: allocate CPU staging pool"); + } + int ret = registerLocalMemory(staging_pool_base_, staging_pool_size_, + kWildcardLocation, false, true); + if (ret) { + ub_free_memory(staging_pool_base_); + staging_pool_base_ = nullptr; + staging_pool_size_ = 0; + return Status::Context( + "UbTransport: register CPU staging pool failed"); + } + LOG(INFO) << "UbTransport: registered CPU staging pool base=" + << staging_pool_base_ << " size=" << staging_pool_size_; + } + + for (auto it = staging_free_list_.begin(); it != staging_free_list_.end(); + ++it) { + if (it->second < aligned_size) continue; + lease.host_ptr = it->first; + lease.size = it->second; + staging_free_list_.erase(it); + return Status::OK(); + } + + if (staging_pool_offset_ + aligned_size > staging_pool_size_) { + return Status::Memory("UbTransport: CPU staging pool exhausted"); + } + lease.host_ptr = static_cast(staging_pool_base_) + + staging_pool_offset_; + lease.size = aligned_size; + staging_pool_offset_ += aligned_size; + return Status::OK(); +} + +void UbTransport::releaseStaging(const StagingLease& lease) { + if (!lease.host_ptr || lease.size == 0) return; + std::lock_guard lock(staging_pool_mutex_); + staging_free_list_.push_back({lease.host_ptr, lease.size}); +} + +void UbTransport::attachStaging(TransferTask* task, + std::shared_ptr state) { + if (!task || !state) return; + std::lock_guard lock(staging_state_mutex_); + task_staging_map_[task] = state; +} + +void UbTransport::attachStagingSlice( + Slice* slice, const std::shared_ptr& state) { + if (!slice || !state) return; + std::lock_guard lock(staging_state_mutex_); + slice_staging_map_[slice] = state; +} + +void UbTransport::detachStagingSlice(Slice* slice) { + if (!slice) return; + std::lock_guard lock(staging_state_mutex_); + slice_staging_map_.erase(slice); +} + +std::shared_ptr UbTransport::stagingStateForSlice( + Slice* slice) { + std::lock_guard lock(staging_state_mutex_); + auto it = slice_staging_map_.find(slice); + return it == slice_staging_map_.end() ? nullptr : it->second; +} + +bool UbTransport::isStagedSlice(Slice* slice) { + return stagingStateForSlice(slice) != nullptr; +} + +bool UbTransport::shouldDeferSuccess(Slice* slice) { + auto state = stagingStateForSlice(slice); + return state && state->opcode == TransferRequest::READ; +} + +void UbTransport::cleanupStagingForTask(TransferTask* task, + bool detach_all_slices) { + std::shared_ptr state; + { + std::lock_guard lock(staging_state_mutex_); + auto task_it = task_staging_map_.find(task); + if (task_it == task_staging_map_.end()) return; + state = task_it->second; + task_staging_map_.erase(task_it); + if (detach_all_slices) { + for (auto* slice : task->slice_list) { + slice_staging_map_.erase(slice); + } + } + } + releaseStaging(StagingLease{state->staging_ptr, state->lease_size}); +} + +void UbTransport::onStagedSliceSuccess(Slice* slice) { + auto state = stagingStateForSlice(slice); + if (!state) { + slice->markSuccess(); + return; + } + + if (state->opcode == TransferRequest::WRITE) { + auto completed = + state->completed_slices.fetch_add(1, std::memory_order_acq_rel) + + 1; + if (completed == state->total_slices) { + cleanupStagingForTask(slice->task); + } + detachStagingSlice(slice); + slice->markSuccess(); + return; + } + + auto completed = + state->completed_slices.fetch_add(1, std::memory_order_acq_rel) + 1; + { + std::lock_guard lock(state->deferred_mutex); + state->deferred_success_slices.push_back(slice); + } + if (completed != state->total_slices) return; + + if (state->failed.load(std::memory_order_acquire)) { + std::vector deferred; + { + std::lock_guard lock(state->deferred_mutex); + deferred.swap(state->deferred_success_slices); + } + cleanupStagingForTask(slice->task); + for (auto* deferred_slice : deferred) { + detachStagingSlice(deferred_slice); + deferred_slice->markFailed(); + } + return; + } + + if (!copyHostToDevice(state->original_device_ptr, state->staging_ptr, + state->size)) { + LOG(ERROR) << "UbTransport: H2D staging copy failed for READ, size=" + << state->size << " dst=" << state->original_device_ptr; + state->failed.store(true, std::memory_order_release); + std::vector deferred; + { + std::lock_guard lock(state->deferred_mutex); + deferred.swap(state->deferred_success_slices); + } + cleanupStagingForTask(slice->task); + for (auto* deferred_slice : deferred) { + detachStagingSlice(deferred_slice); + deferred_slice->markFailed(); + } + return; + } + + std::vector deferred; + { + std::lock_guard lock(state->deferred_mutex); + deferred.swap(state->deferred_success_slices); + } + cleanupStagingForTask(slice->task); + for (auto* deferred_slice : deferred) { + detachStagingSlice(deferred_slice); + deferred_slice->markSuccess(); + } +} + +void UbTransport::onStagedSliceFinalFailure(Slice* slice) { + auto state = stagingStateForSlice(slice); + if (!state) { + slice->markFailed(); + return; + } + state->failed.store(true, std::memory_order_release); + auto completed = + state->completed_slices.fetch_add(1, std::memory_order_acq_rel) + 1; + detachStagingSlice(slice); + if (completed != state->total_slices) { + slice->markFailed(); + return; + } + + std::vector deferred; + if (state->opcode == TransferRequest::READ) { + std::lock_guard lock(state->deferred_mutex); + deferred.swap(state->deferred_success_slices); + } + cleanupStagingForTask(slice->task); + slice->markFailed(); + for (auto* deferred_slice : deferred) { + detachStagingSlice(deferred_slice); + deferred_slice->markFailed(); + } +} + int UbTransport::install(std::string& local_server_name, std::shared_ptr meta, std::shared_ptr topo) { @@ -90,7 +441,17 @@ int UbTransport::registerLocalMemory(void* addr, size_t length, const std::string& name, bool remote_accessible, bool update_metadata) { - (void)remote_accessible; + if (isDevicePointer(addr)) { + if (!stagingEnabled()) { + LOG(ERROR) << "UbTransport: refusing to register device memory " + "while CPU staging is disabled, addr=" + << addr << " length=" << length; + return ERR_INVALID_ARGUMENT; + } + return registerLogicalDeviceRegion(addr, length, name, + remote_accessible); + } + BufferDesc buffer_desc; for (auto& context : context_list_) { int ret = context->registerMemoryRegion((uint64_t)addr, length); @@ -135,6 +496,9 @@ int UbTransport::registerLocalMemory(void* addr, size_t length, } int UbTransport::unregisterLocalMemory(void* addr, bool update_metadata) { + int logical_rc = unregisterLogicalDeviceRegion(addr); + if (logical_rc == 0) return 0; + int rc = metadata_->removeLocalMemoryBuffer(addr, update_metadata); if (rc) return rc; for (auto& context : context_list_) @@ -225,7 +589,6 @@ Status UbTransport::submitTransferTask( const std::vector& task_list) { std::unordered_map, std::vector> slices_to_post; - auto local_segment_desc = metadata_->getSegmentDescByID(LOCAL_SEGMENT_ID); const size_t kBlockSize = globalConfig().slice_size; const int kMaxRetryCount = globalConfig().retry_cnt; const size_t kFragmentSize = globalConfig().fragment_limit; @@ -238,9 +601,58 @@ Status UbTransport::submitTransferTask( nr_slices = 0; assert(task.request); auto& request = *task.request; + void* effective_source = request.source; + std::shared_ptr staging_state; + StagingLease staging_lease; + bool staged_request = false; + + if (isDevicePointer(request.source)) { + if (!stagingEnabled()) { + return Status::InvalidArgument( + "UbTransport: device pointer requires CPU staging"); + } + if (!isLogicalDeviceRange(request.source, request.length)) { + return Status::AddressNotRegistered( + "UbTransport: device pointer is not registered as a " + "logical UB device region, address: " + + std::to_string( + reinterpret_cast(request.source))); + } + auto staging_status = acquireStaging(request.length, staging_lease); + if (!staging_status.ok()) return staging_status; + + staging_state = std::make_shared(); + if (!staging_state) { + releaseStaging(staging_lease); + return Status::Memory("UbTransport: allocate staging state"); + } + staging_state->original_device_ptr = request.source; + staging_state->staging_ptr = staging_lease.host_ptr; + staging_state->size = request.length; + staging_state->lease_size = staging_lease.size; + staging_state->opcode = request.opcode; + staging_state->total_slices = + countUbSlices(request.length, kBlockSize, kFragmentSize); + effective_source = staging_state->staging_ptr; + staged_request = true; + + if (request.opcode == TransferRequest::WRITE && + !copyDeviceToHost(staging_state->staging_ptr, request.source, + request.length)) { + LOG(ERROR) + << "UbTransport: D2H staging copy failed for WRITE, size=" + << request.length << " src=" << request.source; + releaseStaging(staging_lease); + return Status::Memory("UbTransport: D2H staging copy failed"); + } + attachStaging(&task, staging_state); + } + + auto local_segment_desc = + metadata_->getSegmentDescByID(LOCAL_SEGMENT_ID); auto request_buffer_id = -1, request_device_id = -1; - if (selectDevice(local_segment_desc.get(), (uint64_t)request.source, + if (selectDevice(local_segment_desc.get(), (uint64_t)effective_source, request.length, request_buffer_id, request_device_id)) { request_buffer_id = -1; @@ -258,7 +670,7 @@ Status UbTransport::submitTransferTask( slice->dest_rkeys.clear(); bool merge_final_slice = request.length - offset <= kBlockSize + kFragmentSize; - slice->source_addr = (char*)request.source + offset; + slice->source_addr = (char*)effective_source + offset; slice->length = merge_final_slice ? request.length - offset : kBlockSize; slice->opcode = request.opcode; @@ -275,6 +687,7 @@ Status UbTransport::submitTransferTask( slice->ub.src_chip_id = INVALID_CHIP_ID; slice->ub.dst_chip_id = INVALID_CHIP_ID; task.slice_list.push_back(slice); + if (staged_request) attachStagingSlice(slice, staging_state); int buffer_id = -1, device_id = -1, retry_cnt = request.advise_retry_cnt; @@ -314,6 +727,7 @@ Status UbTransport::submitTransferTask( LOG(ERROR) << "UbTransport: Address not registered by any device(s) " << source_addr; + if (staged_request) cleanupStagingForTask(&task, true); return Status::AddressNotRegistered( "UbTransport: not registered by any device(s), " "address: " + @@ -323,6 +737,7 @@ Status UbTransport::submitTransferTask( auto& context = context_list_[device_id]; if (!context->active()) { LOG(ERROR) << "Device " << device_id << " is not active"; + if (staged_request) cleanupStagingForTask(&task, true); return Status::InvalidArgument( "Device " + std::to_string(device_id) + " is not active"); } @@ -360,7 +775,7 @@ Status UbTransport::submitTransferTask( slices_to_post[context].push_back(slice); task.total_bytes += slice->length; __sync_fetch_and_add(&task.slice_count, 1); - if (nr_slices >= kSubmitWatermark) { + if (!staged_request && nr_slices >= kSubmitWatermark) { for (auto& entry : slices_to_post) entry.first->submitPostSend(entry.second); slices_to_post.clear(); @@ -371,6 +786,9 @@ Status UbTransport::submitTransferTask( break; } } + if (staged_request && task.slice_count == 0) { + cleanupStagingForTask(&task, true); + } } for (auto& entry : slices_to_post) entry.first->submitPostSend(entry.second); diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/urma_endpoint.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/urma_endpoint.cpp index 607453645d..fa69ac3f01 100644 --- a/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/urma_endpoint.cpp +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/urma_endpoint.cpp @@ -635,6 +635,7 @@ bool UrmaContext::transEidFromString(const std::string& eid_str, int UrmaContext::poll(int num_entries, Transport::Slice** failed_slices, int& num_failed, + std::vector& deferred_success_slices, std::unordered_map& jetty_depth_set, int jfc_index) { num_failed = 0; @@ -662,9 +663,17 @@ int UrmaContext::poll(int num_entries, Transport::Slice** failed_slices, jetty_depth_set[depth] = 1; if (cr[i].status == URMA_CR_SUCCESS) { + if (engine().shouldDeferSuccess(slice)) { + deferred_success_slices.push_back(slice); + continue; + } // Safe to publish here — we are done with this slice and do not // return it to the caller, so no one else will deref it. - slice->markSuccess(); + if (engine().isStagedSlice(slice)) { + engine().onStagedSliceSuccess(slice); + } else { + slice->markSuccess(); + } continue; }