Cut fixed per-request costs out of the reshard control plane

tpu_sync/kv_cache/reshard/framed_rpc.cc

@@ -17,12 +17,15 @@
 #include <arpa/inet.h>
 #include <netdb.h>
 #include <netinet/in.h>
+#include <netinet/tcp.h>
 #include <sys/socket.h>
 #include <unistd.h>
 
+#include <algorithm>
 #include <cerrno>
 #include <cstdint>
 #include <cstring>
+#include <memory>
 #include <string>
 #include <thread>  // NOLINT(build/c++11)
 #include <utility>
@@ -60,6 +63,33 @@ bool SendAll(int fd, const char* data, size_t n) {
   return true;
 }
 
+// Sends the 4-byte length prefix and the body as ONE buffer. Two separate
+// send() calls under Nagle stall the body until the peer ACKs the prefix
+// segment (delayed-ACK interaction), putting a tens-of-ms floor on small
+// RPCs.
+bool SendFramed(int fd, absl::string_view payload) {
+  std::string framed;
+  framed.reserve(sizeof(uint32_t) + payload.size());
+  uint32_t net_len = htonl(static_cast<uint32_t>(payload.size()));
+  framed.append(reinterpret_cast<const char*>(&net_len), sizeof(net_len));
+  framed.append(payload.data(), payload.size());
+  return SendAll(fd, framed.data(), framed.size());
+}
+
+// Control-plane RPCs are short request/response exchanges: disable Nagle so
+// each frame goes out immediately, and enable keepalive plus
+// TCP_USER_TIMEOUT so a black-holed peer surfaces as a socket error within
+// the I/O timeout instead of only at the full receive deadline.
+void ConfigureControlSocket(int fd, absl::Duration io_timeout) {
+  int one = 1;
+  setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &one, sizeof(one));
+  setsockopt(fd, SOL_SOCKET, SO_KEEPALIVE, &one, sizeof(one));
+  unsigned int user_timeout_ms = static_cast<unsigned int>(
+      absl::ToInt64Milliseconds(io_timeout));
+  setsockopt(fd, IPPROTO_TCP, TCP_USER_TIMEOUT, &user_timeout_ms,
+             sizeof(user_timeout_ms));
+}
+
 // Splits "host:port" at the last colon; strips IPv6 brackets, mirroring
 // raiden_controller.connect_socket.
 absl::Status SplitAddress(absl::string_view address, std::string* host,
@@ -104,6 +134,7 @@ int TryConnectOnce(const std::string& host, int port,
     tv.tv_usec = 0;
     setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
     setsockopt(fd, SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof(tv));
+    ConfigureControlSocket(fd, io_timeout);
     if (connect(fd, res->ai_addr, res->ai_addrlen) == 0) break;
     close(fd);
     fd = -1;
@@ -140,10 +171,7 @@ absl::StatusOr<std::string> SocketFramedTransport::Call(
 
   std::string response;
   {
-    uint32_t net_len = htonl(static_cast<uint32_t>(payload.size()));
-    if (!SendAll(fd, reinterpret_cast<const char*>(&net_len),
-                 sizeof(net_len)) ||
-        !SendAll(fd, payload.data(), payload.size())) {
+    if (!SendFramed(fd, payload)) {
       close(fd);
       return absl::UnavailableError(
           absl::StrCat("Failed to send framed payload to ", address, ": ",
@@ -242,10 +270,10 @@ void FramedServer::Stop() {
     server_fd_ = -1;
   }
   if (accept_thread_.joinable()) accept_thread_.join();
-  for (std::thread& t : connection_threads_) {
-    if (t.joinable()) t.join();
+  for (const std::unique_ptr<Connection>& connection : connections_) {
+    if (connection->thread.joinable()) connection->thread.join();
   }
-  connection_threads_.clear();
+  connections_.clear();
 }
 
 void FramedServer::AcceptLoop() {
@@ -262,12 +290,29 @@ void FramedServer::AcceptLoop() {
       close(client_fd);
       break;
     }
-    connection_threads_.emplace_back(&FramedServer::ServeConnection, this,
-                                     client_fd);
+    connections_.erase(
+        std::remove_if(connections_.begin(), connections_.end(),
+                       [](const std::unique_ptr<Connection>& connection) {
+                         if (!connection->done.load()) return false;
+                         if (connection->thread.joinable()) {
+                           connection->thread.join();
+                         }
+                         return true;
+                       }),
+        connections_.end());
+    auto connection = std::make_unique<Connection>();
+    Connection* raw = connection.get();
+    raw->thread = std::thread([this, raw, client_fd]() {
+      ServeConnection(client_fd);
+      raw->done.store(true);
+    });
+    connections_.push_back(std::move(connection));
   }
 }
 
 void FramedServer::ServeConnection(int client_fd) {
+  int one = 1;
+  setsockopt(client_fd, IPPROTO_TCP, TCP_NODELAY, &one, sizeof(one));
   uint32_t net_len = 0;
   if (!ReadExactly(client_fd, reinterpret_cast<char*>(&net_len),
                    sizeof(net_len))) {
@@ -281,10 +326,7 @@ void FramedServer::ServeConnection(int client_fd) {
     return;
   }
   std::string response = handler_(request);
-  uint32_t resp_net_len = htonl(static_cast<uint32_t>(response.size()));
-  SendAll(client_fd, reinterpret_cast<const char*>(&resp_net_len),
-          sizeof(resp_net_len));
-  SendAll(client_fd, response.data(), response.size());
+  SendFramed(client_fd, response);
   close(client_fd);
 }
 

tpu_sync/kv_cache/reshard/framed_rpc.h

@@ -17,6 +17,7 @@
 
 #include <atomic>
 #include <functional>
+#include <memory>
 #include <string>
 #include <thread>  // NOLINT(build/c++11)
 #include <vector>
@@ -79,6 +80,14 @@ class FramedServer final {
   int port() const { return port_; }
 
  private:
+  // One handler thread per accepted connection; `done` lets the accept loop
+  // reap finished threads so a long-lived server does not accumulate one
+  // un-joined thread (and its stack) per request.
+  struct Connection {
+    std::thread thread;
+    std::atomic<bool> done{false};
+  };
+
   void AcceptLoop();
   void ServeConnection(int client_fd);
 
@@ -88,7 +97,7 @@ class FramedServer final {
   Handler handler_;
   std::atomic<bool> stopping_{false};
   std::thread accept_thread_;
-  std::vector<std::thread> connection_threads_;
+  std::vector<std::unique_ptr<Connection>> connections_;
 };
 
 }  // namespace reshard

tpu_sync/kv_cache/reshard/reshard_coordinator.cc

@@ -328,22 +328,65 @@ absl::Status ReshardCoordinator::ExecutePoolReshard(
         "(the destination-side relay is retired)");
   }
 
-  const int64_t controller_start_ns = MonotonicNs();
-  const int64_t plan_build_start_ns = controller_start_ns;
-
-  auto src_metadata = directory_->LocalMetadata(args.src_units);
-  if (!src_metadata.ok()) return src_metadata.status();
   std::vector<tpu_sync::rpc::RegisterWorkUnitRequest> dst_metadata;
+  bool used_cache = false;
   if (!args.dst_controller_address.empty()) {
-    auto remote = QueryRemoteMetadata(args.dst_controller_address);
-    if (!remote.ok()) return remote.status();
-    dst_metadata = *std::move(remote);
+    {
+      absl::MutexLock lock(metadata_cache_mu_);
+      auto it = remote_metadata_cache_.find(args.dst_controller_address);
+      if (it != remote_metadata_cache_.end()) {
+        dst_metadata = it->second;
+        used_cache = true;
+      }
+    }
+    if (!used_cache) {
+      auto remote = QueryRemoteMetadata(args.dst_controller_address);
+      if (!remote.ok()) return remote.status();
+      dst_metadata = *std::move(remote);
+      absl::MutexLock lock(metadata_cache_mu_);
+      remote_metadata_cache_[args.dst_controller_address] = dst_metadata;
+    }
   } else {
     auto local = directory_->LocalMetadata(args.dst_units);
     if (!local.ok()) return local.status();
     dst_metadata = *std::move(local);
   }
 
+  bool receiver_armed = false;
+  absl::Status status =
+      ExecutePoolReshardAttempt(args, dst_metadata, &receiver_armed);
+  if (status.ok() || !used_cache) return status;
+  // Cached destination metadata can be stale after an engine replacement,
+  // so a failed attempt drops the entry and the next request re-queries.
+  // Planning and arming are side-effect-free beyond the abandoned claim
+  // until a receiver acknowledges its arm, so the attempt is replayed on
+  // fresh metadata only while no receiver has acknowledged.
+  {
+    absl::MutexLock lock(metadata_cache_mu_);
+    remote_metadata_cache_.erase(args.dst_controller_address);
+  }
+  if (receiver_armed) return status;
+  auto remote = QueryRemoteMetadata(args.dst_controller_address);
+  if (!remote.ok()) return remote.status();
+  dst_metadata = *std::move(remote);
+  {
+    absl::MutexLock lock(metadata_cache_mu_);
+    remote_metadata_cache_[args.dst_controller_address] = dst_metadata;
+  }
+  receiver_armed = false;
+  return ExecutePoolReshardAttempt(args, dst_metadata, &receiver_armed);
+}
+
+absl::Status ReshardCoordinator::ExecutePoolReshardAttempt(
+    const PoolReshardArgs& args,
+    const std::vector<tpu_sync::rpc::RegisterWorkUnitRequest>& dst_metadata,
+    bool* receiver_armed) {
+  const int64_t controller_start_ns = MonotonicNs();
+  const int64_t plan_build_start_ns = controller_start_ns;
+
+  auto src_metadata = directory_->LocalMetadata(args.src_units);
+  if (!src_metadata.ok()) return src_metadata.status();
+
   int64_t uuid = args.uuid;
   if (uuid <= 0) {
     // Python: random.randint(1, 2**63 - 1) when the wire carries no uuid.
@@ -355,7 +398,7 @@ absl::Status ReshardCoordinator::ExecutePoolReshard(
   plan_request.src_units = args.src_units;
   plan_request.dst_units = args.dst_units;
   plan_request.src_metadata = *std::move(src_metadata);
-  plan_request.dst_metadata = std::move(dst_metadata);
+  plan_request.dst_metadata = dst_metadata;
   plan_request.req_id = args.req_id;
   plan_request.uuid = uuid;
   plan_request.dst_device_block_ids = args.dst_device_block_ids;
@@ -401,6 +444,9 @@ absl::Status ReshardCoordinator::ExecutePoolReshard(
       });
     }
     for (std::thread& t : armers) t.join();
+    for (const absl::Status& status : arm_status) {
+      if (status.ok()) *receiver_armed = true;
+    }
     for (const absl::Status& status : arm_status) {
       if (!status.ok()) {
         registry_->AbandonClaim(args.req_id, uuid, claim_owner);

tpu_sync/kv_cache/reshard/reshard_coordinator.h

@@ -112,6 +112,15 @@ class ReshardCoordinator {
  private:
   absl::Status ExecutePoolReshard(const PoolReshardArgs& args);
 
+  // One planning/arming/dispatch pass over the given destination metadata.
+  // Sets *receiver_armed once any destination acknowledges its arm; past
+  // that point the transfer has side effects beyond the abandoned claim and
+  // must not be replayed.
+  absl::Status ExecutePoolReshardAttempt(
+      const PoolReshardArgs& args,
+      const std::vector<tpu_sync::rpc::RegisterWorkUnitRequest>& dst_metadata,
+      bool* receiver_armed);
+
   // GET_METADATA against the destination controller (dst_controller_address
   // path), recorded shape-identical to Python's _query_remote_metadata.
   absl::StatusOr<std::vector<tpu_sync::rpc::RegisterWorkUnitRequest>>
@@ -124,6 +133,16 @@ class ReshardCoordinator {
 
   mutable absl::Mutex status_mu_;
   std::map<std::string, Status> transfer_status_ ABSL_GUARDED_BY(status_mu_);
+
+  // Destination work-unit metadata rarely changes (units register once per
+  // engine lifetime), so the per-request GET_METADATA round trip is served
+  // from this per-address cache. Staleness surfaces as a plan-build or
+  // receiver-arm failure; a failed attempt drops the entry, and
+  // ExecutePoolReshard replays once on fresh metadata while no receiver
+  // has acknowledged its arm.
+  mutable absl::Mutex metadata_cache_mu_;
+  std::map<std::string, std::vector<tpu_sync::rpc::RegisterWorkUnitRequest>>
+      remote_metadata_cache_ ABSL_GUARDED_BY(metadata_cache_mu_);
 };
 
 }  // namespace reshard