From 5627ce3a6622db1376b82d49647abfd32b70c02b Mon Sep 17 00:00:00 2001
From: Aman Sharma <amsharma@meta.com>
Date: Wed, 22 Jan 2025 16:27:07 -0800
Subject: [PATCH] Make a writemGSO API with just iovecs [3/n]

Summary:
**Background for the diff stack**: We'd like to excise the usage of folly from a library that we maintain (quic), which is why we'd like to have a version of `writemGSO` that takes in iovecs instead of IOBufs.

**What this diff is doing**: Previously, `writemGSO` called a helper function `writeImpl` that takes in `IOBuf`s. In this diff, I'm having two separate helper functions `writeImpl` and `writeImplIOBufs`. The latter uses `IOBuf`s, while the former doesn't. The version of `writemGSO` that uses `IOBuf`s will call `writeImplIOBufs`, while the version that just uses iovecs will call `writeImpl`.

Reviewed By: jbeshay

Differential Revision: D68459292

fbshipit-source-id: 954fca1b805399c9716a73206c1782a1f03c6e13
---
 folly/io/async/AsyncUDPSocket.cpp | 58 ++++++++++++++++---------------
 folly/io/async/AsyncUDPSocket.h   | 11 +++++-
 2 files changed, 40 insertions(+), 29 deletions(-)

diff --git a/folly/io/async/AsyncUDPSocket.cpp b/folly/io/async/AsyncUDPSocket.cpp
index 86cb3b2f613..25d29375588 100644
--- a/folly/io/async/AsyncUDPSocket.cpp
+++ b/folly/io/async/AsyncUDPSocket.cpp
@@ -847,7 +847,7 @@ int AsyncUDPSocket::writemGSO(
     controlPtr = control;
 #endif
     FOLLY_POP_WARNING
-    ret = writeImpl(addrs, bufs, count, vec, options, controlPtr);
+    ret = writeImplIOBufs(addrs, bufs, count, vec, options, controlPtr);
   } else {
     std::unique_ptr<mmsghdr[]> vec(new mmsghdr[count]);
 #ifdef FOLLY_HAVE_MSG_ERRQUEUE
@@ -859,7 +859,7 @@ int AsyncUDPSocket::writemGSO(
     memset(control.get(), 0, controlBufSize);
     controlPtr = control.get();
 #endif
-    ret = writeImpl(addrs, bufs, count, vec.get(), options, controlPtr);
+    ret = writeImplIOBufs(addrs, bufs, count, vec.get(), options, controlPtr);
   }
 
   return ret;
@@ -988,29 +988,18 @@ void AsyncUDPSocket::fillMsgVec(
   }
 }
 
-int AsyncUDPSocket::writeImpl(
+int AsyncUDPSocket::writeImplIOBufs(
     Range<SocketAddress const*> addrs,
     const std::unique_ptr<folly::IOBuf>* bufs,
     size_t count,
     struct mmsghdr* msgvec,
     const WriteOptions* options,
     char* control) {
-  // most times we have a single destination addr
-  auto addr_count = addrs.size();
-  constexpr size_t kAddrCountMax = 1;
-  small_vector<full_sockaddr_storage, kAddrCountMax> addrStorage(addr_count);
-
-  for (size_t i = 0; i < addr_count; i++) {
-    addrs[i].getAddress(&addrStorage[i].storage);
-    addrStorage[i].len = folly::to_narrow(addrs[i].getActualSize());
-  }
-
   size_t iov_count = 0;
   for (size_t i = 0; i < count; i++) {
     iov_count += bufs[i]->countChainElements();
   }
 
-  int ret;
   constexpr size_t kSmallSizeMax = 8;
   if (iov_count <= kSmallSizeMax) {
     // suppress "warning: variable length array 'vec' is used [-Wvla]"
@@ -1020,31 +1009,44 @@ int AsyncUDPSocket::writeImpl(
     size_t messageIovLens[BOOST_PP_IF(FOLLY_HAVE_VLA_01, count, kSmallSizeMax)];
     FOLLY_POP_WARNING
     fillIoVec(bufs, iov, messageIovLens, count, iov_count);
-    fillMsgVec(
-        range(addrStorage),
-        messageIovLens,
-        count,
-        msgvec,
-        iov,
-        options,
-        control);
-    ret = sendmmsg(fd_, msgvec, count, 0);
+    return writeImpl(
+        addrs, messageIovLens, iov, count, msgvec, options, control);
   } else {
     std::unique_ptr<iovec[]> iov(new iovec[iov_count]);
     std::unique_ptr<size_t[]> messageIovLens(new size_t[count]);
     fillIoVec(bufs, iov.get(), messageIovLens.get(), count, iov_count);
-    fillMsgVec(
-        range(addrStorage),
+    return writeImpl(
+        addrs,
         messageIovLens.get(),
+        iov.get(),
         count,
         msgvec,
-        iov.get(),
         options,
         control);
-    ret = sendmmsg(fd_, msgvec, count, 0);
+  }
+}
+
+int AsyncUDPSocket::writeImpl(
+    Range<SocketAddress const*> addrs,
+    size_t* messageIovLens,
+    struct iovec* iov,
+    size_t count,
+    struct mmsghdr* msgvec,
+    const WriteOptions* options,
+    char* control) {
+  // most times we have a single destination addr
+  auto addr_count = addrs.size();
+  constexpr size_t kAddrCountMax = 1;
+  small_vector<full_sockaddr_storage, kAddrCountMax> addrStorage(addr_count);
+
+  for (size_t i = 0; i < addr_count; i++) {
+    addrs[i].getAddress(&addrStorage[i].storage);
+    addrStorage[i].len = folly::to_narrow(addrs[i].getActualSize());
   }
 
-  return ret;
+  fillMsgVec(
+      range(addrStorage), messageIovLens, count, msgvec, iov, options, control);
+  return sendmmsg(fd_, msgvec, count, 0);
 }
 
 ssize_t AsyncUDPSocket::recvmsg(struct msghdr* msg, int flags) {
diff --git a/folly/io/async/AsyncUDPSocket.h b/folly/io/async/AsyncUDPSocket.h
index 274c174361e..7f97b7e8a7e 100644
--- a/folly/io/async/AsyncUDPSocket.h
+++ b/folly/io/async/AsyncUDPSocket.h
@@ -557,7 +557,7 @@ class AsyncUDPSocket : public EventHandler {
       const WriteOptions* options,
       char* control);
 
-  virtual int writeImpl(
+  virtual int writeImplIOBufs(
       Range<SocketAddress const*> addrs,
       const std::unique_ptr<folly::IOBuf>* bufs,
       size_t count,
@@ -565,6 +565,15 @@ class AsyncUDPSocket : public EventHandler {
       const WriteOptions* options,
       char* control);
 
+  virtual int writeImpl(
+      Range<SocketAddress const*> addrs,
+      size_t* messageIovLens,
+      struct iovec* iov,
+      size_t count,
+      struct mmsghdr* msgvec,
+      const WriteOptions* options,
+      char* control);
+
   virtual ssize_t writevImpl(
       netops::Msgheader* msg, [[maybe_unused]] WriteOptions options);