From d1fb3348a3404875e6e58a8bf0735903efa05b14 Mon Sep 17 00:00:00 2001 From: Husam Date: Wed, 30 Sep 2026 14:51:59 +0000 Subject: [PATCH 01/15] fix(replication): respond to kGetObj over existing client socket --- .../replication/transfer_service/task.cpp | 2 +- .../replication/transfer_service/task.h | 11 ++- .../transfer_service/transfer_service.cpp | 42 ++++------- .../transfer_service_p2p_test.cpp | 72 +++++++++++++++++++ 4 files changed, 92 insertions(+), 35 deletions(-) diff --git a/src/ml_flashpoint/replication/transfer_service/task.cpp b/src/ml_flashpoint/replication/transfer_service/task.cpp index 1e34f51..437aff0 100644 --- a/src/ml_flashpoint/replication/transfer_service/task.cpp +++ b/src/ml_flashpoint/replication/transfer_service/task.cpp @@ -32,7 +32,7 @@ void GetTask::Execute(TransferService* service) { } void RespondToGetTask::Execute(TransferService* service) { - LOG(INFO) << "Executing RespondToGetTask to " << dest_addr_; + LOG(INFO) << "Executing RespondToGetTask on fd=" << client_fd_; service->ExecuteRespondToGetTask(this); } diff --git a/src/ml_flashpoint/replication/transfer_service/task.h b/src/ml_flashpoint/replication/transfer_service/task.h index 97e9786..daa7f56 100644 --- a/src/ml_flashpoint/replication/transfer_service/task.h +++ b/src/ml_flashpoint/replication/transfer_service/task.h @@ -167,13 +167,12 @@ class RespondToGetTask : public Task { public: RespondToGetTask( std::string task_id, std::string source_obj_id, std::string dest_obj_id, - std::string source_addr, std::string dest_addr, + int client_fd, std::shared_ptr metric_container) : Task(std::move(task_id)), source_obj_id_(std::move(source_obj_id)), dest_obj_id_(std::move(dest_obj_id)), - source_addr_(std::move(source_addr)), - dest_addr_(std::move(dest_addr)), + client_fd_(client_fd), metric_container_(std::move(metric_container)) {} void Execute(TransferService* service) override; @@ -183,14 +182,12 @@ class RespondToGetTask : public Task { const std::string& GetSourceObjId() const { return source_obj_id_; } const std::string& GetDestObjId() const { return dest_obj_id_; } - const std::string& GetSourceAddr() const { return source_addr_; } - const std::string& GetDestAddr() const { return dest_addr_; } + int GetClientFd() const { return client_fd_; } private: std::string source_obj_id_; std::string dest_obj_id_; - std::string source_addr_; - std::string dest_addr_; + int client_fd_; std::shared_ptr metric_container_; }; diff --git a/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp b/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp index 9d0d7b6..3a18a69 100644 --- a/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp +++ b/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp @@ -812,7 +812,7 @@ void TransferService::ExecuteGetTask(GetTask* task) { } metric_container.header_sent_time = absl::Now(); - // Wait for the immediate response (ACK or ERROR) + // Wait for the response (RespondToGetObj or ERROR) ObjInfoHeader resp_header; if (!RecvHeader(conn.fd(), resp_header).ok()) { ReportResult(task->GetTaskId(), false, @@ -821,15 +821,15 @@ void TransferService::ExecuteGetTask(GetTask* task) { } switch (resp_header.type) { - case MessageType::kAck: - LOG(INFO) << "Received ACK for GET request. Waiting for data transfer."; + case MessageType::kRespondToGetObj: + HandleDataReceive(conn.fd(), resp_header, true); break; case MessageType::kError: ReportResult(task->GetTaskId(), false, "Received error message"); break; + case MessageType::kAck: case MessageType::kPutObj: case MessageType::kGetObj: - case MessageType::kRespondToGetObj: ReportResult(task->GetTaskId(), false, "Received unexpected response for GET request"); break; @@ -842,16 +842,14 @@ void TransferService::ExecuteRespondToGetTask(RespondToGetTask* task) { LOG(INFO) << "Executing RespondToGetTask for task_id=" << task->GetTaskId() << ", source_obj_id=" << task->GetSourceObjId() << ", dest_obj_id=" << task->GetDestObjId() - << ", source_addr=" << task->GetSourceAddr() - << ", dest_addr=" << task->GetDestAddr(); + << ", client_fd=" << task->GetClientFd(); - auto conn_opt = GetConnectionFromPool(task->GetDestAddr()); - if (!conn_opt) { - LOG(ERROR) << "Failed to get connection!"; - ReportResult(task->GetTaskId(), false, "Failed to get connection"); + int sockfd = task->GetClientFd(); + if (sockfd < 0) { + LOG(ERROR) << "Invalid client_fd!"; + ReportResult(task->GetTaskId(), false, "Invalid client_fd"); return; } - ScopedConnection conn = std::move(conn_opt.value()); metric_container.connection_acquired_time = absl::Now(); // Open file as buffer object @@ -862,6 +860,8 @@ void TransferService::ExecuteRespondToGetTask(RespondToGetTask* task) { if (buffer_data_ptr == nullptr) { LOG(ERROR) << "RespondToGetTask failed: Could not open buffer object for '" << task->GetSourceObjId() << "'"; + SendErrorResponse(sockfd, task->GetTaskId().c_str(), + task->GetSourceObjId().c_str()); ReportResult(task->GetTaskId(), false, "Failed to create buffer object"); return; } @@ -883,7 +883,6 @@ void TransferService::ExecuteRespondToGetTask(RespondToGetTask* task) { header.type = MessageType::kRespondToGetObj; header.obj_size = size; - int sockfd = conn.fd(); if (!SendAll(sockfd, &header, kHeaderSize).ok()) { LOG(ERROR) << "Failed to send kRespondToGetObj header"; @@ -987,16 +986,6 @@ void TransferService::HandleGetObjRequest(int client_fd, return; } - // Send ACK to confirm request is accepted, before async data transfer - ObjInfoHeader ack_header; - ack_header.type = MessageType::kAck; - snprintf(ack_header.task_id, sizeof(ack_header.task_id), "%s", - header.task_id); - if (!SendAll(client_fd, &ack_header, kHeaderSize).ok()) { - LOG(ERROR) << "Failed to send ACK for GET request " << header.source_obj_id; - return; // Don't proceed if we can't even ACK - } - auto metric_container = std::make_shared(); metric_container->task_id = header.task_id; metric_container->task_type = TaskMetricContainer::TaskType::kRespondToGet; @@ -1009,18 +998,17 @@ void TransferService::HandleGetObjRequest(int client_fd, } metric_container->submit_time = absl::Now(); - auto task = std::make_unique( - header.task_id, header.source_obj_id, header.dest_obj_id, - header.source_address, header.dest_address, metric_container); + RespondToGetTask task(header.task_id, header.source_obj_id, + header.dest_obj_id, client_fd, metric_container); { std::lock_guard lock(pending_tasks_mutex_); // We don't have a promise for RespondToGetTask as it's triggered // remotely, but we want to track it. - pending_tasks_[task->GetTaskId()] = {nullptr, metric_container}; + pending_tasks_[task.GetTaskId()] = {nullptr, metric_container}; } - task_queue_.enqueue(std::move(task)); + task.Execute(this); } void TransferService::UpdateTaskMetrics( diff --git a/tests/replication/transfer_service/transfer_service_p2p_test.cpp b/tests/replication/transfer_service/transfer_service_p2p_test.cpp index 33f5cd9..76e4f93 100644 --- a/tests/replication/transfer_service/transfer_service_p2p_test.cpp +++ b/tests/replication/transfer_service/transfer_service_p2p_test.cpp @@ -30,6 +30,7 @@ #include "absl/log/log_sink.h" #include "absl/log/log_sink_registry.h" #include "gtest/gtest.h" +#include "net_util.h" #include "transfer_service.h" namespace ml_flashpoint::replication::transfer_service { @@ -916,5 +917,76 @@ TEST(TransferServiceP2PTest, TimestampsAreRecorded) { EXPECT_EQ(timing_logs_count, 3) << "timing_logs_count: " << timing_logs_count; } +// Verifies that handling a kGetObj request responds over the existing client +// socket and never initiates an outbound connection to header.dest_address. +TEST(TransferServiceP2PTest, GetObjDoesNotConnectToDestAddress) { + // Given + TransferService service1("127.0.0.1"); + int port1 = service1.Initialize(); + ASSERT_GT(port1, 0); + + std::string expected_data = "ssrf_test_payload"; + std::string source_obj_id = "ssrf_source_obj"; + std::ofstream(source_obj_id) << expected_data; + + int target_listener_fd = socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK, 0); + ASSERT_GE(target_listener_fd, 0); + sockaddr_in target_addr{}; + target_addr.sin_family = AF_INET; + inet_pton(AF_INET, "127.0.0.1", &target_addr.sin_addr); + target_addr.sin_port = 0; + ASSERT_EQ(bind(target_listener_fd, reinterpret_cast(&target_addr), + sizeof(target_addr)), + 0); + socklen_t target_len = sizeof(target_addr); + ASSERT_EQ(getsockname(target_listener_fd, + reinterpret_cast(&target_addr), &target_len), + 0); + int target_port = ntohs(target_addr.sin_port); + ASSERT_EQ(listen(target_listener_fd, 1), 0); + + int client_fd = socket(AF_INET, SOCK_STREAM, 0); + ASSERT_GE(client_fd, 0); + sockaddr_in serv_addr{}; + serv_addr.sin_family = AF_INET; + serv_addr.sin_port = htons(port1); + inet_pton(AF_INET, "127.0.0.1", &serv_addr.sin_addr); + ASSERT_EQ(connect(client_fd, reinterpret_cast(&serv_addr), + sizeof(serv_addr)), + 0); + + // When + ObjInfoHeader req_header; + req_header.type = MessageType::kGetObj; + snprintf(req_header.task_id, sizeof(req_header.task_id), "test_task_id"); + snprintf(req_header.source_obj_id, sizeof(req_header.source_obj_id), "%s", + source_obj_id.c_str()); + snprintf(req_header.dest_obj_id, sizeof(req_header.dest_obj_id), + "ssrf_dest_obj"); + snprintf(req_header.dest_address, sizeof(req_header.dest_address), + "127.0.0.1:%d", target_port); + ASSERT_TRUE(SendAll(client_fd, &req_header, kHeaderSize).ok()); + + ObjInfoHeader resp_header; + ASSERT_TRUE(RecvHeader(client_fd, resp_header).ok()); + std::string actual_data(resp_header.obj_size, '\0'); + ASSERT_TRUE( + RecvAll(client_fd, actual_data.data(), resp_header.obj_size).ok()); + + ObjInfoHeader ack_header; + ack_header.type = MessageType::kAck; + ASSERT_TRUE(SendAll(client_fd, &ack_header, kHeaderSize).ok()); + + // Then + EXPECT_EQ(resp_header.type, MessageType::kRespondToGetObj); + EXPECT_EQ(actual_data, expected_data); + EXPECT_EQ(accept(target_listener_fd, nullptr, nullptr), -1); + + close(client_fd); + close(target_listener_fd); + std::remove(source_obj_id.c_str()); + service1.Shutdown(); +} + } // namespace } // namespace ml_flashpoint::replication::transfer_service From 711b1e111da7f02c654253873e2f843ccfb79152 Mon Sep 17 00:00:00 2001 From: Husam Date: Wed, 30 Sep 2026 17:46:44 +0000 Subject: [PATCH 02/15] test(replication): add single-conn reuse and bidirectional Get/Put stress tests --- .../transfer_service/transfer_service.cpp | 16 ++- .../transfer_service/transfer_service.h | 1 + tests/replication/test_transer_service.py | 58 +++++++++ .../transfer_service_p2p_test.cpp | 121 +++++++++++++++++- 4 files changed, 189 insertions(+), 7 deletions(-) diff --git a/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp b/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp index 3a18a69..adf0601 100644 --- a/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp +++ b/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp @@ -123,8 +123,9 @@ int TransferService::Initialize(int listen_port, int threads, absl::StrFormat("%s:%d", local_ip_address.value(), listen_port_); LOG(INFO) << "Listening on " << local_address_; - // Create the thread pool. + // Create the thread pools for outbound tasks and inbound epoll events. thread_pool_ = std::make_unique(threads_); + epoll_thread_pool_ = std::make_unique(threads_); // Create epoll instance. epoll_fd_ = epoll_create1(0); @@ -222,10 +223,13 @@ void TransferService::Shutdown() { connection_pools_.clear(); } - // 7. Stop the thread pool. + // 7. Stop the thread pools. if (thread_pool_) { thread_pool_->stop(); } + if (epoll_thread_pool_) { + epoll_thread_pool_->stop(); + } // 8. Clean up epoll fd. if (epoll_fd_ != -1) { @@ -428,7 +432,7 @@ void TransferService::ProcessEpollEventsLoop() { running_.store(false); break; } - thread_pool_->enqueue([this]() { this->HandleNewConnection(); }); + epoll_thread_pool_->enqueue([this]() { this->HandleNewConnection(); }); } else { if ((current_events & EPOLLERR) || (current_events & EPOLLHUP)) { if (current_events & EPOLLERR) { @@ -439,7 +443,7 @@ void TransferService::ProcessEpollEventsLoop() { LOG(INFO) << "Epoll HUP event (EPOLLHUP) on client fd=" << current_fd << ". Closing."; } - thread_pool_->enqueue( + epoll_thread_pool_->enqueue( [this, current_fd]() { this->RemoveClient(current_fd); }); continue; } @@ -448,8 +452,8 @@ void TransferService::ProcessEpollEventsLoop() { LOG(INFO) << "Epoll loop: Dispatching incoming data " "processing for fd=" << current_fd; - thread_pool_->enqueue(&TransferService::ProcessIncomingData, this, - current_fd); + epoll_thread_pool_->enqueue(&TransferService::ProcessIncomingData, + this, current_fd); } else { LOG(WARNING) << "Epoll loop: Unhandled event flags=" << current_events << " on client fd=" << current_fd; diff --git a/src/ml_flashpoint/replication/transfer_service/transfer_service.h b/src/ml_flashpoint/replication/transfer_service/transfer_service.h index ddcdd10..20fdf7d 100644 --- a/src/ml_flashpoint/replication/transfer_service/transfer_service.h +++ b/src/ml_flashpoint/replication/transfer_service/transfer_service.h @@ -162,6 +162,7 @@ class TransferService final { int global_rank_ = -1; std::unique_ptr thread_pool_; + std::unique_ptr epoll_thread_pool_; std::thread epoll_thread_; std::thread task_queue_thread_; TaskQueue task_queue_; diff --git a/tests/replication/test_transer_service.py b/tests/replication/test_transer_service.py index 9cb3890..b776e62 100644 --- a/tests/replication/test_transer_service.py +++ b/tests/replication/test_transer_service.py @@ -351,3 +351,61 @@ def test_async_get_non_existent_object( # Then with pytest.raises(RuntimeError, match="Received error message"): get_future.result(timeout=10) + + +def test_concurrent_and_repeated_async_get( + transfer_services: tuple[ + transfer_service_ext.TransferService, + transfer_service_ext.TransferService, + str, + str, + ], + tmp_path: pathlib.Path, +) -> None: + """Verifies that concurrent and repeated async_get calls succeed and recover cleanly after an error.""" + # Given + _, service2, addr1, _ = transfer_services + num_objects = 10 + source_paths = [] + dest_paths = [] + expected_payloads = [] + + for i in range(num_objects): + src_path = str(tmp_path / f"concurrent_src_{i}.object") + dst_path = str(tmp_path / f"concurrent_dst_{i}.object") + expected_data = np.arange(i * 100, (i + 1) * 100, dtype=np.int64).tobytes() + with open(src_path, "wb") as f: + f.write(expected_data) + source_paths.append(src_path) + dest_paths.append(dst_path) + expected_payloads.append(expected_data) + + # When + first_round_futures = [service2.async_get(source_paths[i], addr1, dest_paths[i]) for i in range(num_objects)] + missing_future = service2.async_get( + str(tmp_path / "missing_src.object"), + addr1, + str(tmp_path / "missing_dst.object"), + ) + with pytest.raises(RuntimeError, match="Received error message"): + missing_future.result(timeout=10) + + second_round_dest_paths = [str(tmp_path / f"concurrent_dst_round2_{i}.object") for i in range(num_objects)] + second_round_futures = [ + service2.async_get(source_paths[i], addr1, second_round_dest_paths[i]) for i in range(num_objects) + ] + + # Then + for i in range(num_objects): + actual_result_1 = first_round_futures[i].result(timeout=10) + actual_result_2 = second_round_futures[i].result(timeout=10) + assert actual_result_1.success is True + assert actual_result_2.success is True + + with open(dest_paths[i], "rb") as f: + actual_content_1 = f.read() + with open(second_round_dest_paths[i], "rb") as f: + actual_content_2 = f.read() + + assert actual_content_1 == expected_payloads[i] + assert actual_content_2 == expected_payloads[i] diff --git a/tests/replication/transfer_service/transfer_service_p2p_test.cpp b/tests/replication/transfer_service/transfer_service_p2p_test.cpp index 76e4f93..19e9698 100644 --- a/tests/replication/transfer_service/transfer_service_p2p_test.cpp +++ b/tests/replication/transfer_service/transfer_service_p2p_test.cpp @@ -185,7 +185,8 @@ TEST(TransferServiceP2PTest, ShutdownInterruptsTransfer) { LOG(INFO) << "Transfer failed as expected after shutdown."; } catch (const std::runtime_error& e) { EXPECT_THAT(e.what(), testing::HasSubstr("Service is shutting down")); - LOG(INFO) << "Transfer threw exception as expected after shutdown: " << e.what(); + LOG(INFO) << "Transfer threw exception as expected after shutdown: " + << e.what(); } // Cleanup file if it was partially created @@ -988,5 +989,123 @@ TEST(TransferServiceP2PTest, GetObjDoesNotConnectToDestAddress) { service1.Shutdown(); } +// Verifies that a single pooled connection (conn_pool_per_peer=1) is cleanly +// reused across sequential AsyncGet calls, including after a kError response. +TEST(TransferServiceP2PTest, + GetReusesSinglePooledConnectionAfterSuccessAndError) { + // Given + TransferService service1("127.0.0.1"); + int port1 = service1.Initialize(0, /*threads=*/4, /*conn_pool_per_peer=*/1); + ASSERT_GT(port1, 0); + + TransferService service2("127.0.0.1"); + int port2 = service2.Initialize(0, /*threads=*/4, /*conn_pool_per_peer=*/1); + ASSERT_GT(port2, 0); + + std::string peer1_addr = "127.0.0.1:" + std::to_string(port1); + std::string expected_data_1 = "first_payload_over_single_conn"; + std::string expected_data_2 = "second_payload_after_error_over_single_conn"; + std::string source_obj_1 = "single_conn_source_1"; + std::string source_obj_2 = "single_conn_source_2"; + std::string dest_obj_1 = "single_conn_dest_1"; + std::string dest_obj_2 = "single_conn_dest_2"; + + std::ofstream(source_obj_1) << expected_data_1; + std::ofstream(source_obj_2) << expected_data_2; + + // When + TransferResult actual_result_1 = + service2.AsyncGet(source_obj_1, peer1_addr, dest_obj_1).get(); + auto missing_future = + service2.AsyncGet("non_existent_single_conn_obj", peer1_addr, + "non_existent_single_conn_dest"); + EXPECT_THROW(missing_future.get(), std::runtime_error); + TransferResult actual_result_2 = + service2.AsyncGet(source_obj_2, peer1_addr, dest_obj_2).get(); + + // Then + EXPECT_TRUE(actual_result_1.success); + EXPECT_TRUE(actual_result_2.success); + VerifyFileContentAndRemove(dest_obj_1, expected_data_1); + VerifyFileContentAndRemove(dest_obj_2, expected_data_2); + + std::remove(source_obj_1.c_str()); + std::remove(source_obj_2.c_str()); + service1.Shutdown(); + service2.Shutdown(); +} + +// Verifies that simultaneous bidirectional AsyncGet and AsyncPut calls between +// two peers complete without deadlock or stream corruption. +TEST(TransferServiceP2PTest, BidirectionalConcurrentGetAndPut) { + // Given + TransferService service1("127.0.0.1"); + int port1 = service1.Initialize(); + ASSERT_GT(port1, 0); + + TransferService service2("127.0.0.1"); + int port2 = service2.Initialize(); + ASSERT_GT(port2, 0); + + std::string addr1 = "127.0.0.1:" + std::to_string(port1); + std::string addr2 = "127.0.0.1:" + std::to_string(port2); + + const int kNumOps = 8; + std::vector s1_get_sources(kNumOps); + std::vector s2_get_sources(kNumOps); + std::vector s1_get_dests(kNumOps); + std::vector s2_get_dests(kNumOps); + std::vector s1_put_dests(kNumOps); + std::vector s2_put_dests(kNumOps); + std::vector expected_payloads(kNumOps); + + for (int i = 0; i < kNumOps; ++i) { + expected_payloads[i] = + "Bidirectional payload " + std::to_string(i) + std::string(1024, 'X'); + s1_get_sources[i] = "bidir_s1_src_" + std::to_string(i); + s2_get_sources[i] = "bidir_s2_src_" + std::to_string(i); + s1_get_dests[i] = "bidir_s1_get_dst_" + std::to_string(i); + s2_get_dests[i] = "bidir_s2_get_dst_" + std::to_string(i); + s1_put_dests[i] = "bidir_s1_put_dst_" + std::to_string(i); + s2_put_dests[i] = "bidir_s2_put_dst_" + std::to_string(i); + + std::ofstream(s1_get_sources[i]) << expected_payloads[i]; + std::ofstream(s2_get_sources[i]) << expected_payloads[i]; + } + + // When + std::vector> futures; + for (int i = 0; i < kNumOps; ++i) { + futures.push_back( + service1.AsyncGet(s2_get_sources[i], addr2, s1_get_dests[i])); + futures.push_back( + service2.AsyncGet(s1_get_sources[i], addr1, s2_get_dests[i])); + futures.push_back( + service1.AsyncPut(const_cast(expected_payloads[i].data()), + expected_payloads[i].size(), addr2, s2_put_dests[i])); + futures.push_back( + service2.AsyncPut(const_cast(expected_payloads[i].data()), + expected_payloads[i].size(), addr1, s1_put_dests[i])); + } + + // Then + for (auto& fut : futures) { + TransferResult actual_result = fut.get(); + EXPECT_TRUE(actual_result.success); + } + + for (int i = 0; i < kNumOps; ++i) { + VerifyFileContentAndRemove(s1_get_dests[i], expected_payloads[i]); + VerifyFileContentAndRemove(s2_get_dests[i], expected_payloads[i]); + VerifyFileContentAndRemove(s1_put_dests[i], expected_payloads[i]); + VerifyFileContentAndRemove(s2_put_dests[i], expected_payloads[i]); + std::remove(s1_get_sources[i].c_str()); + std::remove(s2_get_sources[i].c_str()); + } + + service1.Shutdown(); + service2.Shutdown(); +} + } // namespace } // namespace ml_flashpoint::replication::transfer_service From a27785c4b01840ef1f3e59bcd3b5ffa65d628839 Mon Sep 17 00:00:00 2001 From: Husam Date: Fri, 2 Oct 2026 18:17:03 +0000 Subject: [PATCH 03/15] fix(replication): harden same-socket kGetObj response path Follow-up hardening for the b/565101350 fix (respond to kGetObj over the request's own connection instead of connecting back to header.dest_address). Wire-protocol change (intentional): kGetObj is no longer acknowledged with kAck; the responder streams kRespondToGetObj + payload on the request socket and the requester replies kAck there. Old and new builds do not interoperate for Gets. This is acceptable because the transport is scoped to a single training job in which all ranks run the same build; the on-disk checkpoint format is unchanged. Bump the minor version at release. - Shutdown(): skip the promise-less RespondToGetTask entries in pending_tasks_ (was a null dereference whenever a remote Get was in flight), and shutdown(2) every accepted client socket before joining the thread pools so a responder blocked on a stalled peer cannot hang Shutdown(). Accepted sockets are tracked in client_fds_ and closed once. - ExecuteRespondToGetTask(): BufferObject throws for empty files, directories and unreadable paths; the old "nullptr -> kError" branch was dead and the escaping exception orphaned the connection and hung the requester forever. Catch it and send kError on the request socket. - ProcessIncomingData(): close the connection on an unsolicited kRespondToGetObj (it was an arbitrary file write sink), on an unknown message type, and on any handler exception, instead of leaving the fd registered but never re-armed. - ExecuteGetTask(): report a failure (and shut the pooled connection down so it is not reused mid-payload) if the destination BufferObject cannot be created. - Document that `threads` sizes each of the two pools, why two pools are required, and that RespondToGetTask must run synchronously on a non-owning fd. Tests: strengthen GetObjDoesNotConnectToDestAddress (poll the decoy listener instead of an instantaneous accept()), add a 12-case SpoofedDestAddressTest, unterminated address fields, unopenable object -> kError on the same socket, unsolicited kRespondToGetObj / unknown type -> connection closed, Shutdown() while a responder waits for ACK and while a peer is not reading (16 MiB, 4 KiB SO_RCVBUF), and the ported GetSucceedsWhenRequesterAdvertisesUnreachableIp. --- .../replication/transfer_service/task.h | 8 + .../transfer_service/transfer_service.cpp | 142 ++++-- .../transfer_service/transfer_service.h | 20 +- .../transfer_service_p2p_test.cpp | 452 +++++++++++++++++- 4 files changed, 585 insertions(+), 37 deletions(-) diff --git a/src/ml_flashpoint/replication/transfer_service/task.h b/src/ml_flashpoint/replication/transfer_service/task.h index daa7f56..b6796b2 100644 --- a/src/ml_flashpoint/replication/transfer_service/task.h +++ b/src/ml_flashpoint/replication/transfer_service/task.h @@ -163,6 +163,14 @@ class GetTask : public Task { std::shared_ptr metric_container_; }; +// Serves a kGetObj request by streaming the object back over the connection +// the request arrived on. +// +// `client_fd` is NOT owned by the task: it belongs to the epoll loop. The task +// must therefore be executed synchronously from within ProcessIncomingData for +// that fd (as HandleGetObjRequest does), i.e. while the fd is still disarmed +// under EPOLLONESHOT and before any RemoveClient() can close and recycle it. +// Never enqueue this task for deferred execution. class RespondToGetTask : public Task { public: RespondToGetTask( diff --git a/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp b/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp index adf0601..2ce5610 100644 --- a/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp +++ b/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp @@ -181,6 +181,8 @@ void TransferService::Shutdown() { { std::lock_guard lock(pending_tasks_mutex_); for (auto const& [task_id, context] : pending_tasks_) { + // Remotely triggered tasks (RespondToGetTask) have no promise. + if (!context.promise) continue; try { context.promise->set_exception(std::make_exception_ptr( std::runtime_error("Service is shutting down"))); @@ -212,7 +214,19 @@ void TransferService::Shutdown() { epoll_thread_.join(); } - // 6. Clean up connection pools. + // 6. Unblock any worker that is blocked on an accepted client socket (e.g. a + // RespondToGetTask whose peer stopped reading, or a kPutObj receive whose + // sender stalled). shutdown() makes their pending send/recv fail so the + // thread pools below can be joined. The fds are closed in step 9, once no + // worker can be using them anymore. + { + std::lock_guard lock(client_fds_mutex_); + for (int fd : client_fds_) { + shutdown(fd, SHUT_RDWR); + } + } + + // 7. Clean up connection pools. { std::unique_lock write_lock(connection_pools_mutex_); for (auto const& [peer_addr, pool] : connection_pools_) { @@ -223,7 +237,7 @@ void TransferService::Shutdown() { connection_pools_.clear(); } - // 7. Stop the thread pools. + // 8. Stop the thread pools. if (thread_pool_) { thread_pool_->stop(); } @@ -231,7 +245,16 @@ void TransferService::Shutdown() { epoll_thread_pool_->stop(); } - // 8. Clean up epoll fd. + // 9. Close any accepted client sockets that are still open. + { + std::lock_guard lock(client_fds_mutex_); + for (int fd : client_fds_) { + close(fd); + } + client_fds_.clear(); + } + + // 10. Clean up epoll fd. if (epoll_fd_ != -1) { if (close(epoll_fd_) == -1) { PLOG(WARNING) << "Failed to close epoll_fd"; @@ -551,6 +574,13 @@ void TransferService::HandleNewConnection() { continue; } + // Register the client before arming it in epoll, so that a RemoveClient() + // triggered by an early event cannot race ahead of the registration. + { + std::lock_guard lock(client_fds_mutex_); + client_fds_.insert(conn_fd); + } + // Add the new client to epoll. struct epoll_event event; event.events = @@ -558,7 +588,7 @@ void TransferService::HandleNewConnection() { event.data.fd = conn_fd; if (epoll_ctl(epoll_fd_, EPOLL_CTL_ADD, conn_fd, &event) == -1) { PLOG(ERROR) << "Failed to add client fd to epoll"; - close(conn_fd); + RemoveClient(conn_fd); continue; } @@ -571,6 +601,14 @@ void TransferService::HandleNewConnection() { void TransferService::RemoveClient(int client_fd) { if (client_fd < 0) return; + { + // Only the caller that removes the fd from the registry closes it. This + // makes RemoveClient idempotent and prevents a double close (and closing + // an unrelated, reused fd) when Shutdown() and a worker both tear down + // the same connection. + std::lock_guard lock(client_fds_mutex_); + if (client_fds_.erase(client_fd) == 0) return; + } if (epoll_fd_ >= 0) { epoll_ctl(epoll_fd_, EPOLL_CTL_DEL, client_fd, nullptr); } @@ -594,26 +632,47 @@ void TransferService::ProcessIncomingData(int client_fd) { << ", dest_obj_id=" << header.dest_obj_id; bool keep_connection = true; - switch (header.type) { - case MessageType::kPutObj: - HandleDataReceive(client_fd, header, false); - break; - case MessageType::kRespondToGetObj: - HandleDataReceive(client_fd, header, true); - break; - case MessageType::kGetObj: - HandleGetObjRequest(client_fd, header); - break; - case MessageType::kAck: - LOG(ERROR) << "Worker fd=" << client_fd - << ": Unexpected kAck message received for task " - << header.task_id; - break; - case MessageType::kError: - LOG(ERROR) << "Worker fd=" << client_fd - << ": Error message received for task " << header.task_id; - ReportResult(header.task_id, false, "Received error message"); - break; + try { + switch (header.type) { + case MessageType::kPutObj: + HandleDataReceive(client_fd, header, false); + break; + case MessageType::kGetObj: + HandleGetObjRequest(client_fd, header); + break; + case MessageType::kRespondToGetObj: + // Responses to kGetObj are read by ExecuteGetTask on the connection + // that sent the request. An unsolicited one arriving on the listener + // is a protocol violation; do not write its payload to disk. + LOG(ERROR) << "Worker fd=" << client_fd + << ": Unsolicited kRespondToGetObj received for task " + << header.task_id << ". Closing connection."; + keep_connection = false; + break; + case MessageType::kAck: + LOG(ERROR) << "Worker fd=" << client_fd + << ": Unexpected kAck message received for task " + << header.task_id; + break; + case MessageType::kError: + LOG(ERROR) << "Worker fd=" << client_fd + << ": Error message received for task " << header.task_id; + ReportResult(header.task_id, false, "Received error message"); + break; + default: + // The stream is desynced if we cannot interpret the header. + LOG(ERROR) << "Worker fd=" << client_fd << ": Unknown message type " + << static_cast(header.type) << ". Closing connection."; + keep_connection = false; + break; + } + } catch (const std::exception& e) { + // A handler that throws leaves the connection in an unknown state; close + // it rather than leaving it registered in epoll but never re-armed. + LOG(ERROR) << "Worker fd=" << client_fd + << ": Exception while handling message type=" + << static_cast(header.type) << ": " << e.what(); + keep_connection = false; } if (!keep_connection) { @@ -826,7 +885,21 @@ void TransferService::ExecuteGetTask(GetTask* task) { switch (resp_header.type) { case MessageType::kRespondToGetObj: - HandleDataReceive(conn.fd(), resp_header, true); + try { + HandleDataReceive(conn.fd(), resp_header, true); + } catch (const std::exception& e) { + // Typically BufferObject failing to create the destination file. The + // payload is still in flight on this connection, so it must not be + // reused as-is: shut it down so the next user fails fast instead of + // parsing stale payload bytes as a header. + // TODO: Let ScopedConnection discard a broken connection instead of + // returning it to the pool. + LOG(ERROR) << "GetTask " << task->GetTaskId() + << " failed while receiving data: " << e.what(); + shutdown(conn.fd(), SHUT_RDWR); + ReportResult(task->GetTaskId(), false, + std::string("Failed to receive data: ") + e.what()); + } break; case MessageType::kError: ReportResult(task->GetTaskId(), false, "Received error message"); @@ -856,19 +929,22 @@ void TransferService::ExecuteRespondToGetTask(RespondToGetTask* task) { } metric_container.connection_acquired_time = absl::Now(); - // Open file as buffer object - BufferObject buffer_obj(task->GetSourceObjId()); - void* buffer_data_ptr = buffer_obj.get_data_ptr(); - size_t size = buffer_obj.get_capacity(); - - if (buffer_data_ptr == nullptr) { + // Open file as buffer object. The constructor throws if the object cannot be + // opened (unreadable, a directory, or an empty file), so the requester must + // be told explicitly; it is blocked waiting for our response on this socket. + std::optional buffer_obj; + try { + buffer_obj.emplace(task->GetSourceObjId()); + } catch (const std::exception& e) { LOG(ERROR) << "RespondToGetTask failed: Could not open buffer object for '" - << task->GetSourceObjId() << "'"; + << task->GetSourceObjId() << "': " << e.what(); SendErrorResponse(sockfd, task->GetTaskId().c_str(), task->GetSourceObjId().c_str()); - ReportResult(task->GetTaskId(), false, "Failed to create buffer object"); + ReportResult(task->GetTaskId(), false, "Failed to open buffer object"); return; } + void* buffer_data_ptr = buffer_obj->get_data_ptr(); + size_t size = buffer_obj->get_capacity(); ObjInfoHeader header; diff --git a/src/ml_flashpoint/replication/transfer_service/transfer_service.h b/src/ml_flashpoint/replication/transfer_service/transfer_service.h index 20fdf7d..f30c6ae 100644 --- a/src/ml_flashpoint/replication/transfer_service/transfer_service.h +++ b/src/ml_flashpoint/replication/transfer_service/transfer_service.h @@ -37,6 +37,7 @@ #include #include #include +#include #include "absl/status/statusor.h" #include "connection_pool.h" @@ -66,7 +67,9 @@ class TransferService final { // // Args: // listen_port: The port to listen on. If 0, an ephemeral port is chosen. - // threads: The number of worker threads in the thread pool. + // threads: The number of worker threads in *each* of the two thread pools + // (outbound tasks and inbound epoll events), i.e. 2 * threads + // workers in total. // conn_pool_per_peer: The size of the connection pool for each peer. // global_rank: The global rank of the process. // @@ -161,8 +164,13 @@ class TransferService final { std::string local_address_; // Local address used in data transfer. int global_rank_ = -1; - std::unique_ptr thread_pool_; - std::unique_ptr epoll_thread_pool_; + // Two separate pools are required. Outbound tasks (Put/Get) block a worker + // for the entire transfer, including waiting for the remote peer to respond. + // The peer can only respond if *its* inbound work is being serviced, so if + // inbound epoll work shared a pool with outbound tasks, two peers that + // saturate their pools with Gets to each other would deadlock. + std::unique_ptr thread_pool_; // Outbound tasks. + std::unique_ptr epoll_thread_pool_; // Inbound epoll events. std::thread epoll_thread_; std::thread task_queue_thread_; TaskQueue task_queue_; @@ -173,6 +181,12 @@ class TransferService final { std::map> connection_pools_; mutable std::shared_mutex connection_pools_mutex_; // Guard connection_pools_ + // Accepted (inbound) client sockets. Tracked so that Shutdown() can unblock + // workers that are blocked on one of them (e.g. a RespondToGetTask whose + // peer stopped reading) and so the sockets are closed on shutdown. + std::unordered_set client_fds_; + std::mutex client_fds_mutex_; // Guard client_fds_ + struct PendingTaskContext { std::shared_ptr> promise; std::shared_ptr metric_container; diff --git a/tests/replication/transfer_service/transfer_service_p2p_test.cpp b/tests/replication/transfer_service/transfer_service_p2p_test.cpp index 19e9698..b834d60 100644 --- a/tests/replication/transfer_service/transfer_service_p2p_test.cpp +++ b/tests/replication/transfer_service/transfer_service_p2p_test.cpp @@ -12,25 +12,34 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include #include #include #include +#include +#include #include +#include #include +#include #include +#include +#include #include #include #include #include #include #include +#include #include "absl/log/log.h" #include "absl/log/log_sink.h" #include "absl/log/log_sink_registry.h" #include "gtest/gtest.h" #include "net_util.h" +#include "protocol.h" #include "transfer_service.h" namespace ml_flashpoint::replication::transfer_service { @@ -918,6 +927,144 @@ TEST(TransferServiceP2PTest, TimestampsAreRecorded) { EXPECT_EQ(timing_logs_count, 3) << "timing_logs_count: " << timing_logs_count; } +// --------------------------------------------------------------------------- +// Helpers for driving a TransferService with a raw client socket, used by the +// kGetObj / dest_address (SSRF) hardening tests below. +// --------------------------------------------------------------------------- + +// Connects a blocking TCP socket to 127.0.0.1:port. Returns -1 on failure. +// A positive `rcvbuf_bytes` shrinks SO_RCVBUF before connecting, so that a +// peer writing to this socket fills the pipe quickly and blocks in send(). +int ConnectToLocalPort(int port, int rcvbuf_bytes = 0) { + int fd = socket(AF_INET, SOCK_STREAM, 0); + if (fd < 0) return -1; + if (rcvbuf_bytes > 0) { + setsockopt(fd, SOL_SOCKET, SO_RCVBUF, &rcvbuf_bytes, sizeof(rcvbuf_bytes)); + } + sockaddr_in addr{}; + addr.sin_family = AF_INET; + addr.sin_port = htons(port); + inet_pton(AF_INET, "127.0.0.1", &addr.sin_addr); + if (connect(fd, reinterpret_cast(&addr), sizeof(addr)) != 0) { + close(fd); + return -1; + } + return fd; +} + +// Creates a non-blocking listener on 127.0.0.1 with an ephemeral port. Returns +// the listener fd and sets `port` to the bound port, or returns -1 on failure. +int ListenOnLocalEphemeralPort(int* port) { + int fd = socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK, 0); + if (fd < 0) return -1; + sockaddr_in addr{}; + addr.sin_family = AF_INET; + addr.sin_port = 0; + inet_pton(AF_INET, "127.0.0.1", &addr.sin_addr); + socklen_t addr_len = sizeof(addr); + if (bind(fd, reinterpret_cast(&addr), addr_len) != 0 || + getsockname(fd, reinterpret_cast(&addr), &addr_len) != 0 || + listen(fd, 16) != 0) { + close(fd); + return -1; + } + *port = ntohs(addr.sin_port); + return fd; +} + +void WriteTestFile(const std::string& path, const std::string& data) { + std::ofstream out_file(path); + out_file << data; + out_file.close(); +} + +ObjInfoHeader BuildGetRequest(const std::string& task_id, + const std::string& source_obj_id, + const std::string& dest_obj_id, + const std::string& dest_address) { + ObjInfoHeader request; + request.type = MessageType::kGetObj; + snprintf(request.task_id, sizeof(request.task_id), "%s", task_id.c_str()); + snprintf(request.source_obj_id, sizeof(request.source_obj_id), "%s", + source_obj_id.c_str()); + snprintf(request.dest_obj_id, sizeof(request.dest_obj_id), "%s", + dest_obj_id.c_str()); + snprintf(request.dest_address, sizeof(request.dest_address), "%s", + dest_address.c_str()); + return request; +} + +// Sends `request` on `client_fd` and returns the response header. +ObjInfoHeader SendRequestAndRecvResponse(int client_fd, + const ObjInfoHeader& request) { + ObjInfoHeader response; + EXPECT_TRUE(SendAll(client_fd, &request, kHeaderSize).ok()); + EXPECT_TRUE(RecvHeader(client_fd, response).ok()); + return response; +} + +// Reads the payload announced by a kRespondToGetObj `response` from `fd` and +// acknowledges it so the responder completes. Returns the payload, or an empty +// string if `response` is not a kRespondToGetObj with a payload. +std::string RecvGetPayloadAndAck(int fd, const ObjInfoHeader& response) { + if (response.type != MessageType::kRespondToGetObj || + response.obj_size <= 0) { + return ""; + } + std::string payload(response.obj_size, '\0'); + EXPECT_TRUE(RecvAll(fd, payload.data(), response.obj_size).ok()); + ObjInfoHeader ack; + ack.type = MessageType::kAck; + EXPECT_TRUE(SendAll(fd, &ack, kHeaderSize).ok()); + return payload; +} + +// Returns true if no connection attempt reaches `listener` within `timeout_ms`. +bool ListenerStaysIdle(int listener, int timeout_ms) { + pollfd pfd{listener, POLLIN, 0}; + return poll(&pfd, 1, timeout_ms) == 0; +} + +// Returns true if the peer closes `fd` (EOF or reset) within `timeout_ms`. +// Any data still arriving before the close is drained and ignored. +bool PeerClosed(int fd, int timeout_ms) { + const auto deadline = + std::chrono::steady_clock::now() + std::chrono::milliseconds(timeout_ms); + char drain[256]; + while (true) { + const auto remaining = + std::chrono::duration_cast( + deadline - std::chrono::steady_clock::now()); + if (remaining.count() < 0) return false; + pollfd pfd{fd, POLLIN, 0}; + if (poll(&pfd, 1, static_cast(remaining.count())) <= 0) return false; + if (recv(fd, drain, sizeof(drain), 0) <= 0) return true; + } +} + +// Runs `service.Shutdown()` on a helper thread and returns true if it completes +// within `timeout`. If it does not, `*blocking_fd` is closed (and set to -1) to +// release whichever worker Shutdown() is stuck behind, so that the thread can +// still be joined and the test fails cleanly instead of hanging. +bool ShutdownCompletesWithin(TransferService& service, + std::chrono::milliseconds timeout, + int* blocking_fd) { + std::promise done; + std::future done_future = done.get_future(); + std::thread shutdown_thread([&service, &done]() { + service.Shutdown(); + done.set_value(); + }); + const bool completed = + done_future.wait_for(timeout) == std::future_status::ready; + if (!completed) { + close(*blocking_fd); + *blocking_fd = -1; + } + shutdown_thread.join(); + return completed; +} + // Verifies that handling a kGetObj request responds over the existing client // socket and never initiates an outbound connection to header.dest_address. TEST(TransferServiceP2PTest, GetObjDoesNotConnectToDestAddress) { @@ -981,7 +1128,52 @@ TEST(TransferServiceP2PTest, GetObjDoesNotConnectToDestAddress) { // Then EXPECT_EQ(resp_header.type, MessageType::kRespondToGetObj); EXPECT_EQ(actual_data, expected_data); - EXPECT_EQ(accept(target_listener_fd, nullptr, nullptr), -1); + EXPECT_TRUE(ListenerStaysIdle(target_listener_fd, /*timeout_ms=*/200)); + + close(client_fd); + close(target_listener_fd); + std::remove(source_obj_id.c_str()); + service1.Shutdown(); +} + +class SpoofedDestAddressTest : public ::testing::TestWithParam {}; + +// Whatever dest_address a kGetObj request carries (other hosts, metadata IP, +// garbage, invalid or missing port), the object is streamed back over the +// request's own socket and no outbound connection is ever attempted. +TEST_P(SpoofedDestAddressTest, GetResponseGoesBackOnRequestSocket) { + // Given + TransferService service1("127.0.0.1"); + int port1 = service1.Initialize(/*listen_port=*/0, /*threads=*/4, + /*conn_pool_per_peer=*/1); + ASSERT_GT(port1, 0); + const std::string expected_data = "spoofed_dest_payload"; + const std::string source_obj_id = "spoofed_dest_source_obj"; + WriteTestFile(source_obj_id, expected_data); + int target_port = 0; + int target_listener_fd = ListenOnLocalEphemeralPort(&target_port); + ASSERT_GE(target_listener_fd, 0); + int client_fd = ConnectToLocalPort(port1); + ASSERT_GE(client_fd, 0); + // "{port}" stands in for a real, listening port so that a regression to + // connecting back would be observable on target_listener_fd. + std::string dest_address = GetParam(); + const size_t placeholder = dest_address.find("{port}"); + if (placeholder != std::string::npos) { + dest_address.replace(placeholder, 6, std::to_string(target_port)); + } + + // When + ObjInfoHeader actual_response = SendRequestAndRecvResponse( + client_fd, BuildGetRequest("spoofed_task", source_obj_id, + "spoofed_dest_obj", dest_address)); + std::string actual_data = RecvGetPayloadAndAck(client_fd, actual_response); + + // Then + EXPECT_EQ(actual_response.type, MessageType::kRespondToGetObj); + EXPECT_STREQ(actual_response.task_id, "spoofed_task"); + EXPECT_EQ(actual_data, expected_data); + EXPECT_TRUE(ListenerStaysIdle(target_listener_fd, /*timeout_ms=*/200)); close(client_fd); close(target_listener_fd); @@ -989,6 +1181,264 @@ TEST(TransferServiceP2PTest, GetObjDoesNotConnectToDestAddress) { service1.Shutdown(); } +INSTANTIATE_TEST_SUITE_P( + TransferServiceP2PTest, SpoofedDestAddressTest, + ::testing::Values("127.0.0.1:{port}", // the requester itself + "203.0.113.1:{port}", // TEST-NET-3 + "10.0.0.1:{port}", // private + "169.254.169.254:{port}", // metadata server + "0.0.0.0:{port}", "255.255.255.255:{port}", + "not-a-host:{port}", "127.0.0.1:0", "127.0.0.1:65536", + "127.0.0.1:abc", ":", "")); + +// Address fields that fill their whole buffer without a NUL terminator are +// handled safely (no over-read) and the request is served like any other. +TEST(TransferServiceP2PTest, GetRequestWithUnterminatedAddressFieldsIsServed) { + // Given + TransferService service1("127.0.0.1"); + int port1 = service1.Initialize(/*listen_port=*/0, /*threads=*/4, + /*conn_pool_per_peer=*/1); + ASSERT_GT(port1, 0); + const std::string expected_data = "unterminated_payload"; + const std::string source_obj_id = "unterminated_source_obj"; + WriteTestFile(source_obj_id, expected_data); + int client_fd = ConnectToLocalPort(port1); + ASSERT_GE(client_fd, 0); + ObjInfoHeader request = BuildGetRequest("unterminated_task", source_obj_id, + "unterminated_dest_obj", ""); + std::memset(request.dest_address, 'A', sizeof(request.dest_address)); + std::memset(request.source_address, 'B', sizeof(request.source_address)); + + // When + ObjInfoHeader actual_response = + SendRequestAndRecvResponse(client_fd, request); + std::string actual_data = RecvGetPayloadAndAck(client_fd, actual_response); + + // Then + EXPECT_EQ(actual_response.type, MessageType::kRespondToGetObj); + EXPECT_STREQ(actual_response.task_id, "unterminated_task"); + EXPECT_EQ(actual_data, expected_data); + + close(client_fd); + std::remove(source_obj_id.c_str()); + service1.Shutdown(); +} + +// An object that exists but cannot be opened (empty file, directory) yields a +// kError on the request socket rather than an exception that would orphan the +// connection and hang the requester, and the socket stays usable afterwards. +TEST(TransferServiceP2PTest, + GetRequestForUnopenableObjectReturnsErrorOnSameSocket) { + // Given + TransferService service1("127.0.0.1"); + int port1 = service1.Initialize(/*listen_port=*/0, /*threads=*/4, + /*conn_pool_per_peer=*/1); + ASSERT_GT(port1, 0); + const std::string empty_obj_id = "unopenable_empty_obj"; + const std::string dir_obj_id = "unopenable_dir_obj"; + const std::string valid_obj_id = "unopenable_valid_obj"; + const std::string expected_data = "payload_after_errors"; + WriteTestFile(empty_obj_id, ""); + std::filesystem::create_directory(dir_obj_id); + WriteTestFile(valid_obj_id, expected_data); + int client_fd = ConnectToLocalPort(port1); + ASSERT_GE(client_fd, 0); + + // When + ObjInfoHeader actual_empty_response = SendRequestAndRecvResponse( + client_fd, + BuildGetRequest("empty_task", empty_obj_id, "unopenable_empty_dest", "")); + ObjInfoHeader actual_dir_response = SendRequestAndRecvResponse( + client_fd, + BuildGetRequest("dir_task", dir_obj_id, "unopenable_dir_dest", "")); + ObjInfoHeader actual_valid_response = SendRequestAndRecvResponse( + client_fd, + BuildGetRequest("valid_task", valid_obj_id, "unopenable_valid_dest", "")); + std::string actual_data = + RecvGetPayloadAndAck(client_fd, actual_valid_response); + + // Then + EXPECT_EQ(actual_empty_response.type, MessageType::kError); + EXPECT_STREQ(actual_empty_response.task_id, "empty_task"); + EXPECT_EQ(actual_dir_response.type, MessageType::kError); + EXPECT_STREQ(actual_dir_response.task_id, "dir_task"); + EXPECT_EQ(actual_valid_response.type, MessageType::kRespondToGetObj); + EXPECT_STREQ(actual_valid_response.task_id, "valid_task"); + EXPECT_EQ(actual_data, expected_data); + + close(client_fd); + std::remove(empty_obj_id.c_str()); + std::filesystem::remove(dir_obj_id); + std::remove(valid_obj_id.c_str()); + service1.Shutdown(); +} + +// An inbound kRespondToGetObj that no GetTask is waiting for must not act as +// an arbitrary file write: the connection is closed and nothing is written. +TEST(TransferServiceP2PTest, + UnsolicitedRespondToGetObjClosesConnectionWithoutWritingFile) { + // Given + TransferService service1("127.0.0.1"); + int port1 = service1.Initialize(/*listen_port=*/0, /*threads=*/4, + /*conn_pool_per_peer=*/1); + ASSERT_GT(port1, 0); + const std::string dest_obj_id = "unsolicited_dest_obj"; + std::remove(dest_obj_id.c_str()); + int client_fd = ConnectToLocalPort(port1); + ASSERT_GE(client_fd, 0); + const std::string payload = "pwnd"; + ObjInfoHeader header; + header.type = MessageType::kRespondToGetObj; + snprintf(header.task_id, sizeof(header.task_id), "unsolicited_task"); + snprintf(header.dest_obj_id, sizeof(header.dest_obj_id), "%s", + dest_obj_id.c_str()); + header.obj_size = payload.size(); + + // When + ASSERT_TRUE(SendAll(client_fd, &header, kHeaderSize).ok()); + // The service may already have closed the socket, so this send may fail. + send(client_fd, payload.data(), payload.size(), MSG_NOSIGNAL); + + // Then + EXPECT_TRUE(PeerClosed(client_fd, /*timeout_ms=*/5000)); + EXPECT_FALSE(std::filesystem::exists(dest_obj_id)); + EXPECT_FALSE(std::filesystem::exists(dest_obj_id + ".tmp")); + + close(client_fd); + service1.Shutdown(); +} + +// A header with an unknown message type leaves the stream desynchronised, so +// the service closes the connection instead of guessing. +TEST(TransferServiceP2PTest, UnknownMessageTypeClosesConnection) { + // Given + TransferService service1("127.0.0.1"); + int port1 = service1.Initialize(/*listen_port=*/0, /*threads=*/4, + /*conn_pool_per_peer=*/1); + ASSERT_GT(port1, 0); + int client_fd = ConnectToLocalPort(port1); + ASSERT_GE(client_fd, 0); + ObjInfoHeader header; + header.type = static_cast(42); + snprintf(header.task_id, sizeof(header.task_id), "unknown_type_task"); + + // When + ASSERT_TRUE(SendAll(client_fd, &header, kHeaderSize).ok()); + + // Then + EXPECT_TRUE(PeerClosed(client_fd, /*timeout_ms=*/5000)); + + close(client_fd); + service1.Shutdown(); +} + +// Shutdown() must neither crash on the promise-less RespondToGetTask nor hang +// behind a responder that is blocked waiting for the requester's final ACK on +// an accepted socket. +TEST(TransferServiceP2PTest, ShutdownWhileResponderWaitsForAckIsSafe) { + // Given + TransferService service1("127.0.0.1"); + int port1 = service1.Initialize(/*listen_port=*/0, /*threads=*/4, + /*conn_pool_per_peer=*/1); + ASSERT_GT(port1, 0); + const std::string expected_data = "shutdown_no_ack_payload"; + const std::string source_obj_id = "shutdown_no_ack_source_obj"; + WriteTestFile(source_obj_id, expected_data); + int client_fd = ConnectToLocalPort(port1); + ASSERT_GE(client_fd, 0); + ObjInfoHeader actual_response = SendRequestAndRecvResponse( + client_fd, BuildGetRequest("no_ack_task", source_obj_id, + "shutdown_no_ack_dest_obj", "")); + ASSERT_EQ(actual_response.type, MessageType::kRespondToGetObj); + std::string actual_data(actual_response.obj_size, '\0'); + ASSERT_TRUE( + RecvAll(client_fd, actual_data.data(), actual_response.obj_size).ok()); + // Deliberately no ACK: the responder is now blocked in RecvHeader. + + // When + const bool actual_completed = + ShutdownCompletesWithin(service1, std::chrono::seconds(10), &client_fd); + + // Then + EXPECT_TRUE(actual_completed) + << "Shutdown() hung behind a responder waiting for an ACK"; + EXPECT_EQ(actual_data, expected_data); + if (client_fd >= 0) { + EXPECT_TRUE(PeerClosed(client_fd, /*timeout_ms=*/1000)); + close(client_fd); + } + TransferService service2("127.0.0.1"); + EXPECT_GT(service2.Initialize(), 0); + service2.Shutdown(); + std::remove(source_obj_id.c_str()); +} + +// Shutdown() must not hang behind a responder whose peer has stopped reading, +// i.e. a worker blocked in send() on an accepted socket with a full buffer. +TEST(TransferServiceP2PTest, ShutdownWhilePeerIsNotReadingGetResponseIsSafe) { + // Given + TransferService service1("127.0.0.1"); + int port1 = service1.Initialize(/*listen_port=*/0, /*threads=*/4, + /*conn_pool_per_peer=*/1); + ASSERT_GT(port1, 0); + // Large enough to exceed any socket buffer autotuning on loopback. + const std::string large_data(16 * 1024 * 1024, 'L'); + const std::string source_obj_id = "shutdown_not_reading_source_obj"; + WriteTestFile(source_obj_id, large_data); + int client_fd = ConnectToLocalPort(port1, /*rcvbuf_bytes=*/4096); + ASSERT_GE(client_fd, 0); + ObjInfoHeader request = BuildGetRequest("not_reading_task", source_obj_id, + "shutdown_not_reading_dest_obj", ""); + ASSERT_TRUE(SendAll(client_fd, &request, kHeaderSize).ok()); + // Wait until the response starts arriving, then let the responder fill the + // socket buffers and block, without ever reading from client_fd. + pollfd pfd{client_fd, POLLIN, 0}; + ASSERT_GT(poll(&pfd, 1, /*timeout_ms=*/5000), 0); + std::this_thread::sleep_for(std::chrono::milliseconds(300)); + + // When + const bool actual_completed = + ShutdownCompletesWithin(service1, std::chrono::seconds(10), &client_fd); + + // Then + EXPECT_TRUE(actual_completed) + << "Shutdown() hung behind a responder blocked in send()"; + if (client_fd >= 0) close(client_fd); + std::remove(source_obj_id.c_str()); +} + +// Happy path through the public API: because the response travels back over +// the request's own connection, a Get succeeds even when the requesting +// service advertises a local address that is not reachable. +TEST(TransferServiceP2PTest, GetSucceedsWhenRequesterAdvertisesUnreachableIp) { + // Given + TransferService service1("127.0.0.1"); + int port1 = service1.Initialize(); + ASSERT_GT(port1, 0); + TransferService service2("203.0.113.5"); + int port2 = service2.Initialize(); + ASSERT_GT(port2, 0); + const std::string expected_data = "reachable after all"; + const std::string source_obj_id = "unreachable_requester_source_obj"; + const std::string dest_obj_id = "unreachable_requester_dest_obj"; + WriteTestFile(source_obj_id, expected_data); + + // When + TransferResult actual_result = + service2 + .AsyncGet(source_obj_id, "127.0.0.1:" + std::to_string(port1), + dest_obj_id) + .get(); + + // Then + EXPECT_TRUE(actual_result.success); + VerifyFileContentAndRemove(dest_obj_id, expected_data); + + std::remove(source_obj_id.c_str()); + service1.Shutdown(); + service2.Shutdown(); +} + // Verifies that a single pooled connection (conn_pool_per_peer=1) is cleanly // reused across sequential AsyncGet calls, including after a kError response. TEST(TransferServiceP2PTest, From 827519b1199924d446caaec886b087e11d7bb875 Mon Sep 17 00:00:00 2001 From: Husam Date: Mon, 5 Oct 2026 14:22:45 +0000 Subject: [PATCH 04/15] fix(replication): keep connections usable and futures resolvable on GET/PUT failures Follow-up hardening for the same-socket kGetObj response path. * HandleDataReceive: when the destination cannot be created, drain the payload the peer is already streaming (new RecvAndDiscard) and reply kError instead of letting the exception escape. Pooled connections are never replaced once closed, so unread payload bytes left on the stream would have poisoned the connection for every later task on it. * ExecuteGetTask: complete the task and write the payload using the local task's task_id and dest_obj_id. The values echoed by the peer were used before, so a peer could leave the AsyncGet future pending forever or choose where the payload was written. Drops the shutdown() stopgap, which would have permanently poisoned a pooled connection. * HandleGetObjRequest: use the non-throwing std::filesystem::exists so that stat errors other than ENOENT (e.g. ENAMETOOLONG for an id that fills its field) answer kError instead of dropping the connection. --- .../replication/transfer_service/net_util.cpp | 18 +++++ .../replication/transfer_service/net_util.h | 15 ++++ .../transfer_service/transfer_service.cpp | 68 +++++++++++++++---- 3 files changed, 86 insertions(+), 15 deletions(-) diff --git a/src/ml_flashpoint/replication/transfer_service/net_util.cpp b/src/ml_flashpoint/replication/transfer_service/net_util.cpp index 7994394..ccf2a28 100644 --- a/src/ml_flashpoint/replication/transfer_service/net_util.cpp +++ b/src/ml_flashpoint/replication/transfer_service/net_util.cpp @@ -22,8 +22,10 @@ #include #include +#include #include #include +#include #include "absl/cleanup/cleanup.h" #include "absl/log/log.h" @@ -181,6 +183,22 @@ absl::Status RecvAll(int sockfd, void* data_ptr, ssize_t data_size) { return absl::OkStatus(); } +absl::Status RecvAndDiscard(int sockfd, ssize_t data_size) { + constexpr ssize_t kDiscardChunkSize = 64 * 1024; + std::vector scratch(static_cast( + std::min(std::max(data_size, 1), kDiscardChunkSize))); + ssize_t remaining_bytes = data_size; + while (remaining_bytes > 0) { + ssize_t chunk = std::min(remaining_bytes, kDiscardChunkSize); + absl::Status status = RecvAll(sockfd, scratch.data(), chunk); + if (!status.ok()) { + return status; + } + remaining_bytes -= chunk; + } + return absl::OkStatus(); +} + absl::Status RecvHeader(int sockfd, ObjInfoHeader& header) { absl::Status status = RecvAll(sockfd, &header, kHeaderSize); if (status.ok()) { diff --git a/src/ml_flashpoint/replication/transfer_service/net_util.h b/src/ml_flashpoint/replication/transfer_service/net_util.h index ec6aa0c..7ef4630 100644 --- a/src/ml_flashpoint/replication/transfer_service/net_util.h +++ b/src/ml_flashpoint/replication/transfer_service/net_util.h @@ -52,6 +52,21 @@ absl::Status SendAll(int sockfd, const void* data_ptr, ssize_t data_size); // absl::OkStatus() on success, or an error status on failure. absl::Status RecvAll(int sockfd, void* data_ptr, ssize_t data_size); +// Receives and discards exactly data_size bytes from the socket. +// +// Used when a receiver cannot store a payload that the peer is already +// sending (e.g. the destination file cannot be created): draining it keeps the +// connection's message stream aligned so the connection stays usable, which +// matters because pooled connections are never replaced once closed. +// +// Args: +// sockfd: The socket file descriptor. +// data_size: The number of bytes to discard. +// +// Returns: +// absl::OkStatus() on success, or an error status on failure. +absl::Status RecvAndDiscard(int sockfd, ssize_t data_size); + // Receives an ObjInfoHeader from the socket and ensures its string fields are // null-terminated. absl::Status RecvHeader(int sockfd, ObjInfoHeader& header); diff --git a/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp b/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp index 2ce5610..cd2dddc 100644 --- a/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp +++ b/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp @@ -29,6 +29,7 @@ #include #include #include +#include #include #include #include @@ -785,10 +786,29 @@ void TransferService::HandleDataReceive(int client_fd, } std::string tmp_obj_id = std::string(header.dest_obj_id) + std::string(kTempFileSuffix); - BufferObject buffer_obj(tmp_obj_id, header.obj_size, - /*overwrite=*/true); + std::optional buffer_obj; + try { + buffer_obj.emplace(tmp_obj_id, header.obj_size, /*overwrite=*/true); + } catch (const std::exception& e) { + LOG(ERROR) << "Failed to create buffer object for " << header.dest_obj_id + << ": " << e.what(); + // The sender is already streaming obj_size bytes on this connection. Drain + // them so the connection stays usable (pooled connections are never + // replaced), then tell the sender we failed so it does not wait for an ACK. + absl::Status drain_status = RecvAndDiscard(client_fd, header.obj_size); + if (!drain_status.ok()) { + LOG(ERROR) << "Failed to drain undeliverable payload for " + << header.dest_obj_id << ": " << drain_status; + } + if (is_respond_get_task) { + ReportResult(header.task_id, false, + std::string("Failed to create buffer object: ") + e.what()); + } + SendErrorResponse(client_fd, header.task_id, header.dest_obj_id); + return; + } LOG(INFO) << "Successfully created buffer object"; - void* receiver_data_ptr = buffer_obj.get_data_ptr(); + void* receiver_data_ptr = buffer_obj->get_data_ptr(); if (!RecvAll(client_fd, receiver_data_ptr, header.obj_size).ok()) { LOG(ERROR) << "Failed to receive data for " << header.dest_obj_id; @@ -801,7 +821,7 @@ void TransferService::HandleDataReceive(int client_fd, // Close the buffer object to ensure data is flushed and the file descriptor // is released before renaming. - buffer_obj.close(); + buffer_obj->close(); // Rename the temporary file to the final destination. if (rename(tmp_obj_id.c_str(), header.dest_obj_id) != 0) { @@ -884,23 +904,36 @@ void TransferService::ExecuteGetTask(GetTask* task) { } switch (resp_header.type) { - case MessageType::kRespondToGetObj: + case MessageType::kRespondToGetObj: { + // Only obj_size is new information. The task already knows its own id + // and destination, so neither is taken from the peer: a mismatched + // task_id would otherwise leave this task pending forever, and a + // different dest_obj_id would let the peer choose where we write. + if (task->GetTaskId() != resp_header.task_id) { + LOG(WARNING) << "GetTask " << task->GetTaskId() + << ": response carries task_id '" << resp_header.task_id + << "'; completing the local task regardless."; + } + ObjInfoHeader local_header = resp_header; + snprintf(local_header.task_id, sizeof(local_header.task_id), "%s", + task->GetTaskId().c_str()); + snprintf(local_header.dest_obj_id, sizeof(local_header.dest_obj_id), "%s", + task->GetDestObjId().c_str()); try { - HandleDataReceive(conn.fd(), resp_header, true); + HandleDataReceive(conn.fd(), local_header, true); } catch (const std::exception& e) { - // Typically BufferObject failing to create the destination file. The - // payload is still in flight on this connection, so it must not be - // reused as-is: shut it down so the next user fails fast instead of - // parsing stale payload bytes as a header. - // TODO: Let ScopedConnection discard a broken connection instead of - // returning it to the pool. + // HandleDataReceive reports expected failures (e.g. the destination + // file cannot be created) itself and keeps the connection aligned. + // This is a last-resort net: an exception escaping a worker thread is + // swallowed by the thread pool, which would leave the caller's future + // pending forever. LOG(ERROR) << "GetTask " << task->GetTaskId() << " failed while receiving data: " << e.what(); - shutdown(conn.fd(), SHUT_RDWR); ReportResult(task->GetTaskId(), false, std::string("Failed to receive data: ") + e.what()); } break; + } case MessageType::kError: ReportResult(task->GetTaskId(), false, "Received error message"); break; @@ -1060,8 +1093,13 @@ void TransferService::HandleGetObjRequest(int client_fd, LOG(INFO) << "Handling get object request for requested_obj_id: " << header.source_obj_id; - if (!std::filesystem::exists(header.source_obj_id)) { - PLOG(ERROR) << "Object not found: " << header.source_obj_id; + // Non-throwing overload: stat errors other than "not found" (e.g. a name + // that is too long, or an unreadable path component) must also produce a + // kError for the requester rather than an exception that drops the socket. + std::error_code ec; + if (!std::filesystem::exists(header.source_obj_id, ec)) { + LOG(ERROR) << "Object not found: " << header.source_obj_id + << (ec ? ": " + ec.message() : ""); SendErrorResponse(client_fd, header.task_id, header.source_obj_id); return; } From d62ace9ef8a31aabfaa5e9860df88ce625516766 Mon Sep 17 00:00:00 2001 From: Husam Date: Mon, 5 Oct 2026 14:22:45 +0000 Subject: [PATCH 05/15] test(replication): cover kGetObj happy paths and failure edge cases * net_util_test: RecvAndDiscard byte accounting, zero-length no-op, multi-chunk payloads and EOF reporting. * transfer_service_p2p_test: - sequential and over-subscribed concurrent raw GETs (8 clients on 2 inbound workers) are all served in order with the right task_id; - the responder survives a client that disconnects mid-response or closes instead of ACKing (single inbound worker, so a pinned worker would be visible); - invalid GET requests (empty, unknown and unterminated source_obj_id) get kError and the same socket then serves a valid GET; - uncreatable destinations are drained and answered with kError over a raw Put, AsyncPut and AsyncGet, and the single pooled connection is reused afterwards; - a fake peer returning malformed GET responses (zero size, wrong type, bare ACK, truncated payload, close) never hangs the requester's future, and a peer naming a different dest_obj_id cannot redirect where the payload is written. All waits are bounded so a regression fails the test instead of hanging the suite. Full suite green on 3 runs; the new and timing-sensitive tests pass --gtest_repeat=10. --- .../transfer_service/net_util_test.cpp | 87 +++ .../transfer_service_p2p_test.cpp | 577 ++++++++++++++++++ 2 files changed, 664 insertions(+) diff --git a/tests/replication/transfer_service/net_util_test.cpp b/tests/replication/transfer_service/net_util_test.cpp index a726426..6e9ae70 100644 --- a/tests/replication/transfer_service/net_util_test.cpp +++ b/tests/replication/transfer_service/net_util_test.cpp @@ -121,6 +121,93 @@ TEST_F(SendRecvTest, RecvAll_IncompleteRead) { EXPECT_EQ(recv_status.code(), absl::StatusCode::kUnavailable); } +// RecvAndDiscard consumes exactly the requested byte count, leaving whatever +// follows on the stream intact and in order. +TEST_F(SendRecvTest, RecvAndDiscard_SkipsExactlyRequestedBytes) { + // Given + const std::string discarded_data = "discard-me"; + const std::string expected_data = "keep-me"; + ASSERT_TRUE( + SendAll(fds_[0], discarded_data.data(), discarded_data.size()).ok()); + ASSERT_TRUE( + SendAll(fds_[0], expected_data.data(), expected_data.size()).ok()); + std::string actual_data(expected_data.size(), '\0'); + + // When + absl::Status discard_status = RecvAndDiscard(fds_[1], discarded_data.size()); + absl::Status recv_status = + RecvAll(fds_[1], actual_data.data(), actual_data.size()); + + // Then + EXPECT_TRUE(discard_status.ok()); + EXPECT_TRUE(recv_status.ok()); + EXPECT_EQ(actual_data, expected_data); +} + +// Discarding zero bytes succeeds without consuming anything from the stream. +TEST_F(SendRecvTest, RecvAndDiscard_ZeroBytesIsNoOp) { + // Given + const std::string expected_data = "untouched"; + ASSERT_TRUE( + SendAll(fds_[0], expected_data.data(), expected_data.size()).ok()); + std::string actual_data(expected_data.size(), '\0'); + + // When + absl::Status discard_status = RecvAndDiscard(fds_[1], 0); + absl::Status recv_status = + RecvAll(fds_[1], actual_data.data(), actual_data.size()); + + // Then + EXPECT_TRUE(discard_status.ok()); + EXPECT_TRUE(recv_status.ok()); + EXPECT_EQ(actual_data, expected_data); +} + +// Payloads larger than the internal scratch buffer are drained in multiple +// chunks, still stopping exactly at the requested byte count. +TEST_F(SendRecvTest, RecvAndDiscard_DrainsPayloadLargerThanChunkSize) { + // Given + const std::string discarded_data(1024 * 1024, 'd'); + const std::string expected_data = "trailer"; + // Sent from a helper thread: the payload exceeds the socketpair buffers, so + // SendAll only completes once the receiver starts draining. + std::thread writer_thread([this, &discarded_data, &expected_data]() { + ASSERT_TRUE( + SendAll(fds_[0], discarded_data.data(), discarded_data.size()).ok()); + ASSERT_TRUE( + SendAll(fds_[0], expected_data.data(), expected_data.size()).ok()); + }); + std::string actual_data(expected_data.size(), '\0'); + + // When + absl::Status discard_status = RecvAndDiscard(fds_[1], discarded_data.size()); + absl::Status recv_status = + RecvAll(fds_[1], actual_data.data(), actual_data.size()); + writer_thread.join(); + + // Then + EXPECT_TRUE(discard_status.ok()); + EXPECT_TRUE(recv_status.ok()); + EXPECT_EQ(actual_data, expected_data); +} + +// If the peer closes before the requested byte count arrives, RecvAndDiscard +// reports the same kUnavailable error as RecvAll instead of spinning. +TEST_F(SendRecvTest, RecvAndDiscard_FailsOnEof) { + // Given + const std::string partial_data = "short"; + ASSERT_TRUE(SendAll(fds_[0], partial_data.data(), partial_data.size()).ok()); + close(fds_[0]); + + // When + absl::Status discard_status = + RecvAndDiscard(fds_[1], partial_data.size() + 5); + + // Then + EXPECT_FALSE(discard_status.ok()); + EXPECT_EQ(discard_status.code(), absl::StatusCode::kUnavailable); +} + TEST_F(SendRecvTest, RecvHeader_Basic) { // Given ObjInfoHeader sent_header; diff --git a/tests/replication/transfer_service/transfer_service_p2p_test.cpp b/tests/replication/transfer_service/transfer_service_p2p_test.cpp index b834d60..ff472e6 100644 --- a/tests/replication/transfer_service/transfer_service_p2p_test.cpp +++ b/tests/replication/transfer_service/transfer_service_p2p_test.cpp @@ -22,11 +22,13 @@ #include #include +#include #include #include #include #include #include +#include #include #include #include @@ -1065,6 +1067,76 @@ bool ShutdownCompletesWithin(TransferService& service, return completed; } +// Reads a header from `fd` if one starts arriving within `timeout_ms`. Returns +// false on timeout or read failure, so that a stalled or crashed worker fails +// the test instead of hanging it. +bool RecvHeaderWithin(int fd, ObjInfoHeader* header, int timeout_ms) { + pollfd pfd{fd, POLLIN, 0}; + if (poll(&pfd, 1, timeout_ms) <= 0) return false; + return RecvHeader(fd, *header).ok(); +} + +// Returns true if `future` becomes ready within `timeout`. +template +bool FutureReadyWithin(const std::future& future, + std::chrono::milliseconds timeout) { + return future.wait_for(timeout) == std::future_status::ready; +} + +ObjInfoHeader BuildPutRequest(const std::string& task_id, + const std::string& dest_obj_id, + ssize_t obj_size) { + ObjInfoHeader request; + request.type = MessageType::kPutObj; + snprintf(request.task_id, sizeof(request.task_id), "%s", task_id.c_str()); + snprintf(request.dest_obj_id, sizeof(request.dest_obj_id), "%s", + dest_obj_id.c_str()); + request.obj_size = obj_size; + return request; +} + +// Stands in for a misbehaving peer: accepts one connection, reads one kGetObj +// request and answers it with whatever `respond` writes to the socket. Every +// wait is bounded so the helper thread always terminates. +class FakeGetResponder { + public: + explicit FakeGetResponder(std::function respond) + : respond_(std::move(respond)) { + listener_fd_ = ListenOnLocalEphemeralPort(&port_); + thread_ = std::thread([this]() { Serve(); }); + } + + ~FakeGetResponder() { + thread_.join(); + if (listener_fd_ >= 0) close(listener_fd_); + } + + std::string address() const { return "127.0.0.1:" + std::to_string(port_); } + bool received_request() const { return received_request_; } + + private: + void Serve() { + if (listener_fd_ < 0) return; + pollfd pfd{listener_fd_, POLLIN, 0}; + if (poll(&pfd, 1, /*timeout_ms=*/5000) <= 0) return; + int client_fd = accept(listener_fd_, nullptr, nullptr); + if (client_fd < 0) return; + ObjInfoHeader request; + if (RecvHeaderWithin(client_fd, &request, /*timeout_ms=*/5000) && + request.type == MessageType::kGetObj) { + received_request_ = true; + respond_(client_fd); + } + close(client_fd); + } + + std::function respond_; + int listener_fd_ = -1; + int port_ = 0; + std::atomic received_request_{false}; + std::thread thread_; +}; + // Verifies that handling a kGetObj request responds over the existing client // socket and never initiates an outbound connection to header.dest_address. TEST(TransferServiceP2PTest, GetObjDoesNotConnectToDestAddress) { @@ -1439,6 +1511,511 @@ TEST(TransferServiceP2PTest, GetSucceedsWhenRequesterAdvertisesUnreachableIp) { service2.Shutdown(); } +// Several kGetObj requests issued back-to-back on one accepted socket are each +// answered on that socket, in order, with the matching task_id and payload. +TEST(TransferServiceP2PTest, SequentialGetRequestsOnOneSocketAreServedInOrder) { + // Given + TransferService service1("127.0.0.1"); + int port1 = service1.Initialize(/*listen_port=*/0, /*threads=*/4, + /*conn_pool_per_peer=*/1); + ASSERT_GT(port1, 0); + const int kNumRequests = 3; + std::vector source_obj_ids; + std::vector expected_payloads; + for (int i = 0; i < kNumRequests; ++i) { + source_obj_ids.push_back("sequential_source_obj_" + std::to_string(i)); + expected_payloads.push_back("sequential_payload_" + std::to_string(i) + + std::string(4096 * (i + 1), 'a' + i)); + WriteTestFile(source_obj_ids[i], expected_payloads[i]); + } + int client_fd = ConnectToLocalPort(port1); + ASSERT_GE(client_fd, 0); + + // When + std::vector actual_responses; + std::vector actual_payloads; + for (int i = 0; i < kNumRequests; ++i) { + actual_responses.push_back(SendRequestAndRecvResponse( + client_fd, + BuildGetRequest("sequential_task_" + std::to_string(i), + source_obj_ids[i], "sequential_dest_obj", ""))); + actual_payloads.push_back( + RecvGetPayloadAndAck(client_fd, actual_responses.back())); + } + + // Then + for (int i = 0; i < kNumRequests; ++i) { + EXPECT_EQ(actual_responses[i].type, MessageType::kRespondToGetObj); + EXPECT_EQ(std::string(actual_responses[i].task_id), + "sequential_task_" + std::to_string(i)); + EXPECT_EQ(actual_payloads[i], expected_payloads[i]); + } + + close(client_fd); + for (const auto& source_obj_id : source_obj_ids) { + std::remove(source_obj_id.c_str()); + } + service1.Shutdown(); +} + +// Because a kGetObj response now occupies an inbound worker for the whole +// transfer, more simultaneous requests than inbound workers must still all be +// served (queued behind each other) rather than dropped or deadlocked. +TEST(TransferServiceP2PTest, ConcurrentGetRequestsBeyondThreadCountAreServed) { + // Given + TransferService service1("127.0.0.1"); + int port1 = service1.Initialize(/*listen_port=*/0, /*threads=*/2, + /*conn_pool_per_peer=*/1); + ASSERT_GT(port1, 0); + const int kNumClients = 8; + std::vector source_obj_ids(kNumClients); + std::vector expected_payloads(kNumClients); + for (int i = 0; i < kNumClients; ++i) { + source_obj_ids[i] = "oversubscribed_source_obj_" + std::to_string(i); + expected_payloads[i] = "oversubscribed_payload_" + std::to_string(i) + + std::string(64 * 1024, 'A' + i); + WriteTestFile(source_obj_ids[i], expected_payloads[i]); + } + std::vector actual_types(kNumClients, MessageType::kError); + std::vector actual_payloads(kNumClients); + // Not vector: its packed elements cannot be written concurrently. + std::vector actual_responded(kNumClients, 0); + + // When + std::vector clients; + for (int i = 0; i < kNumClients; ++i) { + clients.emplace_back([&, i]() { + int client_fd = ConnectToLocalPort(port1); + if (client_fd < 0) return; + ObjInfoHeader request = + BuildGetRequest("oversubscribed_task_" + std::to_string(i), + source_obj_ids[i], "oversubscribed_dest_obj", ""); + ObjInfoHeader response; + if (SendAll(client_fd, &request, kHeaderSize).ok() && + RecvHeaderWithin(client_fd, &response, /*timeout_ms=*/10000)) { + actual_responded[i] = 1; + actual_types[i] = response.type; + actual_payloads[i] = RecvGetPayloadAndAck(client_fd, response); + } + close(client_fd); + }); + } + for (auto& client : clients) client.join(); + + // Then + for (int i = 0; i < kNumClients; ++i) { + EXPECT_TRUE(actual_responded[i]) << "client " << i << " got no response"; + EXPECT_EQ(actual_types[i], MessageType::kRespondToGetObj) << "client " << i; + EXPECT_EQ(actual_payloads[i], expected_payloads[i]) << "client " << i; + } + + for (const auto& source_obj_id : source_obj_ids) { + std::remove(source_obj_id.c_str()); + } + service1.Shutdown(); +} + +// A requester that disconnects while its response is still being streamed must +// only fail its own request: with a single inbound worker, the next client's +// request proves the worker was released promptly and the service is healthy. +TEST(TransferServiceP2PTest, ResponderSurvivesClientDisconnectMidResponse) { + // Given + TransferService service1("127.0.0.1"); + int port1 = service1.Initialize(/*listen_port=*/0, /*threads=*/1, + /*conn_pool_per_peer=*/1); + ASSERT_GT(port1, 0); + // Large enough to exceed any socket buffer autotuning on loopback. + const std::string large_data(16 * 1024 * 1024, 'L'); + const std::string large_obj_id = "mid_response_large_source_obj"; + const std::string expected_data = "served after the disconnect"; + const std::string small_obj_id = "mid_response_small_source_obj"; + WriteTestFile(large_obj_id, large_data); + WriteTestFile(small_obj_id, expected_data); + int leaving_fd = ConnectToLocalPort(port1, /*rcvbuf_bytes=*/4096); + ASSERT_GE(leaving_fd, 0); + ObjInfoHeader large_request = BuildGetRequest( + "mid_response_leaving_task", large_obj_id, "mid_response_dest_obj", ""); + ASSERT_TRUE(SendAll(leaving_fd, &large_request, kHeaderSize).ok()); + // Let the response start flowing and the responder block in send() on the + // full socket buffers, then walk away with most of the payload unread. + pollfd pfd{leaving_fd, POLLIN, 0}; + ASSERT_GT(poll(&pfd, 1, /*timeout_ms=*/5000), 0); + std::this_thread::sleep_for(std::chrono::milliseconds(200)); + + // When + close(leaving_fd); + int client_fd = ConnectToLocalPort(port1); + ASSERT_GE(client_fd, 0); + ObjInfoHeader request = BuildGetRequest( + "mid_response_next_task", small_obj_id, "mid_response_dest_obj", ""); + ASSERT_TRUE(SendAll(client_fd, &request, kHeaderSize).ok()); + ObjInfoHeader actual_response; + const bool actual_responded = + RecvHeaderWithin(client_fd, &actual_response, /*timeout_ms=*/10000); + + // Then + ASSERT_TRUE(actual_responded) + << "the single inbound worker is still stuck on the departed client"; + EXPECT_EQ(actual_response.type, MessageType::kRespondToGetObj); + EXPECT_EQ(RecvGetPayloadAndAck(client_fd, actual_response), expected_data); + + close(client_fd); + std::remove(large_obj_id.c_str()); + std::remove(small_obj_id.c_str()); + service1.Shutdown(); +} + +// A requester that reads the whole response but closes instead of sending the +// final ACK must not pin the inbound worker: the next client is still served. +TEST(TransferServiceP2PTest, ResponderSurvivesClientClosingInsteadOfAck) { + // Given + TransferService service1("127.0.0.1"); + int port1 = service1.Initialize(/*listen_port=*/0, /*threads=*/1, + /*conn_pool_per_peer=*/1); + ASSERT_GT(port1, 0); + const std::string expected_data = "no_ack_then_next_payload"; + const std::string source_obj_id = "no_ack_next_source_obj"; + WriteTestFile(source_obj_id, expected_data); + int leaving_fd = ConnectToLocalPort(port1); + ASSERT_GE(leaving_fd, 0); + ObjInfoHeader first_response = SendRequestAndRecvResponse( + leaving_fd, + BuildGetRequest("no_ack_leaving_task", source_obj_id, "no_ack_dest", "")); + ASSERT_EQ(first_response.type, MessageType::kRespondToGetObj); + std::string first_payload(first_response.obj_size, '\0'); + ASSERT_TRUE( + RecvAll(leaving_fd, first_payload.data(), first_response.obj_size).ok()); + + // When + close(leaving_fd); // Deliberately no ACK. + int client_fd = ConnectToLocalPort(port1); + ASSERT_GE(client_fd, 0); + ObjInfoHeader request = + BuildGetRequest("no_ack_next_task", source_obj_id, "no_ack_dest", ""); + ASSERT_TRUE(SendAll(client_fd, &request, kHeaderSize).ok()); + ObjInfoHeader actual_response; + const bool actual_responded = + RecvHeaderWithin(client_fd, &actual_response, /*timeout_ms=*/10000); + + // Then + EXPECT_EQ(first_payload, expected_data); + ASSERT_TRUE(actual_responded) + << "the single inbound worker is still waiting for the departed ACK"; + EXPECT_EQ(actual_response.type, MessageType::kRespondToGetObj); + EXPECT_EQ(RecvGetPayloadAndAck(client_fd, actual_response), expected_data); + + close(client_fd); + std::remove(source_obj_id.c_str()); + service1.Shutdown(); +} + +class InvalidGetRequestTest : public ::testing::TestWithParam {}; + +// A kGetObj naming an object that cannot be served (empty id, unknown id, or an +// id that fills its whole field without a terminator) is answered with a kError +// carrying the request's task_id, and the same socket then serves a valid GET. +TEST_P(InvalidGetRequestTest, RepliesWithErrorAndKeepsSocketUsable) { + // Given + TransferService service1("127.0.0.1"); + int port1 = service1.Initialize(/*listen_port=*/0, /*threads=*/4, + /*conn_pool_per_peer=*/1); + ASSERT_GT(port1, 0); + const std::string expected_data = "valid_after_invalid_request"; + const std::string valid_obj_id = "invalid_request_valid_source_obj"; + WriteTestFile(valid_obj_id, expected_data); + int client_fd = ConnectToLocalPort(port1); + ASSERT_GE(client_fd, 0); + ObjInfoHeader invalid_request = + BuildGetRequest("invalid_task", GetParam(), "invalid_dest_obj", ""); + if (GetParam() == "{unterminated}") { + std::memset(invalid_request.source_obj_id, 'A', + sizeof(invalid_request.source_obj_id)); + } + + // When + ObjInfoHeader actual_error_response = + SendRequestAndRecvResponse(client_fd, invalid_request); + ObjInfoHeader actual_valid_response = SendRequestAndRecvResponse( + client_fd, + BuildGetRequest("valid_task", valid_obj_id, "invalid_request_dest", "")); + std::string actual_data = + RecvGetPayloadAndAck(client_fd, actual_valid_response); + + // Then + EXPECT_EQ(actual_error_response.type, MessageType::kError); + EXPECT_STREQ(actual_error_response.task_id, "invalid_task"); + EXPECT_EQ(actual_valid_response.type, MessageType::kRespondToGetObj); + EXPECT_STREQ(actual_valid_response.task_id, "valid_task"); + EXPECT_EQ(actual_data, expected_data); + + close(client_fd); + std::remove(valid_obj_id.c_str()); + service1.Shutdown(); +} + +INSTANTIATE_TEST_SUITE_P(TransferServiceP2PTest, InvalidGetRequestTest, + ::testing::Values("", "non_existent_source_obj", + "{unterminated}")); + +// When the receiver cannot create the destination of an inbound kPutObj, it +// drains the payload the sender is already streaming and answers kError, so the +// connection stays aligned and the next request on it is served normally. +TEST(TransferServiceP2PTest, + PutToUncreatableDestinationRepliesWithErrorAndKeepsSocketUsable) { + // Given + TransferService service1("127.0.0.1"); + int port1 = service1.Initialize(/*listen_port=*/0, /*threads=*/4, + /*conn_pool_per_peer=*/1); + ASSERT_GT(port1, 0); + const std::string undeliverable_payload(256 * 1024, 'U'); + // A regular file as a path component makes the destination uncreatable + // (missing directories alone would simply be created). + const std::string not_a_dir = "put_not_a_dir"; + WriteTestFile(not_a_dir, "regular file"); + const std::string uncreatable_dest = not_a_dir + "/uncreatable_dest_obj"; + const std::string expected_data = "delivered_after_undeliverable_put"; + const std::string dest_obj_id = "uncreatable_then_valid_dest_obj"; + std::remove(dest_obj_id.c_str()); + int client_fd = ConnectToLocalPort(port1); + ASSERT_GE(client_fd, 0); + + // When + ObjInfoHeader uncreatable_request = BuildPutRequest( + "uncreatable_task", uncreatable_dest, undeliverable_payload.size()); + ASSERT_TRUE(SendAll(client_fd, &uncreatable_request, kHeaderSize).ok()); + ASSERT_TRUE(SendAll(client_fd, undeliverable_payload.data(), + undeliverable_payload.size()) + .ok()); + ObjInfoHeader actual_error_response; + const bool actual_error_responded = + RecvHeaderWithin(client_fd, &actual_error_response, /*timeout_ms=*/10000); + ObjInfoHeader valid_request = + BuildPutRequest("valid_put_task", dest_obj_id, expected_data.size()); + ASSERT_TRUE(SendAll(client_fd, &valid_request, kHeaderSize).ok()); + ASSERT_TRUE( + SendAll(client_fd, expected_data.data(), expected_data.size()).ok()); + ObjInfoHeader actual_ack_response; + const bool actual_ack_responded = + RecvHeaderWithin(client_fd, &actual_ack_response, /*timeout_ms=*/10000); + + // Then + ASSERT_TRUE(actual_error_responded); + EXPECT_EQ(actual_error_response.type, MessageType::kError); + EXPECT_STREQ(actual_error_response.task_id, "uncreatable_task"); + ASSERT_TRUE(actual_ack_responded); + EXPECT_EQ(actual_ack_response.type, MessageType::kAck); + EXPECT_TRUE(std::filesystem::is_regular_file(not_a_dir)); + VerifyFileContentAndRemove(dest_obj_id, expected_data); + + close(client_fd); + std::remove(not_a_dir.c_str()); + service1.Shutdown(); +} + +// Through the public API, an AsyncPut whose destination cannot be created on +// the peer fails with an error (not a hang), and the single pooled connection +// it used is still good for the next AsyncPut. +TEST(TransferServiceP2PTest, + PutFailsCleanlyAndReusesConnectionWhenDestinationCannotBeCreated) { + // Given + TransferService service1("127.0.0.1"); + int port1 = service1.Initialize(/*listen_port=*/0, /*threads=*/4, + /*conn_pool_per_peer=*/1); + ASSERT_GT(port1, 0); + TransferService service2("127.0.0.1"); + int port2 = service2.Initialize(/*listen_port=*/0, /*threads=*/4, + /*conn_pool_per_peer=*/1); + ASSERT_GT(port2, 0); + const std::string peer1_addr = "127.0.0.1:" + std::to_string(port1); + std::string undeliverable_payload(256 * 1024, 'U'); + std::string expected_data = "put_after_uncreatable_destination"; + const std::string not_a_dir = "api_put_not_a_dir"; + WriteTestFile(not_a_dir, "regular file"); + const std::string dest_obj_id = "api_put_uncreatable_then_valid_dest_obj"; + + // When + std::future failed_future = service2.AsyncPut( + undeliverable_payload.data(), undeliverable_payload.size(), peer1_addr, + not_a_dir + "/uncreatable_dest_obj"); + const bool actual_failed_ready = + FutureReadyWithin(failed_future, std::chrono::seconds(10)); + std::future valid_future = service2.AsyncPut( + expected_data.data(), expected_data.size(), peer1_addr, dest_obj_id); + const bool actual_valid_ready = + FutureReadyWithin(valid_future, std::chrono::seconds(10)); + + // Then + ASSERT_TRUE(actual_failed_ready) << "AsyncPut hung instead of failing"; + EXPECT_THROW(failed_future.get(), std::runtime_error); + ASSERT_TRUE(actual_valid_ready) << "pooled connection unusable after error"; + EXPECT_TRUE(valid_future.get().success); + EXPECT_TRUE(std::filesystem::is_regular_file(not_a_dir)); + VerifyFileContentAndRemove(dest_obj_id, expected_data); + + std::remove(not_a_dir.c_str()); + service1.Shutdown(); + service2.Shutdown(); +} + +// Through the public API, an AsyncGet whose local destination cannot be created +// fails with an error (not a hang), the responder is told so it does not wait +// for an ACK, and the single pooled connection is still good for the next Get. +TEST(TransferServiceP2PTest, + GetFailsCleanlyAndReusesConnectionWhenDestinationCannotBeCreated) { + // Given + TransferService service1("127.0.0.1"); + int port1 = service1.Initialize(/*listen_port=*/0, /*threads=*/4, + /*conn_pool_per_peer=*/1); + ASSERT_GT(port1, 0); + TransferService service2("127.0.0.1"); + int port2 = service2.Initialize(/*listen_port=*/0, /*threads=*/4, + /*conn_pool_per_peer=*/1); + ASSERT_GT(port2, 0); + const std::string peer1_addr = "127.0.0.1:" + std::to_string(port1); + const std::string large_obj_id = "api_get_uncreatable_large_source_obj"; + const std::string source_obj_id = "api_get_uncreatable_valid_source_obj"; + const std::string not_a_dir = "api_get_not_a_dir"; + const std::string dest_obj_id = "api_get_uncreatable_then_valid_dest_obj"; + const std::string expected_data = "get_after_uncreatable_destination"; + WriteTestFile(large_obj_id, std::string(256 * 1024, 'U')); + WriteTestFile(source_obj_id, expected_data); + WriteTestFile(not_a_dir, "regular file"); + + // When + std::future failed_future = service2.AsyncGet( + large_obj_id, peer1_addr, not_a_dir + "/uncreatable_dest_obj"); + const bool actual_failed_ready = + FutureReadyWithin(failed_future, std::chrono::seconds(10)); + std::future valid_future = + service2.AsyncGet(source_obj_id, peer1_addr, dest_obj_id); + const bool actual_valid_ready = + FutureReadyWithin(valid_future, std::chrono::seconds(10)); + + // Then + ASSERT_TRUE(actual_failed_ready) << "AsyncGet hung instead of failing"; + EXPECT_THROW(failed_future.get(), std::runtime_error); + ASSERT_TRUE(actual_valid_ready) << "pooled connection unusable after error"; + EXPECT_TRUE(valid_future.get().success); + EXPECT_TRUE(std::filesystem::is_regular_file(not_a_dir)); + VerifyFileContentAndRemove(dest_obj_id, expected_data); + + std::remove(large_obj_id.c_str()); + std::remove(source_obj_id.c_str()); + std::remove(not_a_dir.c_str()); + service1.Shutdown(); + service2.Shutdown(); +} + +class MalformedGetResponseTest : public ::testing::TestWithParam { +}; + +// Whatever a misbehaving peer sends back for a GET (zero-size response, wrong +// message type, a bare ACK, a payload shorter than announced, or nothing at +// all), the requester's future resolves with an error instead of hanging. +TEST_P(MalformedGetResponseTest, RequesterFutureFailsInsteadOfHanging) { + // Given + const std::string scenario = GetParam(); + const std::string dest_obj_id = "malformed_response_dest_obj"; + FakeGetResponder responder([&scenario](int fd) { + ObjInfoHeader response; + if (scenario == "close_without_response") return; + if (scenario == "zero_size_response") { + response.type = MessageType::kRespondToGetObj; + response.obj_size = 0; + } else if (scenario == "wrong_type_response") { + response.type = MessageType::kPutObj; + } else if (scenario == "ack_response") { + response.type = MessageType::kAck; + } else if (scenario == "truncated_payload") { + response.type = MessageType::kRespondToGetObj; + response.obj_size = 1024; + } + snprintf(response.task_id, sizeof(response.task_id), "ignored"); + SendAll(fd, &response, kHeaderSize).IgnoreError(); + if (scenario == "truncated_payload") { + const std::string partial(16, 'p'); + SendAll(fd, partial.data(), partial.size()).IgnoreError(); + } + }); + TransferService service2("127.0.0.1"); + int port2 = service2.Initialize(/*listen_port=*/0, /*threads=*/2, + /*conn_pool_per_peer=*/1); + ASSERT_GT(port2, 0); + + // When + std::future actual_future = + service2.AsyncGet("any_source_obj", responder.address(), dest_obj_id); + const bool actual_ready = + FutureReadyWithin(actual_future, std::chrono::seconds(10)); + + // Then + ASSERT_TRUE(actual_ready) << "AsyncGet hung on scenario " << scenario; + EXPECT_THROW(actual_future.get(), std::runtime_error); + EXPECT_TRUE(responder.received_request()); + EXPECT_FALSE(std::filesystem::exists(dest_obj_id)); + + std::remove((dest_obj_id + ".tmp").c_str()); + service2.Shutdown(); +} + +INSTANTIATE_TEST_SUITE_P(TransferServiceP2PTest, MalformedGetResponseTest, + ::testing::Values("zero_size_response", + "wrong_type_response", + "ack_response", "truncated_payload", + "close_without_response"), + [](const ::testing::TestParamInfo& info) { + return info.param; + }); + +// A peer that answers a GET correctly but names a different dest_obj_id (and +// task_id) in its response must not redirect where the payload is written nor +// which task it completes: the file lands at the requester's own destination. +TEST(TransferServiceP2PTest, GetWritesToLocalDestinationNotToPeerNamedPath) { + // Given + const std::string expected_data = "payload written where the caller asked"; + const std::string dest_obj_id = "peer_redirect_requested_dest_obj"; + const std::string redirected_obj_id = "peer_redirect_peer_named_obj"; + std::remove(dest_obj_id.c_str()); + std::remove(redirected_obj_id.c_str()); + std::promise ack_promise; + std::future ack_future = ack_promise.get_future(); + FakeGetResponder responder([&](int fd) { + ObjInfoHeader response; + response.type = MessageType::kRespondToGetObj; + response.obj_size = expected_data.size(); + snprintf(response.task_id, sizeof(response.task_id), "peer_chosen_task"); + snprintf(response.dest_obj_id, sizeof(response.dest_obj_id), "%s", + redirected_obj_id.c_str()); + SendAll(fd, &response, kHeaderSize).IgnoreError(); + SendAll(fd, expected_data.data(), expected_data.size()).IgnoreError(); + ObjInfoHeader ack; + ack_promise.set_value(RecvHeaderWithin(fd, &ack, /*timeout_ms=*/5000) + ? ack.type + : MessageType::kError); + }); + TransferService service2("127.0.0.1"); + int port2 = service2.Initialize(/*listen_port=*/0, /*threads=*/2, + /*conn_pool_per_peer=*/1); + ASSERT_GT(port2, 0); + + // When + std::future actual_future = + service2.AsyncGet("any_source_obj", responder.address(), dest_obj_id); + const bool actual_ready = + FutureReadyWithin(actual_future, std::chrono::seconds(10)); + + // Then + ASSERT_TRUE(actual_ready) << "AsyncGet hung on a redirecting responder"; + EXPECT_TRUE(actual_future.get().success); + ASSERT_TRUE(FutureReadyWithin(ack_future, std::chrono::seconds(5))); + EXPECT_EQ(ack_future.get(), MessageType::kAck); + VerifyFileContentAndRemove(dest_obj_id, expected_data); + EXPECT_FALSE(std::filesystem::exists(redirected_obj_id)); + EXPECT_FALSE(std::filesystem::exists(redirected_obj_id + ".tmp")); + + service2.Shutdown(); +} + // Verifies that a single pooled connection (conn_pool_per_peer=1) is cleanly // reused across sequential AsyncGet calls, including after a kError response. TEST(TransferServiceP2PTest, From b48c472bfa320cd3306952996525956ae850b55e Mon Sep 17 00:00:00 2001 From: Husam Date: Mon, 5 Oct 2026 15:42:58 +0000 Subject: [PATCH 06/15] docs(replication): rewrite kGetObj comments in plain English Comment-only change to the code and tests added in the previous three commits: shorter declarative sentences, no clefts or roundabout phrasing. The only non-comment change is a shorter warning message in ExecuteGetTask for a response whose task_id does not match the request. --- .../replication/transfer_service/net_util.h | 8 +- .../replication/transfer_service/task.h | 13 +- .../transfer_service/transfer_service.cpp | 74 +++++----- .../transfer_service/transfer_service.h | 16 +- .../transfer_service/net_util_test.cpp | 17 ++- .../transfer_service_p2p_test.cpp | 139 +++++++++--------- 6 files changed, 136 insertions(+), 131 deletions(-) diff --git a/src/ml_flashpoint/replication/transfer_service/net_util.h b/src/ml_flashpoint/replication/transfer_service/net_util.h index 7ef4630..c786aa4 100644 --- a/src/ml_flashpoint/replication/transfer_service/net_util.h +++ b/src/ml_flashpoint/replication/transfer_service/net_util.h @@ -54,10 +54,10 @@ absl::Status RecvAll(int sockfd, void* data_ptr, ssize_t data_size); // Receives and discards exactly data_size bytes from the socket. // -// Used when a receiver cannot store a payload that the peer is already -// sending (e.g. the destination file cannot be created): draining it keeps the -// connection's message stream aligned so the connection stays usable, which -// matters because pooled connections are never replaced once closed. +// Call this when the peer is already sending a payload that the receiver +// cannot store (for example, the destination file cannot be created). Reading +// the payload keeps the message stream in sync, so the connection can be +// reused. This matters because the pool never replaces a closed connection. // // Args: // sockfd: The socket file descriptor. diff --git a/src/ml_flashpoint/replication/transfer_service/task.h b/src/ml_flashpoint/replication/transfer_service/task.h index b6796b2..b4c9b86 100644 --- a/src/ml_flashpoint/replication/transfer_service/task.h +++ b/src/ml_flashpoint/replication/transfer_service/task.h @@ -163,14 +163,13 @@ class GetTask : public Task { std::shared_ptr metric_container_; }; -// Serves a kGetObj request by streaming the object back over the connection -// the request arrived on. +// Serves a kGetObj request by sending the object back on the connection the +// request arrived on. // -// `client_fd` is NOT owned by the task: it belongs to the epoll loop. The task -// must therefore be executed synchronously from within ProcessIncomingData for -// that fd (as HandleGetObjRequest does), i.e. while the fd is still disarmed -// under EPOLLONESHOT and before any RemoveClient() can close and recycle it. -// Never enqueue this task for deferred execution. +// The task does not own `client_fd`; the epoll loop does. Run the task +// synchronously inside ProcessIncomingData for that fd, as HandleGetObjRequest +// does. At that point EPOLLONESHOT has disarmed the fd, so RemoveClient() +// cannot close or recycle it. Never queue this task to run later. class RespondToGetTask : public Task { public: RespondToGetTask( diff --git a/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp b/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp index cd2dddc..0cf70b7 100644 --- a/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp +++ b/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp @@ -215,11 +215,11 @@ void TransferService::Shutdown() { epoll_thread_.join(); } - // 6. Unblock any worker that is blocked on an accepted client socket (e.g. a - // RespondToGetTask whose peer stopped reading, or a kPutObj receive whose - // sender stalled). shutdown() makes their pending send/recv fail so the - // thread pools below can be joined. The fds are closed in step 9, once no - // worker can be using them anymore. + // 6. Unblock workers that are stuck on an accepted client socket, such as a + // RespondToGetTask whose peer stopped reading or a kPutObj receive whose + // sender stalled. shutdown() makes their pending send/recv fail, so the + // thread pools below can be joined. Step 9 closes the fds, after all + // workers have stopped. { std::lock_guard lock(client_fds_mutex_); for (int fd : client_fds_) { @@ -575,8 +575,8 @@ void TransferService::HandleNewConnection() { continue; } - // Register the client before arming it in epoll, so that a RemoveClient() - // triggered by an early event cannot race ahead of the registration. + // Register the client before arming it in epoll. Otherwise an early event + // could trigger RemoveClient() before the fd is in the registry. { std::lock_guard lock(client_fds_mutex_); client_fds_.insert(conn_fd); @@ -604,9 +604,9 @@ void TransferService::RemoveClient(int client_fd) { if (client_fd < 0) return; { // Only the caller that removes the fd from the registry closes it. This - // makes RemoveClient idempotent and prevents a double close (and closing - // an unrelated, reused fd) when Shutdown() and a worker both tear down - // the same connection. + // makes RemoveClient idempotent. Without it, Shutdown() and a worker could + // both close the same fd, and the second close could hit an unrelated fd + // that reused the number. std::lock_guard lock(client_fds_mutex_); if (client_fds_.erase(client_fd) == 0) return; } @@ -642,9 +642,9 @@ void TransferService::ProcessIncomingData(int client_fd) { HandleGetObjRequest(client_fd, header); break; case MessageType::kRespondToGetObj: - // Responses to kGetObj are read by ExecuteGetTask on the connection - // that sent the request. An unsolicited one arriving on the listener - // is a protocol violation; do not write its payload to disk. + // ExecuteGetTask reads kGetObj responses on the connection that sent + // the request. A response arriving on the listener is a protocol + // violation. Do not write its payload to disk. LOG(ERROR) << "Worker fd=" << client_fd << ": Unsolicited kRespondToGetObj received for task " << header.task_id << ". Closing connection."; @@ -661,15 +661,15 @@ void TransferService::ProcessIncomingData(int client_fd) { ReportResult(header.task_id, false, "Received error message"); break; default: - // The stream is desynced if we cannot interpret the header. + // We cannot parse this header, so the stream is out of sync. LOG(ERROR) << "Worker fd=" << client_fd << ": Unknown message type " << static_cast(header.type) << ". Closing connection."; keep_connection = false; break; } } catch (const std::exception& e) { - // A handler that throws leaves the connection in an unknown state; close - // it rather than leaving it registered in epoll but never re-armed. + // A handler that throws leaves the connection in an unknown state. Close + // it; otherwise it would stay registered in epoll but never be re-armed. LOG(ERROR) << "Worker fd=" << client_fd << ": Exception while handling message type=" << static_cast(header.type) << ": " << e.what(); @@ -792,9 +792,10 @@ void TransferService::HandleDataReceive(int client_fd, } catch (const std::exception& e) { LOG(ERROR) << "Failed to create buffer object for " << header.dest_obj_id << ": " << e.what(); - // The sender is already streaming obj_size bytes on this connection. Drain - // them so the connection stays usable (pooled connections are never - // replaced), then tell the sender we failed so it does not wait for an ACK. + // The sender is already streaming obj_size bytes on this connection. Read + // and discard them so the connection stays in sync (the pool never + // replaces a closed connection). Then send kError so the sender does not + // wait for an ACK. absl::Status drain_status = RecvAndDiscard(client_fd, header.obj_size); if (!drain_status.ok()) { LOG(ERROR) << "Failed to drain undeliverable payload for " @@ -905,14 +906,14 @@ void TransferService::ExecuteGetTask(GetTask* task) { switch (resp_header.type) { case MessageType::kRespondToGetObj: { - // Only obj_size is new information. The task already knows its own id - // and destination, so neither is taken from the peer: a mismatched - // task_id would otherwise leave this task pending forever, and a - // different dest_obj_id would let the peer choose where we write. + // Use the task's own task_id and dest_obj_id, not the peer's. Only + // obj_size is new information. If we used the peer's task_id, a + // mismatch would leave this task pending forever. If we used the peer's + // dest_obj_id, the peer could choose where we write. if (task->GetTaskId() != resp_header.task_id) { LOG(WARNING) << "GetTask " << task->GetTaskId() - << ": response carries task_id '" << resp_header.task_id - << "'; completing the local task regardless."; + << ": response has task_id '" << resp_header.task_id + << "'; ignoring it."; } ObjInfoHeader local_header = resp_header; snprintf(local_header.task_id, sizeof(local_header.task_id), "%s", @@ -922,11 +923,11 @@ void TransferService::ExecuteGetTask(GetTask* task) { try { HandleDataReceive(conn.fd(), local_header, true); } catch (const std::exception& e) { - // HandleDataReceive reports expected failures (e.g. the destination - // file cannot be created) itself and keeps the connection aligned. - // This is a last-resort net: an exception escaping a worker thread is - // swallowed by the thread pool, which would leave the caller's future - // pending forever. + // HandleDataReceive handles expected failures itself (for example, + // the destination file cannot be created) and keeps the connection in + // sync. This catch is a last resort: the thread pool swallows any + // exception that escapes a worker, which would leave the caller's + // future pending forever. LOG(ERROR) << "GetTask " << task->GetTaskId() << " failed while receiving data: " << e.what(); ReportResult(task->GetTaskId(), false, @@ -962,9 +963,10 @@ void TransferService::ExecuteRespondToGetTask(RespondToGetTask* task) { } metric_container.connection_acquired_time = absl::Now(); - // Open file as buffer object. The constructor throws if the object cannot be - // opened (unreadable, a directory, or an empty file), so the requester must - // be told explicitly; it is blocked waiting for our response on this socket. + // Open the file as a buffer object. The constructor throws if the object + // cannot be opened (unreadable, a directory, or an empty file). The + // requester is blocked waiting for our response on this socket, so we must + // tell it about the failure. std::optional buffer_obj; try { buffer_obj.emplace(task->GetSourceObjId()); @@ -1093,9 +1095,9 @@ void TransferService::HandleGetObjRequest(int client_fd, LOG(INFO) << "Handling get object request for requested_obj_id: " << header.source_obj_id; - // Non-throwing overload: stat errors other than "not found" (e.g. a name - // that is too long, or an unreadable path component) must also produce a - // kError for the requester rather than an exception that drops the socket. + // Use the non-throwing overload. A stat error other than "not found" (for + // example, a name that is too long) must also send kError to the requester + // instead of throwing and dropping the socket. std::error_code ec; if (!std::filesystem::exists(header.source_obj_id, ec)) { LOG(ERROR) << "Object not found: " << header.source_obj_id diff --git a/src/ml_flashpoint/replication/transfer_service/transfer_service.h b/src/ml_flashpoint/replication/transfer_service/transfer_service.h index f30c6ae..5b714cb 100644 --- a/src/ml_flashpoint/replication/transfer_service/transfer_service.h +++ b/src/ml_flashpoint/replication/transfer_service/transfer_service.h @@ -164,11 +164,11 @@ class TransferService final { std::string local_address_; // Local address used in data transfer. int global_rank_ = -1; - // Two separate pools are required. Outbound tasks (Put/Get) block a worker - // for the entire transfer, including waiting for the remote peer to respond. - // The peer can only respond if *its* inbound work is being serviced, so if - // inbound epoll work shared a pool with outbound tasks, two peers that - // saturate their pools with Gets to each other would deadlock. + // Two separate pools are required. An outbound task (Put/Get) blocks a + // worker for the whole transfer, including the wait for the peer's + // response. The peer can only respond if its own inbound work is running. + // If inbound and outbound work shared one pool, two peers that fill their + // pools with Gets to each other would deadlock. std::unique_ptr thread_pool_; // Outbound tasks. std::unique_ptr epoll_thread_pool_; // Inbound epoll events. std::thread epoll_thread_; @@ -181,9 +181,9 @@ class TransferService final { std::map> connection_pools_; mutable std::shared_mutex connection_pools_mutex_; // Guard connection_pools_ - // Accepted (inbound) client sockets. Tracked so that Shutdown() can unblock - // workers that are blocked on one of them (e.g. a RespondToGetTask whose - // peer stopped reading) and so the sockets are closed on shutdown. + // Accepted (inbound) client sockets. Shutdown() uses this set to unblock + // workers stuck on one of them (for example, a RespondToGetTask whose peer + // stopped reading) and to close the sockets. std::unordered_set client_fds_; std::mutex client_fds_mutex_; // Guard client_fds_ diff --git a/tests/replication/transfer_service/net_util_test.cpp b/tests/replication/transfer_service/net_util_test.cpp index 6e9ae70..4d8f2b0 100644 --- a/tests/replication/transfer_service/net_util_test.cpp +++ b/tests/replication/transfer_service/net_util_test.cpp @@ -121,8 +121,8 @@ TEST_F(SendRecvTest, RecvAll_IncompleteRead) { EXPECT_EQ(recv_status.code(), absl::StatusCode::kUnavailable); } -// RecvAndDiscard consumes exactly the requested byte count, leaving whatever -// follows on the stream intact and in order. +// RecvAndDiscard consumes exactly the requested number of bytes. The bytes +// after them stay on the stream, intact and in order. TEST_F(SendRecvTest, RecvAndDiscard_SkipsExactlyRequestedBytes) { // Given const std::string discarded_data = "discard-me"; @@ -163,14 +163,15 @@ TEST_F(SendRecvTest, RecvAndDiscard_ZeroBytesIsNoOp) { EXPECT_EQ(actual_data, expected_data); } -// Payloads larger than the internal scratch buffer are drained in multiple -// chunks, still stopping exactly at the requested byte count. +// A payload larger than the internal scratch buffer is read in several +// chunks and still stops exactly at the requested byte count. TEST_F(SendRecvTest, RecvAndDiscard_DrainsPayloadLargerThanChunkSize) { // Given const std::string discarded_data(1024 * 1024, 'd'); const std::string expected_data = "trailer"; - // Sent from a helper thread: the payload exceeds the socketpair buffers, so - // SendAll only completes once the receiver starts draining. + // A helper thread sends the data because the payload is larger than the + // socketpair buffers: SendAll cannot finish until the receiver starts + // reading. std::thread writer_thread([this, &discarded_data, &expected_data]() { ASSERT_TRUE( SendAll(fds_[0], discarded_data.data(), discarded_data.size()).ok()); @@ -191,8 +192,8 @@ TEST_F(SendRecvTest, RecvAndDiscard_DrainsPayloadLargerThanChunkSize) { EXPECT_EQ(actual_data, expected_data); } -// If the peer closes before the requested byte count arrives, RecvAndDiscard -// reports the same kUnavailable error as RecvAll instead of spinning. +// If the peer closes before all requested bytes arrive, RecvAndDiscard +// returns the same kUnavailable error as RecvAll instead of looping forever. TEST_F(SendRecvTest, RecvAndDiscard_FailsOnEof) { // Given const std::string partial_data = "short"; diff --git a/tests/replication/transfer_service/transfer_service_p2p_test.cpp b/tests/replication/transfer_service/transfer_service_p2p_test.cpp index ff472e6..3f7024b 100644 --- a/tests/replication/transfer_service/transfer_service_p2p_test.cpp +++ b/tests/replication/transfer_service/transfer_service_p2p_test.cpp @@ -935,8 +935,8 @@ TEST(TransferServiceP2PTest, TimestampsAreRecorded) { // --------------------------------------------------------------------------- // Connects a blocking TCP socket to 127.0.0.1:port. Returns -1 on failure. -// A positive `rcvbuf_bytes` shrinks SO_RCVBUF before connecting, so that a -// peer writing to this socket fills the pipe quickly and blocks in send(). +// A positive `rcvbuf_bytes` shrinks SO_RCVBUF before connecting, so a peer +// that writes to this socket fills the buffers quickly and blocks in send(). int ConnectToLocalPort(int port, int rcvbuf_bytes = 0) { int fd = socket(AF_INET, SOCK_STREAM, 0); if (fd < 0) return -1; @@ -1005,9 +1005,9 @@ ObjInfoHeader SendRequestAndRecvResponse(int client_fd, return response; } -// Reads the payload announced by a kRespondToGetObj `response` from `fd` and -// acknowledges it so the responder completes. Returns the payload, or an empty -// string if `response` is not a kRespondToGetObj with a payload. +// Reads the payload that a kRespondToGetObj `response` announced from `fd`, +// then sends kAck so the responder can finish. Returns the payload, or an +// empty string if `response` is not a kRespondToGetObj with a payload. std::string RecvGetPayloadAndAck(int fd, const ObjInfoHeader& response) { if (response.type != MessageType::kRespondToGetObj || response.obj_size <= 0) { @@ -1028,7 +1028,7 @@ bool ListenerStaysIdle(int listener, int timeout_ms) { } // Returns true if the peer closes `fd` (EOF or reset) within `timeout_ms`. -// Any data still arriving before the close is drained and ignored. +// Data that arrives before the close is read and ignored. bool PeerClosed(int fd, int timeout_ms) { const auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(timeout_ms); @@ -1044,10 +1044,10 @@ bool PeerClosed(int fd, int timeout_ms) { } } -// Runs `service.Shutdown()` on a helper thread and returns true if it completes -// within `timeout`. If it does not, `*blocking_fd` is closed (and set to -1) to -// release whichever worker Shutdown() is stuck behind, so that the thread can -// still be joined and the test fails cleanly instead of hanging. +// Runs `service.Shutdown()` on a helper thread and returns true if it finishes +// within `timeout`. If it does not, closes `*blocking_fd` (and sets it to -1) +// to release the worker that Shutdown() is waiting for, so the thread can be +// joined and the test fails instead of hanging. bool ShutdownCompletesWithin(TransferService& service, std::chrono::milliseconds timeout, int* blocking_fd) { @@ -1067,9 +1067,9 @@ bool ShutdownCompletesWithin(TransferService& service, return completed; } -// Reads a header from `fd` if one starts arriving within `timeout_ms`. Returns -// false on timeout or read failure, so that a stalled or crashed worker fails -// the test instead of hanging it. +// Reads a header from `fd` if one starts to arrive within `timeout_ms`. +// Returns false on timeout or read failure, so a stalled or crashed worker +// fails the test instead of hanging it. bool RecvHeaderWithin(int fd, ObjInfoHeader* header, int timeout_ms) { pollfd pfd{fd, POLLIN, 0}; if (poll(&pfd, 1, timeout_ms) <= 0) return false; @@ -1095,9 +1095,9 @@ ObjInfoHeader BuildPutRequest(const std::string& task_id, return request; } -// Stands in for a misbehaving peer: accepts one connection, reads one kGetObj -// request and answers it with whatever `respond` writes to the socket. Every -// wait is bounded so the helper thread always terminates. +// A fake peer. It accepts one connection, reads one kGetObj request, and +// answers with whatever `respond` writes to the socket. Every wait has a +// timeout, so the helper thread always exits. class FakeGetResponder { public: explicit FakeGetResponder(std::function respond) @@ -1210,9 +1210,9 @@ TEST(TransferServiceP2PTest, GetObjDoesNotConnectToDestAddress) { class SpoofedDestAddressTest : public ::testing::TestWithParam {}; -// Whatever dest_address a kGetObj request carries (other hosts, metadata IP, -// garbage, invalid or missing port), the object is streamed back over the -// request's own socket and no outbound connection is ever attempted. +// The service sends the object back on the request's own socket and never +// opens an outbound connection, whatever dest_address the kGetObj request +// carries (other hosts, metadata IP, garbage, invalid or missing port). TEST_P(SpoofedDestAddressTest, GetResponseGoesBackOnRequestSocket) { // Given TransferService service1("127.0.0.1"); @@ -1227,8 +1227,8 @@ TEST_P(SpoofedDestAddressTest, GetResponseGoesBackOnRequestSocket) { ASSERT_GE(target_listener_fd, 0); int client_fd = ConnectToLocalPort(port1); ASSERT_GE(client_fd, 0); - // "{port}" stands in for a real, listening port so that a regression to - // connecting back would be observable on target_listener_fd. + // "{port}" is replaced with a real listening port, so if the service ever + // connects back again, target_listener_fd will see it. std::string dest_address = GetParam(); const size_t placeholder = dest_address.find("{port}"); if (placeholder != std::string::npos) { @@ -1263,8 +1263,8 @@ INSTANTIATE_TEST_SUITE_P( "not-a-host:{port}", "127.0.0.1:0", "127.0.0.1:65536", "127.0.0.1:abc", ":", "")); -// Address fields that fill their whole buffer without a NUL terminator are -// handled safely (no over-read) and the request is served like any other. +// Address fields that fill their whole buffer with no NUL terminator do not +// cause an over-read, and the request is served like any other. TEST(TransferServiceP2PTest, GetRequestWithUnterminatedAddressFieldsIsServed) { // Given TransferService service1("127.0.0.1"); @@ -1296,9 +1296,9 @@ TEST(TransferServiceP2PTest, GetRequestWithUnterminatedAddressFieldsIsServed) { service1.Shutdown(); } -// An object that exists but cannot be opened (empty file, directory) yields a -// kError on the request socket rather than an exception that would orphan the -// connection and hang the requester, and the socket stays usable afterwards. +// An object that exists but cannot be opened (empty file, directory) gets a +// kError on the request socket. Before, the exception orphaned the connection +// and hung the requester. The socket stays usable afterwards. TEST(TransferServiceP2PTest, GetRequestForUnopenableObjectReturnsErrorOnSameSocket) { // Given @@ -1345,8 +1345,8 @@ TEST(TransferServiceP2PTest, service1.Shutdown(); } -// An inbound kRespondToGetObj that no GetTask is waiting for must not act as -// an arbitrary file write: the connection is closed and nothing is written. +// An inbound kRespondToGetObj that no GetTask is waiting for must not write a +// file. The service closes the connection and writes nothing. TEST(TransferServiceP2PTest, UnsolicitedRespondToGetObjClosesConnectionWithoutWritingFile) { // Given @@ -1380,7 +1380,7 @@ TEST(TransferServiceP2PTest, service1.Shutdown(); } -// A header with an unknown message type leaves the stream desynchronised, so +// A header with an unknown message type leaves the stream out of sync, so // the service closes the connection instead of guessing. TEST(TransferServiceP2PTest, UnknownMessageTypeClosesConnection) { // Given @@ -1404,9 +1404,9 @@ TEST(TransferServiceP2PTest, UnknownMessageTypeClosesConnection) { service1.Shutdown(); } -// Shutdown() must neither crash on the promise-less RespondToGetTask nor hang -// behind a responder that is blocked waiting for the requester's final ACK on -// an accepted socket. +// Shutdown() must not crash on a RespondToGetTask (which has no promise) and +// must not hang while a responder waits for the requester's final ACK on an +// accepted socket. TEST(TransferServiceP2PTest, ShutdownWhileResponderWaitsForAckIsSafe) { // Given TransferService service1("127.0.0.1"); @@ -1425,7 +1425,7 @@ TEST(TransferServiceP2PTest, ShutdownWhileResponderWaitsForAckIsSafe) { std::string actual_data(actual_response.obj_size, '\0'); ASSERT_TRUE( RecvAll(client_fd, actual_data.data(), actual_response.obj_size).ok()); - // Deliberately no ACK: the responder is now blocked in RecvHeader. + // No ACK on purpose: the responder is now blocked in RecvHeader. // When const bool actual_completed = @@ -1445,8 +1445,9 @@ TEST(TransferServiceP2PTest, ShutdownWhileResponderWaitsForAckIsSafe) { std::remove(source_obj_id.c_str()); } -// Shutdown() must not hang behind a responder whose peer has stopped reading, -// i.e. a worker blocked in send() on an accepted socket with a full buffer. +// Shutdown() must not hang while a responder's peer has stopped reading, that +// is, while a worker is blocked in send() on an accepted socket with a full +// buffer. TEST(TransferServiceP2PTest, ShutdownWhilePeerIsNotReadingGetResponseIsSafe) { // Given TransferService service1("127.0.0.1"); @@ -1462,8 +1463,8 @@ TEST(TransferServiceP2PTest, ShutdownWhilePeerIsNotReadingGetResponseIsSafe) { ObjInfoHeader request = BuildGetRequest("not_reading_task", source_obj_id, "shutdown_not_reading_dest_obj", ""); ASSERT_TRUE(SendAll(client_fd, &request, kHeaderSize).ok()); - // Wait until the response starts arriving, then let the responder fill the - // socket buffers and block, without ever reading from client_fd. + // Wait until the response starts to arrive, then let the responder fill the + // socket buffers and block. Never read from client_fd. pollfd pfd{client_fd, POLLIN, 0}; ASSERT_GT(poll(&pfd, 1, /*timeout_ms=*/5000), 0); std::this_thread::sleep_for(std::chrono::milliseconds(300)); @@ -1479,9 +1480,9 @@ TEST(TransferServiceP2PTest, ShutdownWhilePeerIsNotReadingGetResponseIsSafe) { std::remove(source_obj_id.c_str()); } -// Happy path through the public API: because the response travels back over -// the request's own connection, a Get succeeds even when the requesting -// service advertises a local address that is not reachable. +// Happy path through the public API. The response comes back on the request's +// own connection, so a Get succeeds even when the requesting service +// advertises a local address that is not reachable. TEST(TransferServiceP2PTest, GetSucceedsWhenRequesterAdvertisesUnreachableIp) { // Given TransferService service1("127.0.0.1"); @@ -1511,8 +1512,9 @@ TEST(TransferServiceP2PTest, GetSucceedsWhenRequesterAdvertisesUnreachableIp) { service2.Shutdown(); } -// Several kGetObj requests issued back-to-back on one accepted socket are each -// answered on that socket, in order, with the matching task_id and payload. +// Several kGetObj requests sent back-to-back on one accepted socket are +// answered on that socket, in order, each with the matching task_id and +// payload. TEST(TransferServiceP2PTest, SequentialGetRequestsOnOneSocketAreServedInOrder) { // Given TransferService service1("127.0.0.1"); @@ -1558,9 +1560,9 @@ TEST(TransferServiceP2PTest, SequentialGetRequestsOnOneSocketAreServedInOrder) { service1.Shutdown(); } -// Because a kGetObj response now occupies an inbound worker for the whole -// transfer, more simultaneous requests than inbound workers must still all be -// served (queued behind each other) rather than dropped or deadlocked. +// A kGetObj response now holds an inbound worker for the whole transfer. When +// there are more concurrent requests than inbound workers, all of them must +// still be served (they queue behind each other), not dropped or deadlocked. TEST(TransferServiceP2PTest, ConcurrentGetRequestsBeyondThreadCountAreServed) { // Given TransferService service1("127.0.0.1"); @@ -1615,9 +1617,9 @@ TEST(TransferServiceP2PTest, ConcurrentGetRequestsBeyondThreadCountAreServed) { service1.Shutdown(); } -// A requester that disconnects while its response is still being streamed must -// only fail its own request: with a single inbound worker, the next client's -// request proves the worker was released promptly and the service is healthy. +// A requester that disconnects while its response is still streaming must only +// fail its own request. With a single inbound worker, the next client's +// request proves the worker was released quickly and the service is healthy. TEST(TransferServiceP2PTest, ResponderSurvivesClientDisconnectMidResponse) { // Given TransferService service1("127.0.0.1"); @@ -1636,8 +1638,8 @@ TEST(TransferServiceP2PTest, ResponderSurvivesClientDisconnectMidResponse) { ObjInfoHeader large_request = BuildGetRequest( "mid_response_leaving_task", large_obj_id, "mid_response_dest_obj", ""); ASSERT_TRUE(SendAll(leaving_fd, &large_request, kHeaderSize).ok()); - // Let the response start flowing and the responder block in send() on the - // full socket buffers, then walk away with most of the payload unread. + // Let the response start to flow and the responder block in send() on the + // full socket buffers. Then close with most of the payload unread. pollfd pfd{leaving_fd, POLLIN, 0}; ASSERT_GT(poll(&pfd, 1, /*timeout_ms=*/5000), 0); std::this_thread::sleep_for(std::chrono::milliseconds(200)); @@ -1666,7 +1668,7 @@ TEST(TransferServiceP2PTest, ResponderSurvivesClientDisconnectMidResponse) { } // A requester that reads the whole response but closes instead of sending the -// final ACK must not pin the inbound worker: the next client is still served. +// final ACK must not hold the inbound worker. The next client is still served. TEST(TransferServiceP2PTest, ResponderSurvivesClientClosingInsteadOfAck) { // Given TransferService service1("127.0.0.1"); @@ -1711,9 +1713,9 @@ TEST(TransferServiceP2PTest, ResponderSurvivesClientClosingInsteadOfAck) { class InvalidGetRequestTest : public ::testing::TestWithParam {}; -// A kGetObj naming an object that cannot be served (empty id, unknown id, or an -// id that fills its whole field without a terminator) is answered with a kError -// carrying the request's task_id, and the same socket then serves a valid GET. +// A kGetObj for an object that cannot be served (empty id, unknown id, or an id +// that fills its whole field with no terminator) gets a kError with the +// request's task_id. The same socket then serves a valid GET. TEST_P(InvalidGetRequestTest, RepliesWithErrorAndKeepsSocketUsable) { // Given TransferService service1("127.0.0.1"); @@ -1758,8 +1760,8 @@ INSTANTIATE_TEST_SUITE_P(TransferServiceP2PTest, InvalidGetRequestTest, "{unterminated}")); // When the receiver cannot create the destination of an inbound kPutObj, it -// drains the payload the sender is already streaming and answers kError, so the -// connection stays aligned and the next request on it is served normally. +// reads and discards the payload the sender is already streaming and answers +// kError. The connection stays in sync and serves the next request normally. TEST(TransferServiceP2PTest, PutToUncreatableDestinationRepliesWithErrorAndKeepsSocketUsable) { // Given @@ -1768,8 +1770,8 @@ TEST(TransferServiceP2PTest, /*conn_pool_per_peer=*/1); ASSERT_GT(port1, 0); const std::string undeliverable_payload(256 * 1024, 'U'); - // A regular file as a path component makes the destination uncreatable - // (missing directories alone would simply be created). + // A regular file as a path component makes the destination impossible to + // create. A missing directory would not: BufferObject creates those. const std::string not_a_dir = "put_not_a_dir"; WriteTestFile(not_a_dir, "regular file"); const std::string uncreatable_dest = not_a_dir + "/uncreatable_dest_obj"; @@ -1813,8 +1815,8 @@ TEST(TransferServiceP2PTest, } // Through the public API, an AsyncPut whose destination cannot be created on -// the peer fails with an error (not a hang), and the single pooled connection -// it used is still good for the next AsyncPut. +// the peer fails with an error instead of hanging. The single pooled +// connection it used still works for the next AsyncPut. TEST(TransferServiceP2PTest, PutFailsCleanlyAndReusesConnectionWhenDestinationCannotBeCreated) { // Given @@ -1858,8 +1860,9 @@ TEST(TransferServiceP2PTest, } // Through the public API, an AsyncGet whose local destination cannot be created -// fails with an error (not a hang), the responder is told so it does not wait -// for an ACK, and the single pooled connection is still good for the next Get. +// fails with an error instead of hanging. The responder is told, so it does +// not wait for an ACK, and the single pooled connection still works for the +// next Get. TEST(TransferServiceP2PTest, GetFailsCleanlyAndReusesConnectionWhenDestinationCannotBeCreated) { // Given @@ -1909,9 +1912,9 @@ TEST(TransferServiceP2PTest, class MalformedGetResponseTest : public ::testing::TestWithParam { }; -// Whatever a misbehaving peer sends back for a GET (zero-size response, wrong -// message type, a bare ACK, a payload shorter than announced, or nothing at -// all), the requester's future resolves with an error instead of hanging. +// The requester's future fails with an error instead of hanging, whatever a bad +// peer sends back for a GET: a zero-size response, the wrong message type, a +// bare ACK, a payload shorter than announced, or nothing at all. TEST_P(MalformedGetResponseTest, RequesterFutureFailsInsteadOfHanging) { // Given const std::string scenario = GetParam(); @@ -1967,9 +1970,9 @@ INSTANTIATE_TEST_SUITE_P(TransferServiceP2PTest, MalformedGetResponseTest, return info.param; }); -// A peer that answers a GET correctly but names a different dest_obj_id (and -// task_id) in its response must not redirect where the payload is written nor -// which task it completes: the file lands at the requester's own destination. +// A peer that answers a GET correctly but puts a different dest_obj_id (and +// task_id) in its response must not change where the payload is written or +// which task completes. The file lands at the requester's own destination. TEST(TransferServiceP2PTest, GetWritesToLocalDestinationNotToPeerNamedPath) { // Given const std::string expected_data = "payload written where the caller asked"; From 31109e77e7f4882ffe315e0addfa075d22aa3dec Mon Sep 17 00:00:00 2001 From: Husam Date: Mon, 5 Oct 2026 16:47:03 +0000 Subject: [PATCH 07/15] fix(replication): do not create connection pools once Shutdown() has started ThreadPool::stop() runs every task that is still queued before the workers exit. Shutdown() has already shut down and cleared the connection pools by then, so a queued Put or Get re-created a pool that nothing would ever shut down. If the peer completed the TCP handshake but never answered (for example, a wedged rank), the task blocked in recv() forever and Shutdown() never returned, which hangs the Python thread that called shutdown(). GetOrCreateConnectionPool now returns nullptr once running_ is false, so drained tasks fail fast with 'Failed to get connection'. This predates the kGetObj change and affects Put as well. New test ShutdownWithQueuedTasksAgainstSilentPeerCompletes reproduces the hang (10 s timeout before the fix) and passes 5/5 with it. Full transfer_service_test 114/114; pytest tests/replication 46 passed. --- .../transfer_service/transfer_service.cpp | 10 +++++ .../transfer_service_p2p_test.cpp | 38 +++++++++++++++++++ 2 files changed, 48 insertions(+) diff --git a/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp b/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp index 0cf70b7..48b9df3 100644 --- a/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp +++ b/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp @@ -1077,6 +1077,16 @@ std::shared_ptr TransferService::GetOrCreateConnectionPool( return it->second; // Double check } + // Shutdown() sets running_ to false before it shuts down and clears the + // pools. ThreadPool::stop() then runs every task that is still queued. If one + // of those tasks created a pool here, nothing would ever shut it down, and a + // task blocked on a peer that never answers would block Shutdown() forever. + if (!running_.load()) { + LOG(WARNING) << "Not creating connection pool for " << peer_addr + << ": service is shutting down"; + return nullptr; + } + // Pool doesn't exist, create it dynamically auto new_pool = std::make_shared(peer_host, peer_port, conn_pool_size_per_peer_); diff --git a/tests/replication/transfer_service/transfer_service_p2p_test.cpp b/tests/replication/transfer_service/transfer_service_p2p_test.cpp index 3f7024b..5456465 100644 --- a/tests/replication/transfer_service/transfer_service_p2p_test.cpp +++ b/tests/replication/transfer_service/transfer_service_p2p_test.cpp @@ -2137,5 +2137,43 @@ TEST(TransferServiceP2PTest, BidirectionalConcurrentGetAndPut) { service2.Shutdown(); } +// A peer whose kernel completes TCP handshakes but whose service never answers +// (for example, a wedged rank) must not make Shutdown() hang. Shutdown() runs +// the tasks that are still queued behind the workers, so those tasks must fail +// fast instead of opening new connections to that peer. +TEST(TransferServiceP2PTest, + ShutdownWithQueuedTasksAgainstSilentPeerCompletes) { + // Given + int silent_port = 0; + int silent_listener = ListenOnLocalEphemeralPort(&silent_port); + ASSERT_GE(silent_listener, 0); + TransferService service1("127.0.0.1"); + int port1 = service1.Initialize(/*listen_port=*/0, /*threads=*/2, + /*conn_pool_per_peer=*/1); + ASSERT_GT(port1, 0); + const std::string silent_addr = "127.0.0.1:" + std::to_string(silent_port); + constexpr int kNumTasks = 6; // More than the 2 workers, so 4 stay queued. + std::vector> futures; + for (int i = 0; i < kNumTasks; ++i) { + futures.push_back(service1.AsyncGet("silent_source_obj", silent_addr, + "silent_dest_" + std::to_string(i))); + } + // Let the first task connect and block waiting for a reply that never comes. + std::this_thread::sleep_for(std::chrono::milliseconds(300)); + + // When + const bool actual_completed = ShutdownCompletesWithin( + service1, std::chrono::seconds(10), &silent_listener); + + // Then + EXPECT_TRUE(actual_completed) + << "Shutdown() hung running queued tasks against a silent peer"; + for (auto& future : futures) { + ASSERT_TRUE(FutureReadyWithin(future, std::chrono::seconds(5))); + EXPECT_THROW(future.get(), std::runtime_error); + } + if (silent_listener >= 0) close(silent_listener); +} + } // namespace } // namespace ml_flashpoint::replication::transfer_service From c171e2a6e43111e02631c2259a9b5afa3dba5d33 Mon Sep 17 00:00:00 2001 From: Husam Date: Mon, 5 Oct 2026 20:10:25 +0000 Subject: [PATCH 08/15] refactor(replication): share the failure tail in HandleDataReceive The four failure branches before the ACK (bad obj_size, cannot create the buffer, receive failed, rename failed) all ended with the same three steps: resolve the local Get future if there is one, reply kError, return. A local fail() lambda now does that, so each branch is one line. The separate drain-failure log is gone: if the drain fails the peer is gone, and SendErrorResponse already logs the failed kError send. No behavior change. Error messages now carry the obj_size and the rename errno text; the prefixes tests match on are unchanged. Full transfer_service_test 114/114, the failure-path tests 5/5 under --gtest_repeat, pytest tests/replication 46 passed. --- .../transfer_service/transfer_service.cpp | 54 +++++++++---------- 1 file changed, 25 insertions(+), 29 deletions(-) diff --git a/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp b/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp index 48b9df3..d01a863 100644 --- a/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp +++ b/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp @@ -22,6 +22,7 @@ #include #include +#include #include #include #include @@ -767,12 +768,24 @@ void TransferService::HandleDataReceive(int client_fd, const ObjInfoHeader& header, bool is_respond_get_task) { LOG(INFO) << "Handling data receive for filename: " << header.dest_obj_id; - if (header.obj_size <= 0) { - LOG(ERROR) << "Invalid obj_size received: " << header.obj_size; + + // Every failure before the ACK ends the same way: reply kError so the sender + // stops waiting for an ACK, and, when this is the receiving half of a local + // Get, resolve the caller's future. The connection is left open and stays + // usable: the branches before the payload either have nothing to read or + // drain it first, and RecvAll only fails once the connection is already + // broken. The caller decides what happens to the connection next. + auto fail = [&](const std::string& message) { + LOG(ERROR) << "HandleDataReceive for " << header.dest_obj_id << ": " + << message; if (is_respond_get_task) { - ReportResult(header.task_id, false, "Invalid obj_size received"); + ReportResult(header.task_id, false, message); } SendErrorResponse(client_fd, header.task_id, header.dest_obj_id); + }; + + if (header.obj_size <= 0) { + fail("Invalid obj_size received: " + std::to_string(header.obj_size)); return; } if (is_respond_get_task) { @@ -786,37 +799,24 @@ void TransferService::HandleDataReceive(int client_fd, } std::string tmp_obj_id = std::string(header.dest_obj_id) + std::string(kTempFileSuffix); + // BufferObject's constructor throws when it cannot create the file. std::optional buffer_obj; try { buffer_obj.emplace(tmp_obj_id, header.obj_size, /*overwrite=*/true); } catch (const std::exception& e) { - LOG(ERROR) << "Failed to create buffer object for " << header.dest_obj_id - << ": " << e.what(); // The sender is already streaming obj_size bytes on this connection. Read - // and discard them so the connection stays in sync (the pool never - // replaces a closed connection). Then send kError so the sender does not - // wait for an ACK. - absl::Status drain_status = RecvAndDiscard(client_fd, header.obj_size); - if (!drain_status.ok()) { - LOG(ERROR) << "Failed to drain undeliverable payload for " - << header.dest_obj_id << ": " << drain_status; - } - if (is_respond_get_task) { - ReportResult(header.task_id, false, - std::string("Failed to create buffer object: ") + e.what()); - } - SendErrorResponse(client_fd, header.task_id, header.dest_obj_id); + // and discard them so the connection stays in sync and can be reused (the + // pool never replaces a closed connection). If the drain fails, the peer + // is gone, and the kError send inside fail() logs that. + RecvAndDiscard(client_fd, header.obj_size).IgnoreError(); + fail(std::string("Failed to create buffer object: ") + e.what()); return; } LOG(INFO) << "Successfully created buffer object"; void* receiver_data_ptr = buffer_obj->get_data_ptr(); if (!RecvAll(client_fd, receiver_data_ptr, header.obj_size).ok()) { - LOG(ERROR) << "Failed to receive data for " << header.dest_obj_id; - if (is_respond_get_task) { - ReportResult(header.task_id, false, "Failed to receive data"); - } - SendErrorResponse(client_fd, header.task_id, header.dest_obj_id); + fail("Failed to receive data"); return; } @@ -826,12 +826,8 @@ void TransferService::HandleDataReceive(int client_fd, // Rename the temporary file to the final destination. if (rename(tmp_obj_id.c_str(), header.dest_obj_id) != 0) { - PLOG(ERROR) << "Failed to rename temporary file " << tmp_obj_id << " to " - << header.dest_obj_id; - if (is_respond_get_task) { - ReportResult(header.task_id, false, "Failed to rename temporary file"); - } - SendErrorResponse(client_fd, header.task_id, header.dest_obj_id); + fail("Failed to rename temporary file " + tmp_obj_id + " to " + + header.dest_obj_id + ": " + std::strerror(errno)); return; } if (is_respond_get_task) { From 64d30bff0cbfe144f9516d118ad0df29a579a52c Mon Sep 17 00:00:00 2001 From: Husam Date: Tue, 6 Oct 2026 14:27:55 +0000 Subject: [PATCH 09/15] refactor(replication): drop the default case from the message dispatch switch Move the switch on header.type into DispatchMessage. Every enumerator returns, so the switch has no default and -Wswitch reports a new MessageType that is not handled. The type byte comes from the network and can hold any uint8_t value, so the unknown-type check stays, but after the switch where it cannot mask that warning. No behavior change. UnknownMessageTypeClosesConnection covers the unknown-type path. --- .../transfer_service/transfer_service.cpp | 75 ++++++++++--------- .../transfer_service/transfer_service.h | 4 + 2 files changed, 44 insertions(+), 35 deletions(-) diff --git a/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp b/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp index d01a863..5409fc6 100644 --- a/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp +++ b/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp @@ -633,48 +633,15 @@ void TransferService::ProcessIncomingData(int client_fd) { << ": Received header type=" << static_cast(header.type) << ", dest_obj_id=" << header.dest_obj_id; - bool keep_connection = true; + bool keep_connection = false; try { - switch (header.type) { - case MessageType::kPutObj: - HandleDataReceive(client_fd, header, false); - break; - case MessageType::kGetObj: - HandleGetObjRequest(client_fd, header); - break; - case MessageType::kRespondToGetObj: - // ExecuteGetTask reads kGetObj responses on the connection that sent - // the request. A response arriving on the listener is a protocol - // violation. Do not write its payload to disk. - LOG(ERROR) << "Worker fd=" << client_fd - << ": Unsolicited kRespondToGetObj received for task " - << header.task_id << ". Closing connection."; - keep_connection = false; - break; - case MessageType::kAck: - LOG(ERROR) << "Worker fd=" << client_fd - << ": Unexpected kAck message received for task " - << header.task_id; - break; - case MessageType::kError: - LOG(ERROR) << "Worker fd=" << client_fd - << ": Error message received for task " << header.task_id; - ReportResult(header.task_id, false, "Received error message"); - break; - default: - // We cannot parse this header, so the stream is out of sync. - LOG(ERROR) << "Worker fd=" << client_fd << ": Unknown message type " - << static_cast(header.type) << ". Closing connection."; - keep_connection = false; - break; - } + keep_connection = DispatchMessage(client_fd, header); } catch (const std::exception& e) { // A handler that throws leaves the connection in an unknown state. Close // it; otherwise it would stay registered in epoll but never be re-armed. LOG(ERROR) << "Worker fd=" << client_fd << ": Exception while handling message type=" << static_cast(header.type) << ": " << e.what(); - keep_connection = false; } if (!keep_connection) { @@ -694,6 +661,44 @@ void TransferService::ProcessIncomingData(int client_fd) { } } +bool TransferService::DispatchMessage(int client_fd, + const ObjInfoHeader& header) { + // No default case on purpose: with -Wswitch the compiler reports a new + // MessageType that is not handled here. + switch (header.type) { + case MessageType::kPutObj: + HandleDataReceive(client_fd, header, false); + return true; + case MessageType::kGetObj: + HandleGetObjRequest(client_fd, header); + return true; + case MessageType::kRespondToGetObj: + // ExecuteGetTask reads kGetObj responses on the connection that sent + // the request. A response arriving on the listener is a protocol + // violation. Do not write its payload to disk. + LOG(ERROR) << "Worker fd=" << client_fd + << ": Unsolicited kRespondToGetObj received for task " + << header.task_id << ". Closing connection."; + return false; + case MessageType::kAck: + LOG(ERROR) << "Worker fd=" << client_fd + << ": Unexpected kAck message received for task " + << header.task_id; + return true; + case MessageType::kError: + LOG(ERROR) << "Worker fd=" << client_fd + << ": Error message received for task " << header.task_id; + ReportResult(header.task_id, false, "Received error message"); + return true; + } + // The type byte comes from the network, so it can hold any uint8_t value, + // not only the enumerators above. We cannot parse this header, so the + // stream is out of sync. + LOG(ERROR) << "Worker fd=" << client_fd << ": Unknown message type " + << static_cast(header.type) << ". Closing connection."; + return false; +} + void TransferService::ExecutePutTask(PutTask* task) { auto& metric_container = task->GetMetricContainer(); metric_container.start_execution_time = absl::Now(); diff --git a/src/ml_flashpoint/replication/transfer_service/transfer_service.h b/src/ml_flashpoint/replication/transfer_service/transfer_service.h index 5b714cb..81e3226 100644 --- a/src/ml_flashpoint/replication/transfer_service/transfer_service.h +++ b/src/ml_flashpoint/replication/transfer_service/transfer_service.h @@ -135,6 +135,10 @@ class TransferService final { void ProcessEpollEventsLoop(); void HandleNewConnection(); void ProcessIncomingData(int client_fd); + // Routes one received header to its handler. Returns true when the + // connection can carry another message afterwards, false when it must be + // closed. + bool DispatchMessage(int client_fd, const ObjInfoHeader& header); void RemoveClient(int client_fd); std::shared_ptr GetOrCreateConnectionPool( const std::string& peer_addr); From 3c0c914e4cd736758bbf069f2e5c6b4e2d528055 Mon Sep 17 00:00:00 2001 From: Husam Date: Tue, 6 Oct 2026 15:21:18 +0000 Subject: [PATCH 10/15] fix(replication): close the connection on any reply that reaches the listener The listener only receives requests. Every reply (kRespondToGetObj, kAck, kError) is read by the worker that owns the request socket, so a reply arriving on the listener is a protocol violation. Treat all three the same way: log and close the connection. Before, kAck was logged and the connection kept, and kError called ReportResult(task_id, false). That let any peer that could reach the listener fail an in-flight local task by naming its task_id; the responder learns that id from the request header. Tests: UnsolicitedAckOnListenerClosesConnection and UnsolicitedErrorOnListenerDoesNotFailLocalTask. The second one forges a kError for a real in-flight Get and checks that the Get still completes on the request socket. Both fail against the previous code. --- .../transfer_service/transfer_service.cpp | 25 +++--- .../transfer_service_p2p_test.cpp | 86 +++++++++++++++++++ 2 files changed, 96 insertions(+), 15 deletions(-) diff --git a/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp b/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp index 5409fc6..18997d5 100644 --- a/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp +++ b/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp @@ -673,23 +673,18 @@ bool TransferService::DispatchMessage(int client_fd, HandleGetObjRequest(client_fd, header); return true; case MessageType::kRespondToGetObj: - // ExecuteGetTask reads kGetObj responses on the connection that sent - // the request. A response arriving on the listener is a protocol - // violation. Do not write its payload to disk. - LOG(ERROR) << "Worker fd=" << client_fd - << ": Unsolicited kRespondToGetObj received for task " - << header.task_id << ". Closing connection."; - return false; case MessageType::kAck: - LOG(ERROR) << "Worker fd=" << client_fd - << ": Unexpected kAck message received for task " - << header.task_id; - return true; case MessageType::kError: - LOG(ERROR) << "Worker fd=" << client_fd - << ": Error message received for task " << header.task_id; - ReportResult(header.task_id, false, "Received error message"); - return true; + // The listener only receives requests. Replies are read by the worker + // that owns the request socket (ExecutePutTask, ExecuteGetTask, and + // ExecuteRespondToGetTask for the final ACK). A reply arriving here is a + // protocol violation, so close the connection and do not act on it: a + // kRespondToGetObj payload must not be written to disk, and a kError + // must not fail a local task by task_id. + LOG(ERROR) << "Worker fd=" << client_fd << ": Unsolicited reply (type " + << static_cast(header.type) << ") received for task " + << header.task_id << ". Closing connection."; + return false; } // The type byte comes from the network, so it can hold any uint8_t value, // not only the enumerators above. We cannot parse this header, so the diff --git a/tests/replication/transfer_service/transfer_service_p2p_test.cpp b/tests/replication/transfer_service/transfer_service_p2p_test.cpp index 5456465..c0f4d76 100644 --- a/tests/replication/transfer_service/transfer_service_p2p_test.cpp +++ b/tests/replication/transfer_service/transfer_service_p2p_test.cpp @@ -1404,6 +1404,92 @@ TEST(TransferServiceP2PTest, UnknownMessageTypeClosesConnection) { service1.Shutdown(); } +// The listener only receives requests. A bare kAck is a reply that no worker +// on this side is waiting for, so the service closes the connection. +TEST(TransferServiceP2PTest, UnsolicitedAckOnListenerClosesConnection) { + // Given + TransferService service1("127.0.0.1"); + int port1 = service1.Initialize(/*listen_port=*/0, /*threads=*/4, + /*conn_pool_per_peer=*/1); + ASSERT_GT(port1, 0); + int client_fd = ConnectToLocalPort(port1); + ASSERT_GE(client_fd, 0); + ObjInfoHeader header; + header.type = MessageType::kAck; + snprintf(header.task_id, sizeof(header.task_id), "unsolicited_ack_task"); + + // When + ASSERT_TRUE(SendAll(client_fd, &header, kHeaderSize).ok()); + + // Then + EXPECT_TRUE(PeerClosed(client_fd, /*timeout_ms=*/5000)); + + close(client_fd); + service1.Shutdown(); +} + +// A kError that arrives on the listener must not fail a local task, even when +// it names the task_id of a Get that is in flight. Only the reply read on the +// request socket may complete a task. The service closes the rogue connection +// and the Get still completes. +TEST(TransferServiceP2PTest, UnsolicitedErrorOnListenerDoesNotFailLocalTask) { + // Given + const std::string expected_data = "payload delivered after a forged kError"; + const std::string dest_obj_id = "forged_error_dest_obj"; + std::remove(dest_obj_id.c_str()); + int peer_port = 0; + int peer_listener = ListenOnLocalEphemeralPort(&peer_port); + ASSERT_GE(peer_listener, 0); + TransferService service1("127.0.0.1"); + int port1 = service1.Initialize(/*listen_port=*/0, /*threads=*/2, + /*conn_pool_per_peer=*/1); + ASSERT_GT(port1, 0); + std::future actual_future = service1.AsyncGet( + "any_source_obj", "127.0.0.1:" + std::to_string(peer_port), dest_obj_id); + // Act as the peer: accept the request and learn its task_id. + pollfd peer_pfd{peer_listener, POLLIN, 0}; + ASSERT_GT(poll(&peer_pfd, 1, /*timeout_ms=*/5000), 0); + int request_fd = accept(peer_listener, nullptr, nullptr); + ASSERT_GE(request_fd, 0); + ObjInfoHeader request; + ASSERT_TRUE(RecvHeaderWithin(request_fd, &request, /*timeout_ms=*/5000)); + ASSERT_EQ(request.type, MessageType::kGetObj); + + // When: send kError for that task on a new connection to the listener. + int rogue_fd = ConnectToLocalPort(port1); + ASSERT_GE(rogue_fd, 0); + ObjInfoHeader forged_error; + forged_error.type = MessageType::kError; + snprintf(forged_error.task_id, sizeof(forged_error.task_id), "%s", + request.task_id); + ASSERT_TRUE(SendAll(rogue_fd, &forged_error, kHeaderSize).ok()); + + // Then: the rogue connection is closed and the task is still pending. + EXPECT_TRUE(PeerClosed(rogue_fd, /*timeout_ms=*/5000)); + EXPECT_FALSE( + FutureReadyWithin(actual_future, std::chrono::milliseconds(300))); + + // The real reply on the request socket completes the task. + ObjInfoHeader response; + response.type = MessageType::kRespondToGetObj; + response.obj_size = expected_data.size(); + snprintf(response.task_id, sizeof(response.task_id), "%s", request.task_id); + ASSERT_TRUE(SendAll(request_fd, &response, kHeaderSize).ok()); + ASSERT_TRUE( + SendAll(request_fd, expected_data.data(), expected_data.size()).ok()); + ObjInfoHeader ack; + ASSERT_TRUE(RecvHeaderWithin(request_fd, &ack, /*timeout_ms=*/5000)); + EXPECT_EQ(ack.type, MessageType::kAck); + ASSERT_TRUE(FutureReadyWithin(actual_future, std::chrono::seconds(5))); + EXPECT_TRUE(actual_future.get().success); + VerifyFileContentAndRemove(dest_obj_id, expected_data); + + close(rogue_fd); + close(request_fd); + close(peer_listener); + service1.Shutdown(); +} + // Shutdown() must not crash on a RespondToGetTask (which has no promise) and // must not hang while a responder waits for the requester's final ACK on an // accepted socket. From f13b95c2000a236aa44d9f977be5d10a74d51e7d Mon Sep 17 00:00:00 2001 From: Husam Date: Wed, 7 Oct 2026 16:54:17 +0000 Subject: [PATCH 11/15] fix(replication): destroy connection pools only after the workers are joined Shutdown() cleared connection_pools_ before it stopped the thread pools. A worker that was still inside a Get or Put held a ScopedConnection, which returns its socket to the pool through a raw pointer when it goes out of scope. With the pool already destroyed, that release was a heap-use-after-free. Shutdown() still shuts the pools down first, so blocked workers wake up and give up, but it now destroys the pools only after both thread pools have joined their workers. Verified with an ASAN build of transfer_service_p2p_test.cpp (-fsanitize=address): ShutdownWithQueuedTasksAgainstSilentPeerCompletes reports the use-after-free in ConnectionPool::ReleaseConnection before this change; all four *Shutdown* tests are clean for 10 iterations after it. Without ASAN the bug has no deterministic symptom, so there is no new plain test. --- .../transfer_service/transfer_service.cpp | 26 +++++++++++++------ 1 file changed, 18 insertions(+), 8 deletions(-) diff --git a/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp b/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp index 18997d5..1d82e3b 100644 --- a/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp +++ b/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp @@ -228,7 +228,11 @@ void TransferService::Shutdown() { } } - // 7. Clean up connection pools. + // 7. Shut down the connection pools. This wakes workers that are waiting + // for a connection and shuts down the sockets that workers are using, so + // their pending send/recv fail. The pools must stay alive until step 8 has + // joined the workers: a worker that holds a ScopedConnection returns it to + // its pool through a raw pointer when it finishes. Step 9 destroys them. { std::unique_lock write_lock(connection_pools_mutex_); for (auto const& [peer_addr, pool] : connection_pools_) { @@ -236,10 +240,10 @@ void TransferService::Shutdown() { pool->Shutdown(); } } - connection_pools_.clear(); } - // 8. Stop the thread pools. + // 8. Stop the thread pools. This runs the tasks that are still queued and + // joins every worker. if (thread_pool_) { thread_pool_->stop(); } @@ -247,7 +251,13 @@ void TransferService::Shutdown() { epoll_thread_pool_->stop(); } - // 9. Close any accepted client sockets that are still open. + // 9. Destroy the connection pools. No worker can touch them anymore. + { + std::unique_lock write_lock(connection_pools_mutex_); + connection_pools_.clear(); + } + + // 10. Close any accepted client sockets that are still open. { std::lock_guard lock(client_fds_mutex_); for (int fd : client_fds_) { @@ -256,7 +266,7 @@ void TransferService::Shutdown() { client_fds_.clear(); } - // 10. Clean up epoll fd. + // 11. Clean up epoll fd. if (epoll_fd_ != -1) { if (close(epoll_fd_) == -1) { PLOG(WARNING) << "Failed to close epoll_fd"; @@ -1073,9 +1083,9 @@ std::shared_ptr TransferService::GetOrCreateConnectionPool( return it->second; // Double check } - // Shutdown() sets running_ to false before it shuts down and clears the - // pools. ThreadPool::stop() then runs every task that is still queued. If one - // of those tasks created a pool here, nothing would ever shut it down, and a + // Shutdown() sets running_ to false before it shuts down the pools. + // ThreadPool::stop() then runs every task that is still queued. If one of + // those tasks created a pool here, nothing would ever shut it down, and a // task blocked on a peer that never answers would block Shutdown() forever. if (!running_.load()) { LOG(WARNING) << "Not creating connection pool for " << peer_addr From 88e4cef838eb02008ae388b0ede9d5a258248994 Mon Sep 17 00:00:00 2001 From: Husam Date: Wed, 7 Oct 2026 18:07:31 +0000 Subject: [PATCH 12/15] fix(replication): keep remotely triggered work out of pending_tasks_ HandleGetObjRequest registered the RespondToGetTask in pending_tasks_ under the task_id from the request header, and ExecuteRespondToGetTask reported its outcome through ReportResult with that id. Both are the peer's choice, and the peer learns our task ids from the kGetObj requests we send it. A peer that sent kGetObj naming the id of a local Get in flight overwrote that task's entry, which dropped its promise: the caller's future failed with std::future_error (broken_promise) and the real reply later found no task to complete. The same path also broke a service that fetched an object from itself. RespondToGetTask now stays out of pending_tasks_ entirely. It logs its outcome from its own metric container through LogTaskResult, which is the same log line ReportResult writes, so the timing logs keep their format (TimestampsAreRecorded still parses three of them). The null-promise guard in Shutdown() goes away with the only null entry. Test: GetRequestNamingLocalTaskIdDoesNotTouchLocalTask sends a kGetObj for an existing object naming an in-flight Get's task_id, checks that the request is served, that the Get's future stays pending, and that the real reply then completes it. Fails against the previous code with "Broken promise". --- .../transfer_service/transfer_service.cpp | 59 +++++++-------- .../transfer_service_p2p_test.cpp | 72 ++++++++++++++++++- 2 files changed, 101 insertions(+), 30 deletions(-) diff --git a/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp b/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp index 1d82e3b..1777828 100644 --- a/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp +++ b/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp @@ -183,8 +183,6 @@ void TransferService::Shutdown() { { std::lock_guard lock(pending_tasks_mutex_); for (auto const& [task_id, context] : pending_tasks_) { - // Remotely triggered tasks (RespondToGetTask) have no promise. - if (!context.promise) continue; try { context.promise->set_exception(std::make_exception_ptr( std::runtime_error("Service is shutting down"))); @@ -503,21 +501,28 @@ void TransferService::ProcessEpollEventsLoop() { LOG(INFO) << "Epoll processing loop finished."; } +// Logs the outcome of a task. Local tasks (Put, Get) and remotely triggered +// ones (RespondToGet) share this format, so log consumers parse them alike. +static void LogTaskResult(const std::string& task_id, + const TaskMetricContainer& metrics, bool success, + const std::string& message) { + LOG(INFO) << "TransferService::report_result: task_id=" << task_id + << ", task_type=" << TaskTypeToString(metrics.task_type) + << ", data_size=" << metrics.data_size << ", success=" << success + << ", message=" << message << ", timing=" << metrics.ToString(); +} + void TransferService::ReportResult(std::string task_id, bool success, const std::string& message) { std::shared_ptr> promise; - std::string timing_message = ""; - std::string task_type = ""; - size_t data_size = 0; + std::shared_ptr metric_container; { std::lock_guard lock(pending_tasks_mutex_); auto it = pending_tasks_.find(task_id); if (it != pending_tasks_.end()) { promise = it->second.promise; - timing_message = it->second.metric_container->ToString(); - task_type = TaskTypeToString(it->second.metric_container->task_type); - data_size = it->second.metric_container->data_size; + metric_container = it->second.metric_container; pending_tasks_.erase(it); } else { LOG(WARNING) << "No promise found for task_id: " << task_id; @@ -525,10 +530,7 @@ void TransferService::ReportResult(std::string task_id, bool success, } } - LOG(INFO) << "TransferService::report_result: task_id=" << task_id - << ", task_type=" << task_type << ", data_size=" << data_size - << ", success=" << success << ", message=" << message - << ", timing=" << timing_message; + LogTaskResult(task_id, *metric_container, success, message); if (promise) { TransferResult result; @@ -961,10 +963,18 @@ void TransferService::ExecuteRespondToGetTask(RespondToGetTask* task) { << ", dest_obj_id=" << task->GetDestObjId() << ", client_fd=" << task->GetClientFd(); + // The task_id is the peer's, and the peer learned ours from the request + // header, so it may be the id of a local Get that is in flight. This task + // therefore has no entry in pending_tasks_ and must not call ReportResult: + // that would complete, or drop, the local task's promise. + auto finish = [&](bool success, const std::string& message) { + LogTaskResult(task->GetTaskId(), metric_container, success, message); + }; + int sockfd = task->GetClientFd(); if (sockfd < 0) { LOG(ERROR) << "Invalid client_fd!"; - ReportResult(task->GetTaskId(), false, "Invalid client_fd"); + finish(false, "Invalid client_fd"); return; } metric_container.connection_acquired_time = absl::Now(); @@ -981,7 +991,7 @@ void TransferService::ExecuteRespondToGetTask(RespondToGetTask* task) { << task->GetSourceObjId() << "': " << e.what(); SendErrorResponse(sockfd, task->GetTaskId().c_str(), task->GetSourceObjId().c_str()); - ReportResult(task->GetTaskId(), false, "Failed to open buffer object"); + finish(false, "Failed to open buffer object"); return; } void* buffer_data_ptr = buffer_obj->get_data_ptr(); @@ -1007,14 +1017,14 @@ void TransferService::ExecuteRespondToGetTask(RespondToGetTask* task) { if (!SendAll(sockfd, &header, kHeaderSize).ok()) { LOG(ERROR) << "Failed to send kRespondToGetObj header"; - ReportResult(task->GetTaskId(), false, "Failed to send header"); + finish(false, "Failed to send header"); return; } metric_container.header_sent_time = absl::Now(); if (!SendAll(sockfd, buffer_data_ptr, size).ok()) { LOG(ERROR) << "Failed to send buffer data"; - ReportResult(task->GetTaskId(), false, "Failed to send data"); + finish(false, "Failed to send data"); return; } metric_container.data_sent_time = absl::Now(); @@ -1022,18 +1032,17 @@ void TransferService::ExecuteRespondToGetTask(RespondToGetTask* task) { ObjInfoHeader ack_header; if (!RecvHeader(sockfd, ack_header).ok()) { LOG(ERROR) << "Failed to receive ACK"; - ReportResult(task->GetTaskId(), false, "Failed to receive ACK"); + finish(false, "Failed to receive ACK"); return; } if (ack_header.type != MessageType::kAck) { LOG(ERROR) << "Failed to receive ACK for RespondToGetObj"; - ReportResult(task->GetTaskId(), false, "Received unexpected ACK"); + finish(false, "Received unexpected ACK"); return; } metric_container.finish_time = absl::Now(); - ReportResult(task->GetTaskId(), true, - "RespondToGetTask completed successfully"); + finish(true, "RespondToGetTask completed successfully"); } std::optional TransferService::GetConnectionFromPool( @@ -1136,14 +1145,8 @@ void TransferService::HandleGetObjRequest(int client_fd, RespondToGetTask task(header.task_id, header.source_obj_id, header.dest_obj_id, client_fd, metric_container); - - { - std::lock_guard lock(pending_tasks_mutex_); - // We don't have a promise for RespondToGetTask as it's triggered - // remotely, but we want to track it. - pending_tasks_[task.GetTaskId()] = {nullptr, metric_container}; - } - + // Not registered in pending_tasks_: the task_id is the peer's and could + // name a local task. ExecuteRespondToGetTask logs its own outcome. task.Execute(this); } diff --git a/tests/replication/transfer_service/transfer_service_p2p_test.cpp b/tests/replication/transfer_service/transfer_service_p2p_test.cpp index c0f4d76..ef75656 100644 --- a/tests/replication/transfer_service/transfer_service_p2p_test.cpp +++ b/tests/replication/transfer_service/transfer_service_p2p_test.cpp @@ -1490,8 +1490,76 @@ TEST(TransferServiceP2PTest, UnsolicitedErrorOnListenerDoesNotFailLocalTask) { service1.Shutdown(); } -// Shutdown() must not crash on a RespondToGetTask (which has no promise) and -// must not hang while a responder waits for the requester's final ACK on an +// A kGetObj that reaches the listener naming the task_id of a local Get in +// flight is served like any other request and must not touch the local task. +// The responder runs under the peer's task_id, so it has to stay out of the +// local task bookkeeping: the local future stays pending until the real reply +// arrives on the request socket. +TEST(TransferServiceP2PTest, GetRequestNamingLocalTaskIdDoesNotTouchLocalTask) { + // Given + const std::string expected_data = "payload delivered after a colliding Get"; + const std::string expected_served_data = "object served to the rogue peer"; + const std::string dest_obj_id = "colliding_get_dest_obj"; + const std::string served_obj_id = "colliding_get_served_source_obj"; + std::remove(dest_obj_id.c_str()); + WriteTestFile(served_obj_id, expected_served_data); + int peer_port = 0; + int peer_listener = ListenOnLocalEphemeralPort(&peer_port); + ASSERT_GE(peer_listener, 0); + TransferService service1("127.0.0.1"); + int port1 = service1.Initialize(/*listen_port=*/0, /*threads=*/2, + /*conn_pool_per_peer=*/1); + ASSERT_GT(port1, 0); + std::future actual_future = service1.AsyncGet( + "any_source_obj", "127.0.0.1:" + std::to_string(peer_port), dest_obj_id); + // Act as the peer: accept the request and learn its task_id. + pollfd peer_pfd{peer_listener, POLLIN, 0}; + ASSERT_GT(poll(&peer_pfd, 1, /*timeout_ms=*/5000), 0); + int request_fd = accept(peer_listener, nullptr, nullptr); + ASSERT_GE(request_fd, 0); + ObjInfoHeader request; + ASSERT_TRUE(RecvHeaderWithin(request_fd, &request, /*timeout_ms=*/5000)); + ASSERT_EQ(request.type, MessageType::kGetObj); + + // When: send a kGetObj for that task_id on a new connection to the listener. + int rogue_fd = ConnectToLocalPort(port1); + ASSERT_GE(rogue_fd, 0); + ObjInfoHeader actual_response = SendRequestAndRecvResponse( + rogue_fd, BuildGetRequest(request.task_id, served_obj_id, + "colliding_get_rogue_dest_obj", "")); + std::string actual_served_data = + RecvGetPayloadAndAck(rogue_fd, actual_response); + + // Then: the rogue request is served and the local task is still pending. + EXPECT_EQ(actual_response.type, MessageType::kRespondToGetObj); + EXPECT_EQ(actual_served_data, expected_served_data); + EXPECT_FALSE( + FutureReadyWithin(actual_future, std::chrono::milliseconds(300))); + + // The real reply on the request socket completes the task. + ObjInfoHeader response; + response.type = MessageType::kRespondToGetObj; + response.obj_size = expected_data.size(); + snprintf(response.task_id, sizeof(response.task_id), "%s", request.task_id); + ASSERT_TRUE(SendAll(request_fd, &response, kHeaderSize).ok()); + ASSERT_TRUE( + SendAll(request_fd, expected_data.data(), expected_data.size()).ok()); + ObjInfoHeader ack; + ASSERT_TRUE(RecvHeaderWithin(request_fd, &ack, /*timeout_ms=*/5000)); + EXPECT_EQ(ack.type, MessageType::kAck); + ASSERT_TRUE(FutureReadyWithin(actual_future, std::chrono::seconds(5))); + EXPECT_TRUE(actual_future.get().success); + VerifyFileContentAndRemove(dest_obj_id, expected_data); + + close(rogue_fd); + close(request_fd); + close(peer_listener); + std::remove(served_obj_id.c_str()); + service1.Shutdown(); +} + +// Shutdown() must not crash while a RespondToGetTask is in flight and must +// not hang while that responder waits for the requester's final ACK on an // accepted socket. TEST(TransferServiceP2PTest, ShutdownWhileResponderWaitsForAckIsSafe) { // Given From a47146d5d8e8f1b15a78f6bd230460955d267781 Mon Sep 17 00:00:00 2001 From: Husam Date: Wed, 7 Oct 2026 18:39:32 +0000 Subject: [PATCH 13/15] fix(replication): zero-initialize every byte of ObjInfoHeader The header constructor set only the first byte of each char array. snprintf writes up to the terminator, and SendAll ships all of kHeaderSize, so the bytes after each terminator went on the wire as whatever the worker's stack held before: pointers, and fragments of headers from other peers. Every kError, kAck and kGetObj header and the address fields of every kRespondToGetObj header leaked this way. Default member initializers now zero the whole struct. It stays packed, and its size does not change. Test: ReplyHeadersCarryNoBytesAfterTerminators reads a kError and a kRespondToGetObj reply off a raw socket and checks that every byte after each field's terminator is zero. Fails against the previous code on all five fields of the kError reply and on both address fields of the kRespondToGetObj reply. --- .../replication/transfer_service/protocol.h | 25 ++++----- .../transfer_service_p2p_test.cpp | 51 +++++++++++++++++++ 2 files changed, 62 insertions(+), 14 deletions(-) diff --git a/src/ml_flashpoint/replication/transfer_service/protocol.h b/src/ml_flashpoint/replication/transfer_service/protocol.h index 2806aa4..07ff9a8 100644 --- a/src/ml_flashpoint/replication/transfer_service/protocol.h +++ b/src/ml_flashpoint/replication/transfer_service/protocol.h @@ -41,22 +41,19 @@ enum class MessageType : uint8_t { // string null terminator character. constexpr char kNullTerm = '\0'; +// The whole struct goes on the wire, so every byte is zero-initialized. The +// bytes after each string's terminator are sent too; if they were left +// uninitialized, every message would leak the worker's stack to the peer. struct ObjInfoHeader { - MessageType type; - char source_obj_id[1024]; - char dest_obj_id[1024]; - char source_address[64]; - char dest_address[64]; - char task_id[64]; - ssize_t obj_size; + MessageType type = MessageType::kAck; + char source_obj_id[1024] = {}; + char dest_obj_id[1024] = {}; + char source_address[64] = {}; + char dest_address[64] = {}; + char task_id[64] = {}; + ssize_t obj_size = 0; - ObjInfoHeader() : type(MessageType::kAck), obj_size(0) { - source_obj_id[0] = kNullTerm; - dest_obj_id[0] = kNullTerm; - source_address[0] = kNullTerm; - dest_address[0] = kNullTerm; - task_id[0] = kNullTerm; - } + ObjInfoHeader() = default; // Ensure all char arrays are null-terminated to prevent buffer over-read // segfaults when they are converted to std::string or used in logging. diff --git a/tests/replication/transfer_service/transfer_service_p2p_test.cpp b/tests/replication/transfer_service/transfer_service_p2p_test.cpp index ef75656..5f6ff34 100644 --- a/tests/replication/transfer_service/transfer_service_p2p_test.cpp +++ b/tests/replication/transfer_service/transfer_service_p2p_test.cpp @@ -1095,6 +1095,16 @@ ObjInfoHeader BuildPutRequest(const std::string& task_id, return request; } +// Returns true if every byte after the first NUL in `field` is zero, that is, +// if the field carries nothing but its string. +template +bool TailIsZero(const char (&field)[N]) { + for (size_t i = strnlen(field, N); i < N; ++i) { + if (field[i] != '\0') return false; + } + return true; +} + // A fake peer. It accepts one connection, reads one kGetObj request, and // answers with whatever `respond` writes to the socket. Every wait has a // timeout, so the helper thread always exits. @@ -1913,6 +1923,47 @@ INSTANTIATE_TEST_SUITE_P(TransferServiceP2PTest, InvalidGetRequestTest, ::testing::Values("", "non_existent_source_obj", "{unterminated}")); +// Headers go on the wire whole, so the bytes after each string's terminator +// must be zero in every reply. Otherwise a reply would carry whatever the +// worker's stack held before, such as fragments of other peers' headers. +TEST(TransferServiceP2PTest, ReplyHeadersCarryNoBytesAfterTerminators) { + // Given + TransferService service1("127.0.0.1"); + int port1 = service1.Initialize(/*listen_port=*/0, /*threads=*/2, + /*conn_pool_per_peer=*/1); + ASSERT_GT(port1, 0); + const std::string source_obj_id = "clean_header_source_obj"; + WriteTestFile(source_obj_id, "clean header payload"); + int client_fd = ConnectToLocalPort(port1); + ASSERT_GE(client_fd, 0); + + // When: one request that fails (kError) and one that succeeds + // (kRespondToGetObj), on the same socket. + ObjInfoHeader actual_error = SendRequestAndRecvResponse( + client_fd, + BuildGetRequest("clean_header_task", "missing_clean_header_obj", + "clean_header_dest_obj", "")); + ObjInfoHeader actual_response = SendRequestAndRecvResponse( + client_fd, BuildGetRequest("clean_header_task", source_obj_id, + "clean_header_dest_obj", "")); + RecvGetPayloadAndAck(client_fd, actual_response); + + // Then + ASSERT_EQ(actual_error.type, MessageType::kError); + ASSERT_EQ(actual_response.type, MessageType::kRespondToGetObj); + for (const ObjInfoHeader* header : {&actual_error, &actual_response}) { + EXPECT_TRUE(TailIsZero(header->source_obj_id)); + EXPECT_TRUE(TailIsZero(header->dest_obj_id)); + EXPECT_TRUE(TailIsZero(header->source_address)); + EXPECT_TRUE(TailIsZero(header->dest_address)); + EXPECT_TRUE(TailIsZero(header->task_id)); + } + + close(client_fd); + std::remove(source_obj_id.c_str()); + service1.Shutdown(); +} + // When the receiver cannot create the destination of an inbound kPutObj, it // reads and discards the payload the sender is already streaming and answers // kError. The connection stays in sync and serves the next request normally. From 0bff102826efaba7f13343c0ae066941410e0d87 Mon Sep 17 00:00:00 2001 From: Husam Date: Wed, 7 Oct 2026 18:45:54 +0000 Subject: [PATCH 14/15] fix(replication): remove the temporary file when an inbound transfer fails HandleDataReceive writes an incoming object to .tmp and renames it at the end. When the receive or the rename failed, or when BufferObject's constructor threw after it had created the file, the temporary file was left behind. obj_size comes from the peer, so a kPutObj that announces a large size and then stops leaves a sparse file of that apparent size on disk, and every failed Get left its destination's .tmp behind. The failure paths after the file is created now close and remove it. The rename error text is captured before the removal so errno is not overwritten. Tests: TruncatedPutLeavesNoTemporaryFile (kPutObj announcing 16 MiB, 16 bytes sent, then EOF; expects kError and no file). The cleanup lines in PutFailsInRenameWhenTargetExistAsADirectory, GetFailsInRenameWhenTargetExistAsADirectory and MalformedGetResponseTest become assertions that the temporary file is gone. All of these fail against the previous code. --- .../transfer_service/transfer_service.cpp | 15 +++++- .../transfer_service_p2p_test.cpp | 48 +++++++++++++++++-- 2 files changed, 58 insertions(+), 5 deletions(-) diff --git a/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp b/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp index 1777828..b6ef084 100644 --- a/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp +++ b/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp @@ -813,9 +813,19 @@ void TransferService::HandleDataReceive(int client_fd, std::string(header.dest_obj_id) + std::string(kTempFileSuffix); // BufferObject's constructor throws when it cannot create the file. std::optional buffer_obj; + // Every failure from here on removes the temporary file. obj_size is the + // peer's number, so a transfer that stops early would otherwise leave a + // sparse file of that apparent size behind. The constructor can also throw + // after it created the file (for example when mmap fails), so this runs on + // that path too; removing a file that was never created is harmless. + auto remove_tmp = [&]() { + buffer_obj.reset(); + std::remove(tmp_obj_id.c_str()); + }; try { buffer_obj.emplace(tmp_obj_id, header.obj_size, /*overwrite=*/true); } catch (const std::exception& e) { + remove_tmp(); // The sender is already streaming obj_size bytes on this connection. Read // and discard them so the connection stays in sync and can be reused (the // pool never replaces a closed connection). If the drain fails, the peer @@ -828,6 +838,7 @@ void TransferService::HandleDataReceive(int client_fd, void* receiver_data_ptr = buffer_obj->get_data_ptr(); if (!RecvAll(client_fd, receiver_data_ptr, header.obj_size).ok()) { + remove_tmp(); fail("Failed to receive data"); return; } @@ -838,8 +849,10 @@ void TransferService::HandleDataReceive(int client_fd, // Rename the temporary file to the final destination. if (rename(tmp_obj_id.c_str(), header.dest_obj_id) != 0) { + const std::string rename_error = std::strerror(errno); + remove_tmp(); fail("Failed to rename temporary file " + tmp_obj_id + " to " + - header.dest_obj_id + ": " + std::strerror(errno)); + header.dest_obj_id + ": " + rename_error); return; } if (is_respond_get_task) { diff --git a/tests/replication/transfer_service/transfer_service_p2p_test.cpp b/tests/replication/transfer_service/transfer_service_p2p_test.cpp index 5f6ff34..758a571 100644 --- a/tests/replication/transfer_service/transfer_service_p2p_test.cpp +++ b/tests/replication/transfer_service/transfer_service_p2p_test.cpp @@ -711,8 +711,8 @@ TEST(TransferServiceP2PTest, PutFailsInRenameWhenTargetExistAsADirectory) { testing::HasSubstr("Received error from destination")); } - // Cleanup - remove((obj_id + ".tmp").c_str()); // Remove temporary file if created + // The receiver removes its temporary file when the rename fails. + EXPECT_FALSE(std::filesystem::exists(obj_id + ".tmp")); rmdir(target_dir.c_str()); service1.Shutdown(); @@ -845,7 +845,8 @@ TEST(TransferServiceP2PTest, GetFailsInRenameWhenTargetExistAsADirectory) { // Cleanup std::remove(obj_id.c_str()); - remove((dest_obj_id + ".tmp").c_str()); + // The requester removes its temporary file when the rename fails. + EXPECT_FALSE(std::filesystem::exists(dest_obj_id + ".tmp")); rmdir(target_dir.c_str()); service1.Shutdown(); @@ -1964,6 +1965,44 @@ TEST(TransferServiceP2PTest, ReplyHeadersCarryNoBytesAfterTerminators) { service1.Shutdown(); } +// A kPutObj whose sender stops before obj_size bytes arrived must not leave +// the temporary file behind. obj_size is the peer's number, so without this +// cleanup a peer could leave a sparse file of any apparent size on disk. +TEST(TransferServiceP2PTest, TruncatedPutLeavesNoTemporaryFile) { + // Given + TransferService service1("127.0.0.1"); + int port1 = service1.Initialize(/*listen_port=*/0, /*threads=*/2, + /*conn_pool_per_peer=*/1); + ASSERT_GT(port1, 0); + const std::string dest_obj_id = "truncated_put_dest_obj"; + const std::string tmp_obj_id = dest_obj_id + ".tmp"; + int client_fd = ConnectToLocalPort(port1); + ASSERT_GE(client_fd, 0); + ObjInfoHeader request = BuildPutRequest("truncated_put_task", dest_obj_id, + /*obj_size=*/16 * 1024 * 1024); + const std::string partial(16, 'p'); + + // When: announce 16 MiB, send 16 bytes, then signal EOF. The read side + // stays open so the receiver's kError can be observed. + ASSERT_TRUE(SendAll(client_fd, &request, kHeaderSize).ok()); + ASSERT_TRUE(SendAll(client_fd, partial.data(), partial.size()).ok()); + ASSERT_EQ(shutdown(client_fd, SHUT_WR), 0); + ObjInfoHeader actual_response; + const bool actual_replied = + RecvHeaderWithin(client_fd, &actual_response, /*timeout_ms=*/5000); + + // Then: the receiver reports the failure and leaves no file behind. The + // temporary file is removed before kError is sent, so no waiting is needed. + ASSERT_TRUE(actual_replied); + EXPECT_EQ(actual_response.type, MessageType::kError); + EXPECT_FALSE(std::filesystem::exists(tmp_obj_id)); + EXPECT_FALSE(std::filesystem::exists(dest_obj_id)); + + close(client_fd); + std::remove(tmp_obj_id.c_str()); // Keep the directory clean on failure. + service1.Shutdown(); +} + // When the receiver cannot create the destination of an inbound kPutObj, it // reads and discards the payload the sender is already streaming and answers // kError. The connection stays in sync and serves the next request normally. @@ -2161,8 +2200,9 @@ TEST_P(MalformedGetResponseTest, RequesterFutureFailsInsteadOfHanging) { EXPECT_THROW(actual_future.get(), std::runtime_error); EXPECT_TRUE(responder.received_request()); EXPECT_FALSE(std::filesystem::exists(dest_obj_id)); + // No temporary file either, including after a truncated payload. + EXPECT_FALSE(std::filesystem::exists(dest_obj_id + ".tmp")); - std::remove((dest_obj_id + ".tmp").c_str()); service2.Shutdown(); } From 50df04452abb2de86f5d53c5715f5bb8f3198828 Mon Sep 17 00:00:00 2001 From: Husam Date: Wed, 7 Oct 2026 18:47:39 +0000 Subject: [PATCH 15/15] docs(replication): correct the -Wswitch note and document the worker pools - DispatchMessage: the comment implied the compiler reports a new MessageType by default. -Wswitch is part of -Wall, which this build does not enable yet. Say so. - transfer_service.h: state the trade-off of the inbound pool (a worker is held for the whole transfer it serves and socket calls have no timeout) and that the service runs 2 * threads workers. - bindings.cpp: the initialize() docstring says the same, since Python callers only see the single `threads` argument. - net_util_test.cpp: use EXPECT_TRUE instead of ASSERT_TRUE inside the writer thread, as gtest requires off the main thread. No behavior change. --- .../replication/transfer_service/bindings.cpp | 5 ++++- .../replication/transfer_service/transfer_service.cpp | 5 +++-- .../replication/transfer_service/transfer_service.h | 6 ++++++ tests/replication/transfer_service/net_util_test.cpp | 7 ++++--- 4 files changed, 17 insertions(+), 6 deletions(-) diff --git a/src/ml_flashpoint/replication/transfer_service/bindings.cpp b/src/ml_flashpoint/replication/transfer_service/bindings.cpp index 258db89..0031d09 100644 --- a/src/ml_flashpoint/replication/transfer_service/bindings.cpp +++ b/src/ml_flashpoint/replication/transfer_service/bindings.cpp @@ -38,7 +38,10 @@ PYBIND11_MODULE(transfer_service_ext, m) { .def("initialize", &TransferService::Initialize, py::arg("listen_port"), py::arg("threads") = 16, py::arg("conn_pool_per_peer") = 16, py::arg("global_rank") = -1, - "Initializes and starts the C++ transfer service.") + "Initializes and starts the C++ transfer service. `threads` is the " + "size of each of the two worker pools, one for outbound tasks " + "(async_put/async_get) and one for inbound requests from peers, so " + "the service runs 2 * threads worker threads.") .def("shutdown", &TransferService::Shutdown, "Stops the C++ transer service.") diff --git a/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp b/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp index b6ef084..6293de4 100644 --- a/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp +++ b/src/ml_flashpoint/replication/transfer_service/transfer_service.cpp @@ -675,8 +675,9 @@ void TransferService::ProcessIncomingData(int client_fd) { bool TransferService::DispatchMessage(int client_fd, const ObjInfoHeader& header) { - // No default case on purpose: with -Wswitch the compiler reports a new - // MessageType that is not handled here. + // No default case on purpose, so that a new MessageType that is not handled + // here is reported by -Wswitch (part of -Wall; the project does not enable + // it yet) instead of being silently routed to the unknown-type path. switch (header.type) { case MessageType::kPutObj: HandleDataReceive(client_fd, header, false); diff --git a/src/ml_flashpoint/replication/transfer_service/transfer_service.h b/src/ml_flashpoint/replication/transfer_service/transfer_service.h index 81e3226..85ca4fc 100644 --- a/src/ml_flashpoint/replication/transfer_service/transfer_service.h +++ b/src/ml_flashpoint/replication/transfer_service/transfer_service.h @@ -173,6 +173,12 @@ class TransferService final { // response. The peer can only respond if its own inbound work is running. // If inbound and outbound work shared one pool, two peers that fill their // pools with Gets to each other would deadlock. + // + // Trade-off: an inbound worker is also held for the whole transfer it + // serves, including the wait for the requester's final ACK, and the socket + // calls have no timeout. A peer that stops reading or writing holds that + // worker until it disconnects or Shutdown() runs. Both pools have `threads` + // workers, so the service runs twice that many worker threads. std::unique_ptr thread_pool_; // Outbound tasks. std::unique_ptr epoll_thread_pool_; // Inbound epoll events. std::thread epoll_thread_; diff --git a/tests/replication/transfer_service/net_util_test.cpp b/tests/replication/transfer_service/net_util_test.cpp index 4d8f2b0..ca766d8 100644 --- a/tests/replication/transfer_service/net_util_test.cpp +++ b/tests/replication/transfer_service/net_util_test.cpp @@ -171,11 +171,12 @@ TEST_F(SendRecvTest, RecvAndDiscard_DrainsPayloadLargerThanChunkSize) { const std::string expected_data = "trailer"; // A helper thread sends the data because the payload is larger than the // socketpair buffers: SendAll cannot finish until the receiver starts - // reading. + // reading. EXPECT rather than ASSERT: gtest's fatal assertions only return + // from the current function, which off the main thread is the lambda. std::thread writer_thread([this, &discarded_data, &expected_data]() { - ASSERT_TRUE( + EXPECT_TRUE( SendAll(fds_[0], discarded_data.data(), discarded_data.size()).ok()); - ASSERT_TRUE( + EXPECT_TRUE( SendAll(fds_[0], expected_data.data(), expected_data.size()).ok()); }); std::string actual_data(expected_data.size(), '\0');