From: Kyle Nekritz <knekritz@fb.com>
Date: Thu, 3 Nov 2016 22:53:52 +0000 (-0700)
Subject: Add splitAtMost to IOBufQueue.
X-Git-Tag: v2016.11.07.00~5
X-Git-Url: http://demsky.eecs.uci.edu/git/?a=commitdiff_plain;h=4b4183cec9704344fa139d8eb3a962b28aad3f87;p=folly.git

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
---

diff --git a/folly/io/IOBufQueue.cpp b/folly/io/IOBufQueue.cpp
index 6faf5c79..4147ea2d 100644
--- a/folly/io/IOBufQueue.cpp
+++ b/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) {
-      throw std::underflow_error(
-          "Attempt to remove more bytes than are present in IOBufQueue");
+      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();
diff --git a/folly/io/IOBufQueue.h b/folly/io/IOBufQueue.h
index bbf28fe9..30b7fed3 100644
--- a/folly/io/IOBufQueue.h
+++ b/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;
diff --git a/folly/io/test/IOBufQueueTest.cpp b/folly/io/test/IOBufQueueTest.cpp
index 737c4bf4..a5e8988c 100644
--- a/folly/io/test/IOBufQueueTest.cpp
+++ b/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"));