folly/io_uring: deliver writeStarting() exactly once via shared WriteCallbackWithState

Summary:
What: Deliver `writeStarting()` to the `WriteCallback` exactly once per write on the native `AsyncSocket` + io_uring send path. `IoUringSend`'s `SendRequest` now holds the `WriteCallbackWithState`, and `onSendStarted()` calls `notifyOnWrite()` (which dedups via `writeInProgress_`); `AsyncSocket::writeImpl` passes its `callbackWithState` into `iouSendHandle_->write()` so an immediate write arrives already-notified.

Why: The io_uring send handle fired `writeStarting()` on every SQE submit while `AsyncSocket` also fired it via `notifyOnWrite()`, so a single write got two calls (and more on partial / event-re-registration re-submits). `RocketServerConnection::writeStarting()` DCHECKs that `startRawByteOffset` is unset, so the duplicate aborts the process (`RocketServerConnection.cpp:718`); in opt builds it silently corrupts the byte offset. This crashed graphstore once it served on native `AsyncSocket` over io_uring with zero-copy (D111176951).

```
Rocket flushWrites
  -> AsyncSocket::writeImpl
       |- notifyOnWrite() --------------> writeStarting() #1   (sets startRawByteOffset)
       `- iouSendHandle_->write -> submit
            `- onSendStarted() ---------> writeStarting() facebook#2   -> DCHECK abort
                                          (+ again on every partial / re-arm re-submit)

  fix: both paths funnel through one WriteCallbackWithState.notifyOnWrite() -> fires once
```

Rendered flow: https://pxl.cl/bBVG9

Reviewed By: spikeh

Differential Revision: D111316678

fbshipit-source-id: 037ff58303d0fbff4e8e77af6b3921f9463a7196

Author: jlhe97

third-party/folly/src/folly/io/async/AsyncSocket.cpp

