[Store] Dual rdma with coroutine te submit by G0-KU · Pull Request #3355 · kvcache-ai/Mooncake

G0-KU · 2026-08-10T07:44:39Z

Description

Offload TransferEngine wait (WaitAllTransferBatches / poll + sleep) from coroutine/RPC workers onto a dedicated te_poll thread pool, so multi-key P2P coroutines can yield during TE completion instead of occupying business threads.

This is a store-side bridge, not true TE async completion callbacks. Submit stays on the caller; only wait is offloaded when te_async_poll_worker_num > 0 (default 32; 0 keeps the previous sync-wait behavior).

What this adds

DataManager TE wait bridge
- New te_poll_executor_ (AsyncMemcpyExecutor) sized by te_async_poll_worker_num
- WaitAllTransferBatchesAsync: when the pool is enabled, runs only WaitAllTransferBatches on poll workers and returns a Future; when te_async_poll_worker_num == 0, wait stays synchronous on the caller (legacy behavior)
- TransferDataAsync: validate + submit on the caller, then returns that wait Future (submit is not offloaded)
Forward (accessor TE)
- Remote forward write/read already use 3-phase peer RPCs (PreWrite/Pin → TE → WriteCommit/UnPin)
- TE wait is now co_await TransferDataAsync / TeTransferFn Future instead of blocking the coroutine worker
- After wait, AwaitTeExpectedFuture resumes on CurrentExecutor, or coro_io::get_global_executor() if none is bound—so Commit/Revoke/UnPin/metrics do not run on te_poll threads
- Forward write path catches poll shutdown/setException so the outer write Promise always completes
Local TE Put/Get
- BuildDataCopierViaTe: offloads wait via FutureHandle, keeping TeSubmitResult (handle / temp buffer) alive until wait finishes
- PutViaTe: PreWrite + submit on caller; only wait goes to the pool; DRAM staging copy + WriteCommit / revoke run after wait on the caller / resume executor (PutViaTeTaskHandle), not inside the poll task
Reverse (owner TE)
- ClientRpcService::{Read,Write}RemoteData become Lazy handlers
- Delegate to ReadRemoteDataAsync / WriteRemoteDataAsync: submit on the RPC path, co_await wait offload, then copy/commit (write) or return (read)
- Sync DataManager::{Read,Write}RemoteData remain as syncAwait wrappers for existing unit tests
Config / wiring
- te_async_poll_worker_num on P2PClientConfig / LocalTransferConfig, exposed via client config builder (incl. dict), real_client_main gflag, and Python MooncakeStore ctor
- Default 32; set 0 to disable the poll pool

InflightTracker change (needed for reverse transfer)

Reverse Lazy handlers hold peer_tracker_ Guard across co_await. The old Guard kept a std::shared_lock, which must unlock on the same OS thread that locked—unsafe if the coroutine resumes on a te_poll (or other) thread.

Change: drain by atomic inflight_ + CV instead of r/w lock.

Close() rejects new Enter()
Wait() blocks until inflight_ == 0 (does not by itself reject new enters—always Close() then Wait())
Guard may be destroyed on any thread

peer_tracker_ and local_inflight_tracker_ share this implementation.

Module

Type of Change

How Has This Been Tested?

Unit

inflight_tracker_test — admit/close/drain; cross-thread Guard destroy; Close+Wait
data_manager_test — poll executor create/disable; TransferDataAsync / WaitAllTransferBatchesAsync; FutureHandle; shutdown cancel surfaces as exception
client_config_builder_test — te_async_poll_worker_num default / args / dict
client_rpc_service_test — Lazy Read/WriteRemoteData via syncAwait

Integration

p2p_client_integration_test
- TeAsyncPollForwardRemoteBatchPutAndGet (te_async={0,4})
- TeAsyncPollReverseRemoteBatchPutAndGet (te_async={0,4})

Test commands:

# adjust binary/ctest names to local build
ctest -R 'inflight_tracker_test|data_manager_test|client_config_builder_test|client_rpc_service_test|p2p_client_integration_test' --output-on-failure

Test results:

Unit tests pass
Integration tests pass (if applicable)

Checklist

I have performed a self-review of my own code
I have formatted my code using ./scripts/code_format.sh
I have run pre-commit run --all-files and all hooks pass
I have updated the documentation (if applicable)
I have added tests to prove my changes are effective
For changes >500 LOC: I have filed an RFC issue

AI Assistance Disclosure

No AI tools were used
AI tools were used (specify below)

Cursor agent assisted implementation and PR text; author reviewed and owns the change.

codecov-commenter · 2026-08-10T14:15:50Z

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 80.82437% with 107 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
mooncake-store/src/data_manager.cpp	65.15%	77 Missing ⚠️
mooncake-store/include/task_handle.h	40.90%	13 Missing ⚠️
...oncake-store/tests/p2p_client_integration_test.cpp	93.27%	8 Missing ⚠️
mooncake-store/tests/data_manager_test.cpp	96.66%	3 Missing ⚠️
mooncake-integration/store/store_py.cpp	0.00%	2 Missing ⚠️
mooncake-store/include/inflight_tracker.h	84.61%	2 Missing ⚠️
mooncake-store/src/p2p_client_service.cpp	88.23%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!

