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..6e546e80d2 --- /dev/null +++ b/mooncake-common/tests/rpc_drain_guard_test.cpp @@ -0,0 +1,92 @@ +// 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(); +} + + +// 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/include/client_service.h b/mooncake-store/include/client_service.h index d7e5eb7d47..26f81e7d0d 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&)>; @@ -933,6 +940,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/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 76f1c03f7b..e676d3de0a 100644 --- a/mooncake-store/include/master_client.h +++ b/mooncake-store/include/master_client.h @@ -104,6 +104,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(); void EnableHaConnectionPolicy(); @@ -725,6 +726,7 @@ class MasterClient { invoke_batch_rpc(size_t input_size, Args&&... args); RpcClientPool client_accessor_; + RpcDrainGuard rpc_drain_; RpcClientPool ha_control_client_accessor_; RpcClientPool ha_probe_client_accessor_; 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/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 f422497177..e1761f0393 100644 --- a/mooncake-store/src/client_service.cpp +++ b/mooncake-store/src/client_service.cpp @@ -426,7 +426,25 @@ 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"; + return false; + } + return true; +} + Client::~Client() { + // Never free TransferEngine / master pool under an in-flight Get/Put. + // 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()) { task_poll_thread_.join(); @@ -1111,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 @@ -1315,6 +1339,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); @@ -1392,6 +1421,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) { @@ -1587,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; @@ -1875,6 +1914,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( @@ -3534,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_; @@ -4090,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/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 151325432d..f0c06f1381 100644 --- a/mooncake-store/src/master_client.cpp +++ b/mooncake-store/src/master_client.cpp @@ -343,6 +343,10 @@ template tl::expected MasterClient::invoke_rpc_with_client_pool( const std::shared_ptr& client_pool, Args&&... args) { + RpcDrainGuard::ScopedCall inflight(rpc_drain_); + if (!inflight.ok()) { + return tl::make_unexpected(ErrorCode::RPC_FAIL); + } // Increment RPC counter if (metrics_) { metrics_->rpc_count.inc({RpcNameTraits::value}); @@ -398,6 +402,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 @@ -446,7 +455,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"; + } +} void MasterClient::EnableHaConnectionPolicy() { MutexLocker lock(&connect_mutex_); diff --git a/mooncake-store/src/real_client.cpp b/mooncake-store/src/real_client.cpp index c03be72b48..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, @@ -2944,13 +2961,22 @@ 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); }); } @@ -7303,6 +7329,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, @@ -7338,6 +7371,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 c0fb45a899..3040d726e8 100644 --- a/mooncake-store/tests/rpc_timeout_test.cpp +++ b/mooncake-store/tests/rpc_timeout_test.cpp @@ -26,8 +26,11 @@ #include #include +#include #include #include +#include +#include #include #include @@ -154,6 +157,44 @@ 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(); + } + } // During failover the heartbeat can still be connecting to the deleted // leader's pod IP. It must return promptly so the HA loop can use the newly // published view instead of spending the default retry budget on a stale