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,6 +679,9 @@ void AsyncUDPSocket::handleRead() noexcept {
     return readCallback_->onNotifyDataAvailable(*this);
   }
 
+  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(
@@ -771,6 +774,9 @@ void AsyncUDPSocket::handleRead() noexcept {
 
       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_;