wanyue-wy · 2026-08-12T12:10:20Z

+// Await a TE-poll Future and always resume off the poll worker.
+// Prefer the Lazy's CurrentExecutor; if none is bound (e.g. syncAwait without
+// via), fall back to coro_io's global executor so copy/commit/RPC never
+// continue on te_poll threads.
+template <typename V>
+inline async_simple::coro::Lazy<tl::expected<V, ErrorCode>>
+AwaitTeExpectedFuture(async_simple::Future<tl::expected<V, ErrorCode>> fut) {
+    async_simple::Executor* ex = co_await async_simple::CurrentExecutor{};
+    if (ex == nullptr) {
+        ex = coro_io::get_global_executor();
+        LOG(WARNING) << "AwaitTeExpectedFuture: no CurrentExecutor; "
+                        "resuming via coro_io::get_global_executor()";
+    }
+    try {
+        co_return co_await std::move(fut).via(ex);
+    } catch (const std::exception& e) {
+        LOG(ERROR) << "TE wait future exception: " << e.what();
+        co_return tl::unexpected(ErrorCode::INTERNAL_ERROR);
+    } catch (...) {
+        LOG(ERROR) << "TE wait future unknown exception";
+        co_return tl::unexpected(ErrorCode::INTERNAL_ERROR);
+    }
+}


If this is intended to be a general‑purpose coroutine component, the naming and logging should not contain TE‑specific information, and it should be placed in task_handle.h.

If data_manager removes te_poll_executor_ and WaitAllTransferBatchesAsync uses an independent executor for coroutine logic, would AwaitTeExpectedFuture still be necessary?

Yes — still needed (now AwaitExpectedFuture). Wait runs on te_wait_pool_, so the Future completes there. Without hop-back, copy/commit/UnPin would resume on that pool. We don’t run wait on the caller’s executor: many callers are the RPC pool, and TE polling would crowd PreWrite/PinKey/Heartbeat off those threads. Hop-back is the cost of keeping wait off RPC.

…ines can yield

…o coroutine te_poll

…ng pin and holding the transfer Future in std::optional

…e P2P via from RPC Co-authored-by: Cursor <cursoragent@cursor.com>

Co-authored-by: Cursor <cursoragent@cursor.com>

G0-KU requested review from XucSh, YiXR, stmatengss and ykwd as code owners August 10, 2026 07:44

github-actions Bot added run-ci Store Integration labels Aug 10, 2026

G0-KU force-pushed the dual-rdma-p2p-te-coroutine-0808 branch from c9947fc to bd0ebe8 Compare August 11, 2026 02:07

wanyue-wy reviewed Aug 12, 2026

View reviewed changes

shichangzhang064 and others added 7 commits August 22, 2026 21:28

[Store] feat: Offload forward/local TE wait to te_poll pool so corout…

b6c272c

…ines can yield

[Store] Drain InflightTracker by counter and modify reverse TE wait t…

9ac9294

…o coroutine te_poll

fix: code format fix

23296b3

[Store] fix gcc-11/12 coroutine ICE in RunForwardReadOnRoute by scopi…

9981c3b

…ng pin and holding the transfer Future in std::optional

[Store] Address TE-coroutine PR review: generalize handles and isolat…

af2892f

…e P2P via from RPC Co-authored-by: Cursor <cursoragent@cursor.com>

[Store] Replace te_poll blocking wait with coro_io polling coroutines

281e519

Co-authored-by: Cursor <cursoragent@cursor.com>

code format fix

6ec9319

G0-KU force-pushed the dual-rdma-p2p-te-coroutine-0808 branch from bd0ebe8 to 6ec9319 Compare August 22, 2026 14:55

coro miss csginal include

f74716f

Aionw mentioned this pull request Aug 24, 2026

[Store] Migrate async Python client to unified package #3624

Merged

12 of 31 tasks

wanyue-wy merged commit 4db1c02 into kvcache-ai:P2P-Mooncake-Store Aug 24, 2026
23 of 25 checks passed

Navigation Menu

Repository navigation

[Store] Dual rdma with coroutine te submit - #3355

[Store] Dual rdma with coroutine te submit#3355
wanyue-wy merged 8 commits into
kvcache-ai:P2P-Mooncake-Storefrom
G0-KU:dual-rdma-p2p-te-coroutine-0808

G0-KU commented Aug 10, 2026 •

edited

Loading

Uh oh!

codecov-commenter commented Aug 10, 2026 •

edited

Loading

Uh oh!

wanyue-wy Aug 12, 2026

Uh oh!

G0-KU Aug 22, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

G0-KU commented Aug 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

What this adds

InflightTracker change (needed for reverse transfer)

Module

Type of Change

How Has This Been Tested?

Checklist

AI Assistance Disclosure

Uh oh!

codecov-commenter commented Aug 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

wanyue-wy Aug 12, 2026

Choose a reason for hiding this comment

Uh oh!

G0-KU Aug 22, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

G0-KU commented Aug 10, 2026 •

edited

Loading

codecov-commenter commented Aug 10, 2026 •

edited

Loading