io pattern实现

mooncake-common/etcd/etcd_wrapper.go

@@ -811,8 +811,49 @@ func EtcdStoreBatchCreateWrapper(keys **C.char, values **C.char, count C.int, er
 	return 0
 }
 
+//export EtcdStoreBatchPutWithLeaseWrapper
+func EtcdStoreBatchPutWithLeaseWrapper(keys **C.char, keySizes *C.int, values **C.char, valueSizes *C.int, count C.int, leaseId int64, errMsg **C.char) int {
+	cli := getStoreClient()
+	if cli == nil {
+		*errMsg = C.CString("etcd client not initialized")
+		return -1
+	}
+
+	n := int(count)
+	if n == 0 {
+		return 0
+	}
+
+	// Unsafe casting to access C arrays as Go slices; key/value buffers are
+	// binary-safe (carry explicit sizes).
+	keyPtrs := (*[1 << 28]*C.char)(unsafe.Pointer(keys))[:n:n]
+	keySizeList := (*[1 << 28]C.int)(unsafe.Pointer(keySizes))[:n:n]
+	valPtrs := (*[1 << 28]*C.char)(unsafe.Pointer(values))[:n:n]
+	valSizeList := (*[1 << 28]C.int)(unsafe.Pointer(valueSizes))[:n:n]
+
+	ops := make([]clientv3.Op, 0, n)
+	for i := 0; i < n; i++ {
+		k := C.GoStringN(keyPtrs[i], keySizeList[i])
+		v := C.GoStringN(valPtrs[i], valSizeList[i])
+		// Bind all keys to the single master lease. Caller is responsible for
+		// fencing (the keys must already be absent / owned by self); this is an
+		// unconditional atomic batch, so it can also be reused for reaffirm.
+		ops = append(ops, clientv3.OpPut(k, v, clientv3.WithLease(clientv3.LeaseID(leaseId))))
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	_, err := cli.Txn(ctx).Then(ops...).Commit()
+	if err != nil {
+		*errMsg = C.CString(err.Error())
+		return -1
+	}
+	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 {
+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, deleteKeys **C.char, deleteKeySizes *C.int, deleteCount C.int, errMsg **C.char) int {
 	cli := getStoreClient()
 	if cli == nil {
 		*errMsg = C.CString("etcd client not initialized")
@@ -821,6 +862,7 @@ func EtcdStoreTxnCompareAndPutWrapper(compareKeys **C.char, compareKeySizes *C.i
 
 	cmpN := int(compareCount)
 	putN := int(putCount)
+	deleteN := int(deleteCount)
 
 	cmps := make([]clientv3.Cmp, 0, cmpN)
 	if cmpN > 0 {
@@ -844,7 +886,7 @@ func EtcdStoreTxnCompareAndPutWrapper(compareKeys **C.char, compareKeySizes *C.i
 		}
 	}
 
-	ops := make([]clientv3.Op, 0, putN)
+	ops := make([]clientv3.Op, 0, putN+deleteN)
 	if putN > 0 {
 		putKeyPtrs := (*[1 << 28]*C.char)(unsafe.Pointer(putKeys))[:putN:putN]
 		putKeySizeList := (*[1 << 28]C.int)(unsafe.Pointer(putKeySizes))[:putN:putN]
@@ -856,6 +898,14 @@ func EtcdStoreTxnCompareAndPutWrapper(compareKeys **C.char, compareKeySizes *C.i
 			ops = append(ops, clientv3.OpPut(k, v))
 		}
 	}
+	if deleteN > 0 {
+		deleteKeyPtrs := (*[1 << 28]*C.char)(unsafe.Pointer(deleteKeys))[:deleteN:deleteN]
+		deleteKeySizeList := (*[1 << 28]C.int)(unsafe.Pointer(deleteKeySizes))[:deleteN:deleteN]
+		for i := 0; i < deleteN; i++ {
+			k := C.GoStringN(deleteKeyPtrs[i], deleteKeySizeList[i])
+			ops = append(ops, clientv3.OpDelete(k))
+		}
+	}
 
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()

mooncake-common/include/rpc_client_io_context.h

@@ -8,6 +8,7 @@
 #include <shared_mutex>
 #include <string>
 #include <string_view>
+#include <unordered_map>
 #include <utility>
 
 #include <ylt/coro_io/client_pool.hpp>
@@ -26,9 +27,10 @@ coro_io::io_context_pool& GetRpcClientIoContextPool(uint32_t thread_count) {
 }
 
 /**
- * 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 client pool accessor that caches one client pool per target address, so
+ * callers that alternate between several targets (e.g. submaster routing)
+ * reuse existing connections instead of recreating a pool on every switch.
+ * The most recently selected pool is also exposed via GetClientPool().
  */
 class RpcClientPool {
    public:
@@ -38,18 +40,22 @@ 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.
+        // Explicit target selection supersedes background recovery of an old
+        // host; the pool's own connect/retry still applies per request.
         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_);
-            address_ = address;
+        std::string addr(address);
+        auto it = pools_.find(addr);
+        if (it == pools_.end()) {
+            auto pool = ClientPool::create(addr, config_, io_context_pool_);
+            it = pools_.emplace(addr, std::move(pool)).first;
         }
+        address_ = std::move(addr);
+        client_pool_ = it->second;
         return client_pool_;
     }
 
@@ -58,12 +64,21 @@ class RpcClientPool {
         return client_pool_;
     }
 
+    std::string GetAddress() const {
+        std::shared_lock<std::shared_mutex> lock(mutex_);
+        return address_;
+    }
+
    private:
     mutable std::shared_mutex mutex_;
     coro_io::io_context_pool& io_context_pool_;
     PoolConfig config_;
     std::string address_;
     std::shared_ptr<ClientPool> client_pool_;
+    // Cache of client pools keyed by target address, so switching back and
+    // forth between submaster addresses reuses existing connections instead of
+    // recreating the pool on every switch.
+    std::unordered_map<std::string, std::shared_ptr<ClientPool>> pools_;
 };
 
 }  // namespace mooncake

mooncake-common/tests/rpc_client_io_context_test.cpp

@@ -41,18 +41,19 @@ TEST(RpcClientIoContextPoolTest, UsesConfiguredSizeAndReusesPool) {
     EXPECT_NE(&first_pool, &second_pool);
 }
 
-TEST(RpcClientIoContextPoolTest, ReplacesPoolWhenTargetChanges) {
+TEST(RpcClientIoContextPoolTest, CachesPoolsPerAddress) {
     RpcClientPool pools(GetFirstTestRpcClientIoContextPool());
 
     auto first = pools.GetOrCreateClientPool("127.0.0.1:10001");
-    std::weak_ptr<RpcClientPool::ClientPool> old_pool = first;
     EXPECT_EQ(pools.GetOrCreateClientPool("127.0.0.1:10001"), first);
 
     auto second = pools.GetOrCreateClientPool("127.0.0.1:10002");
     EXPECT_NE(first, second);
-    first.reset();
-    EXPECT_TRUE(old_pool.expired());
     EXPECT_EQ(pools.GetClientPool(), second);
+
+    // Switching back reuses the cached pool instead of recreating it.
+    EXPECT_EQ(pools.GetOrCreateClientPool("127.0.0.1:10001"), first);
+    EXPECT_EQ(pools.GetClientPool(), first);
 }
 
 TEST(RpcClientIoContextPoolTest, SendsToNewAddressAfterSwitch) {

mooncake-store/benchmarks/CMakeLists.txt

@@ -38,6 +38,13 @@ target_link_libraries(
   stress_cluster_bench PRIVATE mooncake_store transfer_engine asio_shared
                                gflags::gflags glog::glog pthread)
 
+# CFM client benchmark. Simulates vLLM inference requests as KV-cache block
+# accesses and reports both client performance and IO Pattern observability.
+add_executable(cfm_client_bench cfm_client_bench.cpp)
+target_link_libraries(
+  cfm_client_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).
@@ -75,3 +82,10 @@ if(STORE_USE_ETCD)
     oplog_batch_bench PRIVATE mooncake_store glog::glog gflags::gflags
                               JsonCpp::JsonCpp)
 endif()
+
+if(STORE_USE_ETCD)
+  add_executable(vchunk_distributed_bench vchunk_distributed_bench.cpp)
+  target_link_libraries(
+    vchunk_distributed_bench PRIVATE mooncake_store transfer_engine asio_shared
+                                     gflags::gflags glog::glog pthread)
+endif()