@@ -2102,7 +2102,7 @@ void AsyncSocket::writeImpl(
       }
     }
     iouSendHandle_->write(
-        callback,
+        callbackWithState,
         vec + countWritten,
         count - countWritten,
         partialWritten,

third-party/folly/src/folly/io/async/IoUringSend.cpp

@@ -39,7 +39,7 @@ class IoUringSendHandle::SendRequest : public IoSqeBase {
   }
 
   explicit SendRequest(
-      AsyncWriter::WriteCallback* callback,
+      WriteCallbackWithState callback,
       const struct iovec* iov,
       size_t iovCount,
       size_t partialWritten,
@@ -48,8 +48,11 @@ class IoUringSendHandle::SendRequest : public IoSqeBase {
       WriteFlags flags,
       NetworkSocket fd)
       : IoSqeBase(IoSqeBase::Type::Write),
-        callback_(callback),
-        releaseCb_(callback ? callback->getReleaseIOBufCallback() : nullptr),
+        callbackWithState_(callback),
+        releaseCb_(
+            callback.getCallback()
+                ? callback.getCallback()->getReleaseIOBufCallback()
+                : nullptr),
         iovRemaining_(iovCount),
         bytesWritten_(bytesWritten),
         data_(std::move(data)),
@@ -81,7 +84,10 @@ class IoUringSendHandle::SendRequest : public IoSqeBase {
   }
   SendRequest* getNext() { return next_; }
   void append(SendRequest* request) { next_ = request; }
-  AsyncWriter::WriteCallback* getCallback() { return callback_; }
+  AsyncWriter::WriteCallback* getCallback() {
+    return callbackWithState_.getCallback();
+  }
+  void notifyOnWrite() { callbackWithState_.notifyOnWrite(); }
   size_t getTotalBytesWritten() { return bytesWritten_; }
   folly::IOBuf* getData() const { return data_.get(); }
   bool notifPending() const { return refs_ > 1; }
@@ -107,7 +113,7 @@ class IoUringSendHandle::SendRequest : public IoSqeBase {
     CHECK(handle_ == nullptr);
     void* buf = alloc(msg_.msg_iovlen);
     auto clone = new (buf) SendRequest(
-        callback_,
+        callbackWithState_,
         msg_.msg_iov,
         msg_.msg_iovlen,
         0,
@@ -222,7 +228,7 @@ class IoUringSendHandle::SendRequest : public IoSqeBase {
     msg_.msg_iovlen = std::min<size_t>(iovRemaining_, kIovMax);
   }
 
-  AsyncWriter::WriteCallback* callback_;
+  WriteCallbackWithState callbackWithState_;
   AsyncWriter::ReleaseIOBufCallback* releaseCb_;
   size_t iovRemaining_;
   size_t bytesWritten_;
@@ -338,7 +344,7 @@ bool IoUringSendHandle::update(uint16_t eventFlags) {
 }
 
 void IoUringSendHandle::write(
-    AsyncWriter::WriteCallback* callback,
+    WriteCallbackWithState callback,
     const struct iovec* iov,
     size_t iovCount,
     size_t partialWritten,
@@ -412,9 +418,7 @@ void IoUringSendHandle::trySubmit() {
 }
 
 void IoUringSendHandle::onSendStarted() {
-  if (auto* cb = requestHead_->getCallback()) {
-    cb->writeStarting();
-  }
+  requestHead_->notifyOnWrite();
 }
 
 void IoUringSendHandle::onSendPartial(size_t bytesWritten) {
@@ -516,7 +520,7 @@ bool IoUringSendHandle::update(uint16_t /*eventFlags*/) {
 }
 
 void IoUringSendHandle::write(
-    AsyncWriter::WriteCallback* /*callback*/,
+    WriteCallbackWithState /*callback*/,
     const struct iovec* /*iov*/,
     size_t /*iovCount*/,
     size_t /*partialWritten*/,

third-party/folly/src/folly/io/async/IoUringSend.h

@@ -23,6 +23,7 @@
 #include <folly/io/async/AsyncTransport.h>
 #include <folly/io/async/DelayedDestruction.h>
 #include <folly/io/async/IoUringBase.h>
+#include <folly/io/async/WriteCallbackWithState.h>
 #include <folly/net/NetworkSocket.h>
 
 namespace folly {
@@ -59,7 +60,7 @@ class IoUringSendHandle : public DelayedDestruction {
 
   bool update(uint16_t eventFlags);
   void write(
-      AsyncWriter::WriteCallback* callback,
+      WriteCallbackWithState callback,
       const struct iovec* iov,
       size_t iovCount,
       size_t partialWritten,

third-party/folly/src/folly/io/async/test/AsyncSocketTest2.cpp

@@ -1127,6 +1127,78 @@ TEST_P(AsyncSocketTest, MoveEventBaseWithInflightZeroCopyWrite) {
   server.verifyConnection(buf, kLen);
 }
 
+/**
+ * A partially-written buffer is handed to IoUringSendHandle and re-submitted to
+ * io_uring as the socket buffer drains. writeStarting() must reach the
+ * WriteCallback exactly once, not once per (re-)submission. Regression test for
+ * the double writeStarting() that aborted RocketServerConnection's DCHECK on
+ * the native AsyncSocket + io_uring send path.
+ */
+TEST_P(AsyncSocketTest, PartialWriteFiresWriteStartingOnce) {
+  if (GetParam() != BackendType::IO_URING) {
+    GTEST_SKIP() << "IoUringSendHandle send path is io_uring-only";
+  }
+
+  // Small socket buffers so a large write can't complete in one shot.
+  constexpr size_t kSockBufSize = 8 * 1024;
+  TestServer server(false, kSockBufSize);
+
+  SocketOptionMap options{
+      {{SOL_SOCKET, SO_SNDBUF}, int(kSockBufSize)},
+      {{SOL_SOCKET, SO_RCVBUF}, int(kSockBufSize)},
+      {{IPPROTO_TCP, TCP_NODELAY}, 1},
+  };
+
+  // The receiver drains on this thread; the sender gets its own EVB thread.
+  EventBase& senderEvb = getEventBase();
+  std::thread senderThread([&]() { senderEvb.loopForever(); });
+
+  ConnCallback ccb;
+  WriteCallback wcb;
+  std::shared_ptr<AsyncSocket> socket;
+
+  senderEvb.runInEventBaseThreadAndWait([&]() {
+    socket = AsyncSocket::newSocket(&senderEvb);
+    socket->connect(&ccb, server.getAddress(), 30, options);
+  });
+
+  std::shared_ptr<BlockingSocket> acceptedSocket = server.accept();
+
+  // Completion is signalled via an atomic set on the sender thread; wcb's
+  // non-atomic fields are only read after that thread is joined.
+  std::atomic<bool> writeDone{false};
+  wcb.successCallback = [&writeDone]() { writeDone = true; };
+
+  // Big enough to overflow the send+recv buffers, so the first write is partial
+  // and its remainder is re-submitted through IoUringSendHandle as we drain.
+  constexpr size_t kSendSize = 100 * 1024;
+  auto const sendBuf = std::vector<char>(kSendSize, 'a');
+
+  senderEvb.runInEventBaseThreadAndWait([&]() {
+    socket->write(&wcb, sendBuf.data(), kSendSize);
+  });
+
+  // Drain everything so the write ultimately succeeds.
+  std::vector<uint8_t> recvBuf(kSendSize);
+  auto bytesRead = acceptedSocket->readAll(recvBuf.data(), recvBuf.size());
+  ASSERT_EQ(kSendSize, bytesRead);
+  EXPECT_EQ(0, memcmp(recvBuf.data(), sendBuf.data(), bytesRead));
+
+  using clock = std::chrono::steady_clock;
+  auto const deadline = clock::now() + std::chrono::seconds(30);
+  while (!writeDone.load() && clock::now() < deadline) {
+    std::this_thread::yield();
+  }
+
+  senderEvb.terminateLoopSoon();
+  senderThread.join();
+  socket.reset();
+
+  EXPECT_EQ(STATE_SUCCEEDED, wcb.state);
+  // The core assertion: exactly one writeStarting despite multiple submits.
+  EXPECT_EQ(1, wcb.writeStartingInvocations);
+}
+
 /**
  * Test calling close() immediately after connect()
  */