[anranxia]

Description

The TCP transfer path currently runs a single asio::io_context on one worker thread (TcpTransport::worker → io_context.run()). Every per-peer lane, socket completion, admission timer, and reconnect handler executes on that one thread. On a fast cross-node TCP path (~100 GbE) that single thread saturates at 100% CPU and caps throughput well below line rate — adding more lanes stops helping once they all funnel through the one thread.

This PR introduces TcpIoPool, a bounded pool of io_context threads, and spreads I/O across them while preserving the #2974 bounded-lane model, admission control, reconnect, and exactly-once completion.

Design:

• New TcpIoPool owns N io_context instances, each with its own thread and a work-guard; threads are named mc-tcp-io-<i>. N = MC_NUM_TCP_IO_THREADS (default 1 = unchanged behavior, bounded [1, 32]).
• Each PeerConnectionGroup's coordination state (queue/admission/retry/retirement timers, pump) is pinned to a single shard chosen by hash of the connection key, so its mutex-guarded state stays serialized on one executor.
• Each ConnectionLane's socket, and each accepted server socket, are round-robin bound to a shard (not hashed per-peer), so a single hot peer uses multiple cores. runGroupPump dispatches a lane's connect/session onto that lane's shard via asio::dispatch (runs inline when N=1, so the single-thread path is byte-for-byte the previous behavior).
• Shutdown stops every shard and joins every thread; that join-all is the quiescence barrier.

Correctness rests on the machinery that already existed for the single-thread path — PeerConnectionGroup::mutex, the operation_epoch/pump_epoch/retry_epoch/admission_epoch stale-handler guards, and the move-only TerminalAction + ClientSession::finalize exactly-once path. These are mutex/atomic based; they are now genuinely exercised concurrently. Because each socket is pinned to exactly one shard's single-threaded io_context, asio's per-socket handler serialization holds without adding strands.

This PR also fixes a pre-existing data race (found while adding TSan coverage): TcpTransport::getTransferStatus read task.transferred_bytes / success_slice_count / failed_slice_count with plain loads while the I/O thread updates them via __atomic_fetch_add. It now uses __atomic_load_n(..., __ATOMIC_ACQUIRE), matching MultiTransport::getTransferStatus. This race predates this change (it exists between the single io thread and the polling submission thread).

Relates to RFC #1749; builds on the bounded-lane model from #2974.

Module

• Transfer Engine (mooncake-transfer-engine)
• Mooncake Store (mooncake-store)
• Reshard (mooncake-reshard)
• Mooncake EP (mooncake-ep)
• Mooncake PG (mooncake-pg)
• Integration (mooncake-integration)
• P2P Store (mooncake-p2p-store)
• Python Wheel (mooncake-wheel)
• Common (mooncake-common)
• Mooncake RL (mooncake-rl)
• CI/CD
• Docs
• Other

Type of Change

• Bug fix
• New feature
• Refactor
• Breaking change
• Documentation update
• Performance improvement
• Other

How Has This Been Tested?

Two cross-node hosts over a ~100 GbE TCP path (data-transfer bytes confirmed to cross the NIC via /proc/net/dev). Built with -DUSE_CUDA=ON -DBUILD_EXAMPLES=ON. MC_TCP_ENABLE_CONNECTION_POOL=1 (default).

Regression — MC_NUM_TCP_IO_THREADS=1 is unchanged: cross-node 1 MiB read = 2.49 GB/s, identical to the pre-change baseline; loopback matches too (see below).

Cross-node scaling (1 MiB, DRAM read):

MC_NUM_TCP_IO_THREADS	lanes	submit threads	GB/s
1	8	1	2.49
2	8	1	4.61
4	8	1	7.51
8	16	4	12.28 (~5×, ≈98 Gbit/s)

WRITE, N=8 lanes=16 threads=4: 12.40 GB/s. Per-thread CPU spreads across mc-tcp-io-* instead of one thread at 100%.

Loopback, reproducing the #2974 reviewer config (transfer_engine_bench, threads=4/batch=32/64 KiB/10s):

config	READ GB/s	WRITE GB/s
4-lane pool (baseline, this repo)	1.70 / 1.72	2.11 / 1.12
io-pool N=1, 4-lane	1.45 / 1.53	1.78 / 1.07
io-pool N=4, 4-lane	2.96 / 2.99	2.80 / 2.96
io-pool N=8, 8-lane	3.11 / 3.28	3.25 / 3.25

(The baseline 4-lane numbers line up with the #2974 "main 4-lane" measurements, confirming the setup is comparable.)

Correctness — transfer_engine_validator, cross-node, 12 workers: all 12 workers report Data validation passed at both MC_NUM_TCP_IO_THREADS=4 and =8.

ThreadSanitizer: built -fsanitize=thread and ran the validator (loopback, N=4 and N=8). No data races on the new shared structures (TcpIoPool / PeerConnectionGroup / ConnectionLane) — the existing mutex+epoch model holds under real concurrency. The only transport race TSan reported is the pre-existing getTransferStatus counter read (reproducible at N=1), which this PR fixes; after the fix no transport hot-path races remain. (Remaining reports are in the benchmark harness — a running stop-flag and buffer reuse — and a condition_variable shutdown-tracker report where both sides hold the same mutex.)

Test commands:

cmake -S mooncake-transfer-engine -B build -G Ninja \
  -DUSE_ETCD=OFF -DUSE_CUDA=ON -DBUILD_EXAMPLES=ON -DCMAKE_BUILD_TYPE=Release
cmake --build build -j

# target host
MC_NUM_TCP_IO_THREADS=8 MC_TCP_LANES_PER_PEER=16 \
  ./build/example/transfer_engine_bench --mode=target \
  --metadata_server=P2PHANDSHAKE --protocol=tcp --local_server_name=<ip>:17001
# initiator host (segment_id = the target's "listening on <ip:port>")
MC_NUM_TCP_IO_THREADS=8 MC_TCP_LANES_PER_PEER=16 \
  ./build/example/transfer_engine_bench --mode=initiator \
  --metadata_server=P2PHANDSHAKE --protocol=tcp --segment_id=<ip>:<port> \
  --operation=read --block_size=1048576 --batch_size=32 --threads=4 --duration=10

Test results:

• Unit tests pass
• Integration tests pass (if applicable)
• Manual testing done (cross-node bench + validator + TSan, described above)

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 on the files changed in this PR 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)

AI (Claude Code) assisted with codebase exploration, drafting the TcpIoPool implementation, and running the benchmark/validator/TSan sweeps. The submitter has reviewed every changed line and can defend the change end-to-end.

[staryxchen]

LGTM

[anranxia]

LGTM

All checks have passed. Whenever you have a moment, could you help merge this? Thanks a lot for the reviews!

[staryxchen]

LGTM

All checks have passed. Whenever you have a moment, could you help merge this? Thanks a lot for the reviews!

Merged. Thanks for your contribution.