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/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..c786aa4 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. +// +// 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. +// 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/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/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..b4c9b86 100644 --- a/src/ml_flashpoint/replication/transfer_service/task.h +++ b/src/ml_flashpoint/replication/transfer_service/task.h @@ -163,17 +163,23 @@ class GetTask : public Task { std::shared_ptr metric_container_; }; +// Serves a kGetObj request by sending the object back on the connection the +// request arrived on. +// +// 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( 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 +189,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..6293de4 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 @@ -29,6 +30,7 @@ #include #include #include +#include #include #include #include @@ -123,8 +125,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); @@ -211,7 +214,23 @@ void TransferService::Shutdown() { epoll_thread_.join(); } - // 6. Clean up connection pools. + // 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_) { + shutdown(fd, SHUT_RDWR); + } + } + + // 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_) { @@ -219,15 +238,33 @@ void TransferService::Shutdown() { pool->Shutdown(); } } - connection_pools_.clear(); } - // 7. Stop the thread pool. + // 8. Stop the thread pools. This runs the tasks that are still queued and + // joins every worker. if (thread_pool_) { thread_pool_->stop(); } + if (epoll_thread_pool_) { + epoll_thread_pool_->stop(); + } + + // 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_) { + close(fd); + } + client_fds_.clear(); + } - // 8. 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"; @@ -428,7 +465,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 +476,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 +485,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; @@ -464,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; @@ -486,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; @@ -547,6 +588,13 @@ void TransferService::HandleNewConnection() { continue; } + // 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); + } + // Add the new client to epoll. struct epoll_event event; event.events = @@ -554,7 +602,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; } @@ -567,6 +615,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. 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; + } if (epoll_fd_ >= 0) { epoll_ctl(epoll_fd_, EPOLL_CTL_DEL, client_fd, nullptr); } @@ -589,27 +645,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; - 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; + bool keep_connection = false; + try { + 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(); } if (!keep_connection) { @@ -629,6 +673,40 @@ void TransferService::ProcessIncomingData(int client_fd) { } } +bool TransferService::DispatchMessage(int client_fd, + const ObjInfoHeader& header) { + // 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); + return true; + case MessageType::kGetObj: + HandleGetObjRequest(client_fd, header); + return true; + case MessageType::kRespondToGetObj: + case MessageType::kAck: + case MessageType::kError: + // 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 + // 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(); @@ -703,12 +781,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) { @@ -722,32 +812,48 @@ 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); + // 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 + // 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(); + 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); + remove_tmp(); + fail("Failed to receive data"); return; } // 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) { - 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); + const std::string rename_error = std::strerror(errno); + remove_tmp(); + fail("Failed to rename temporary file " + tmp_obj_id + " to " + + header.dest_obj_id + ": " + rename_error); return; } if (is_respond_get_task) { @@ -812,7 +918,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 +927,42 @@ 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: { + // 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 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", + 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(), local_header, true); + } catch (const std::exception& e) { + // 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, + std::string("Failed to receive data: ") + e.what()); + } 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,29 +975,41 @@ 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(); - - auto conn_opt = GetConnectionFromPool(task->GetDestAddr()); - if (!conn_opt) { - LOG(ERROR) << "Failed to get connection!"; - ReportResult(task->GetTaskId(), false, "Failed to get connection"); + << ", 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!"; + finish(false, "Invalid client_fd"); return; } - ScopedConnection conn = std::move(conn_opt.value()); 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 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()); + } catch (const std::exception& e) { LOG(ERROR) << "RespondToGetTask failed: Could not open buffer object for '" - << task->GetSourceObjId() << "'"; - ReportResult(task->GetTaskId(), false, "Failed to create buffer object"); + << task->GetSourceObjId() << "': " << e.what(); + SendErrorResponse(sockfd, task->GetTaskId().c_str(), + task->GetSourceObjId().c_str()); + finish(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; @@ -883,18 +1028,17 @@ 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"; - 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(); @@ -902,18 +1046,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( @@ -963,6 +1106,16 @@ std::shared_ptr TransferService::GetOrCreateConnectionPool( return it->second; // Double check } + // 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 + << ": 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_); @@ -981,22 +1134,17 @@ 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; + // 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 + << (ec ? ": " + ec.message() : ""); SendErrorResponse(client_fd, header.task_id, header.source_obj_id); 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 +1157,11 @@ 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); - - { - 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}; - } - - task_queue_.enqueue(std::move(task)); + RespondToGetTask task(header.task_id, header.source_obj_id, + header.dest_obj_id, client_fd, 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); } void TransferService::UpdateTaskMetrics( diff --git a/src/ml_flashpoint/replication/transfer_service/transfer_service.h b/src/ml_flashpoint/replication/transfer_service/transfer_service.h index ddcdd10..85ca4fc 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. // @@ -132,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); @@ -161,7 +168,19 @@ class TransferService final { std::string local_address_; // Local address used in data transfer. int global_rank_ = -1; - std::unique_ptr thread_pool_; + // 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. + // + // 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_; std::thread task_queue_thread_; TaskQueue task_queue_; @@ -172,6 +191,12 @@ class TransferService final { std::map> connection_pools_; mutable std::shared_mutex connection_pools_mutex_; // Guard connection_pools_ + // 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_ + struct PendingTaskContext { std::shared_ptr> promise; std::shared_ptr metric_container; 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/net_util_test.cpp b/tests/replication/transfer_service/net_util_test.cpp index a726426..ca766d8 100644 --- a/tests/replication/transfer_service/net_util_test.cpp +++ b/tests/replication/transfer_service/net_util_test.cpp @@ -121,6 +121,95 @@ TEST_F(SendRecvTest, RecvAll_IncompleteRead) { EXPECT_EQ(recv_status.code(), absl::StatusCode::kUnavailable); } +// 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"; + 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); +} + +// 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"; + // A helper thread sends the data because the payload is larger than the + // socketpair buffers: SendAll cannot finish until the receiver starts + // 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]() { + EXPECT_TRUE( + SendAll(fds_[0], discarded_data.data(), discarded_data.size()).ok()); + EXPECT_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 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"; + 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 33f5cd9..758a571 100644 --- a/tests/replication/transfer_service/transfer_service_p2p_test.cpp +++ b/tests/replication/transfer_service/transfer_service_p2p_test.cpp @@ -12,24 +12,36 @@ // 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 +#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 { @@ -184,7 +196,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 @@ -698,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(); @@ -832,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(); @@ -916,5 +930,1495 @@ 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 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; + 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 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) { + 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`. +// 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); + 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 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) { + 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; +} + +// 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; + 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; +} + +// 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. +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) { + // 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_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 {}; + +// 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"); + 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}" 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) { + 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); + std::remove(source_obj_id.c_str()); + 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 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"); + 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) 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 + 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 write a +// file. The service closes the connection and writes nothing. +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 out of sync, 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(); +} + +// 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(); +} + +// 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 + 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()); + // No ACK on purpose: 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 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"); + 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 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)); + + // 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. 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"); + 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(); +} + +// 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"); + 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(); +} + +// 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"); + 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 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"); + 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 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)); + + // 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 hold 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 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"); + 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}")); + +// 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(); +} + +// 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. +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 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"; + 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 instead of hanging. The single pooled +// connection it used still works 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 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 + 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 { +}; + +// 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(); + 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)); + // No temporary file either, including after a truncated payload. + EXPECT_FALSE(std::filesystem::exists(dest_obj_id + ".tmp")); + + 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 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"; + 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, + 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(); +} + +// 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