comms/uniflow/tcp: fail loudly on the paths that quietly did the wrong thing (meta-pytorch#3921)

Summary:

Folded from two adjacent diffs; each half is stated separately below so the two
arguments stay reviewable on their own terms.

--- outbound: a failed operation must not reach the peer (was D117927060) ----

One invariant, enforced from the two ends it was missing: the peer must not do work for
an operation whose caller has already been told it failed. put() has held that line since
its pre-flight/commit split; these are the two places that did not. Folded together
because they are halves of the same argument and each is small -- the second is the one
that made the first insufficient.

--- Before the first frame is queued (was D117888775) --------------------------

get() looked up each request's remote handle in the loop that also admits and
queues frames, so a multi-request get() whose second request carries no TCP handle
failed after the first request's ReadRequests were already queued. The peer serviced
reads for an operation the caller had already been told had failed.

put() settles everything that can fail before its first frame is queued and says so
at the top of its pre-flight loop. get() did not hold that line. This moves the
lookup into the loop that already walks requests to count chunks, and carries the
resolved segIds forward, so a rejected get() leaves the peer untouched.

Not a correctness fix in the load-bearing sense and not a performance one. The caller
already saw the same InvalidArgument, and the destination buffer was never at risk:
fail() marks the op done, so a late ReadReply fails tryBeginWrite() and skips the
copy rather than writing into a buffer the caller may have released. What changed is
that the peer no longer does work for a rejected operation, and those inflight_ slots
are no longer held for a round trip.

No hot-path effect either way: findRemoteHandle() was already called once per request
rather than once per chunk, so this reorders the same work instead of removing any.
Benchmarked regardless, because the get path is measured rather than argued about.

--- After teardown has swept (was D117927060) ----------------------------------

The three admission points disagreed about connBroken_. admitInflight() and
recvImpl() re-test it under their container mutex, because the caller's entry check is
not enough on its own -- failAllPending() can land in the gap between that check and
the insert. The three enqueue paths listed connBroken_ only in their wait predicate, as
a wake condition, and after waking checked outClosed alone.

failAllPending() sets connBroken_ and clears every lane queue, but never sets outClosed:
the only writer is a sender that has died. handleFrame's exception containment sweeps
without closing the connection on purpose, so the reachable state is connBroken_ set,
reader stopped, connection open, sender alive indefinitely. A frame admitted just before
that sweep then lands in the just-cleared queue and the live sender transmits it.

For put and get that means a Write reaching the peer's segment for an operation whose
caller has already been resolved with ConnectionFailed -- a partial write at offsets
nobody is told about, which is the case put()'s pre-flight/commit split exists to
prevent. For send it means the promise completes successfully on a transport that has
failed everything else and stopped reading.

Now all three check connBroken_ alongside outClosed and route it down the path they
already had for refusal: enqueueFrame/enqueueFrames report false so the caller fails the
op, and enqueueSendFrame takes ownership of the promise and fails it.

senderLoop's own outClosed check is deliberately left alone. It has to keep draining
whatever is already queued; refusing there would abandon frames rather than admit them.

--- inbound: a mismatched reply must not complete an operation (was D117932250) ---

TcpInflight::isRead records whether a chunk came from get() or put(), and the
reply handler only consulted it on the ReadReply branch. The Ack branch did not, so an
Ack naming a get chunk fell through to completeOne(): the chunk resolved Ok with
entry.dst never written, and the caller read back whatever its destination buffer
already held.

That is the one peer-supplied dimension on this path that failed silently. segId,
offset, len and payload size are all checked and all produce an error; a crossed reply
kind produced success with wrong data.

Both directions are now settled in one place, immediately after the entry lookup, and
the check is exhaustive over the three ops that reach it: an Ack must name a write and
a ReadReply must name a read. Error stays exempt because it is kind-agnostic by design.

Consolidating also fixes a diagnostic that was wrong before: !entry.isRead used to
report "read reply size mismatch", so a skewed peer sending a ReadReply for a put chunk
sent the reader chasing a length bug that did not exist. That branch is now purely a
size check and its message is accurate.

Not reachable from a same-version peer -- Ack answers a Write, ReadReply answers a
ReadRequest, and reqIds are unique per chunk -- so this is version skew or a hostile
peer, the same bar as the oversized-ReadRequest check.

Reviewed By: rmahidhar

Differential Revision: D117927060

Author: Peng Chen

comms/uniflow/transport/tcp/TcpTransport.cpp

@@ -988,14 +988,26 @@ std::future<Status> TcpTransport::get(
   auto state = std::make_shared<TcpOpState>();
   auto future = state->promise.get_future();
 
+  // Pre-flight. Same reason as put(): nothing that can fail may run after the
+  // first frame is queued, because a queued ReadRequest is already on its way
+  // to the peer and cannot be recalled. Resolving segIds here rather than in
+  // the emit loop is what makes a rejected get() leave the peer untouched.
   size_t totalChunks = 0;
+  std::vector<uint64_t> segIds;
+  segIds.reserve(requests.size());
   for (const auto& req : requests) {
     if (req.local.size() != req.remote.size()) {
       state->fail(
           Err(ErrCode::InvalidArgument,
               "tcp get: local and remote buffer sizes must match"));
       return future;
     }
+    auto remoteHandle = findRemoteHandle(req.remote);
+    if (!remoteHandle) {
+      state->fail(std::move(remoteHandle).error());
+      return future;
+    }
+    segIds.push_back(remoteHandle.value()->segId());
     const size_t len = req.local.size();
     const size_t chunkSize = adaptiveGetChunk(len, lanes_.size());
     totalChunks += (len == 0) ? 1 : (len + chunkSize - 1) / chunkSize;
@@ -1013,13 +1025,9 @@ std::future<Status> TcpTransport::get(
     }
   }
 
-  for (const auto& req : requests) {
-    auto remoteHandle = findRemoteHandle(req.remote);
-    if (!remoteHandle) {
-      state->fail(std::move(remoteHandle).error());
-      return future;
-    }
-    const uint64_t segId = remoteHandle.value()->segId();
+  for (size_t reqIdx = 0; reqIdx < requests.size(); ++reqIdx) {
+    const auto& req = requests[reqIdx];
+    const uint64_t segId = segIds[reqIdx];
     const uint64_t baseOffset = static_cast<uint64_t>(req.remote.remoteOffset_);
     const size_t len = req.local.size();
     const MemoryType memType = req.local.memType();
@@ -1834,7 +1842,15 @@ bool TcpTransport::enqueueFrame(TcpFrame frame, bool mayBlock) {
     // the connection. It cannot wait either: that stops it draining the socket
     // and reintroduces the mutual-READ deadlock the reader/sender split exists
     // to avoid. What bounds this queue is the drain rate, not a byte cap.
-    if (lane.outClosed) {
+    // Refused on connBroken_ as well as outClosed, so this admission point
+    // gives the same answer as admitInflight() and recvImpl(). failAllPending()
+    // sets connBroken_ and clears this queue but never sets outClosed -- only a
+    // dead sender does that -- and handleFrame's exception containment sweeps
+    // without closing the connection, leaving the sender alive. Checking
+    // outClosed alone therefore lets a frame admitted before the sweep land in
+    // the cleared queue and go out on the wire for an op whose caller has
+    // already been told it failed.
+    if (lane.outClosed || connBroken_.load(std::memory_order_acquire)) {
       return false;
     }
     lane.queue.push_back(TcpOutItem{std::move(frame), nullptr});
@@ -1871,7 +1887,7 @@ bool TcpTransport::enqueueFrames(std::vector<TcpFrame> frames, bool mayBlock) {
             lane.queue.empty() || lane.bytes + bytes <= cap;
       });
     }
-    if (lane.outClosed) {
+    if (lane.outClosed || connBroken_.load(std::memory_order_acquire)) {
       return false;
     }
     for (auto& frame : frames) {
@@ -1913,7 +1929,7 @@ void TcpTransport::enqueueSendFrame(
       return lane.outClosed || connBroken_.load(std::memory_order_acquire) ||
           lane.queue.empty() || lane.bytes + bytes <= cap;
     });
-    if (lane.outClosed) {
+    if (lane.outClosed || connBroken_.load(std::memory_order_acquire)) {
       closed = true;
       // Taken over here so each path has exactly one owner: the queue takes it
       // when the frame is enqueued, this does when it cannot be.
@@ -2354,13 +2370,29 @@ void TcpTransport::handleFrameImpl(
         std::lock_guard<std::mutex> lk(inflightMu_);
         inflight_.erase(header.reqId);
       };
+      // An Ack answers a Write and a ReadReply answers a ReadRequest, so a
+      // reply whose kind disagrees with the request it names is version skew or
+      // a hostile peer. Rejected here rather than per-branch because the Ack
+      // direction is the dangerous one and it used to fall straight through to
+      // completeOne(): the get chunk resolved Ok with entry.dst never written,
+      // handing the caller back whatever its buffer already held. Every other
+      // peer-supplied dimension on this path fails loudly; this one did not
+      // fail at all. Error is exempt because it is kind-agnostic by design.
+      if (op != TcpOp::Error && (op == TcpOp::ReadReply) != entry.isRead) {
+        eraseInflight();
+        entry.state->fail(
+            Err(ErrCode::TransportError,
+                "tcp: reply op does not match the request kind"));
+        break;
+      }
       if (op == TcpOp::Error) {
         eraseInflight();
         entry.state->fail(
             Err(ErrCode::TransportError, "tcp: peer reported an error"));
       } else if (op == TcpOp::ReadReply) {
-        if (!entry.isRead || payload.size() != header.len ||
-            header.len != entry.len) {
+        // Kind is settled above, so this is purely a size check and its message
+        // says so.
+        if (payload.size() != header.len || header.len != entry.len) {
           eraseInflight();
           entry.state->fail(Err(
               ErrCode::TransportError, "tcp get: read reply size mismatch"));

comms/uniflow/transport/tcp/tests/unit/TcpTransportFrameTest.cpp

@@ -206,6 +206,8 @@ class TcpTransportFrameTest : public ::testing::Test {
  protected:
   static constexpr uint64_t kSegId = 42;
   static constexpr size_t kSegLen = 256;
+  /// Destination length for the reply-kind-mismatch tests.
+  static constexpr size_t kAckDstLen = 16;
 
   void SetUp() override {
     evbThread_ = std::make_unique<ScopedEventBaseThread>("tcp-frame-test");
@@ -1179,6 +1181,122 @@ TEST_F(TcpTransportFrameTest, AdmissionIsRefusedOnceTeardownHasSwept) {
       << "admitInflight must not resolve the promise itself";
 }
 
+// The three admission points must agree about connBroken_. admitInflight() and
+// recvImpl() re-test it under their container mutex; the enqueue paths listed
+// it only in the wait predicate, as a wake condition, and then checked
+// outClosed alone. failAllPending() does not set outClosed -- the only writer
+// is the dead sender -- so after a sweep that leaves the connection open
+// (handleFrame's exception containment does exactly that: it stops the reader
+// without closing) a frame admitted before the sweep still lands in the
+// just-cleared queue and the live sender transmits it.
+//
+// For put that means a Write applied to the peer's segment for an operation the
+// caller was already told had failed, which is the partial-write-nobody-knows
+// -about case put()'s pre-flight exists to prevent.
+TEST_F(TcpTransportFrameTest, EnqueueIsRefusedOnceTeardownHasSwept) {
+  setConnected();
+  failAllPending();
+
+  EXPECT_FALSE(enqueueFrame(makeFrame(
+      TcpOp::Write, kSegId, /*offset=*/0, /*len=*/8, /*payloadBytes=*/8)))
+      << "a frame queued after the sweep is transmitted for a failed op";
+  EXPECT_EQ(outQueueDepth(), 0u) << "no frame may remain queued after a sweep";
+}
+
+TEST_F(TcpTransportFrameTest, EnqueueFramesIsRefusedOnceTeardownHasSwept) {
+  setConnected();
+  failAllPending();
+
+  std::vector<TcpFrame> group;
+  group.emplace_back(makeFrame(
+      TcpOp::Write, kSegId, /*offset=*/0, /*len=*/8, /*payloadBytes=*/8));
+
+  EXPECT_FALSE(enqueueFrames(std::move(group), /*mayBlock=*/false))
+      << "a staged wave queued after the sweep is transmitted for a failed op";
+  EXPECT_EQ(outQueueDepth(), 0u);
+}
+
+// send() owns a promise, so refusal has to fail it rather than drop the frame.
+// Reporting success here says a send completed on a transport that has failed
+// everything else and stopped reading.
+TEST_F(TcpTransportFrameTest, SendFrameIsRefusedOnceTeardownHasSwept) {
+  setConnected();
+  failAllPending();
+
+  auto state = std::make_shared<TcpOpState>();
+  state->remaining = 1;
+  auto future = state->promise.get_future();
+
+  enqueueSendFrame(
+      makeFrame(TcpOp::Send, kSegId, /*offset=*/0, /*len=*/8, 8), state);
+
+  EXPECT_EQ(outQueueDepth(), 0u) << "the send frame must not be queued";
+  ASSERT_EQ(future.wait_for(std::chrono::seconds{5}), std::future_status::ready)
+      << "a refused send must fail its promise, not leave the caller waiting";
+  EXPECT_TRUE(future.get().hasError())
+      << "send must not report success on a failed connection";
+}
+
+// Ack answers a Write, ReadReply answers a ReadRequest, and reqIds are unique
+// per chunk, so a same-version peer never crosses them. The ReadReply direction
+// was already rejected; the Ack direction was not, and it failed silently
+// rather than loudly: the get chunk resolved Ok with its destination never
+// written, so the caller read back whatever the buffer held before. Every other
+// peer-supplied dimension on this path -- segId, offset, len, payload size --
+// is checked, and all of those fail with an error.
+//
+// Version-skew or a hostile peer only, which is the same bar as
+// OversizedReadRequestIsRefusedPerRequest.
+TEST_F(TcpTransportFrameTest, AckOnAReadRequestIsRejected) {
+  setConnected();
+
+  // A destination pre-filled with a known pattern, so "never written" is
+  // detectable rather than merely assumed.
+  std::vector<uint8_t> dst(kAckDstLen, uint8_t{0x11});
+  const std::vector<uint8_t> pristineDst = dst;
+
+  auto state = std::make_shared<TcpOpState>();
+  state->remaining = 1;
+  auto future = state->promise.get_future();
+  ASSERT_FALSE(admitInflight(
+                   /*reqId=*/1,
+                   TcpInflight{state, dst.data(), kAckDstLen, /*isRead=*/true})
+                   .hasError());
+
+  feed(makeFrame(
+      TcpOp::Ack, kSegId, /*offset=*/0, /*len=*/0, /*payloadBytes=*/0));
+
+  ASSERT_EQ(future.wait_for(std::chrono::seconds{5}), std::future_status::ready)
+      << "the op must be resolved rather than left outstanding";
+  EXPECT_TRUE(future.get().hasError())
+      << "an Ack cannot complete a get: the destination was never written";
+  EXPECT_EQ(dst, pristineDst) << "the destination buffer must be untouched";
+  EXPECT_EQ(inflightCount(), 0u)
+      << "a rejected reply must not leave its admission slot held";
+}
+
+// The mirror direction, already guarded before this change. Kept as a pair so a
+// future edit cannot close one direction and reopen the other.
+TEST_F(TcpTransportFrameTest, ReadReplyOnAWriteIsRejected) {
+  setConnected();
+
+  auto state = std::make_shared<TcpOpState>();
+  state->remaining = 1;
+  auto future = state->promise.get_future();
+  ASSERT_FALSE(admitInflight(
+                   /*reqId=*/1,
+                   TcpInflight{state, nullptr, kAckDstLen, /*isRead=*/false})
+                   .hasError());
+
+  feed(makeFrame(
+      TcpOp::ReadReply, kSegId, /*offset=*/0, kAckDstLen, kAckDstLen));
+
+  ASSERT_EQ(
+      future.wait_for(std::chrono::seconds{5}), std::future_status::ready);
+  EXPECT_TRUE(future.get().hasError()) << "a ReadReply cannot complete a put";
+  EXPECT_EQ(inflightCount(), 0u);
+}
+
 TEST_F(TcpTransportFrameTest, AdmissionSucceedsOnAHealthyTransport) {
   setConnected();