Let AsyncUDPSocket send to multiple addresses

Summary:
Add support for variable number of addrs so we can send data to multiple destinations.

(Note: this ignores all push blocking failures!)

Reviewed By: kevin-vigor

Differential Revision: D21032857

fbshipit-source-id: f73b98d44f5d7d92f3692dfddb9b1c76ebcc51c5

folly/io/async/AsyncUDPSocket.cpp

@@ -23,6 +23,7 @@
 #include <folly/portability/Fcntl.h>
 #include <folly/portability/Sockets.h>
 #include <folly/portability/Unistd.h>
+#include <folly/small_vector.h>
 
 #include <boost/preprocessor/control/if.hpp>
 #include <cerrno>
@@ -357,14 +358,14 @@ ssize_t AsyncUDPSocket::writev(
  * ::sendmmsg.
  */
 int AsyncUDPSocket::writem(
-    const folly::SocketAddress& address,
+    Range<SocketAddress const*> addrs,
     const std::unique_ptr<folly::IOBuf>* bufs,
     size_t count) {
-  return writemGSO(address, bufs, count, nullptr);
+  return writemGSO(addrs, bufs, count, nullptr);
 }
 
 int AsyncUDPSocket::writemGSO(
-    const folly::SocketAddress& address,
+    Range<SocketAddress const*> addrs,
     const std::unique_ptr<folly::IOBuf>* bufs,
     size_t count,
     const int* gso) {
@@ -390,7 +391,7 @@ int AsyncUDPSocket::writemGSO(
     }
 #endif
     FOLLY_POP_WARNING
-    ret = writeImpl(address, bufs, count, vec, gso, gsoControl);
+    ret = writeImpl(addrs, bufs, count, vec, gso, gsoControl);
   } else {
     std::unique_ptr<mmsghdr[]> vec(new mmsghdr[count]);
 #ifdef FOLLY_HAVE_MSG_ERRQUEUE
@@ -400,15 +401,14 @@ int AsyncUDPSocket::writemGSO(
       gsoControl = control.get();
     }
 #endif
-    ret = writeImpl(address, bufs, count, vec.get(), gso, gsoControl);
+    ret = writeImpl(addrs, bufs, count, vec.get(), gso, gsoControl);
   }
 
   return ret;
 }
 
 void AsyncUDPSocket::fillMsgVec(
-    sockaddr_storage* addr,
-    socklen_t addr_len,
+    Range<full_sockaddr_storage*> addrs,
     const std::unique_ptr<folly::IOBuf>* bufs,
     size_t count,
     struct mmsghdr* msgvec,
@@ -416,6 +416,8 @@ void AsyncUDPSocket::fillMsgVec(
     size_t iov_count,
     const int* gso,
     char* gsoControl) {
+  auto addr_count = addrs.size();
+  DCHECK(addr_count);
   size_t remaining = iov_count;
 
   size_t iov_pos = 0;
@@ -425,8 +427,15 @@ void AsyncUDPSocket::fillMsgVec(
     size_t iovec_len = ret.numIovecs;
     remaining -= iovec_len;
     auto& msg = msgvec[i].msg_hdr;
-    msg.msg_name = reinterpret_cast<void*>(addr);
-    msg.msg_namelen = addr_len;
+    // if we have less addrs compared to count
+    // we use the last addr
+    if (i < addr_count) {
+      msg.msg_name = reinterpret_cast<void*>(&addrs[i].storage);
+      msg.msg_namelen = addrs[i].len;
+    } else {
+      msg.msg_name = reinterpret_cast<void*>(&addrs[addr_count - 1].storage);
+      msg.msg_namelen = addrs[addr_count - 1].len;
+    }
     msg.msg_iov = &iov[iov_pos];
     msg.msg_iovlen = iovec_len;
 #ifdef FOLLY_HAVE_MSG_ERRQUEUE
@@ -459,14 +468,21 @@ void AsyncUDPSocket::fillMsgVec(
 }
 
 int AsyncUDPSocket::writeImpl(
-    const folly::SocketAddress& address,
+    Range<SocketAddress const*> addrs,
     const std::unique_ptr<folly::IOBuf>* bufs,
     size_t count,
     struct mmsghdr* msgvec,
     const int* gso,
     char* gsoControl) {
-  sockaddr_storage addrStorage;
-  address.getAddress(&addrStorage);
+  // 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++) {
@@ -482,8 +498,7 @@ int AsyncUDPSocket::writeImpl(
     iovec iov[BOOST_PP_IF(FOLLY_HAVE_VLA_01, iov_count, kSmallSizeMax)];
     FOLLY_POP_WARNING
     fillMsgVec(
-        &addrStorage,
-        folly::to_narrow(address.getActualSize()),
+        range(addrStorage),
         bufs,
         count,
         msgvec,
@@ -495,8 +510,7 @@ int AsyncUDPSocket::writeImpl(
   } else {
     std::unique_ptr<iovec[]> iov(new iovec[iov_count]);
     fillMsgVec(
-        &addrStorage,
-        folly::to_narrow(address.getActualSize()),
+        range(addrStorage),
         bufs,
         count,
         msgvec,

folly/io/async/AsyncUDPSocket.h

@@ -169,7 +169,7 @@ class AsyncUDPSocket : public EventHandler {
    * of size num
    */
   virtual int writem(
-      const folly::SocketAddress& address,
+      Range<SocketAddress const*> addrs,
       const std::unique_ptr<folly::IOBuf>* bufs,
       size_t count);
 
@@ -197,7 +197,7 @@ class AsyncUDPSocket : public EventHandler {
    *  verify GSO is supported on this platform by calling getGSO
    */
   virtual int writemGSO(
-      const folly::SocketAddress& address,
+      Range<SocketAddress const*> addrs,
       const std::unique_ptr<folly::IOBuf>* bufs,
       size_t count,
       const int* gso);
@@ -365,6 +365,11 @@ class AsyncUDPSocket : public EventHandler {
       SocketOptionKey::ApplyPos pos);
 
  protected:
+  struct full_sockaddr_storage {
+    sockaddr_storage storage;
+    socklen_t len;
+  };
+
   virtual ssize_t
   sendmsg(NetworkSocket socket, const struct msghdr* message, int flags) {
     return netops::sendmsg(socket, message, flags);
@@ -379,8 +384,7 @@ class AsyncUDPSocket : public EventHandler {
   }
 
   void fillMsgVec(
-      sockaddr_storage* addr,
-      socklen_t addr_len,
+      Range<full_sockaddr_storage*> addrs,
       const std::unique_ptr<folly::IOBuf>* bufs,
       size_t count,
       struct mmsghdr* msgvec,
@@ -390,7 +394,7 @@ class AsyncUDPSocket : public EventHandler {
       char* gsoControl);
 
   virtual int writeImpl(
-      const folly::SocketAddress& address,
+      Range<SocketAddress const*> addrs,
       const std::unique_ptr<folly::IOBuf>* bufs,
       size_t count,
       struct mmsghdr* msgvec,

folly/io/async/test/AsyncUDPSocketGSOGROTest.cpp

@@ -336,7 +336,8 @@ class UDPClient : private AsyncUDPSocket::ReadCallback, private AsyncTimeout {
   virtual void writePing(
       const std::vector<std::unique_ptr<folly::IOBuf>>& vec,
       const int* gso) {
-    socket_->writemGSO(server_, vec.data(), vec.size(), gso);
+    socket_->writemGSO(
+        folly::range(&server_, &server_ + 1), vec.data(), vec.size(), gso);
   }
 
   void getReadBuffer(void** buf, size_t* len) noexcept override {

folly/io/async/test/AsyncUDPSocketSendmmsgTest.cpp

@@ -235,7 +235,8 @@ class UDPClient : private AsyncUDPSocket::ReadCallback, private AsyncTimeout {
   }
 
   virtual void writePing(IOBufVec& bufs) {
-    socket_->writem(server_, bufs.data(), bufs.size());
+    socket_->writem(
+        folly::range(&server_, &server_ + 1), bufs.data(), bufs.size());
   }
 
   void getReadBuffer(void** buf, size_t* len) noexcept override {