Add AsyncUDPSocket::maxReadsPerEvent_ support

Summary: Add AsyncUDPSocket::maxReadsPerEvent_ support

Reviewed By: yfeldblum

Differential Revision: D21509991

fbshipit-source-id: 5c4f939c1cfdda8711864df3b5f5c58439f8d7bf

folly/io/async/AsyncUDPSocket.cpp

@@ -679,98 +679,104 @@ void AsyncUDPSocket::handleRead() noexcept {
     return readCallback_->onNotifyDataAvailable(*this);
   }
 
-  readCallback_->getReadBuffer(&buf, &len);
-  if (buf == nullptr || len == 0) {
-    AsyncSocketException ex(
-        AsyncSocketException::BAD_ARGS,
-        "AsyncUDPSocket::getReadBuffer() returned empty buffer");
-
-    auto cob = readCallback_;
-    readCallback_ = nullptr;
-
-    cob->onReadError(ex);
-    updateRegistration();
-    return;
-  }
+  size_t numReads = maxReadsPerEvent_ ? maxReadsPerEvent_ : size_t(-1);
+  EventBase* originalEventBase = eventBase_;
+  while (numReads-- && readCallback_ && eventBase_ == originalEventBase) {
+    readCallback_->getReadBuffer(&buf, &len);
+    if (buf == nullptr || len == 0) {
+      AsyncSocketException ex(
+          AsyncSocketException::BAD_ARGS,
+          "AsyncUDPSocket::getReadBuffer() returned empty buffer");
+
+      auto cob = readCallback_;
+      readCallback_ = nullptr;
+
+      cob->onReadError(ex);
+      updateRegistration();
+      return;
+    }
 
-  struct sockaddr_storage addrStorage;
-  socklen_t addrLen = sizeof(addrStorage);
-  memset(&addrStorage, 0, size_t(addrLen));
-  auto rawAddr = reinterpret_cast<sockaddr*>(&addrStorage);
-  rawAddr->sa_family = localAddress_.getFamily();
+    struct sockaddr_storage addrStorage;
+    socklen_t addrLen = sizeof(addrStorage);
+    memset(&addrStorage, 0, size_t(addrLen));
+    auto rawAddr = reinterpret_cast<sockaddr*>(&addrStorage);
+    rawAddr->sa_family = localAddress_.getFamily();
 
-  ssize_t bytesRead;
-  ReadCallback::OnDataAvailableParams params;
+    ssize_t bytesRead;
+    ReadCallback::OnDataAvailableParams params;
 
 #ifdef FOLLY_HAVE_MSG_ERRQUEUE
-  if (gro_.has_value() && (gro_.value() > 0)) {
-    char control[CMSG_SPACE(sizeof(uint16_t))] = {};
-    struct msghdr msg = {};
-    struct iovec iov = {};
-    struct cmsghdr* cmsg;
-    uint16_t* grosizeptr;
-
-    iov.iov_base = buf;
-    iov.iov_len = len;
-
-    msg.msg_iov = &iov;
-    msg.msg_iovlen = 1;
-
-    msg.msg_name = rawAddr;
-    msg.msg_namelen = addrLen;
-
-    msg.msg_control = control;
-    msg.msg_controllen = sizeof(control);
-
-    bytesRead = netops::recvmsg(fd_, &msg, MSG_TRUNC);
-
-    if (bytesRead >= 0) {
-      addrLen = msg.msg_namelen;
-      for (cmsg = CMSG_FIRSTHDR(&msg); cmsg != nullptr;
-           cmsg = CMSG_NXTHDR(&msg, cmsg)) {
-        if (cmsg->cmsg_level == SOL_UDP && cmsg->cmsg_type == UDP_GRO) {
-          grosizeptr = (uint16_t*)CMSG_DATA(cmsg);
-          params.gro_ = *grosizeptr;
-          break;
+    if (gro_.has_value() && (gro_.value() > 0)) {
+      char control[CMSG_SPACE(sizeof(uint16_t))] = {};
+      struct msghdr msg = {};
+      struct iovec iov = {};
+      struct cmsghdr* cmsg;
+      uint16_t* grosizeptr;
+
+      iov.iov_base = buf;
+      iov.iov_len = len;
+
+      msg.msg_iov = &iov;
+      msg.msg_iovlen = 1;
+
+      msg.msg_name = rawAddr;
+      msg.msg_namelen = addrLen;
+
+      msg.msg_control = control;
+      msg.msg_controllen = sizeof(control);
+
+      bytesRead = netops::recvmsg(fd_, &msg, MSG_TRUNC);
+
+      if (bytesRead >= 0) {
+        addrLen = msg.msg_namelen;
+        for (cmsg = CMSG_FIRSTHDR(&msg); cmsg != nullptr;
+             cmsg = CMSG_NXTHDR(&msg, cmsg)) {
+          if (cmsg->cmsg_level == SOL_UDP && cmsg->cmsg_type == UDP_GRO) {
+            grosizeptr = (uint16_t*)CMSG_DATA(cmsg);
+            params.gro_ = *grosizeptr;
+            break;
+          }
         }
       }
+    } else {
+      bytesRead = netops::recvfrom(fd_, buf, len, MSG_TRUNC, rawAddr, &addrLen);
     }
-  } else {
-    bytesRead = netops::recvfrom(fd_, buf, len, MSG_TRUNC, rawAddr, &addrLen);
-  }
 #else
-  bytesRead = netops::recvfrom(fd_, buf, len, MSG_TRUNC, rawAddr, &addrLen);
+    bytesRead = netops::recvfrom(fd_, buf, len, MSG_TRUNC, rawAddr, &addrLen);
 #endif
-  if (bytesRead >= 0) {
-    clientAddress_.setFromSockaddr(rawAddr, addrLen);
-
-    if (bytesRead > 0) {
-      bool truncated = false;
-      if ((size_t)bytesRead > len) {
-        truncated = true;
-        bytesRead = ssize_t(len);
+    if (bytesRead >= 0) {
+      clientAddress_.setFromSockaddr(rawAddr, addrLen);
+
+      if (bytesRead > 0) {
+        bool truncated = false;
+        if ((size_t)bytesRead > len) {
+          truncated = true;
+          bytesRead = ssize_t(len);
+        }
+
+        readCallback_->onDataAvailable(
+            clientAddress_, size_t(bytesRead), truncated, params);
+      }
+    } else {
+      if (errno == EAGAIN || errno == EWOULDBLOCK) {
+        // No data could be read without blocking the socket
+        return;
       }
 
-      readCallback_->onDataAvailable(
-          clientAddress_, size_t(bytesRead), truncated, params);
-    }
-  } else {
-    if (errno == EAGAIN || errno == EWOULDBLOCK) {
-      // No data could be read without blocking the socket
-      return;
-    }
+      AsyncSocketException ex(
+          AsyncSocketException::INTERNAL_ERROR, "::recvfrom() failed", errno);
 
-    AsyncSocketException ex(
-        AsyncSocketException::INTERNAL_ERROR, "::recvfrom() failed", errno);
+      // In case of UDP we can continue reading from the socket
+      // even if the current request fails. We notify the user
+      // so that he can do some logging/stats collection if he wants.
+      auto cob = readCallback_;
+      readCallback_ = nullptr;
 
-    // In case of UDP we can continue reading from the socket
-    // even if the current request fails. We notify the user
-    // so that he can do some logging/stats collection if he wants.
-    auto cob = readCallback_;
-    readCallback_ = nullptr;
+      cob->onReadError(ex);
+      updateRegistration();
 
-    cob->onReadError(ex);
-    updateRegistration();
+      return;
+    }
   }
 }
 

folly/io/async/AsyncUDPSocket.h

@@ -342,6 +342,30 @@ class AsyncUDPSocket : public EventHandler {
     return readCallback_ != nullptr;
   }
 
+  /**
+   * Set the maximum number of reads to execute from the underlying
+   * socket each time the EventBase detects that new ingress data is
+   * available. The default is kMaxReadsPerEvent
+   *
+   * @param maxReads  Maximum number of reads per data-available event;
+   *                  a value of zero means unlimited.
+   */
+  void setMaxReadsPerEvent(uint16_t maxReads) {
+    maxReadsPerEvent_ = maxReads;
+  }
+
+  /**
+   * Get the maximum number of reads this object will execute from
+   * the underlying socket each time the EventBase detects that new
+   * ingress data is available.
+   *
+   * @returns Maximum number of reads per data-available event; a value
+   *          of zero means unlimited.
+   */
+  uint16_t getMaxReadsPerEvent() const {
+    return maxReadsPerEvent_;
+  }
+
   virtual void detachEventBase();
 
   virtual void attachEventBase(folly::EventBase* evb);
@@ -405,6 +429,10 @@ class AsyncUDPSocket : public EventHandler {
 
   void failErrMessageRead(const AsyncSocketException& ex);
 
+  static auto constexpr kDefaultReadsPerEvent = 1;
+  uint16_t maxReadsPerEvent_{
+      kDefaultReadsPerEvent}; ///< Max reads per event loop iteration
+
   // Non-null only when we are reading
   ReadCallback* readCallback_;