Add splitAtMost to IOBufQueue.

Summary: This allows getting up to n bytes from the queue without first having to check the length.

Reviewed By: yfeldblum

Differential Revision: D4083484

fbshipit-source-id: 2a468992c97f036c22f1a0d9f830e6d5286a4bc2

folly/io/IOBufQueue.cpp

@@ -187,13 +187,16 @@ IOBufQueue::preallocateSlow(uint64_t min, uint64_t newAllocationSize,
                    std::min(max, last->tailroom()));
 }
 
-unique_ptr<IOBuf>
-IOBufQueue::split(size_t n) {
+unique_ptr<IOBuf> IOBufQueue::split(size_t n, bool throwOnUnderflow) {
   unique_ptr<IOBuf> result;
   while (n != 0) {
     if (head_ == nullptr) {
+      if (throwOnUnderflow) {
         throw std::underflow_error(
             "Attempt to remove more bytes than are present in IOBufQueue");
+      } else {
+        break;
+      }
     } else if (head_->length() <= n) {
       n -= head_->length();
       chainLength_ -= head_->length();

folly/io/IOBufQueue.h

@@ -195,7 +195,17 @@ class IOBufQueue {
    * @throws std::underflow_error if n exceeds the number of bytes
    *         in the queue.
    */
-  std::unique_ptr<folly::IOBuf> split(size_t n);
+  std::unique_ptr<folly::IOBuf> split(size_t n) {
+    return split(n, true);
+  }
+
+  /**
+   * Similar to split, but will return the entire queue instead of throwing
+   * if n exceeds the number of bytes in the queue.
+   */
+  std::unique_ptr<folly::IOBuf> splitAtMost(size_t n) {
+    return split(n, false);
+  }
 
   /**
    * Similar to IOBuf::trimStart, but works on the whole queue.  Will
@@ -283,6 +293,8 @@ class IOBufQueue {
   std::pair<void*,uint64_t> preallocateSlow(
     uint64_t min, uint64_t newAllocationSize, uint64_t max);
 
+  std::unique_ptr<folly::IOBuf> split(size_t n, bool throwOnUnderflow);
+
   static const size_t kChainLengthNotCached = (size_t)-1;
   /** Not copyable */
   IOBufQueue(const IOBufQueue&) = delete;

folly/io/test/IOBufQueueTest.cpp

@@ -159,6 +159,15 @@ TEST(IOBufQueue, Split) {
   checkConsistency(queue);
 }
 
+TEST(IOBufQueue, SplitAtMost) {
+  IOBufQueue queue(clOptions);
+  queue.append(stringToIOBuf(SCL("Hello,")));
+  queue.append(stringToIOBuf(SCL(" World")));
+  auto buf = queue.splitAtMost(9999);
+  EXPECT_EQ(buf->computeChainDataLength(), 12);
+  EXPECT_TRUE(queue.empty());
+}
+
 TEST(IOBufQueue, Preallocate) {
   IOBufQueue queue(clOptions);
   queue.append(string("Hello"));