fix(replication): respond to kGetObj over existing client socket

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.")

src/ml_flashpoint/replication/transfer_service/net_util.cpp

@@ -22,8 +22,10 @@
 #include <sys/socket.h>
 #include <unistd.h>
 
+#include <algorithm>
 #include <array>
 #include <memory>
+#include <vector>
 
 #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<char> scratch(static_cast<size_t>(
+      std::min<ssize_t>(std::max<ssize_t>(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()) {

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);

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.

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);
 }
 

src/ml_flashpoint/replication/transfer_service/task.h

@@ -163,17 +163,23 @@ class GetTask : public Task {
   std::shared_ptr<GetTaskMetricContainer> 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<RespondToGetTaskMetricContainer> 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<RespondToGetTaskMetricContainer> metric_container_;
 };