From 1009fded5cdd012a1b0b2c1b761e57442402b6c6 Mon Sep 17 00:00:00 2001 From: jipengtian Date: Thu, 27 Aug 2026 17:48:43 +0800 Subject: [PATCH] [TE] Shard TCP transport I/O across a pool of io_context threads --- .../transport/tcp_transport/tcp_transport.h | 125 ++++++++++- .../transport/tcp_transport/tcp_transport.cpp | 54 ++--- .../tcp_transport/tcp_transport_lane_impl.h | 202 +++++++++++------- .../tcp_transport_session_impl.h | 42 ++-- 4 files changed, 300 insertions(+), 123 deletions(-) diff --git a/mooncake-transfer-engine/include/transport/tcp_transport/tcp_transport.h b/mooncake-transfer-engine/include/transport/tcp_transport/tcp_transport.h index 9d53abc36a..bdb8f308e5 100644 --- a/mooncake-transfer-engine/include/transport/tcp_transport/tcp_transport.h +++ b/mooncake-transfer-engine/include/transport/tcp_transport/tcp_transport.h @@ -35,9 +35,14 @@ #include #include +#include +#include #include #include +#include #include +#include +#include #include "transfer_metadata.h" #include "transport/transport.h" @@ -49,6 +54,101 @@ struct ClientSession; class TcpContext; class TcpTransport; +// A bounded pool of asio::io_context instances, each driven by its own thread. +// TCP lanes and accepted server sockets are round-robin distributed across the +// shards so a single hot peer can use more than one CPU core. With size 1 the +// behavior is identical to the historical single-io_context transport. +class TcpIoPool { + public: + explicit TcpIoPool(size_t num_threads) { + if (num_threads < 1) num_threads = 1; + contexts_.reserve(num_threads); + guards_.reserve(num_threads); + for (size_t i = 0; i < num_threads; ++i) { + contexts_.push_back(std::make_unique()); + guards_.push_back( + std::make_unique(contexts_[i]->get_executor())); + } + } + + ~TcpIoPool() { stop(); } + + size_t size() const { return contexts_.size(); } + + asio::io_context &context(size_t i) { + return *contexts_[i % contexts_.size()]; + } + + asio::io_context::executor_type executor(size_t i) { + return contexts_[i % contexts_.size()]->get_executor(); + } + + // Group coordination (pump/admission/retry/retirement timers) is pinned to + // one shard chosen by hash, so a group's serialized state stays on a single + // executor while its lanes spread across shards. + asio::io_context::executor_type coordinatorExecutor(size_t hash) { + return executor(hash % contexts_.size()); + } + + // Round-robin executor for a new lane / accepted socket. + asio::io_context::executor_type nextLaneExecutor() { + return executor(rr_.fetch_add(1, std::memory_order_relaxed)); + } + + size_t nextShardIndex() { + return rr_.fetch_add(1, std::memory_order_relaxed) % contexts_.size(); + } + + // Spawns one thread per shard. before_run(i) runs on shard i's thread + // before each io_context.run() (used to arm the acceptor on shard 0 and to + // re-arm it after an exception-driven restart). + void start(std::function before_run) { + for (size_t i = 0; i < contexts_.size(); ++i) { + threads_.emplace_back([this, i, before_run] { +#ifdef __linux__ + std::string name = "mc-tcp-io-" + std::to_string(i); + pthread_setname_np(pthread_self(), name.c_str()); +#endif + while (!stopping_.load(std::memory_order_acquire)) { + try { + if (before_run) before_run(i); + contexts_[i]->run(); + break; + } catch (const std::exception &e) { + LOG(ERROR) << "TcpIoPool shard " << i + << " encountered an exception during run: " + << e.what(); + // Do not restart once stop() has published its intent: + // re-arming (e.g. shard 0's async_accept) after stop() + // stopped the context would keep run() alive and hang + // join(). + if (stopping_.load(std::memory_order_acquire)) break; + contexts_[i]->restart(); + } + } + }); + } + } + + void stop() { + stopping_.store(true, std::memory_order_release); + guards_.clear(); + for (auto &c : contexts_) c->stop(); + for (auto &t : threads_) + if (t.joinable()) t.join(); + threads_.clear(); + } + + private: + using work_guard_t = + asio::executor_work_guard; + std::vector> contexts_; + std::vector> guards_; + std::vector threads_; + std::atomic rr_{0}; + std::atomic stopping_{false}; +}; + class TcpTransport : public Transport { public: using BufferDesc = TransferMetadata::BufferDesc; @@ -94,8 +194,6 @@ class TcpTransport : public Transport { int unregisterLocalMemoryBatch( const std::vector &addr_list) override; - void worker(); - Slice *prepareTransfer(TransferTask *task, const TransferRequest &request); void startTransfer(Slice *slice, @@ -111,7 +209,8 @@ class TcpTransport : public Transport { private: TcpContext *context_; std::atomic_bool running_; - std::thread thread_; + std::unique_ptr io_pool_; + size_t num_io_threads_ = 1; bool enable_connection_pool_ = true; // Client-side bounded work queues and fixed connection lanes. @@ -211,10 +310,16 @@ class TcpTransport : public Transport { enum class LaneConnectStage { NONE, RESOLVING, CONNECTING }; struct ConnectionLaneRuntime { - explicit ConnectionLaneRuntime(asio::io_context &io_context) - : executor(io_context.get_executor()) {} + explicit ConnectionLaneRuntime(TcpIoPool &pool_arg) : pool(&pool_arg) {} - asio::io_context::executor_type executor; + asio::io_context::executor_type coordinatorExecutor(size_t hash) { + return pool->coordinatorExecutor(hash); + } + asio::io_context::executor_type nextLaneExecutor() { + return pool->nextLaneExecutor(); + } + + TcpIoPool *pool; }; struct ConnectionLaneState; @@ -222,11 +327,15 @@ class TcpTransport : public Transport { struct ConnectionLane { ConnectionLane(size_t lane_id_arg, - const std::shared_ptr &group_arg) - : lane_id(lane_id_arg), group(group_arg) {} + const std::shared_ptr &group_arg, + const asio::io_context::executor_type &executor_arg) + : lane_id(lane_id_arg), group(group_arg), executor(executor_arg) {} size_t lane_id; std::weak_ptr group; + // Shard executor this lane's socket/resolver/session/timers are bound + // to. Round-robin assigned so a group's lanes spread across io threads. + asio::io_context::executor_type executor; LaneState state = LaneState::DISCONNECTED; LaneConnectStage connect_stage = LaneConnectStage::NONE; uint64_t operation_epoch = 0; diff --git a/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp b/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp index 3ad57acb5d..1618037f8b 100644 --- a/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp +++ b/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp @@ -365,6 +365,12 @@ TcpTransport::TcpTransport() "MC_TCP_ADMISSION_TIMEOUT_MS", getenv("MC_TCP_ADMISSION_TIMEOUT_MS"), kDefaultAdmissionTimeoutMs, 1, kMaxAdmissionTimeoutMs)); + + constexpr size_t kDefaultNumIoThreads = 1; + constexpr size_t kMaxNumIoThreads = 32; + num_io_threads_ = parseBoundedTcpSetting( + "MC_NUM_TCP_IO_THREADS", getenv("MC_NUM_TCP_IO_THREADS"), + kDefaultNumIoThreads, 1, kMaxNumIoThreads); } TcpTransport::~TcpTransport() { @@ -419,15 +425,22 @@ int TcpTransport::install(std::string& local_server_name, close(sockfd); // the above function has opened a socket LOG(INFO) << "TcpTransport: listen on port " << tcp_port; auto metadata = metadata_; - context_ = new TcpContext(tcp_port, [metadata = std::move(metadata)]( - uint64_t addr, uint64_t size) { - return validateTcpAddress(metadata, addr, size); - }); - lane_runtime_ = - std::make_shared(context_->io_context); + io_pool_ = std::make_unique(num_io_threads_); + context_ = new TcpContext( + *io_pool_, tcp_port, + [metadata = std::move(metadata)](uint64_t addr, uint64_t size) { + return validateTcpAddress(metadata, addr, size); + }); + lane_runtime_ = std::make_shared(*io_pool_); lane_state_->runtime = lane_runtime_; running_ = true; - thread_ = std::thread(&TcpTransport::worker, this); + // Shard 0 owns the acceptor; arm (and re-arm after a restart) doAccept on + // its own thread. Other shards just drain their io_context. + io_pool_->start([this](size_t shard) { + if (shard == 0) context_->doAccept(); + }); + LOG(INFO) << "TcpTransport: I/O pool started with " << num_io_threads_ + << " thread(s)"; return 0; } @@ -502,9 +515,16 @@ Status TcpTransport::getTransferStatus(BatchID batch_id, size_t task_id, std::to_string(batch_id)); } auto& task = batch_desc.task_list[task_id]; - status.transferred_bytes = task.transferred_bytes; - uint64_t success_slice_count = task.success_slice_count; - uint64_t failed_slice_count = task.failed_slice_count; + // These counters are updated with __atomic_fetch_add from the I/O + // thread(s) in Slice::markSuccess/markFailed; read them atomically here to + // match MultiTransport::getTransferStatus and avoid a data race with the + // completion path (the read runs on the polling submission thread). + status.transferred_bytes = + __atomic_load_n(&task.transferred_bytes, __ATOMIC_ACQUIRE); + uint64_t success_slice_count = + __atomic_load_n(&task.success_slice_count, __ATOMIC_ACQUIRE); + uint64_t failed_slice_count = + __atomic_load_n(&task.failed_slice_count, __ATOMIC_ACQUIRE); if (success_slice_count + failed_slice_count == task.slice_count) { if (failed_slice_count) { status.s = TransferStatusEnum::FAILED; @@ -665,20 +685,6 @@ void TcpTransport::startTransferSequence(std::vector slices) { (*advance)(); } -void TcpTransport::worker() { - while (running_) { - try { - context_->doAccept(); - context_->io_context.run(); - } catch (std::exception& e) { - LOG(ERROR) << "TcpTransport::worker encountered an exception " - "during doAccept/run: " - << e.what(); - context_->io_context.restart(); - } - } -} - std::shared_ptr TcpTransport::getConnection( const std::string& host, uint16_t port) { // The reusable path is owned by fixed connection lanes. This helper is diff --git a/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport_lane_impl.h b/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport_lane_impl.h index 5aae705b94..09412e1663 100644 --- a/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport_lane_impl.h +++ b/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport_lane_impl.h @@ -398,14 +398,15 @@ void TcpTransport::runGroupRetirement( const std::shared_ptr& group) { std::shared_ptr retry_timer; std::shared_ptr admission_timer; - std::array, kMaxTcpLanesPerPeer> sessions; - std::array, kMaxTcpLanesPerPeer> - resolvers; - std::array, kMaxTcpLanesPerPeer> - sockets; - size_t session_count = 0; - size_t resolver_count = 0; - size_t socket_count = 0; + // Each lane's asio objects live on that lane's shard executor; collect them + // together with the executor so teardown can run on the owning shard. + struct LaneTeardown { + asio::io_context::executor_type executor; + std::shared_ptr session; + std::shared_ptr resolver; + std::shared_ptr socket; + }; + std::vector teardowns; uint64_t pump_epoch = 0; bool retired = false; { @@ -436,12 +437,10 @@ void TcpTransport::runGroupRetirement( for (const auto& lane : group->lanes) { ++lane->operation_epoch; if (lane->operation_epoch == 0) ++lane->operation_epoch; - if (lane->session) - sessions[session_count++] = std::move(lane->session); - if (lane->resolver) - resolvers[resolver_count++] = std::move(lane->resolver); - if (lane->socket) - sockets[socket_count++] = std::move(lane->socket); + if (lane->session || lane->resolver || lane->socket) + teardowns.push_back( + {lane->executor, std::move(lane->session), + std::move(lane->resolver), std::move(lane->socket)}); lane->connect_stage = LaneConnectStage::NONE; lane->state = LaneState::CLOSED; } @@ -457,17 +456,28 @@ void TcpTransport::runGroupRetirement( } if (!retired) return; - for (size_t i = 0; i < session_count; ++i) - if (sessions[i]) sessions[i]->cancel(); - for (size_t i = 0; i < resolver_count; ++i) { - const auto& resolver = resolvers[i]; - if (!resolver) continue; - try { - resolver->cancel(); - } catch (...) { - } + // Cancel/close each lane's socket, session and resolver ON ITS OWN SHARD + // so these asio objects are only touched by their owning io thread. With + // MC_NUM_TCP_IO_THREADS == 1 asio::dispatch runs inline (unchanged + // behavior); with more shards it hops to the lane's io thread, avoiding + // concurrent close()/cancel() against an outstanding async op. + for (auto& teardown : teardowns) { + auto session = std::move(teardown.session); + auto resolver = std::move(teardown.resolver); + auto socket = std::move(teardown.socket); + asio::dispatch(teardown.executor, [session, resolver, socket] { + if (session) session->cancel(); + if (resolver) { + try { + resolver->cancel(); + } catch (...) { + } + } + closeSocketNoThrow(socket); + }); } - for (size_t i = 0; i < socket_count; ++i) closeSocketNoThrow(sockets[i]); + // Retry/admission timers are group->executor objects and we are running on + // that executor, so cancel them here. cancelTimerNoThrow(retry_timer); cancelTimerNoThrow(admission_timer); @@ -546,14 +556,15 @@ void TcpTransport::enqueuePooledTransfer(const std::string& logical_peer, auto group_it = state->groups.find(key); if (group_it == state->groups.end()) { group = std::make_shared( - key, runtime->executor, + key, + runtime->coordinatorExecutor(ConnectionKeyHash{}(key)), state->max_queued_transfers_per_peer, state->max_pending_admissions_per_peer, state->admission_timeout, state->failure_counters); group->lanes.reserve(state->lanes_per_peer); for (size_t i = 0; i < state->lanes_per_peer; ++i) { - group->lanes.push_back( - std::make_shared(i, group)); + group->lanes.push_back(std::make_shared( + i, group, runtime->nextLaneExecutor())); } auto [inserted_it, inserted] = state->groups.emplace(key, group); @@ -858,10 +869,26 @@ void TcpTransport::runGroupPump( invokeLaneObserverHook(kLaneAdmissionTimerArmed, pending_depth, 0, 0, false); #endif - for (size_t i = 0; i < connect_count; ++i) - startLaneConnect(group, connects[i].lane, connects[i].epoch); - for (size_t i = 0; i < session_count; ++i) - startLaneSession(group, sessions[i].lane, sessions[i].epoch); + // Dispatch each lane's connect/session onto that lane's shard executor so + // its socket I/O runs on the owning io thread. asio::dispatch runs inline + // when already on the target executor (the single-shard/N=1 case stays + // identical to the historical inline calls) and hops otherwise. + for (size_t i = 0; i < connect_count; ++i) { + auto connect_lane = connects[i].lane; + uint64_t connect_epoch = connects[i].epoch; + asio::dispatch(connect_lane->executor, + [group, connect_lane, connect_epoch] { + startLaneConnect(group, connect_lane, connect_epoch); + }); + } + for (size_t i = 0; i < session_count; ++i) { + auto session_lane = sessions[i].lane; + uint64_t session_epoch = sessions[i].epoch; + asio::dispatch(session_lane->executor, + [group, session_lane, session_epoch] { + startLaneSession(group, session_lane, session_epoch); + }); + } failWorkItems(std::move(failed), failure_reason, group->failure_counters); failWorkItems(std::move(expired), WorkFailureReason::QUEUE_TIMEOUT, group->failure_counters); @@ -895,9 +922,9 @@ void TcpTransport::startLaneConnect( #endif try { lane->resolver = - std::make_shared(group->executor); + std::make_shared(lane->executor); lane->socket = - std::make_shared(group->executor); + std::make_shared(lane->executor); lane->connect_stage = LaneConnectStage::RESOLVING; lane->resolver->async_resolve( group->key.host, std::to_string(group->key.port), @@ -1403,49 +1430,70 @@ void TcpTransport::shutdownConnectionLanes() { } auto cancellation_posts = std::make_shared(); - if (context_ && running_) { + if (io_pool_ && running_) { + struct LaneTeardown { + asio::io_context::executor_type executor; + std::shared_ptr session; + std::shared_ptr resolver; + std::shared_ptr socket; + }; for (const auto& group : groups) { - cancellation_posts->add(); - try { - asio::post(group->executor, [group, cancellation_posts] { - std::vector> sessions; - std::vector> - resolvers; - std::vector> sockets; - std::shared_ptr retry_timer; - std::shared_ptr admission_timer; - { - std::lock_guard lock(group->mutex); - retry_timer = group->retry_timer; - admission_timer = group->admission_timer; - for (const auto& lane : group->lanes) { - if (lane->session) - sessions.push_back(lane->session); - if (lane->resolver) - resolvers.push_back(lane->resolver); - if (lane->socket) sockets.push_back(lane->socket); - } - } - for (const auto& session : sessions) - if (session) session->cancel(); - for (const auto& resolver : resolvers) { - if (!resolver) continue; - try { - resolver->cancel(); - } catch (...) { - } - } - for (const auto& socket : sockets) - closeSocketNoThrow(socket); - if (retry_timer) { - asio::error_code timer_ec; - retry_timer->cancel(timer_ec); - } - cancelTimerNoThrow(admission_timer); + std::vector teardowns; + std::shared_ptr retry_timer; + std::shared_ptr admission_timer; + { + std::lock_guard lock(group->mutex); + retry_timer = group->retry_timer; + admission_timer = group->admission_timer; + for (const auto& lane : group->lanes) { + if (lane->session || lane->resolver || lane->socket) + teardowns.push_back({lane->executor, lane->session, + lane->resolver, lane->socket}); + } + } + // Cancel/close each lane's asio objects on the lane's own shard; + // the group timers stay on the coordinator (group) executor. Each + // post is tracked so waitUntil() waits for all of them. At N=1 + // asio::dispatch runs inline on the single io thread. + for (auto& teardown : teardowns) { + cancellation_posts->add(); + auto session = teardown.session; + auto resolver = teardown.resolver; + auto socket = teardown.socket; + try { + asio::dispatch( + teardown.executor, + [session, resolver, socket, cancellation_posts] { + if (session) session->cancel(); + if (resolver) { + try { + resolver->cancel(); + } catch (...) { + } + } + closeSocketNoThrow(socket); + cancellation_posts->done(); + }); + } catch (...) { cancellation_posts->done(); - }); - } catch (...) { - cancellation_posts->done(); + } + } + if (retry_timer || admission_timer) { + cancellation_posts->add(); + try { + asio::dispatch( + group->executor, + [retry_timer, admission_timer, cancellation_posts] { + if (retry_timer) { + asio::error_code timer_ec; + retry_timer->cancel(timer_ec); + } + cancelTimerNoThrow(admission_timer); + cancellation_posts->done(); + }); + } catch (...) { + cancellation_posts->done(); + } } } cancellation_posts->waitUntil(std::chrono::steady_clock::now() + @@ -1453,8 +1501,10 @@ void TcpTransport::shutdownConnectionLanes() { } running_ = false; - if (context_) context_->io_context.stop(); - if (thread_.joinable()) thread_.join(); + // Stopping the pool stops every shard io_context and joins every io + // thread; after this returns no handler can run, which is the quiescence + // barrier for the synchronous teardown below. + if (io_pool_) io_pool_->stop(); std::deque deferred; std::vector> sessions; diff --git a/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport_session_impl.h b/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport_session_impl.h index 07219929e6..40a1c4c3ca 100644 --- a/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport_session_impl.h +++ b/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport_session_impl.h @@ -921,8 +921,11 @@ struct ClientSession : public std::enable_shared_from_this { }; struct TcpContext { - TcpContext(short port, ValidateAddrFn validate_addr) - : acceptor(io_context), validate_addr_(std::move(validate_addr)) { + TcpContext(TcpIoPool& pool, short port, ValidateAddrFn validate_addr) + : pool_(pool), + io_context(pool.context(0)), + acceptor(io_context), + validate_addr_(std::move(validate_addr)) { std::error_code ec; asio::ip::tcp::endpoint endpoint(asio::ip::tcp::v6(), port); @@ -950,21 +953,30 @@ struct TcpContext { } void doAccept() { - acceptor.async_accept([this](asio::error_code ec, tcpsocket socket) { - if (!ec) { - asio::error_code nodelay_ec; - socket.set_option(asio::ip::tcp::no_delay(true), nodelay_ec); - auto socket_ptr = - std::make_shared(std::move(socket)); - auto session = - std::make_shared(socket_ptr, validate_addr_); - session->start(); - } - doAccept(); - }); + // Round-robin each accepted connection onto a shard so ServerSession + // I/O spreads across io threads. The accept handler runs on shard 0 + // (the acceptor's context); the accepted socket lives on the target + // shard, so start its session there. + acceptor.async_accept( + pool_.context(pool_.nextShardIndex()), + [this](asio::error_code ec, tcpsocket socket) { + if (!ec) { + asio::error_code nodelay_ec; + socket.set_option(asio::ip::tcp::no_delay(true), + nodelay_ec); + auto socket_ptr = + std::make_shared(std::move(socket)); + auto session = std::make_shared( + socket_ptr, validate_addr_); + asio::post(socket_ptr->get_executor(), + [session] { session->start(); }); + } + doAccept(); + }); } - asio::io_context io_context; + TcpIoPool& pool_; + asio::io_context& io_context; asio::ip::tcp::acceptor acceptor; ValidateAddrFn validate_addr_; };