From 145c7ece3c943fda206cf36e8e2c710a105af561 Mon Sep 17 00:00:00 2001 From: Liqian Lin Date: Wed, 19 Aug 2026 22:09:40 +0800 Subject: [PATCH 1/2] test: reproduce stalled TCP payload without progress deadline --- .../tests/tcp_write_visibility_test.cpp | 288 +++++++++++++++++- 1 file changed, 286 insertions(+), 2 deletions(-) diff --git a/mooncake-transfer-engine/tests/tcp_write_visibility_test.cpp b/mooncake-transfer-engine/tests/tcp_write_visibility_test.cpp index 4af53322fd..c08848df6c 100644 --- a/mooncake-transfer-engine/tests/tcp_write_visibility_test.cpp +++ b/mooncake-transfer-engine/tests/tcp_write_visibility_test.cpp @@ -37,6 +37,7 @@ #include #include #include +#include #include #include #include @@ -674,6 +675,142 @@ class HoldingWriteServer { std::atomic max_accepted_count_{0}; }; +class PayloadStallingWriteServer { + public: + explicit PayloadStallingWriteServer(size_t payload_prefix_to_read) + : payload_prefix_to_read_(payload_prefix_to_read) { + listen_fd_ = socket(AF_INET, SOCK_STREAM, 0); + if (listen_fd_ < 0) return; + int one = 1; + if (setsockopt(listen_fd_, SOL_SOCKET, SO_REUSEADDR, &one, + sizeof(one)) != 0) + return; + + sockaddr_in addr{}; + addr.sin_family = AF_INET; + addr.sin_port = 0; + addr.sin_addr.s_addr = htonl(INADDR_ANY); + if (bind(listen_fd_, reinterpret_cast(&addr), + sizeof(addr)) != 0) + return; + if (listen(listen_fd_, 1) != 0) return; + + socklen_t len = sizeof(addr); + if (getsockname(listen_fd_, reinterpret_cast(&addr), &len) != + 0) + return; + port_ = ntohs(addr.sin_port); + ok_ = true; + thread_ = std::thread([this] { serve(); }); + } + + ~PayloadStallingWriteServer() { release(); } + + bool ok() const { return ok_; } + uint16_t port() const { return port_; } + + bool waitForHeader(std::chrono::seconds timeout) { + std::unique_lock lock(event_mutex_); + return event_cv_.wait_for(lock, timeout, + [this] { return header_received_; }); + } + + bool waitForPayloadPrefix(std::chrono::seconds timeout) { + std::unique_lock lock(event_mutex_); + return event_cv_.wait_for(lock, timeout, + [this] { return payload_prefix_received_; }); + } + + void release() { + { + std::lock_guard lock(event_mutex_); + released_ = true; + } + event_cv_.notify_all(); + if (listen_fd_ >= 0) (void)shutdown(listen_fd_, SHUT_RDWR); + int accepted = -1; + { + std::lock_guard lock(socket_mutex_); + accepted = accepted_fd_; + } + if (accepted >= 0) (void)shutdown(accepted, SHUT_RDWR); + if (thread_.joinable()) thread_.join(); + if (accepted >= 0) close(accepted); + { + std::lock_guard lock(socket_mutex_); + accepted_fd_ = -1; + } + if (listen_fd_ >= 0) { + close(listen_fd_); + listen_fd_ = -1; + } + } + + private: + static bool recvExact(int fd, void* buffer, size_t size) { + char* out = static_cast(buffer); + while (size) { + const ssize_t n = recv(fd, out, size, 0); + if (n <= 0) return false; + out += n; + size -= static_cast(n); + } + return true; + } + + void serve() { + const int fd = accept(listen_fd_, nullptr, nullptr); + if (fd < 0) return; + { + std::lock_guard lock(socket_mutex_); + accepted_fd_ = fd; + } + int receive_buffer = 4096; + if (payload_prefix_to_read_ == 0) { + (void)setsockopt(fd, SOL_SOCKET, SO_RCVBUF, &receive_buffer, + sizeof(receive_buffer)); + } + + TestSessionHeader header{}; + if (!recvExact(fd, &header, sizeof(header))) return; + { + std::lock_guard lock(event_mutex_); + header_received_ = true; + } + event_cv_.notify_all(); + + const size_t prefix = + std::min(payload_prefix_to_read_, le64toh(header.size)); + std::vector payload(prefix); + if (prefix != 0 && !recvExact(fd, payload.data(), prefix)) return; + if (prefix != 0) { + (void)setsockopt(fd, SOL_SOCKET, SO_RCVBUF, &receive_buffer, + sizeof(receive_buffer)); + } + { + std::lock_guard lock(event_mutex_); + payload_prefix_received_ = true; + } + event_cv_.notify_all(); + + std::unique_lock lock(event_mutex_); + event_cv_.wait(lock, [this] { return released_; }); + } + + const size_t payload_prefix_to_read_; + int listen_fd_ = -1; + uint16_t port_ = 0; + bool ok_ = false; + std::thread thread_; + std::mutex socket_mutex_; + int accepted_fd_ = -1; + std::mutex event_mutex_; + std::condition_variable event_cv_; + bool header_received_ = false; + bool payload_prefix_received_ = false; + bool released_ = false; +}; + // Keeps a local TCP port bound but deliberately never listens. Linux rejects // connect attempts deterministically, without blackhole-address timing. class UnavailableTcpPeer { @@ -1015,8 +1152,26 @@ bool waitForBatchTerminal(TransferEngine* engine, Transport::BatchID batch_id, return true; } -void expectEverySliceCompletedExactlyOnceAfterShutdown( - Transport::BatchID batch_id) { +template +bool waitForTaskTerminal(TransferEngine* engine, Transport::BatchID batch_id, + std::chrono::duration timeout, + TransferStatusEnum* final_status = nullptr) { + TransferStatus status; + status.s = TransferStatusEnum::WAITING; + const bool terminal = waitForPredicate( + [&] { + status.s = TransferStatusEnum::WAITING; + if (!engine->getTransferStatus(batch_id, 0, status).ok()) + return false; + return status.s == TransferStatusEnum::COMPLETED || + status.s == TransferStatusEnum::FAILED; + }, + timeout); + if (final_status) *final_status = status.s; + return terminal; +} + +void expectEverySliceCompletedExactlyOnce(Transport::BatchID batch_id) { const auto& batch = Transport::toBatchDesc(batch_id); for (size_t task_id = 0; task_id < batch.task_list.size(); ++task_id) { const auto& task = batch.task_list[task_id]; @@ -1035,6 +1190,11 @@ void expectEverySliceCompletedExactlyOnceAfterShutdown( } } +void expectEverySliceCompletedExactlyOnceAfterShutdown( + Transport::BatchID batch_id) { + expectEverySliceCompletedExactlyOnce(batch_id); +} + void expectEverySliceSucceededExactlyOnceAfterShutdown( Transport::BatchID batch_id) { const auto& batch = Transport::toBatchDesc(batch_id); @@ -2877,4 +3037,128 @@ TEST(TcpWriteVisibilityTest, for (const auto batch_id : batch_ids) reclaimBatchDescAfterEngineShutdownForTest(batch_id); } + +TEST(TcpWriteVisibilityTest, WritePayloadNoProgressFailsWithinBoundedDeadline) { + constexpr size_t kRequestSize = 64ull << 20; + ScopedEnvVar lanes("MC_TCP_LANES_PER_PEER", "1"); + ScopedEnvVar status_timeout("MC_TCP_STATUS_TIMEOUT_SEC", "1"); + ScopedEnvVar progress_timeout("MC_TCP_PROGRESS_TIMEOUT_SEC", "1"); + ScopedLaneHooks hooks; + PayloadStallingWriteServer fake_peer(/*payload_prefix_to_read=*/0); + ASSERT_TRUE(fake_peer.ok()); + + EngineHandle h; + h.init(P2PHANDSHAKE, "127.0.0.2:18044", kRequestSize); + ASSERT_TRUE(h.ok); + pointTcpSegmentAt(h, fake_peer.port()); + + const auto batch_id = h.engine->allocateBatchID(1); + ASSERT_TRUE( + h.engine->submitTransfer(batch_id, {makeWriteRequest(h, kRequestSize)}) + .ok()); + ASSERT_TRUE(fake_peer.waitForHeader(std::chrono::seconds(5))); + ASSERT_TRUE(waitForPredicate( + [] { return lane_busy_count.load(std::memory_order_acquire) >= 1; }, + std::chrono::seconds(5))); + + TransferStatusEnum status = TransferStatusEnum::WAITING; + const bool terminal = waitForTaskTerminal( + h.engine.get(), batch_id, std::chrono::milliseconds(1500), &status); + EXPECT_TRUE(terminal) + << "RED: a connected WRITE that makes no payload progress has no " + "general progress deadline"; + EXPECT_EQ(status, TransferStatusEnum::FAILED); + + fake_peer.release(); + if (waitForTaskTerminal(h.engine.get(), batch_id, + std::chrono::seconds(5))) { + expectEverySliceCompletedExactlyOnce(batch_id); + EXPECT_TRUE(h.engine->freeBatchID(batch_id).ok()); + } else { + h.engine.reset(); + reclaimBatchDescAfterEngineShutdownForTest(batch_id); + } +} + +TEST(TcpWriteVisibilityTest, + WritePayloadPartialProgressThenStallFailsWithinBoundedDeadline) { + constexpr size_t kRequestSize = 64ull << 20; + constexpr size_t kObservedPayloadProgress = 128ull << 10; + ScopedEnvVar lanes("MC_TCP_LANES_PER_PEER", "1"); + ScopedEnvVar status_timeout("MC_TCP_STATUS_TIMEOUT_SEC", "1"); + ScopedEnvVar progress_timeout("MC_TCP_PROGRESS_TIMEOUT_SEC", "1"); + ScopedLaneHooks hooks; + PayloadStallingWriteServer fake_peer(kObservedPayloadProgress); + ASSERT_TRUE(fake_peer.ok()); + + EngineHandle h; + h.init(P2PHANDSHAKE, "127.0.0.2:18045", kRequestSize); + ASSERT_TRUE(h.ok); + pointTcpSegmentAt(h, fake_peer.port()); + + const auto batch_id = h.engine->allocateBatchID(1); + ASSERT_TRUE( + h.engine->submitTransfer(batch_id, {makeWriteRequest(h, kRequestSize)}) + .ok()); + ASSERT_TRUE(fake_peer.waitForHeader(std::chrono::seconds(5))); + ASSERT_TRUE(fake_peer.waitForPayloadPrefix(std::chrono::seconds(5))) + << "the server must observe real payload progress before stalling"; + ASSERT_TRUE(waitForPredicate( + [] { return lane_busy_count.load(std::memory_order_acquire) >= 1; }, + std::chrono::seconds(5))); + + TransferStatusEnum status = TransferStatusEnum::WAITING; + const bool terminal = waitForTaskTerminal( + h.engine.get(), batch_id, std::chrono::milliseconds(1500), &status); + EXPECT_TRUE(terminal) + << "RED: the status-frame timer is not a payload progress deadline " + "after real WRITE progress stops"; + EXPECT_EQ(status, TransferStatusEnum::FAILED); + + fake_peer.release(); + if (waitForTaskTerminal(h.engine.get(), batch_id, + std::chrono::seconds(5))) { + expectEverySliceCompletedExactlyOnce(batch_id); + EXPECT_TRUE(h.engine->freeBatchID(batch_id).ok()); + } else { + h.engine.reset(); + reclaimBatchDescAfterEngineShutdownForTest(batch_id); + } +} + +TEST(TcpWriteVisibilityTest, ShutdownWithArmedProgressDeadlineCompletesOnce) { + constexpr size_t kRequestSize = 64ull << 20; + ScopedEnvVar lanes("MC_TCP_LANES_PER_PEER", "1"); + ScopedEnvVar progress_timeout("MC_TCP_PROGRESS_TIMEOUT_SEC", "10"); + ScopedLaneHooks hooks; + PayloadStallingWriteServer fake_peer(/*payload_prefix_to_read=*/0); + ASSERT_TRUE(fake_peer.ok()); + + EngineHandle h; + h.init(P2PHANDSHAKE, "127.0.0.2:18051", kRequestSize); + ASSERT_TRUE(h.ok); + pointTcpSegmentAt(h, fake_peer.port()); + + const auto batch_id = h.engine->allocateBatchID(1); + ASSERT_TRUE( + h.engine->submitTransfer(batch_id, {makeWriteRequest(h, kRequestSize)}) + .ok()); + ASSERT_TRUE(fake_peer.waitForHeader(std::chrono::seconds(5))); + ASSERT_TRUE(waitForPredicate( + [] { return lane_busy_count.load(std::memory_order_acquire) >= 1; }, + std::chrono::seconds(5))); + + auto engine = std::move(h.engine); + auto destruction = + std::async(std::launch::async, + [engine = std::move(engine)]() mutable { engine.reset(); }); + ASSERT_EQ(destruction.wait_for(std::chrono::seconds(5)), + std::future_status::ready); + destruction.get(); + + expectEverySliceCompletedExactlyOnceAfterShutdown(batch_id); + hooks.reset(); + reclaimBatchDescAfterEngineShutdownForTest(batch_id); + fake_peer.release(); +} #endif From 3606fcb8f0a32eb11ab30c695efabe4aac590c02 Mon Sep 17 00:00:00 2001 From: Liqian Lin Date: Wed, 19 Aug 2026 22:09:40 +0800 Subject: [PATCH 2/2] fix: add session-level TCP progress deadline --- .../transport/tcp_transport/tcp_transport.cpp | 31 ++ .../tcp_transport_session_impl.h | 163 ++++++++++- .../tests/tcp_write_visibility_test.cpp | 271 +++++++++++++++++- 3 files changed, 457 insertions(+), 8 deletions(-) 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 038de3bdc5..5169e766ba 100644 --- a/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp +++ b/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp @@ -56,6 +56,7 @@ using LaneRetryHandlerHook = void (*)() noexcept; using LaneAdmissionHandlerHook = void (*)() noexcept; using LaneObserverHook = void (*)(int, size_t, uint64_t, size_t, bool) noexcept; using LaneFailureReasonHook = void (*)(int) noexcept; +using SessionProgressHook = int (*)(int, bool) noexcept; std::mutex lane_test_hook_mutex; LaneConnectHandlerHook lane_connect_handler_hook = nullptr; @@ -64,6 +65,7 @@ LaneRetryHandlerHook lane_retry_handler_hook = nullptr; LaneAdmissionHandlerHook lane_admission_handler_hook = nullptr; LaneObserverHook lane_observer_hook = nullptr; LaneFailureReasonHook lane_failure_reason_hook = nullptr; +SessionProgressHook session_progress_hook = nullptr; enum LaneTestEvent { kLaneQueueAdmitted = 1, @@ -85,6 +87,20 @@ enum LaneTestEvent { kLaneAdmissionHardRejected = 17, }; +enum SessionProgressTestEvent { + kSessionReadBodySuccess = 1, + kSessionWriteAckSuccess = 2, + kSessionTimeoutCommitted = 3, + kSessionTimeoutStale = 4, + kSessionTerminal = 5, +}; + +enum SessionProgressTestAction { + kSessionProgressNoAction = 0, + kSessionCommitTimeoutBeforeProgress = 1, + kSessionReplayPreviousTimeoutAfterProgress = 2, +}; + void invokeLaneConnectHandlerHook() noexcept { LaneConnectHandlerHook hook; { @@ -142,6 +158,15 @@ void invokeLaneFailureReasonHook(int reason) noexcept { } if (hook) hook(reason); } + +int invokeSessionProgressHook(int event, bool detail) noexcept { + SessionProgressHook hook; + { + std::lock_guard lock(lane_test_hook_mutex); + hook = session_progress_hook; + } + return hook ? hook(event, detail) : kSessionProgressNoAction; +} } // namespace void tcpTransportSetLaneConnectHandlerHookForTest( @@ -179,6 +204,12 @@ void tcpTransportSetLaneFailureReasonHookForTest( lane_failure_reason_hook = hook; } +void tcpTransportSetSessionProgressHookForTest( + SessionProgressHook hook) noexcept { + std::lock_guard lock(lane_test_hook_mutex); + session_progress_hook = hook; +} + bool tcpTransportLaneTypesAreMoveOnlyForTest() noexcept { return std::is_move_constructible::value && !std::is_copy_constructible::value && 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 4cb8907e48..d235b58be9 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 @@ -401,6 +401,21 @@ struct ClientSession : public std::enable_shared_from_this { } std::optional status_timer_; bool status_deadline_disarmed_ = false; + // Unlike the status-frame deadline above, this is a sliding liveness + // deadline for the whole request. It is refreshed only after an Asio + // completion reports actual protocol or payload progress. + static int progressTimeoutSec() { + const char* env = std::getenv("MC_TCP_PROGRESS_TIMEOUT_SEC"); + if (env) { + int v = std::atoi(env); + if (v > 0) return v; + } + return 30; + } + std::optional progress_timer_; + uint64_t progress_epoch_ = 0; + bool progress_deadline_disarmed_ = true; + bool progress_timeout_committed_ = false; bool terminal_reported_ = false; OnTerminal on_terminal_; @@ -416,6 +431,7 @@ struct ClientSession : public std::enable_shared_from_this { void cancel() noexcept { cancelStatusDeadline(); + cancelProgressDeadline(); if (!socket_) return; asio::error_code cancel_ec; socket_->cancel(cancel_ec); @@ -424,10 +440,10 @@ struct ClientSession : public std::enable_shared_from_this { } private: - // All handlers run on the transport's single io thread, so arm/cancel - // and the expiry handler never race. Expiry only closes the socket: the - // pending status read then completes with an error and its handler owns - // the failure path (including source-buffer quiescence for WRITE). + // The status-frame deadline is separate from the sliding progress + // deadline below. Its expiry closes the socket, but a status completion + // already ready in the executor queue retains its original result; the + // status handler still owns the terminal path. void armStatusDeadline() { auto self(shared_from_this()); status_deadline_disarmed_ = false; @@ -461,6 +477,83 @@ struct ClientSession : public std::enable_shared_from_this { } } + void refreshProgressDeadline() { + if (terminal_reported_ || progress_timeout_committed_) return; + auto self(shared_from_this()); + progress_deadline_disarmed_ = false; + ++progress_epoch_; + if (progress_epoch_ == 0) ++progress_epoch_; + const uint64_t progress_epoch = progress_epoch_; + if (!progress_timer_) progress_timer_.emplace(socket_->get_executor()); + progress_timer_->expires_after( + std::chrono::seconds(progressTimeoutSec())); + progress_timer_->async_wait( + [this, self, progress_epoch](const asio::error_code& ec) { + handleProgressDeadline(progress_epoch, ec); + }); + } + + void handleProgressDeadline(uint64_t progress_epoch, + const asio::error_code& ec) { + if (ec == asio::error::operation_aborted || + progress_deadline_disarmed_ || terminal_reported_ || + progress_epoch != progress_epoch_) { +#ifdef MOONCAKE_TCP_TRANSPORT_TEST_HOOKS + (void)invokeSessionProgressHook(kSessionTimeoutStale, false); +#endif + return; + } + progress_timeout_committed_ = true; + progress_deadline_disarmed_ = true; + ++progress_epoch_; + if (progress_epoch_ == 0) ++progress_epoch_; + if ((header_.opcode & ~kOpcodeV2Flag) == + static_cast(TransferRequest::WRITE)) { + write_abort_requested_ = true; + } +#ifdef MOONCAKE_TCP_TRANSPORT_TEST_HOOKS + (void)invokeSessionProgressHook(kSessionTimeoutCommitted, + write_body_in_flight_); +#endif + LOG(ERROR) << "ClientSession: no transfer progress within " + << progressTimeoutSec() << "s; dropping connection"; + if (socket_ && socket_->is_open()) { + asio::error_code cancel_ec; + socket_->cancel(cancel_ec); + asio::error_code close_ec; + socket_->close(close_ec); + } + } + + bool acceptProgress() { + if (progress_timeout_committed_) return false; + refreshProgressDeadline(); + return !progress_timeout_committed_; + } + +#ifdef MOONCAKE_TCP_TRANSPORT_TEST_HOOKS + bool acceptProgressForTest(int event) { + const uint64_t previous_epoch = progress_epoch_; + const int action = invokeSessionProgressHook(event, false); + if (action == kSessionCommitTimeoutBeforeProgress) + handleProgressDeadline(previous_epoch, asio::error_code{}); + if (!acceptProgress()) return false; + if (action == kSessionReplayPreviousTimeoutAfterProgress) + handleProgressDeadline(previous_epoch, asio::error_code{}); + return !progress_timeout_committed_; + } +#endif + + void cancelProgressDeadline() { + progress_deadline_disarmed_ = true; + ++progress_epoch_; + if (progress_epoch_ == 0) ++progress_epoch_; + if (progress_timer_) { + asio::error_code ec; + progress_timer_->cancel(ec); + } + } + // Single terminal path. The invoking Asio operation has already released // its buffer before entering its completion handler. The lane posts any // follow-up pump, so a clean socket cannot be reused inline here. @@ -468,6 +561,10 @@ struct ClientSession : public std::enable_shared_from_this { if (terminal_reported_) return; terminal_reported_ = true; cancelStatusDeadline(); + cancelProgressDeadline(); +#ifdef MOONCAKE_TCP_TRANSPORT_TEST_HOOKS + (void)invokeSessionProgressHook(kSessionTerminal, clean); +#endif auto on_terminal = std::move(on_terminal_); if (on_terminal) on_terminal(status, clean); } @@ -497,6 +594,15 @@ struct ClientSession : public std::enable_shared_from_this { finalize(TransferStatusEnum::FAILED, false); return; } + if (!acceptProgress()) { + if ((header_.opcode & ~kOpcodeV2Flag) == + static_cast(TransferRequest::WRITE)) { + abortWrite(); + } else { + finalize(TransferStatusEnum::FAILED, false); + } + return; + } if ((header_.opcode & ~kOpcodeV2Flag) == (uint8_t)TransferRequest::WRITE) { if (v2_) readWriteAck(); // concurrent with the body @@ -507,6 +613,10 @@ struct ClientSession : public std::enable_shared_from_this { readBody(); } }); + // Asio does not invoke the completion handler inline. Arm only after + // initiation succeeds so a synchronous initiation exception cannot + // leave a timer handler retaining this otherwise-abandoned session. + refreshProgressDeadline(); } // v2 READ: the server prefixes the data with a status frame. @@ -525,6 +635,10 @@ struct ClientSession : public std::enable_shared_from_this { finalize(TransferStatusEnum::FAILED, false); return; } + if (!acceptProgress()) { + finalize(TransferStatusEnum::FAILED, false); + return; + } uint64_t frame = le64toh(status_frame_); if (!statusFrameValid(frame)) { LOG(ERROR) << "ClientSession: malformed READ status " @@ -569,6 +683,17 @@ struct ClientSession : public std::enable_shared_from_this { abortWrite(); return; } + bool progress_accepted = false; +#ifdef MOONCAKE_TCP_TRANSPORT_TEST_HOOKS + progress_accepted = + acceptProgressForTest(kSessionWriteAckSuccess); +#else + progress_accepted = acceptProgress(); +#endif + if (!progress_accepted) { + abortWrite(); + return; + } uint64_t frame = le64toh(status_frame_); if (!statusFrameValid(frame)) { LOG(ERROR) << "ClientSession: malformed WRITE ack frame " @@ -640,6 +765,25 @@ struct ClientSession : public std::enable_shared_from_this { return; } + bool progress_accepted = !progress_timeout_committed_; + if (progress_accepted && transferred_bytes > 0) { +#ifdef MOONCAKE_TCP_TRANSPORT_TEST_HOOKS + progress_accepted = + acceptProgressForTest(kSessionReadBodySuccess); +#else + progress_accepted = acceptProgress(); +#endif + } + if (!progress_accepted) { +#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \ + defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \ + defined(USE_COREX) + if (cuda_device >= 0) delete[] dram_buffer; +#endif + finalize(TransferStatusEnum::FAILED, false); + return; + } + #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \ defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \ defined(USE_COREX) @@ -753,9 +897,14 @@ struct ClientSession : public std::enable_shared_from_this { return; } if (write_abort_requested_) { - // The early ack path closed the socket while this - // operation still owned the caller's source buffer. It is - // safe to publish failure now that the handler has run. + // An early ACK failure or a committed progress timeout + // closed the socket while this operation still owned the + // caller's source buffer. It is safe to publish failure + // now that the handler has run. + finalize(TransferStatusEnum::FAILED, false); + return; + } + if (transferred_bytes > 0 && !acceptProgress()) { finalize(TransferStatusEnum::FAILED, false); return; } diff --git a/mooncake-transfer-engine/tests/tcp_write_visibility_test.cpp b/mooncake-transfer-engine/tests/tcp_write_visibility_test.cpp index c08848df6c..dab35bc84f 100644 --- a/mooncake-transfer-engine/tests/tcp_write_visibility_test.cpp +++ b/mooncake-transfer-engine/tests/tcp_write_visibility_test.cpp @@ -64,6 +64,8 @@ void tcpTransportSetLaneAdmissionHandlerHookForTest( void (*hook)() noexcept) noexcept; void tcpTransportSetLaneFailureReasonHookForTest( void (*hook)(int) noexcept) noexcept; +void tcpTransportSetSessionProgressHookForTest( + int (*hook)(int, bool) noexcept) noexcept; bool tcpTransportLaneTypesAreMoveOnlyForTest() noexcept; } // namespace mooncake #endif @@ -108,6 +110,20 @@ enum WorkFailureReasonForTest { kWorkShutdown = 5, }; +enum SessionProgressTestEvent { + kSessionReadBodySuccess = 1, + kSessionWriteAckSuccess = 2, + kSessionTimeoutCommitted = 3, + kSessionTimeoutStale = 4, + kSessionTerminal = 5, +}; + +enum SessionProgressTestAction { + kSessionProgressNoAction = 0, + kSessionCommitTimeoutBeforeProgress = 1, + kSessionReplayPreviousTimeoutAfterProgress = 2, +}; + std::atomic lane_connect_handler_entered{false}; std::atomic lane_connect_injection_call_count{0}; std::atomic release_lane_connect_handler{false}; @@ -130,6 +146,7 @@ std::atomic queue_rejection_count{0}; std::atomic lane_connecting_count{0}; std::atomic lane_busy_count{0}; std::atomic lane_terminal_count{0}; +std::atomic lane_terminal_socket_count{0}; std::atomic lane_shutdown_clean_count{0}; std::atomic late_lane_handler_count{0}; std::atomic retry_armed_count{0}; @@ -148,6 +165,16 @@ std::atomic connect_failure_count{0}; std::atomic runtime_unavailable_failure_count{0}; std::atomic session_failure_count{0}; std::atomic shutdown_failure_count{0}; +std::atomic session_progress_target_event{0}; +std::atomic session_progress_action{0}; +std::atomic session_timeout_committed_count{0}; +std::atomic session_timeout_stale_count{0}; +std::atomic session_terminal_count{0}; +std::atomic session_dirty_terminal_count{0}; +std::atomic session_timeout_commit_entered{false}; +std::atomic session_timeout_commit_had_write_in_flight{false}; +std::atomic hold_session_timeout_commit{false}; +std::atomic release_session_timeout_commit{false}; template bool waitForPredicate(Predicate&& predicate, @@ -192,6 +219,7 @@ void resetLaneTestState() noexcept { lane_connecting_count.store(0, std::memory_order_release); lane_busy_count.store(0, std::memory_order_release); lane_terminal_count.store(0, std::memory_order_release); + lane_terminal_socket_count.store(0, std::memory_order_release); lane_shutdown_clean_count.store(0, std::memory_order_release); late_lane_handler_count.store(0, std::memory_order_release); retry_armed_count.store(0, std::memory_order_release); @@ -212,6 +240,20 @@ void resetLaneTestState() noexcept { shutdown_failure_count.store(0, std::memory_order_release); } +void resetSessionProgressTestState() noexcept { + session_progress_target_event.store(0, std::memory_order_release); + session_progress_action.store(0, std::memory_order_release); + session_timeout_committed_count.store(0, std::memory_order_release); + session_timeout_stale_count.store(0, std::memory_order_release); + session_terminal_count.store(0, std::memory_order_release); + session_dirty_terminal_count.store(0, std::memory_order_release); + session_timeout_commit_entered.store(false, std::memory_order_release); + session_timeout_commit_had_write_in_flight.store(false, + std::memory_order_release); + hold_session_timeout_commit.store(false, std::memory_order_release); + release_session_timeout_commit.store(false, std::memory_order_release); +} + bool failSecondLaneConnectAttemptOnce(size_t) noexcept { const int attempt = lane_connect_injection_call_count.fetch_add(1) + 1; return attempt == 2; @@ -278,8 +320,11 @@ void observeLaneState(int event, size_t queue_depth, uint64_t, } if (event == kLaneQueueRejected) queue_rejection_count.fetch_add(1, std::memory_order_relaxed); - if (event == kLaneTerminal) + if (event == kLaneTerminal) { lane_terminal_count.fetch_add(1, std::memory_order_relaxed); + lane_terminal_socket_count.store(active_sockets, + std::memory_order_release); + } if (event == kLaneShutdownClean) lane_shutdown_clean_count.fetch_add(1, std::memory_order_relaxed); if (event == kLaneLateHandler) @@ -334,6 +379,36 @@ void observeWorkFailureReason(int reason) noexcept { shutdown_failure_count.fetch_add(1, std::memory_order_relaxed); } +int observeSessionProgress(int event, bool detail) noexcept { + if (event == kSessionTimeoutCommitted) { + session_timeout_committed_count.fetch_add(1, std::memory_order_relaxed); + session_timeout_commit_had_write_in_flight.store( + detail, std::memory_order_release); + if (hold_session_timeout_commit.load(std::memory_order_acquire)) { + session_timeout_commit_entered.store(true, + std::memory_order_release); + while (!release_session_timeout_commit.load( + std::memory_order_acquire)) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + } + } + if (event == kSessionTimeoutStale) + session_timeout_stale_count.fetch_add(1, std::memory_order_relaxed); + if (event == kSessionTerminal) { + session_terminal_count.fetch_add(1, std::memory_order_relaxed); + if (!detail) + session_dirty_terminal_count.fetch_add(1, + std::memory_order_relaxed); + } + if (event != + session_progress_target_event.load(std::memory_order_acquire)) { + return kSessionProgressNoAction; + } + return session_progress_action.exchange(kSessionProgressNoAction, + std::memory_order_acq_rel); +} + class ScopedLaneHooks { public: explicit ScopedLaneHooks(bool block_first_connect_handler = false, @@ -390,6 +465,37 @@ class ScopedLaneHooks { private: bool active_ = true; }; + +class ScopedSessionProgressHooks { + public: + ScopedSessionProgressHooks(int target_event = 0, + int action = kSessionProgressNoAction, + bool block_timeout_commit = false) { + resetSessionProgressTestState(); + session_progress_target_event.store(target_event, + std::memory_order_release); + session_progress_action.store(action, std::memory_order_release); + hold_session_timeout_commit.store(block_timeout_commit, + std::memory_order_release); + tcpTransportSetSessionProgressHookForTest(observeSessionProgress); + } + + ~ScopedSessionProgressHooks() { reset(); } + + void releaseTimeoutCommit() noexcept { + release_session_timeout_commit.store(true, std::memory_order_release); + } + + void reset() noexcept { + if (!active_) return; + releaseTimeoutCommit(); + tcpTransportSetSessionProgressHookForTest(nullptr); + active_ = false; + } + + private: + bool active_ = true; +}; #endif void reclaimBatchDescAfterEngineShutdownForTest(Transport::BatchID batch_id) { @@ -1102,6 +1208,18 @@ TransferRequest makeWriteRequest(const EngineHandle& handle, size_t length, return request; } +TransferRequest makeReadRequest(const EngineHandle& handle, void* destination, + size_t length, uint64_t target_offset = 0) { + TransferRequest request; + request.opcode = TransferRequest::READ; + request.length = length; + request.source = destination; + request.target_id = handle.segment_id; + request.target_offset = + target_offset == 0 ? handle.remote_base : target_offset; + return request; +} + // Submit one request and poll until terminal state; returns final status. TransferStatusEnum runOne(TransferEngine* engine, TransferRequest entry) { auto batch_id = engine->allocateBatchID(1); @@ -3038,6 +3156,157 @@ TEST(TcpWriteVisibilityTest, reclaimBatchDescAfterEngineShutdownForTest(batch_id); } +TEST(TcpWriteVisibilityTest, + ProgressTimeoutWinsAgainstQueuedFinalReadCompletion) { + ScopedEnvVar lanes("MC_TCP_LANES_PER_PEER", "1"); + ScopedEnvVar progress_timeout("MC_TCP_PROGRESS_TIMEOUT_SEC", "30"); + ScopedLaneHooks lane_hooks; + ScopedSessionProgressHooks session_hooks( + kSessionReadBodySuccess, kSessionCommitTimeoutBeforeProgress); + + EngineHandle h; + h.init(P2PHANDSHAKE, "127.0.0.2:18052", 16ull << 20); + ASSERT_TRUE(h.ok); + + char* source = static_cast(h.pool); + char* first_destination = source + kRegionAlign; + for (size_t i = 0; i < kSmallLength; ++i) + source[i] = static_cast(i * 17 + 3); + + const auto first_batch = h.engine->allocateBatchID(1); + ASSERT_TRUE(h.engine + ->submitTransfer( + first_batch, + {makeReadRequest(h, first_destination, kSmallLength)}) + .ok()); + TransferStatusEnum first_status = TransferStatusEnum::WAITING; + ASSERT_TRUE(waitForTaskTerminal(h.engine.get(), first_batch, + std::chrono::seconds(5), &first_status)); + EXPECT_EQ(first_status, TransferStatusEnum::FAILED); + expectEverySliceCompletedExactlyOnce(first_batch); + EXPECT_EQ(session_timeout_committed_count.load(std::memory_order_acquire), + 1); + EXPECT_EQ(session_dirty_terminal_count.load(std::memory_order_acquire), 1); + ASSERT_TRUE(waitForPredicate( + [] { return lane_terminal_count.load(std::memory_order_acquire) >= 1; }, + std::chrono::seconds(5))); + EXPECT_EQ(lane_terminal_socket_count.load(std::memory_order_acquire), 0); + ASSERT_TRUE(h.engine->freeBatchID(first_batch).ok()); + EXPECT_EQ(session_terminal_count.load(std::memory_order_acquire), 1); + EXPECT_EQ(session_dirty_terminal_count.load(std::memory_order_acquire), 1); +} + +TEST(TcpWriteVisibilityTest, + ProgressTimeoutWinsAgainstQueuedFinalWriteAndAckCompletion) { + ScopedEnvVar lanes("MC_TCP_LANES_PER_PEER", "1"); + ScopedEnvVar progress_timeout("MC_TCP_PROGRESS_TIMEOUT_SEC", "30"); + ScopedLaneHooks lane_hooks; + ScopedSessionProgressHooks session_hooks( + kSessionWriteAckSuccess, kSessionCommitTimeoutBeforeProgress); + ReusingWriteServer fake_peer; + ASSERT_TRUE(fake_peer.ok()); + + EngineHandle h; + h.init(P2PHANDSHAKE, "127.0.0.2:18053", 8ull << 20); + ASSERT_TRUE(h.ok); + pointTcpSegmentAt(h, fake_peer.port()); + memset(h.pool, 0x5A, kBigLength); + + const auto first_batch = h.engine->allocateBatchID(1); + ASSERT_TRUE( + h.engine->submitTransfer(first_batch, {makeWriteRequest(h, kBigLength)}) + .ok()); + TransferStatusEnum first_status = TransferStatusEnum::WAITING; + ASSERT_TRUE(waitForTaskTerminal(h.engine.get(), first_batch, + std::chrono::seconds(5), &first_status)); + EXPECT_EQ(first_status, TransferStatusEnum::FAILED); + expectEverySliceCompletedExactlyOnce(first_batch); + EXPECT_EQ(session_timeout_committed_count.load(std::memory_order_acquire), + 1); + EXPECT_EQ(session_dirty_terminal_count.load(std::memory_order_acquire), 1); + ASSERT_TRUE(waitForPredicate( + [] { return lane_terminal_count.load(std::memory_order_acquire) >= 1; }, + std::chrono::seconds(5))); + EXPECT_EQ(lane_terminal_socket_count.load(std::memory_order_acquire), 0); + ASSERT_TRUE(h.engine->freeBatchID(first_batch).ok()); + EXPECT_EQ(fake_peer.acceptedCount(), 1); + EXPECT_EQ(session_terminal_count.load(std::memory_order_acquire), 1); + EXPECT_EQ(session_dirty_terminal_count.load(std::memory_order_acquire), 1); +} + +TEST(TcpWriteVisibilityTest, ProgressBeforeDeadlineMakesOldTimerStale) { + ScopedEnvVar lanes("MC_TCP_LANES_PER_PEER", "1"); + ScopedEnvVar progress_timeout("MC_TCP_PROGRESS_TIMEOUT_SEC", "30"); + ScopedSessionProgressHooks session_hooks( + kSessionReadBodySuccess, kSessionReplayPreviousTimeoutAfterProgress); + + EngineHandle h; + h.init(P2PHANDSHAKE, "127.0.0.2:18054", 8ull << 20); + ASSERT_TRUE(h.ok); + + char* source = static_cast(h.pool); + char* destination = source + kRegionAlign; + for (size_t i = 0; i < kSmallLength; ++i) + source[i] = static_cast(i * 29 + 11); + + EXPECT_EQ( + runOne(h.engine.get(), makeReadRequest(h, destination, kSmallLength)), + TransferStatusEnum::COMPLETED); + EXPECT_EQ(memcmp(destination, source, kSmallLength), 0); + EXPECT_EQ(session_timeout_committed_count.load(std::memory_order_acquire), + 0); + EXPECT_GE(session_timeout_stale_count.load(std::memory_order_acquire), 1); + EXPECT_EQ(session_terminal_count.load(std::memory_order_acquire), 1); + EXPECT_EQ(session_dirty_terminal_count.load(std::memory_order_acquire), 0); +} + +TEST(TcpWriteVisibilityTest, + ProgressTimeoutWaitsForOutstandingWriteBeforeTerminal) { + constexpr size_t kRequestSize = 64ull << 20; + ScopedEnvVar lanes("MC_TCP_LANES_PER_PEER", "1"); + ScopedEnvVar progress_timeout("MC_TCP_PROGRESS_TIMEOUT_SEC", "1"); + ScopedSessionProgressHooks session_hooks( + /*target_event=*/0, kSessionProgressNoAction, + /*block_timeout_commit=*/true); + PayloadStallingWriteServer fake_peer(/*payload_prefix_to_read=*/0); + ASSERT_TRUE(fake_peer.ok()); + + EngineHandle h; + h.init(P2PHANDSHAKE, "127.0.0.2:18055", kRequestSize); + ASSERT_TRUE(h.ok); + pointTcpSegmentAt(h, fake_peer.port()); + + const auto batch_id = h.engine->allocateBatchID(1); + ASSERT_TRUE( + h.engine->submitTransfer(batch_id, {makeWriteRequest(h, kRequestSize)}) + .ok()); + ASSERT_TRUE(fake_peer.waitForHeader(std::chrono::seconds(5))); + ASSERT_TRUE(waitForPredicate( + [] { + return session_timeout_commit_entered.load( + std::memory_order_acquire); + }, + std::chrono::seconds(5))); + EXPECT_TRUE(session_timeout_commit_had_write_in_flight.load( + std::memory_order_acquire)); + + TransferStatus status; + status.s = TransferStatusEnum::WAITING; + ASSERT_TRUE(h.engine->getTransferStatus(batch_id, 0, status).ok()); + EXPECT_EQ(status.s, TransferStatusEnum::WAITING); + EXPECT_EQ(session_terminal_count.load(std::memory_order_acquire), 0); + + session_hooks.releaseTimeoutCommit(); + TransferStatusEnum final_status = TransferStatusEnum::WAITING; + ASSERT_TRUE(waitForTaskTerminal(h.engine.get(), batch_id, + std::chrono::seconds(5), &final_status)); + EXPECT_EQ(final_status, TransferStatusEnum::FAILED); + expectEverySliceCompletedExactlyOnce(batch_id); + EXPECT_EQ(session_dirty_terminal_count.load(std::memory_order_acquire), 1); + ASSERT_TRUE(h.engine->freeBatchID(batch_id).ok()); + fake_peer.release(); +} + TEST(TcpWriteVisibilityTest, WritePayloadNoProgressFailsWithinBoundedDeadline) { constexpr size_t kRequestSize = 64ull << 20; ScopedEnvVar lanes("MC_TCP_LANES_PER_PEER", "1");