comms/uniflow: instrument the receive path across the controller and the transport (meta-pytorch#3905)

Summary:
Pull Request resolved: meta-pytorch#3905

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

 --- controller: recv header-wait and payload-drain timings (was D117632427) ---

TcpConn::syncRecv() populates the existing RecvPhaseStats with the time spent
waiting for the length prefix versus draining the payload, plus frame and byte
counts. Splitting the two phases is what makes a slow receive interpretable: a
large headerWaitNs means we were waiting on the peer, while a large
payloadDrainNs means the socket itself was the limit.

All four counters are relaxed fetch_add on atomics already declared in
Controller.h, so this adds two steady_clock reads per frame and no
synchronisation.

 --- transport: receive-slab hits, misses and vector receives (was D117632428) ---

Adds receiveSlabAttempts_, receiveSlabMisses_, and vectorReceiveCount_ to
TcpTransport and reports them on the existing tcp phases log line. A miss means
the reader could not get a pinned slab and fell back to a vector-backed
receive, which changes the H2D path for that frame, so distinguishing the two
is necessary before drawing conclusions from an aggregate drain number.

The counters are relaxed atomics incremented in readerLoop().

Differential Revision: D117632428

Author: Peng Chen

comms/uniflow/controller/TcpController.cpp

@@ -613,12 +613,16 @@ Result<size_t> TcpConn<IOPolicy>::syncRecv(std::span<uint8_t> buf) {
     return Err(ErrCode::NotConnected, "Socket is not connected");
   }
 
+  auto& stats = recvPhaseStats();
+  const auto tStart = std::chrono::steady_clock::now();
+
   uint32_t rawLen = 0;
   if (!recvAll(&rawLen, sizeof(rawLen))) {
     return Err(
         ErrCode::ConnectionFailed,
         "recv header failed: " + std::system_category().message(errno));
   }
+  const auto tFirstByte = std::chrono::steady_clock::now();
 
   uint32_t len = ntohl(rawLen);
   if (len > kMaxMessageSize) {
@@ -642,6 +646,17 @@ Result<size_t> TcpConn<IOPolicy>::syncRecv(std::span<uint8_t> buf) {
         "recv payload failed: " + std::system_category().message(errno));
   }
 
+  const auto tDone = std::chrono::steady_clock::now();
+  using ns = std::chrono::nanoseconds;
+  stats.headerWaitNs.fetch_add(
+      std::chrono::duration_cast<ns>(tFirstByte - tStart).count(),
+      std::memory_order_relaxed);
+  stats.payloadDrainNs.fetch_add(
+      std::chrono::duration_cast<ns>(tDone - tFirstByte).count(),
+      std::memory_order_relaxed);
+  stats.frames.fetch_add(1, std::memory_order_relaxed);
+  stats.payloadBytes.fetch_add(len, std::memory_order_relaxed);
+
   UNIFLOW_LOG_DEBUG("TcpConn::recv(span): fd={} bytes={}", sock_, len);
   return static_cast<size_t>(len);
 }

comms/uniflow/controller/tests/TcpConnTest.cpp

@@ -120,6 +120,31 @@ TEST_P(TcpConnTest, SendRecvBidirectional) {
   EXPECT_TRUE(recv3.empty());
 }
 
