From e36d8287cbc44096114c0f42049f20be7396373d Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Tue, 8 Sep 2026 14:52:10 +0800 Subject: [PATCH 1/6] [Store] Drain in-flight RPCs and keep pools alive across client teardown Teardown of a read-side client after consecutive RPC timeouts released the client pool while request coroutines were still suspended in it, and ylt only documents send_request as thread-safe, so the resumed coroutine touched freed state and segfaulted in asio's epoll_reactor (#3909). Two layers, both verified against current main: - A shared RpcDrainGuard (mooncake-common) stops admitting new calls once draining and makes the destructor wait for in-flight ones, with a 30s bound that logs loudly rather than silently proceeding. Wired into MasterClient, DummyClient and ClientRequester entry points. - RpcClientPool now hands out pools from a process-wide registry keyed by address instead of freeing them on teardown or address flaps. ylt pools own background reconnect coroutines that reference pool storage, so no amount of request draining makes pool destruction safe mid-retry; pools are one-per-master-address and deliberately process-lifetime. Signed-off-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com> --- .../include/rpc_client_io_context.h | 105 +++++++++++++++++- mooncake-common/tests/CMakeLists.txt | 5 + .../tests/rpc_client_io_context_test.cpp | 6 +- .../tests/rpc_drain_guard_test.cpp | 66 +++++++++++ mooncake-store/include/dummy_client.h | 2 + mooncake-store/include/master_client.h | 2 + mooncake-store/include/pyclient.h | 5 + mooncake-store/src/dummy_client.cpp | 18 ++- mooncake-store/src/master_client.cpp | 18 ++- mooncake-store/src/real_client.cpp | 11 ++ mooncake-store/tests/rpc_timeout_test.cpp | 43 +++++++ 11 files changed, 274 insertions(+), 7 deletions(-) create mode 100644 mooncake-common/tests/rpc_drain_guard_test.cpp diff --git a/mooncake-common/include/rpc_client_io_context.h b/mooncake-common/include/rpc_client_io_context.h index b9842f0c99..f9083199b7 100644 --- a/mooncake-common/include/rpc_client_io_context.h +++ b/mooncake-common/include/rpc_client_io_context.h @@ -1,13 +1,16 @@ #pragma once #include +#include #include +#include #include #include #include #include #include #include +#include #include #include @@ -19,16 +22,106 @@ namespace mooncake { std::shared_ptr CreateRpcClientIoContextPool( uint32_t thread_count); +/** + * Teardown guard for RPC entry points that serve out of a shared pool. + * + * The ylt client pool documents only send_request as thread-safe; closing or + * destroying the pool while a request coroutine is suspended leaves the + * resumed coroutine touching freed connection state (the teardown segfault in + * #3909). Entry points take a ScopedCall; a destructor calls drain_for() so + * the pool is only released once nothing is in flight. + */ +class RpcDrainGuard { + public: + class ScopedCall { + public: + explicit ScopedCall(RpcDrainGuard& guard) + : guard_(guard), ok_(guard.try_enter()) {} + ~ScopedCall() { + if (ok_) guard_.leave(); + } + ScopedCall(const ScopedCall&) = delete; + ScopedCall& operator=(const ScopedCall&) = delete; + bool ok() const { return ok_; } + + private: + RpcDrainGuard& guard_; + bool ok_; + }; + + // Bounded wait for in-flight calls; stops admitting new ones first. + // Returns false on timeout, in which case the caller must not free shared + // pool state without accepting the pre-existing UAF risk. + bool drain_for(std::chrono::milliseconds timeout) { + stopping_.store(true, std::memory_order_release); + std::unique_lock lk(mutex_); + return cv_.wait_for(lk, timeout, [&] { + return inflight_.load(std::memory_order_acquire) == 0; + }); + } + + private: + bool try_enter() { + if (stopping_.load(std::memory_order_acquire)) return false; + inflight_.fetch_add(1, std::memory_order_acq_rel); + // a drain that started between the two reads still sees this call + if (stopping_.load(std::memory_order_acquire)) { + leave(); + return false; + } + return true; + } + + void leave() { + if (inflight_.fetch_sub(1, std::memory_order_acq_rel) == 1) { + std::lock_guard lk(mutex_); + cv_.notify_all(); + } + } + + std::atomic stopping_{false}; + std::atomic inflight_{0}; + std::mutex mutex_; + std::condition_variable cv_; +}; + template coro_io::io_context_pool& GetRpcClientIoContextPool(uint32_t thread_count) { static const auto io_pool = CreateRpcClientIoContextPool(thread_count); return *io_pool; } +namespace detail { + +// Process-wide client-pool registry, keyed by address. A ylt client pool owns +// background reconnect coroutines that hold references into pool storage +// (client_pool.hpp's reconnect loop), so freeing a pool while they are +// suspended is a use-after-free regardless of whether any user request is in +// flight (#3909). Pools are few in practice (one per distinct master address +// per process), so they are deliberately kept alive for the process lifetime. +inline std::shared_ptr> +SharedPoolRegistryImpl( + std::string_view address, + coro_io::client_pool::pool_config config, + coro_io::io_context_pool& io_context_pool) { + using Pool = coro_io::client_pool; + static std::mutex registry_mutex; + static std::unordered_map> registry; + std::lock_guard lock(registry_mutex); + auto& slot = registry[std::string(address)]; + if (!slot) { + slot = Pool::create(std::string(address), std::move(config), + io_context_pool); + } + return slot; +} + +} // namespace detail + /** * A replaceable client pool for callers that communicate with one target at a - * time. Requests retain a shared_ptr to the old pool while they are in flight; - * after an address switch the old pool is destroyed when those requests end. + * time. Pools live in the process-wide registry above; an address switch only + * re-points this holder, it never frees a pool. */ class RpcClientPool { public: @@ -46,8 +139,7 @@ class RpcClientPool { std::string_view address) { std::lock_guard lock(mutex_); if (!client_pool_ || address_ != address) { - client_pool_ = - ClientPool::create(address, config_, io_context_pool_); + client_pool_ = SharedPoolRegistry(address); address_ = address; } return client_pool_; @@ -59,6 +151,11 @@ class RpcClientPool { } private: + std::shared_ptr SharedPoolRegistry(std::string_view address) { + return detail::SharedPoolRegistryImpl(address, config_, + io_context_pool_); + } + mutable std::shared_mutex mutex_; coro_io::io_context_pool& io_context_pool_; PoolConfig config_; diff --git a/mooncake-common/tests/CMakeLists.txt b/mooncake-common/tests/CMakeLists.txt index 6524b4cd3b..9a003bf484 100644 --- a/mooncake-common/tests/CMakeLists.txt +++ b/mooncake-common/tests/CMakeLists.txt @@ -21,6 +21,11 @@ target_link_libraries(rpc_client_io_context_test PUBLIC mooncake_common gtest ibverbs pthread) add_test(NAME rpc_client_io_context_test COMMAND rpc_client_io_context_test) +add_executable(rpc_drain_guard_test rpc_drain_guard_test.cpp) +target_link_libraries(rpc_drain_guard_test PUBLIC mooncake_common gtest + gtest_main pthread) +add_test(NAME rpc_drain_guard_test COMMAND rpc_drain_guard_test) + foreach(parser_test ascii_string bool_parser environment_value_parser integer_parser) add_executable(${parser_test}_test ${parser_test}_test.cpp) diff --git a/mooncake-common/tests/rpc_client_io_context_test.cpp b/mooncake-common/tests/rpc_client_io_context_test.cpp index c3b18c7127..a1c014f635 100644 --- a/mooncake-common/tests/rpc_client_io_context_test.cpp +++ b/mooncake-common/tests/rpc_client_io_context_test.cpp @@ -51,7 +51,11 @@ TEST(RpcClientIoContextPoolTest, ReplacesPoolWhenTargetChanges) { auto second = pools.GetOrCreateClientPool("127.0.0.1:10002"); EXPECT_NE(first, second); first.reset(); - EXPECT_TRUE(old_pool.expired()); + // Pools live in the process-wide registry by design: a ylt pool owns + // background reconnect coroutines that reference pool storage, so freeing + // one mid-retry is a use-after-free regardless of request draining + // (#3909). Switching address only re-points the holder. + EXPECT_FALSE(old_pool.expired()); EXPECT_EQ(pools.GetClientPool(), second); } diff --git a/mooncake-common/tests/rpc_drain_guard_test.cpp b/mooncake-common/tests/rpc_drain_guard_test.cpp new file mode 100644 index 0000000000..c908fcad5e --- /dev/null +++ b/mooncake-common/tests/rpc_drain_guard_test.cpp @@ -0,0 +1,66 @@ +// Tests for RpcDrainGuard: the teardown guard that keeps a shared RPC pool +// from being released under a suspended request coroutine (#3909). + +#include + +#include +#include +#include + +#include "rpc_client_io_context.h" + +namespace mooncake { +namespace { + +TEST(RpcDrainGuardTest, DrainWaitsForInFlightAndStopsNewCalls) { + RpcDrainGuard guard; + + auto slow_call = [&] { + RpcDrainGuard::ScopedCall call(guard); + if (!call.ok()) return false; + std::this_thread::sleep_for(std::chrono::milliseconds(150)); + return true; + }; + + bool call_result = false; + std::thread worker([&] { call_result = slow_call(); }); + // let the worker enter the guard before the drain starts + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + + const auto started = std::chrono::steady_clock::now(); + EXPECT_TRUE(guard.drain_for(std::chrono::seconds(5))); + const auto waited = std::chrono::steady_clock::now() - started; + // the drain actually waited for the in-flight call, not just returned + EXPECT_GE(waited, std::chrono::milliseconds(100)); + worker.join(); + EXPECT_TRUE(call_result); + + // once drained, new calls are refused + EXPECT_FALSE(slow_call()); +} + +TEST(RpcDrainGuardTest, DrainWithNothingInFlightReturnsImmediately) { + RpcDrainGuard guard; + const auto started = std::chrono::steady_clock::now(); + EXPECT_TRUE(guard.drain_for(std::chrono::seconds(5))); + EXPECT_LT(std::chrono::steady_clock::now() - started, + std::chrono::milliseconds(500)); +} + +TEST(RpcDrainGuardTest, DrainTimeoutReportsFalse) { + RpcDrainGuard guard; + std::atomic release{false}; + std::thread worker([&] { + RpcDrainGuard::ScopedCall call(guard); + while (!release.load()) std::this_thread::yield(); + }); + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + + EXPECT_FALSE(guard.drain_for(std::chrono::milliseconds(50))); + + release.store(true); + worker.join(); +} + +} // namespace +} // namespace mooncake diff --git a/mooncake-store/include/dummy_client.h b/mooncake-store/include/dummy_client.h index 0256d3af69..aa972ce233 100644 --- a/mooncake-store/include/dummy_client.h +++ b/mooncake-store/include/dummy_client.h @@ -21,6 +21,7 @@ namespace mooncake { class DummyClient : public PyClient { public: DummyClient(); + // Drains in-flight RPCs before the pool member is released (#3909). ~DummyClient(); int64_t unregister_shm(); @@ -312,6 +313,7 @@ class DummyClient : public PyClient { } RpcClientPool client_accessor_; + RpcDrainGuard rpc_drain_; // The client identification. const UUID client_id_; diff --git a/mooncake-store/include/master_client.h b/mooncake-store/include/master_client.h index 33b3e87058..8834e6e2b3 100644 --- a/mooncake-store/include/master_client.h +++ b/mooncake-store/include/master_client.h @@ -88,6 +88,7 @@ class MasterClient { client_id_(client_id), tenant_id_(std::move(tenant_id)), metrics_(metrics) {} + // Drains in-flight RPCs before the pool member is released (#3909). ~MasterClient(); const std::string& tenant_id() const { return tenant_id_.value(); } @@ -697,6 +698,7 @@ class MasterClient { invoke_batch_rpc(size_t input_size, Args&&... args); RpcClientPool client_accessor_; + RpcDrainGuard rpc_drain_; // The client identification. const UUID client_id_; diff --git a/mooncake-store/include/pyclient.h b/mooncake-store/include/pyclient.h index 5d3d567f41..c083bc54c7 100644 --- a/mooncake-store/include/pyclient.h +++ b/mooncake-store/include/pyclient.h @@ -143,6 +143,10 @@ class ClientRequester { public: ClientRequester(); + // Drains in-flight offload RPCs before the pools member is released + // (#3909). + ~ClientRequester(); + /** * @brief Retrieves multiple objects from a remote Transfer Engine (TE) * @param client_addr Network address (e.g., "ip:port") of the remote @@ -194,6 +198,7 @@ class ClientRequester { mutable std::shared_mutex client_pool_mutex_; std::shared_ptr> client_pools_; + RpcDrainGuard rpc_drain_; /** * @brief Generic RPC invocation helper for single-result operations diff --git a/mooncake-store/src/dummy_client.cpp b/mooncake-store/src/dummy_client.cpp index ad5cbdf971..86be4325fc 100644 --- a/mooncake-store/src/dummy_client.cpp +++ b/mooncake-store/src/dummy_client.cpp @@ -183,6 +183,10 @@ constexpr bool can_invoke_when_disconnected() { template tl::expected DummyClient::invoke_rpc(Args&&... args) { + RpcDrainGuard::ScopedCall inflight(rpc_drain_); + if (!inflight.ok()) { + return tl::make_unexpected(ErrorCode::RPC_FAIL); + } auto pool = client_accessor_.GetClientPool(); if constexpr (!can_invoke_when_disconnected()) { @@ -216,6 +220,11 @@ tl::expected DummyClient::invoke_rpc(Args&&... args) { template std::vector> DummyClient::invoke_batch_rpc( size_t input_size, Args&&... args) { + RpcDrainGuard::ScopedCall inflight(rpc_drain_); + if (!inflight.ok()) { + return std::vector>( + input_size, tl::make_unexpected(ErrorCode::RPC_FAIL)); + } auto pool = client_accessor_.GetClientPool(); if (!connected_.load()) { LOG(ERROR) << "Dummy Client not connected"; @@ -266,7 +275,14 @@ DummyClient::DummyClient() mooncake::init_ylt_log_level(); } -DummyClient::~DummyClient() { tearDownAll(); } +DummyClient::~DummyClient() { + // Never release the pool under a suspended request coroutine (#3909). + if (!rpc_drain_.drain_for(std::chrono::seconds(30))) { + LOG(ERROR) << "DummyClient teardown: RPCs still in flight after 30s " + "drain; releasing the pool regardless"; + } + tearDownAll(); +} void DummyClient::ObserveTransferMetric(TransferOperationKind kind, const char* op_name, size_t bytes, diff --git a/mooncake-store/src/master_client.cpp b/mooncake-store/src/master_client.cpp index 92721836bb..762c8a25a2 100644 --- a/mooncake-store/src/master_client.cpp +++ b/mooncake-store/src/master_client.cpp @@ -341,6 +341,10 @@ struct RpcNameTraits<&WrappedMasterService::PollRemoveAll> { template tl::expected MasterClient::invoke_rpc(Args&&... args) { + RpcDrainGuard::ScopedCall inflight(rpc_drain_); + if (!inflight.ok()) { + return tl::make_unexpected(ErrorCode::RPC_FAIL); + } auto pool = client_accessor_.GetClientPool(); // Increment RPC counter @@ -385,6 +389,11 @@ tl::expected MasterClient::invoke_rpc(Args&&... args) { template std::vector> MasterClient::invoke_batch_rpc( size_t input_size, Args&&... args) { + RpcDrainGuard::ScopedCall inflight(rpc_drain_); + if (!inflight.ok()) { + return std::vector>( + input_size, tl::make_unexpected(ErrorCode::RPC_FAIL)); + } auto pool = client_accessor_.GetClientPool(); // Increment RPC counter @@ -433,7 +442,14 @@ std::vector> MasterClient::invoke_batch_rpc( }()); } -MasterClient::~MasterClient() = default; +MasterClient::~MasterClient() { + // Never release the pool under a suspended request coroutine (#3909). + // 30s is generous: every request carries its own RPC timeout. + if (!rpc_drain_.drain_for(std::chrono::seconds(30))) { + LOG(ERROR) << "MasterClient teardown: RPCs still in flight after " + "30s drain; releasing the pool regardless"; + } +} ErrorCode MasterClient::Connect(const std::string& master_addr) { ScopedVLogTimer timer(1, "MasterClient::Connect"); diff --git a/mooncake-store/src/real_client.cpp b/mooncake-store/src/real_client.cpp index cf01ea1ab6..788b122d48 100644 --- a/mooncake-store/src/real_client.cpp +++ b/mooncake-store/src/real_client.cpp @@ -7038,6 +7038,13 @@ ClientRequester::ClientRequester() { pool_conf, GetStoreRpcClientIoContextPool()); } +ClientRequester::~ClientRequester() { + if (!rpc_drain_.drain_for(std::chrono::seconds(30))) { + LOG(ERROR) << "ClientRequester teardown: offload RPCs still in " + "flight after 30s drain; releasing the pools regardless"; + } +} + tl::expected ClientRequester::batch_get_offload_object(const std::string &client_addr, const std::vector &keys, @@ -7073,6 +7080,10 @@ void ClientRequester::release_offload_buffer(const std::string &client_addr, template tl::expected ClientRequester::invoke_rpc( const std::string &client_addr, Args &&...args) { + RpcDrainGuard::ScopedCall inflight(rpc_drain_); + if (!inflight.ok()) { + return tl::make_unexpected(ErrorCode::RPC_FAIL); + } auto client_pool = client_pools_->at(client_addr); return async_simple::coro::syncAwait( [&]() -> async_simple::coro::Lazy> { diff --git a/mooncake-store/tests/rpc_timeout_test.cpp b/mooncake-store/tests/rpc_timeout_test.cpp index 5ab4adc1cb..0d202f1771 100644 --- a/mooncake-store/tests/rpc_timeout_test.cpp +++ b/mooncake-store/tests/rpc_timeout_test.cpp @@ -20,8 +20,11 @@ #include #include +#include #include #include +#include +#include #include #include @@ -147,6 +150,46 @@ TEST_F(RpcTimeoutEnvTest, RpcTimesOutAgainstUnresponsiveMaster) { "timeout still active?)"; } +// #3909: tearing a client down while timeout-backed requests are in flight +// used to release the pool under suspended coroutines (intermittent segfault). +// The destructor now drains first; this loop used to crash within a few rounds +// when the race lost, and must run clean with the guard. +TEST_F(RpcTimeoutEnvTest, TeardownDrainsInFlightTimeoutRequests) { + // Long request timeout so teardown lands squarely inside the call window. + ASSERT_EQ(::setenv("MC_RPC_TIMEOUT_MS", "3000", /*overwrite=*/1), 0); + ASSERT_EQ(::setenv("MC_RPC_CONNECT_TIMEOUT_MS", "100", /*overwrite=*/1), 0); + + BlackHoleServer server; + + for (int round = 0; round < 10; ++round) { + auto client = std::make_unique(generate_uuid()); + // ServiceReady inside Connect times out against the black hole; the + // pool exists by then either way. + (void)client->Connect(server.address()); + + // Each worker issues exactly one ExistKey and never touches the + // client again, which is the owner's actual shape: a torn-down client + // is never invoked after teardown. `entered` plus the 500ms sleep put + // every call squarely inside its 3s timeout window at reset time. + std::atomic entered{0}; + std::vector workers; + for (int i = 0; i < 16; ++i) { + workers.emplace_back([&] { + entered.fetch_add(1); + (void)client->ExistKey("never-there"); + }); + } + while (entered.load() < 16) { + std::this_thread::yield(); + } + std::this_thread::sleep_for(std::chrono::milliseconds(500)); + client.reset(); // teardown with 16 requests suspended mid-flight + for (auto& w : workers) { + w.join(); + } + } +} + // The offload data path (store->store) builds its own client pool, separate // from the master pool, and used to ignore these variables entirely: its // connect timeout stayed at the built-in 30s, so a read that picked a peer From 033524c9403956b72276442ccdb2fab35262bebe Mon Sep 17 00:00:00 2001 From: zhangyupeng Date: Wed, 9 Sep 2026 13:28:36 +0800 Subject: [PATCH 2/6] Drain Client TE paths before tearDown (#3909) MasterClient rpc_drain alone still UAF when Get is in TransferRead during close. Add Client::api_drain_ and RealClient::client_op_drain_ so tearDown waits for in-flight API/TE work before freeing TE. --- mooncake-store/include/client_service.h | 11 ++++++++++ mooncake-store/include/real_client.h | 4 ++++ mooncake-store/src/client_service.cpp | 28 +++++++++++++++++++++++++ mooncake-store/src/real_client.cpp | 23 +++++++++++++++++--- 4 files changed, 63 insertions(+), 3 deletions(-) diff --git a/mooncake-store/include/client_service.h b/mooncake-store/include/client_service.h index 391cbbadb1..bfb7f3bdc8 100644 --- a/mooncake-store/include/client_service.h +++ b/mooncake-store/include/client_service.h @@ -17,6 +17,7 @@ #include "client_metric.h" #include "ha/leadership/leader_coordinator.h" #include "master_client.h" +#include "rpc_client_io_context.h" #include "storage_backend.h" #include "thread_pool.h" #include "transfer_engine.h" @@ -73,6 +74,12 @@ class Client { public: virtual ~Client(); + // Wait for in-flight Get/Put/... before TE/master teardown (#3909). + // Safe to call more than once; tearDownAll should call this before + // unregister/reset so TransferEngine is not freed under TransferRead. + bool DrainInflightOperations( + std::chrono::seconds timeout = std::chrono::seconds(30)); + using WriteBufferStager = std::function, ErrorCode>( const std::vector&)>; @@ -927,6 +934,10 @@ class Client { // Core components std::shared_ptr transfer_engine_; MasterClient master_client_; + // Covers full Client API including TransferEngine paths (not just master + // RPC). MasterClient::rpc_drain_ alone is insufficient when Get has left + // the master RPC and is still in TransferRead (#3909 concurrent close). + RpcDrainGuard api_drain_; std::unique_ptr transfer_submitter_; // Mutex to protect mounted_segments_ diff --git a/mooncake-store/include/real_client.h b/mooncake-store/include/real_client.h index 6733123702..ca4a7fd149 100644 --- a/mooncake-store/include/real_client.h +++ b/mooncake-store/include/real_client.h @@ -22,6 +22,7 @@ #include "mutex.h" #include "common/network.h" #include "pyclient.h" +#include "rpc_client_io_context.h" #include "rpc_types.h" #if defined(USE_SUNRISE) #include "sunrise_allocator.h" @@ -1063,6 +1064,9 @@ class RealClient : public PyClient { // Ensure cleanup executes at most once across multiple entry points std::atomic closed_{false}; + // Covers get_buffer for the full call including post-Get observe (#3909). + RpcDrainGuard client_op_drain_; + // Counts every LOCAL_DISK read served via peer offload-RPC. std::atomic offload_rpc_read_count_{0}; diff --git a/mooncake-store/src/client_service.cpp b/mooncake-store/src/client_service.cpp index d189166929..5df30f8c73 100644 --- a/mooncake-store/src/client_service.cpp +++ b/mooncake-store/src/client_service.cpp @@ -443,7 +443,20 @@ Client::Client(const std::string& local_hostname, } } +bool Client::DrainInflightOperations(std::chrono::seconds timeout) { + if (!api_drain_.drain_for(timeout)) { + LOG(ERROR) << "Client teardown: API calls still in flight after " + << timeout.count() + << "s drain; continuing teardown (UAF risk)"; + return false; + } + return true; +} + Client::~Client() { + // Never free TransferEngine / master pool under an in-flight Get/Put. + (void)DrainInflightOperations(); + task_poll_running_ = false; if (task_poll_thread_.joinable()) { task_poll_thread_.join(); @@ -1352,6 +1365,11 @@ tl::expected, ErrorCode> Client::BatchReplicaClear( tl::expected Client::Get(const std::string& object_key, const QueryResult& query_result, std::vector& slices) { + RpcDrainGuard::ScopedCall inflight(api_drain_); + if (!inflight.ok()) { + return tl::unexpected(ErrorCode::RPC_FAIL); + } + // Find the first complete replica Replica::Descriptor replica; ErrorCode err = FindFirstCompleteReplica(query_result.replicas, replica); @@ -1424,6 +1442,11 @@ tl::expected Client::Get(const std::string& object_key, const QueryResult& query_result, std::vector& slices, uint64_t src_offset) { + RpcDrainGuard::ScopedCall inflight(api_drain_); + if (!inflight.ok()) { + return tl::unexpected(ErrorCode::RPC_FAIL); + } + Replica::Descriptor replica; ErrorCode err = FindFirstCompleteReplica(query_result.replicas, replica); if (err != ErrorCode::OK) { @@ -1853,6 +1876,11 @@ bool Client::RedirectToHotCache(const std::string& key, tl::expected Client::Put(const ObjectKey& key, std::vector& slices, const ReplicateConfig& config) { + RpcDrainGuard::ScopedCall inflight(api_drain_); + if (!inflight.ok()) { + return tl::unexpected(ErrorCode::RPC_FAIL); + } + std::optional object_checksum; if (object_checksum_enabled_) { auto checksum_result = ComputeObjectChecksumForSlices( diff --git a/mooncake-store/src/real_client.cpp b/mooncake-store/src/real_client.cpp index 788b122d48..9c9a6f5aaf 100644 --- a/mooncake-store/src/real_client.cpp +++ b/mooncake-store/src/real_client.cpp @@ -1375,6 +1375,13 @@ tl::expected RealClient::tearDownAll_internal() { // Not initialized or already cleaned; treat as success for idempotence return {}; } + if (!client_op_drain_.drain_for(std::chrono::seconds(30))) { + LOG(ERROR) << "RealClient teardown: client ops still in flight after " + "30s drain; continuing teardown (UAF risk)"; + } + // Finish Get/Put (incl. TransferEngine) before unregister / client_.reset + // (#3909). MasterClient drain alone does not cover TransferRead. + (void)client_->DrainInflightOperations(); if (client_buffer_allocator_ && client_buffer_allocator_->size() > 0 && protocol != "cxl") { auto unregister_result = client_->unregisterLocalMemory( @@ -2945,16 +2952,26 @@ std::shared_ptr RealClient::get_buffer_internal( // Implementation of get_buffer method std::shared_ptr RealClient::get_buffer(const std::string &key) { + RpcDrainGuard::ScopedCall inflight(client_op_drain_); + if (!inflight.ok()) { + return nullptr; + } + auto client = client_; + if (!client) { + LOG(ERROR) << "Client is not initialized"; + return nullptr; + } return execute_timed_operation>( [&]() { return get_buffer_internal(key, client_buffer_allocator_); }, [](const auto &buffer) { return buffer != nullptr; }, [&](uint64_t latency_us, const auto &buffer) { - client_->ObserveTransferOperation(TransferOperationKind::kRead, - "get_buffer", buffer->size(), - latency_us); + client->ObserveTransferOperation(TransferOperationKind::kRead, + "get_buffer", buffer->size(), + latency_us); }); } + tl::expected, ErrorCode> RealClient::acquire_hot_cache(const std::string &key) { if (!client_ || !client_->IsHotCacheEnabled()) { From 8854b5734e38f93dc9520d65947bb958621779ef Mon Sep 17 00:00:00 2001 From: zhangyupeng Date: Wed, 9 Sep 2026 14:26:54 +0800 Subject: [PATCH 3/6] clang-format: drop extra blank line after get_buffer --- mooncake-store/src/real_client.cpp | 1 - 1 file changed, 1 deletion(-) diff --git a/mooncake-store/src/real_client.cpp b/mooncake-store/src/real_client.cpp index 9c9a6f5aaf..63c1821c4c 100644 --- a/mooncake-store/src/real_client.cpp +++ b/mooncake-store/src/real_client.cpp @@ -2971,7 +2971,6 @@ std::shared_ptr RealClient::get_buffer(const std::string &key) { }); } - tl::expected, ErrorCode> RealClient::acquire_hot_cache(const std::string &key) { if (!client_ || !client_->IsHotCacheEnabled()) { From ff5c6708ee2d81fd916e4b64d8e5252fff7b797f Mon Sep 17 00:00:00 2001 From: zhangyupeng Date: Thu, 10 Sep 2026 09:29:28 +0800 Subject: [PATCH 4/6] ci: retrigger CTest after hosted runner lost communication Co-authored-by: Cursor From 3f39e8ae260349d51ceb300f7a5f310de85cded6 Mon Sep 17 00:00:00 2001 From: zhangyupeng Date: Thu, 10 Sep 2026 11:22:03 +0800 Subject: [PATCH 5/6] ci: retrigger after NVIDIA apt mirror sync flake Co-authored-by: Cursor From 22feb12e6310b5b5c819a7d0bdbe583941335d62 Mon Sep 17 00:00:00 2001 From: zhangyupeng Date: Fri, 11 Sep 2026 18:34:21 +0800 Subject: [PATCH 6/6] Respond to #3969 review: abort teardown on drain timeout, guard batch APIs On RealClient tearDown, fail with RPC_TIMEOUT and keep TE/Client if client_op or API drain times out instead of logging UAF risk and continuing. Client dtor waits until idle after a timed-out drain. Wrap BatchGet / BatchPut / BatchGetOffloadObject with api_drain_ ScopedCall. Add RpcDrainGuard regression that a timed-out drain keeps blocking new calls. --- .../tests/rpc_drain_guard_test.cpp | 26 +++++++++++++++ mooncake-store/src/client_service.cpp | 33 +++++++++++++++++-- mooncake-store/src/real_client.cpp | 24 ++++++++++---- 3 files changed, 73 insertions(+), 10 deletions(-) diff --git a/mooncake-common/tests/rpc_drain_guard_test.cpp b/mooncake-common/tests/rpc_drain_guard_test.cpp index c908fcad5e..6e546e80d2 100644 --- a/mooncake-common/tests/rpc_drain_guard_test.cpp +++ b/mooncake-common/tests/rpc_drain_guard_test.cpp @@ -62,5 +62,31 @@ TEST(RpcDrainGuardTest, DrainTimeoutReportsFalse) { worker.join(); } + +// Timed-out drain must not reopen admission; callers that would free shared +// state on timeout (#3909 review) would still race with the in-flight call. +TEST(RpcDrainGuardTest, TimedOutDrainKeepsBlockingNewCalls) { + RpcDrainGuard guard; + + std::atomic entered{false}; + std::thread worker([&] { + RpcDrainGuard::ScopedCall call(guard); + EXPECT_TRUE(call.ok()); + entered.store(true); + std::this_thread::sleep_for(std::chrono::milliseconds(300)); + }); + while (!entered.load()) { + std::this_thread::yield(); + } + + EXPECT_FALSE(guard.drain_for(std::chrono::milliseconds(50))); + { + RpcDrainGuard::ScopedCall late(guard); + EXPECT_FALSE(late.ok()); + } + worker.join(); + EXPECT_TRUE(guard.drain_for(std::chrono::milliseconds(200))); +} + } // namespace } // namespace mooncake diff --git a/mooncake-store/src/client_service.cpp b/mooncake-store/src/client_service.cpp index fae01e0a45..e1761f0393 100644 --- a/mooncake-store/src/client_service.cpp +++ b/mooncake-store/src/client_service.cpp @@ -429,8 +429,7 @@ Client::Client(const std::string& local_hostname, bool Client::DrainInflightOperations(std::chrono::seconds timeout) { if (!api_drain_.drain_for(timeout)) { LOG(ERROR) << "Client teardown: API calls still in flight after " - << timeout.count() - << "s drain; continuing teardown (UAF risk)"; + << timeout.count() << "s drain"; return false; } return true; @@ -438,7 +437,13 @@ bool Client::DrainInflightOperations(std::chrono::seconds timeout) { Client::~Client() { // Never free TransferEngine / master pool under an in-flight Get/Put. - (void)DrainInflightOperations(); + // On timeout keep waiting rather than destroying under flight (#3909 review). + if (!DrainInflightOperations()) { + LOG(ERROR) << "Client dtor: drain timed out; waiting until idle"; + while (!DrainInflightOperations()) { + LOG(ERROR) << "Client dtor: still waiting for in-flight API calls"; + } + } task_poll_running_ = false; if (task_poll_thread_.joinable()) { @@ -1124,6 +1129,12 @@ tl::expected Client::Get(const std::string& object_key, std::vector> Client::BatchGet( const std::vector& object_keys, std::unordered_map>& slices) { + RpcDrainGuard::ScopedCall inflight(api_drain_); + if (!inflight.ok()) { + return std::vector>( + object_keys.size(), tl::unexpected(ErrorCode::RPC_FAIL)); + } + auto batched_query_results = BatchQuery(object_keys); // If any queries failed, return error results immediately for failed @@ -1610,6 +1621,11 @@ std::vector> Client::BatchGet( const std::vector& query_results, std::unordered_map>& slices, bool prefer_alloc_in_same_node) { + RpcDrainGuard::ScopedCall inflight(api_drain_); + if (!inflight.ok()) { + return std::vector>( + object_keys.size(), tl::unexpected(ErrorCode::RPC_FAIL)); + } if (!transfer_submitter_) { LOG(ERROR) << "TransferSubmitter not initialized"; std::vector> results; @@ -3562,6 +3578,12 @@ std::vector> Client::BatchPut( const std::vector& keys, std::vector>& batched_slices, const ReplicateConfig& config, const WriteBufferStager& stager) { + RpcDrainGuard::ScopedCall inflight(api_drain_); + if (!inflight.ok()) { + return std::vector>( + keys.size(), tl::unexpected(ErrorCode::RPC_FAIL)); + } + ReplicateConfig client_cfg = AttachHostId(config); if (protocol_ == "cxl") { client_cfg.preferred_segment = local_hostname_; @@ -4118,6 +4140,11 @@ tl::expected Client::BatchGetOffloadObject( const std::vector& pointers, const std::unordered_map>& batch_slices, OffloadBufferAccess buffer_access) { + RpcDrainGuard::ScopedCall inflight(api_drain_); + if (!inflight.ok()) { + return tl::make_unexpected(ErrorCode::RPC_FAIL); + } + auto future = transfer_submitter_->submit_batch_get_offload_object( transfer_engine_addr, keys, pointers, batch_slices, buffer_access); if (!future) { diff --git a/mooncake-store/src/real_client.cpp b/mooncake-store/src/real_client.cpp index 1e28952f6f..3bf8db52d6 100644 --- a/mooncake-store/src/real_client.cpp +++ b/mooncake-store/src/real_client.cpp @@ -1359,6 +1359,23 @@ int RealClient::initAll(const std::string &protocol_, } tl::expected RealClient::tearDownAll_internal() { + // Drain before claiming closed_ / freeing TE so a timeout can abort and + // leave the client usable (#3909 review: do not tear down under flight). + if (client_) { + if (!client_op_drain_.drain_for(std::chrono::seconds(30))) { + LOG(ERROR) + << "RealClient teardown: client ops still in flight after 30s " + "drain; aborting teardown (resources kept)"; + return tl::unexpected(ErrorCode::RPC_TIMEOUT); + } + if (!client_->DrainInflightOperations()) { + LOG(ERROR) + << "RealClient teardown: Client API still in flight after " + "drain; aborting teardown (resources kept)"; + return tl::unexpected(ErrorCode::RPC_TIMEOUT); + } + } + // Ensure cleanup executes once across destructor/close/signal paths bool expected = false; if (!closed_.compare_exchange_strong(expected, true, @@ -1374,13 +1391,6 @@ tl::expected RealClient::tearDownAll_internal() { // Not initialized or already cleaned; treat as success for idempotence return {}; } - if (!client_op_drain_.drain_for(std::chrono::seconds(30))) { - LOG(ERROR) << "RealClient teardown: client ops still in flight after " - "30s drain; continuing teardown (UAF risk)"; - } - // Finish Get/Put (incl. TransferEngine) before unregister / client_.reset - // (#3909). MasterClient drain alone does not cover TransferRead. - (void)client_->DrainInflightOperations(); if (client_buffer_allocator_ && client_buffer_allocator_->size() > 0 && protocol != "cxl") { auto unregister_result = client_->unregisterLocalMemory(