[Store] Backport HA RPC recovery fixes

mooncake-common/etcd/etcd_wrapper.go

@@ -791,6 +791,65 @@ func EtcdStoreBatchCreateWrapper(keys **C.char, values **C.char, count C.int, er
 	return 0
 }
 
+//export EtcdStoreTxnCompareAndPutWithLeaseWrapper
+func EtcdStoreTxnCompareAndPutWithLeaseWrapper(compareKeys **C.char, compareKeySizes *C.int, compareKinds *C.int, compareValues **C.char, compareValueSizes *C.int, compareCount C.int, putKeys **C.char, putKeySizes *C.int, putValues **C.char, putValueSizes *C.int, putCount C.int, leaseId int64, errMsg **C.char) int {
+	cli := getStoreClient()
+	if cli == nil {
+		*errMsg = C.CString("etcd client not initialized")
+		return -1
+	}
+
+	cmpN := int(compareCount)
+	putN := int(putCount)
+	cmps := make([]clientv3.Cmp, 0, cmpN)
+	if cmpN > 0 {
+		compareKeyPtrs := (*[1 << 28]*C.char)(unsafe.Pointer(compareKeys))[:cmpN:cmpN]
+		compareKeySizeList := (*[1 << 28]C.int)(unsafe.Pointer(compareKeySizes))[:cmpN:cmpN]
+		compareKindList := (*[1 << 28]C.int)(unsafe.Pointer(compareKinds))[:cmpN:cmpN]
+		compareValuePtrs := (*[1 << 28]*C.char)(unsafe.Pointer(compareValues))[:cmpN:cmpN]
+		compareValueSizeList := (*[1 << 28]C.int)(unsafe.Pointer(compareValueSizes))[:cmpN:cmpN]
+		for i := 0; i < cmpN; i++ {
+			k := C.GoStringN(compareKeyPtrs[i], compareKeySizeList[i])
+			switch int(compareKindList[i]) {
+			case 0:
+				v := C.GoStringN(compareValuePtrs[i], compareValueSizeList[i])
+				cmps = append(cmps, clientv3.Compare(clientv3.Value(k), "=", v))
+			case 1:
+				cmps = append(cmps, clientv3.Compare(clientv3.CreateRevision(k), "=", 0))
+			default:
+				*errMsg = C.CString("unsupported compare kind")
+				return -1
+			}
+		}
+	}
+
+	ops := make([]clientv3.Op, 0, putN)
+	if putN > 0 {
+		putKeyPtrs := (*[1 << 28]*C.char)(unsafe.Pointer(putKeys))[:putN:putN]
+		putKeySizeList := (*[1 << 28]C.int)(unsafe.Pointer(putKeySizes))[:putN:putN]
+		putValuePtrs := (*[1 << 28]*C.char)(unsafe.Pointer(putValues))[:putN:putN]
+		putValueSizeList := (*[1 << 28]C.int)(unsafe.Pointer(putValueSizes))[:putN:putN]
+		for i := 0; i < putN; i++ {
+			k := C.GoStringN(putKeyPtrs[i], putKeySizeList[i])
+			v := C.GoStringN(putValuePtrs[i], putValueSizeList[i])
+			ops = append(ops, clientv3.OpPut(k, v, clientv3.WithLease(clientv3.LeaseID(leaseId))))
+		}
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	resp, err := cli.Txn(ctx).If(cmps...).Then(ops...).Commit()
+	if err != nil {
+		*errMsg = C.CString(err.Error())
+		return -1
+	}
+	if !resp.Succeeded {
+		*errMsg = C.CString("etcd transaction compare failed")
+		return -2
+	}
+	return 0
+}
+
 //export EtcdStoreTxnCompareAndPutWrapper
 func EtcdStoreTxnCompareAndPutWrapper(compareKeys **C.char, compareKeySizes *C.int, compareKinds *C.int, compareValues **C.char, compareValueSizes *C.int, compareCount C.int, putKeys **C.char, putKeySizes *C.int, putValues **C.char, putValueSizes *C.int, putCount C.int, errMsg **C.char) int {
 	cli := getStoreClient()

mooncake-common/include/rpc_client_io_context.h

@@ -1,14 +1,17 @@
 #pragma once
 
-#include <cstdint>
+#include <atomic>
 #include <chrono>
-#include <csignal>
+#include <condition_variable>
+#include <cstdint>
 #include <memory>
 #include <mutex>
 #include <shared_mutex>
 #include <string>
 #include <string_view>
+#include <unordered_map>
 #include <utility>
+#include <vector>
 
 #include <ylt/coro_io/client_pool.hpp>
 #include <ylt/coro_io/io_context_pool.hpp>
@@ -19,16 +22,116 @@ namespace mooncake {
 std::shared_ptr<coro_io::io_context_pool> CreateRpcClientIoContextPool(
     uint32_t thread_count);
 
+class RpcDrainGuard {
+   private:
+    struct State {
+        std::atomic<bool> stopping{false};
+        std::atomic<int> inflight{0};
+        std::mutex mutex;
+        std::condition_variable cv;
+    };
+
+   public:
+    class ScopedCall {
+       public:
+        explicit ScopedCall(const RpcDrainGuard& guard)
+            : state_(guard.state_), ok_(try_enter(*state_)) {}
+        ~ScopedCall() {
+            if (ok_) {
+                leave(*state_);
+            }
+        }
+        ScopedCall(const ScopedCall&) = delete;
+        ScopedCall& operator=(const ScopedCall&) = delete;
+        bool ok() const { return ok_; }
+
+       private:
+        static bool try_enter(State& state) {
+            if (state.stopping.load(std::memory_order_acquire)) {
+                return false;
+            }
+            state.inflight.fetch_add(1, std::memory_order_acq_rel);
+            if (state.stopping.load(std::memory_order_acquire)) {
+                leave(state);
+                return false;
+            }
+            return true;
+        }
+
+        static void leave(State& state) {
+            if (state.inflight.fetch_sub(1, std::memory_order_acq_rel) == 1) {
+                std::lock_guard<std::mutex> lock(state.mutex);
+                state.cv.notify_all();
+            }
+        }
+
+        std::shared_ptr<State> state_;
+        bool ok_;
+    };
+
+    bool drain_for(std::chrono::milliseconds timeout) {
+        state_->stopping.store(true, std::memory_order_release);
+        std::unique_lock<std::mutex> lock(state_->mutex);
+        return state_->cv.wait_for(lock, timeout, [&] {
+            return state_->inflight.load(std::memory_order_acquire) == 0;
+        });
+    }
+
+   private:
+    std::shared_ptr<State> state_ = std::make_shared<State>();
+};
+
 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 {
+
+using MasterClientPool =
+    coro_io::client_pool<coro_rpc::coro_rpc_client>;
+
+inline void KeepMasterPoolAlive(std::shared_ptr<MasterClientPool> pool) {
+    static std::mutex keep_mutex;
+    static std::vector<std::shared_ptr<MasterClientPool>> retired_pools;
+    std::lock_guard<std::mutex> lock(keep_mutex);
+    retired_pools.push_back(std::move(pool));
+}
+
+inline std::shared_ptr<MasterClientPool> SharedPoolRegistryImpl(
+    std::string_view address, MasterClientPool::pool_config config,
+    coro_io::io_context_pool& io_context_pool) {
+    static std::mutex registry_mutex;
+    static std::unordered_map<std::string, std::shared_ptr<MasterClientPool>>
+        registry;
+    std::lock_guard<std::mutex> lock(registry_mutex);
+
+    std::string key(address);
+    auto append = [&key](auto value) {
+        key.push_back('|');
+        key.append(std::to_string(value));
+    };
+    append(config.max_connection);
+    append(config.connect_retry_count);
+    append(config.reconnect_wait_time.count());
+    append(config.client_config.connect_timeout_duration.count());
+    append(config.client_config.request_timeout_duration.count());
+    append(static_cast<size_t>(config.client_config.socket_config.index()));
+
+    auto& pool = registry[key];
+    if (!pool) {
+        pool = MasterClientPool::create(std::string(address), std::move(config),
+                                        io_context_pool);
+    }
+    return pool;
+}
+
+}  // 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.
+ * A replaceable client pool. Pools are retained process-wide because ylt's
+ * background reconnect coroutines may outlive the active holder.
  */
 class RpcClientPool {
    public:
@@ -38,21 +141,30 @@ class RpcClientPool {
     explicit RpcClientPool(coro_io::io_context_pool& io_context_pool,
                            PoolConfig config = {})
         : io_context_pool_(io_context_pool), config_(std::move(config)) {
-        // Address replacement supersedes background recovery of the old host.
         config_.host_alive_detect_duration = std::chrono::seconds(0);
     }
 
     std::shared_ptr<ClientPool> GetOrCreateClientPool(
         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_ = detail::SharedPoolRegistryImpl(
+                address, config_, io_context_pool_);
             address_ = address;
         }
         return client_pool_;
     }
 
+    std::shared_ptr<ClientPool> RecreateClientPool(std::string_view address) {
+        std::lock_guard<std::shared_mutex> lock(mutex_);
+        if (client_pool_) {
+            detail::KeepMasterPoolAlive(client_pool_);
+        }
+        client_pool_ = ClientPool::create(address, config_, io_context_pool_);
+        address_ = address;
+        return client_pool_;
+    }
+
     std::shared_ptr<ClientPool> GetClientPool() const {
         std::shared_lock<std::shared_mutex> lock(mutex_);
         return client_pool_;

mooncake-common/tests/rpc_client_io_context_test.cpp

@@ -51,7 +51,8 @@ TEST(RpcClientIoContextPoolTest, ReplacesPoolWhenTargetChanges) {
     auto second = pools.GetOrCreateClientPool("127.0.0.1:10002");
     EXPECT_NE(first, second);
     first.reset();
-    EXPECT_TRUE(old_pool.expired());
+    // The registry retains retired pools for background reconnect coroutines.
+    EXPECT_FALSE(old_pool.expired());
     EXPECT_EQ(pools.GetClientPool(), second);
 }
 

mooncake-store/include/client_service.h

@@ -811,6 +811,7 @@ class Client {
      * @brief Internal helper functions for initialization and data transfer
      */
     ErrorCode ConnectToMaster(const std::string& master_server_entry);
+    void EnterHaRuntimeMode();
     ErrorCode InitTransferEngine(
         const std::string& local_hostname,
         const std::string& metadata_connstring, const std::string& protocol,

mooncake-store/include/etcd_helper.h

@@ -95,6 +95,9 @@ class EtcdHelper {
      */
     static ErrorCode TxnCompareAndPut(const std::vector<TxnCompare>& compares,
                                       const std::vector<TxnPut>& puts);
+    static ErrorCode TxnCompareAndPutWithLease(
+        const std::vector<TxnCompare>& compares,
+        const std::vector<TxnPut>& puts, EtcdLeaseId lease_id);
 
     /*
      * @brief Grant a lease from the etcd.

mooncake-store/include/ha/leadership/backends/etcd/etcd_leader_coordinator.h

@@ -26,6 +26,9 @@ class EtcdLeaderCoordinator final : public LeaderCoordinator {
     tl::expected<AcquireLeadershipResult, ErrorCode> TryAcquireLeadership(
         const std::string& leader_address) override;
 
+    ErrorCode PublishServiceReady(
+        const LeadershipSession& session) override;
+
     tl::expected<bool, ErrorCode> RenewLeadership(
         const LeadershipSession& session) override;
 

mooncake-store/include/ha/leadership/leader_coordinator.h

@@ -28,6 +28,11 @@ class LeaderCoordinator {
     virtual tl::expected<AcquireLeadershipResult, ErrorCode>
     TryAcquireLeadership(const std::string& leader_address) = 0;
 
+    virtual ErrorCode PublishServiceReady(
+        const LeadershipSession& /*session*/) {
+        return ErrorCode::OK;
+    }
+
     virtual tl::expected<bool, ErrorCode> RenewLeadership(
         const LeadershipSession& session) = 0;
 
@@ -42,5 +47,19 @@ class LeaderCoordinator {
     virtual ErrorCode ReleaseLeadership(const LeadershipSession& session) = 0;
 };
 
+inline tl::expected<std::optional<MasterView>, ErrorCode>
+ReadCurrentViewOrWaitForReady(LeaderCoordinator& coordinator,
+                              std::chrono::milliseconds timeout) {
+    auto current_view = coordinator.ReadCurrentView();
+    if (!current_view || current_view->has_value()) {
+        return current_view;
+    }
+    auto ready_view = coordinator.WaitForViewChange(std::nullopt, timeout);
+    if (!ready_view) {
+        return tl::make_unexpected(ready_view.error());
+    }
+    return ready_view->current_view;
+}
+
 }  // namespace ha
 }  // namespace mooncake