[Store] Drain Client TE paths before tearDown (#3909)

mooncake-common/include/rpc_client_io_context.h

@@ -1,13 +1,16 @@
 #pragma once
 
 #include <cstdint>
+#include <atomic>
 #include <chrono>
+#include <condition_variable>
 #include <csignal>
 #include <memory>
 #include <mutex>
 #include <shared_mutex>
 #include <string>
 #include <string_view>
+#include <unordered_map>
 #include <utility>
 
 #include <ylt/coro_io/client_pool.hpp>
@@ -19,16 +22,106 @@ namespace mooncake {
 std::shared_ptr<coro_io::io_context_pool> 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<bool> stopping_{false};
+    std::atomic<int> inflight_{0};
+    std::mutex mutex_;
+    std::condition_variable cv_;
+};
+
 template <typename PoolTag>
 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<coro_io::client_pool<coro_rpc::coro_rpc_client>>
+SharedPoolRegistryImpl(
+    std::string_view address,
+    coro_io::client_pool<coro_rpc::coro_rpc_client>::pool_config config,
+    coro_io::io_context_pool& io_context_pool) {
+    using Pool = coro_io::client_pool<coro_rpc::coro_rpc_client>;
+    static std::mutex registry_mutex;
+    static std::unordered_map<std::string, std::shared_ptr<Pool>> registry;
+    std::lock_guard<std::mutex> 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<std::shared_mutex> 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<ClientPool> 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_;

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)

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);
 }
 

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 <gtest/gtest.h>
+
+#include <atomic>
+#include <chrono>
+#include <thread>
+
+#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<bool> 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<bool> 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

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<tl::expected<std::vector<Slice>, ErrorCode>(
             const std::vector<Slice>&)>;
@@ -933,6 +940,10 @@ class Client {
     // Core components
     std::shared_ptr<TransferEngine> 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<TransferSubmitter> transfer_submitter_;
 
     // Mutex to protect mounted_segments_

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_;

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_;
 

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<coro_io::client_pools<coro_rpc::coro_rpc_client>>
         client_pools_;
+    RpcDrainGuard rpc_drain_;
 
     /**
      * @brief Generic RPC invocation helper for single-result operations

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<bool> 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<int64_t> offload_rpc_read_count_{0};