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 6faf5c79b2b67895d0788413071eccd98cb757c8..4147ea2dd5de3af298bbb320d897f0a7fc98eb52 100644 (file)
--- 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 bbf28fe9f9a7fee2a600ca27d67c05a9cea4d66d..30b7fed3a6473a230ad5c709d4c09c03fb128ac5 100644 (file)
--- 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 737c4bf4109d1913571bd0c39093b3c63947924f..a5e8988cd64de024c91c83d45c2e6f25ffa3c78c 100644 (file)
--- 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"));