+TEST_P(TcpConnTest, SpanRecvRecordsPhaseStats) {
+  auto [clientConn, serverConn] = connectPair();
+  ASSERT_NE(clientConn, nullptr);
+  ASSERT_NE(serverConn, nullptr);
+
+  std::vector<uint8_t> sent(4096, 0xA5);
+  std::vector<uint8_t> received(sent.size());
+  ASSERT_TRUE(clientConn->send(sent).get().hasValue());
+
+  auto recvResult = serverConn->recv(std::span<uint8_t>(received)).get();
+  ASSERT_TRUE(recvResult.hasValue()) << recvResult.error().toString();
+  EXPECT_EQ(recvResult.value(), sent.size());
+  EXPECT_EQ(received, sent);
+
+  auto& stats = serverConn->recvPhaseStats();
+  EXPECT_EQ(stats.frames.load(std::memory_order_relaxed), 1);
+  EXPECT_EQ(stats.payloadBytes.load(std::memory_order_relaxed), sent.size());
+
+  stats.reset();
+  EXPECT_EQ(stats.headerWaitNs.load(std::memory_order_relaxed), 0);
+  EXPECT_EQ(stats.payloadDrainNs.load(std::memory_order_relaxed), 0);
+  EXPECT_EQ(stats.frames.load(std::memory_order_relaxed), 0);
+  EXPECT_EQ(stats.payloadBytes.load(std::memory_order_relaxed), 0);
+}
+
 TEST_P(TcpConnTest, SendRecvLargeData) {
   auto [clientConn, serverConn] = connectPair();
   ASSERT_NE(clientConn, nullptr);

comms/uniflow/transport/tcp/TcpTransport.cpp

@@ -1527,6 +1527,12 @@ void TcpTransport::logAndResetPhaseStats(std::string_view label) {
       h2dState_->copyNs.load(std::memory_order_relaxed);
   const uint64_t copies = dstCopyCount_.load(std::memory_order_relaxed) +
       h2dState_->copyCount.load(std::memory_order_relaxed);
+  const uint64_t slabAttempts =
+      receiveSlabAttempts_.load(std::memory_order_relaxed);
+  const uint64_t slabMisses =
+      receiveSlabMisses_.load(std::memory_order_relaxed);
+  const uint64_t vectorReceives =
+      vectorReceiveCount_.load(std::memory_order_relaxed);
 
   if (frames == 0) {
     UNIFLOW_LOG_INFO("tcp phases [{}]: no frames", label);
@@ -1546,7 +1552,8 @@ void TcpTransport::logAndResetPhaseStats(std::string_view label) {
     UNIFLOW_LOG_INFO(
         "tcp phases [{}]: frames={} bytes={} | first-byte {:.1f}us/frame "
         "({:.1f}%) | drain {:.1f}us/frame ({:.1f}%, {:.2f} GB/s) | dstcopy "
-        "{:.1f}us x{} ({:.1f}% of wire)",
+        "{:.1f}us x{} ({:.1f}% of wire) | receive_slabs attempts={} misses={} "
+        "vector_recvs={}",
         label,
         frames,
         bytes,
@@ -1557,13 +1564,19 @@ void TcpTransport::logAndResetPhaseStats(std::string_view label) {
         drainGBps,
         copies > 0 ? static_cast<double>(copyNs) / copies / 1000.0 : 0.0,
         copies,
-        pct(copyNs));
+        pct(copyNs),
+        slabAttempts,
+        slabMisses,
+        vectorReceives);
   }
   rs.reset();
   dstCopyNs_.store(0, std::memory_order_relaxed);
   dstCopyCount_.store(0, std::memory_order_relaxed);
   h2dState_->copyNs.store(0, std::memory_order_relaxed);
   h2dState_->copyCount.store(0, std::memory_order_relaxed);
+  receiveSlabAttempts_.store(0, std::memory_order_relaxed);
+  receiveSlabMisses_.store(0, std::memory_order_relaxed);
+  vectorReceiveCount_.store(0, std::memory_order_relaxed);
 }
 
 void TcpTransport::readerLoop() noexcept {
@@ -1576,7 +1589,11 @@ void TcpTransport::readerLoop() noexcept {
     TcpPinnedSlab receiveSlab;
     if (config_.asyncGetH2d) {
       if (auto pool = receivePoolIfCreated()) {
+        receiveSlabAttempts_.fetch_add(1, std::memory_order_relaxed);
         receiveSlab = pool->tryAcquire(/*allowReserved=*/true);
+        if (!receiveSlab) {
+          receiveSlabMisses_.fetch_add(1, std::memory_order_relaxed);
+        }
       }
     }
     if (receiveSlab) {
@@ -1596,6 +1613,7 @@ void TcpTransport::readerLoop() noexcept {
       continue;
     }
 
+    vectorReceiveCount_.fetch_add(1, std::memory_order_relaxed);
     auto result = dataConn_->recv(msg).get();
     if (!result) {
       // Connection closed, errored, or idle-timed-out; stop reading.
@@ -2148,7 +2166,7 @@ std::future<Status> TcpTransport::recv(
 }
 
 void TcpTransport::shutdown() {
-  std::lock_guard<std::mutex> lk(lifecycleMu_);
+  std::lock_guard<std::mutex> lifecycleLock(lifecycleMu_);
   // One-shot, but checked under the mutex rather than before taking it:
   // shutdown() is called twice in the normal flow (MultiTransport::shutdown()
   // then ~TcpTransport), and the second caller must not return while the first
@@ -2159,7 +2177,7 @@ void TcpTransport::shutdown() {
   running_.store(false, std::memory_order_release);
 
   {
-    std::lock_guard<std::mutex> lk(outMu_);
+    std::lock_guard<std::mutex> outLock(outMu_);
     outClosed_ = true;
   }
   outCv_.notify_all();

comms/uniflow/transport/tcp/TcpTransport.h

@@ -848,6 +848,9 @@ class TcpTransport : public Transport {
   // nothing more.
   std::atomic<uint64_t> dstCopyNs_{0};
   std::atomic<uint64_t> dstCopyCount_{0};
+  std::atomic<uint64_t> receiveSlabAttempts_{0};
+  std::atomic<uint64_t> receiveSlabMisses_{0};
+  std::atomic<uint64_t> vectorReceiveCount_{0};
   std::shared_ptr<H2dPollState> h2dState_;
 
   // Serialises bind()/connect()/shutdown(), which together own server_,

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

@@ -367,6 +367,127 @@ TEST_F(TcpReceivePoolTest, VramReadReplyKeepsSlabUntilEventCompletes) {
   EXPECT_TRUE(pool->tryAcquire(/*allowReserved=*/true));
 }
 
+TEST_F(
+    TcpReceivePoolTest,
+    EventRecordFailureCompletesGetWhenStreamSynchronizationSucceeds) {
+  makeTransport();
+  auto pool = createReceivePool();
+  ASSERT_NE(pool, nullptr);
+  auto slab = pool->tryAcquire(/*allowReserved=*/true);
+  ASSERT_TRUE(slab);
+  auto frame = readReplyFrame(/*reqId=*/1, kLen);
+  ASSERT_LE(frame.size(), slab.capacity());
+  std::memcpy(slab.data(), frame.data(), frame.size());
+  const auto bytes = std::span<const uint8_t>{slab.data(), frame.size()};
+
+  std::vector<uint8_t> destination(kLen);
+  auto* const callerStream = reinterpret_cast<void*>(0xCA11);
+  auto future = postVramRead(destination.data(), callerStream);
+
+  EXPECT_CALL(
+      *cudaApi_,
+      memcpyAsync(
+          destination.data(),
+          ::testing::_,
+          kLen,
+          kMockMemcpyH2D,
+          static_cast<MockStream>(callerStream)))
+      .WillOnce([&](void* dst, const void* src, size_t len, auto, auto) {
+        std::memcpy(dst, src, len);
+        return Ok();
+      });
+  EXPECT_CALL(*cudaApi_, eventCreate(::testing::_))
+      .WillOnce([](auto* event) -> Status {
+        *event = {};
+        return Ok();
+      });
+  EXPECT_CALL(
+      *cudaApi_,
+      eventRecord(::testing::_, static_cast<MockStream>(callerStream)))
+      .WillOnce(
+          ::testing::Return(
+              Err(ErrCode::DriverError, "test: event record failed")));
+  EXPECT_CALL(
+      *cudaApi_, streamSynchronize(static_cast<MockStream>(callerStream)))
+      .WillOnce(::testing::Return(Ok()));
+  EXPECT_CALL(*cudaApi_, eventDestroy(::testing::_))
+      .WillOnce(::testing::Return(Ok()));
+  EXPECT_CALL(*cudaApi_, eventQuery(::testing::_)).Times(0);
+
+  handleFrame(bytes, std::move(slab));
+
+  ASSERT_EQ(
+      future.wait_for(std::chrono::seconds(5)), std::future_status::ready);
+  EXPECT_TRUE(future.get().hasValue());
+  EXPECT_TRUE(
+      std::all_of(destination.begin(), destination.end(), [](uint8_t b) {
+        return b == uint8_t{0xCD};
+      }));
+  auto returnedSlabs = pool->acquire(2);
+  EXPECT_TRUE(returnedSlabs.hasValue());
+}
+
+TEST_F(
+    TcpReceivePoolTest,
+    EventQueryFailureCompletesGetWhenStreamSynchronizationSucceeds) {
+  makeTransport();
+  auto pool = createReceivePool();
+  ASSERT_NE(pool, nullptr);
+  auto slab = pool->tryAcquire(/*allowReserved=*/true);
+  ASSERT_TRUE(slab);
+  auto frame = readReplyFrame(/*reqId=*/1, kLen);
+  ASSERT_LE(frame.size(), slab.capacity());
+  std::memcpy(slab.data(), frame.data(), frame.size());
+  const auto bytes = std::span<const uint8_t>{slab.data(), frame.size()};
+
+  std::vector<uint8_t> destination(kLen);
+  auto* const callerStream = reinterpret_cast<void*>(0xCA11);
+  auto future = postVramRead(destination.data(), callerStream);
+
+  EXPECT_CALL(
+      *cudaApi_,
+      memcpyAsync(
+          destination.data(),
+          ::testing::_,
+          kLen,
+          kMockMemcpyH2D,
+          static_cast<MockStream>(callerStream)))
+      .WillOnce([&](void* dst, const void* src, size_t len, auto, auto) {
+        std::memcpy(dst, src, len);
+        return Ok();
+      });
+  EXPECT_CALL(*cudaApi_, eventCreate(::testing::_))
+      .WillOnce([](auto* event) -> Status {
+        *event = {};
+        return Ok();
+      });
+  EXPECT_CALL(
+      *cudaApi_,
+      eventRecord(::testing::_, static_cast<MockStream>(callerStream)))
+      .WillOnce(::testing::Return(Ok()));
+  EXPECT_CALL(*cudaApi_, eventQuery(::testing::_))
+      .WillOnce([](auto) -> Result<bool> {
+        return Err(ErrCode::DriverError, "test: event query failed");
+      });
+  EXPECT_CALL(
+      *cudaApi_, streamSynchronize(static_cast<MockStream>(callerStream)))
+      .WillOnce(::testing::Return(Ok()));
+  EXPECT_CALL(*cudaApi_, eventDestroy(::testing::_))
+      .WillOnce(::testing::Return(Ok()));
+
+  handleFrame(bytes, std::move(slab));
+
+  ASSERT_EQ(
+      future.wait_for(std::chrono::seconds(5)), std::future_status::ready);
+  EXPECT_TRUE(future.get().hasValue());
+  EXPECT_TRUE(
+      std::all_of(destination.begin(), destination.end(), [](uint8_t b) {
+        return b == uint8_t{0xCD};
+      }));
+  auto returnedSlabs = pool->acquire(2);
+  EXPECT_TRUE(returnedSlabs.hasValue());
+}
+
 TEST_F(TcpReceivePoolTest, VectorBackedVramReadReplyRemainsSynchronous) {
   makeTransport();
   auto frame = readReplyFrame(/*reqId=*/1, kLen);