Pre-prune iobufs from clone

Summary: It's pretty common during parsing protocols with cursors
that there will be empty iobufs along the way.

This changes it so that we prune empty IOBufs in the cloned
output before returning it to the user.

Reviewed By: @djwatson

Differential Revision: D2464488

diff --git a/folly/io/Cursor.h b/folly/io/Cursor.h
index 2aa4227f3386ded5e53dffee3e70883de51f563b..5d12b6f89069ff183faa23a1351cd579944dbc93 100644 (file)
--- a/folly/io/Cursor.h
+++ b/folly/io/Cursor.h
@@ -162,6 +162,7 @@ class CursorBase {
     if (LIKELY(length() >= sizeof(T))) {
       val = loadUnaligned<T>(data());
       offset_ += sizeof(T);
+      advanceBufferIfEmpty();
     } else {
       pullSlow(&val, sizeof(T));
     }
@@ -191,6 +192,7 @@ class CursorBase {
     if (LIKELY(length() >= len)) {
       str.append(reinterpret_cast<const char*>(data()), len);
       offset_ += len;
+      advanceBufferIfEmpty();
     } else {
       readFixedStringSlow(&str, len);
     }
@@ -232,16 +234,13 @@ class CursorBase {
       }
 
       skip(i);
-
-      if (UNLIKELY(!tryAdvanceBuffer())) {
-        throw std::out_of_range("string underflow");
-      }
     }
   }
 
   size_t skipAtMost(size_t len) {
     if (LIKELY(length() >= len)) {
       offset_ += len;
+      advanceBufferIfEmpty();
       return len;
     }
     return skipAtMostSlow(len);
@@ -250,6 +249,7 @@ class CursorBase {
   void skip(size_t len) {
     if (LIKELY(length() >= len)) {
       offset_ += len;
+      advanceBufferIfEmpty();
     } else {
       skipSlow(len);
     }
@@ -260,6 +260,7 @@ class CursorBase {
     if (LIKELY(length() >= len)) {
       memcpy(buf, data(), len);
       offset_ += len;
+      advanceBufferIfEmpty();
       return len;
     }
     return pullAtMostSlow(buf, len);
@@ -269,6 +270,7 @@ class CursorBase {
     if (LIKELY(length() >= len)) {
       memcpy(buf, data(), len);
       offset_ += len;
+      advanceBufferIfEmpty();
     } else {
       pullSlow(buf, len);
     }
@@ -321,6 +323,7 @@ class CursorBase {
         }
 
         offset_ += len;
+        advanceBufferIfEmpty();
         return copied + len;
       }
 
@@ -419,6 +422,12 @@ class CursorBase {
     return true;
   }
 
+  void advanceBufferIfEmpty() {
+    if (length() == 0) {
+      tryAdvanceBuffer();
+    }
+  }
+
   BufType* crtBuf_;
   size_t offset_ = 0;
 
@@ -433,6 +442,7 @@ class CursorBase {
     }
     str->append(reinterpret_cast<const char*>(data()), len);
     offset_ += len;
+    advanceBufferIfEmpty();
   }
 
   size_t pullAtMostSlow(void* buf, size_t len) {
@@ -449,6 +459,7 @@ class CursorBase {
     }
     memcpy(p, data(), len);
     offset_ += len;
+    advanceBufferIfEmpty();
     return copied + len;
   }
 
@@ -468,6 +479,7 @@ class CursorBase {
       len -= available;
     }
     offset_ += len;
+    advanceBufferIfEmpty();
     return skipped + len;
   }
 

diff --git a/folly/io/test/IOBufCursorTest.cpp b/folly/io/test/IOBufCursorTest.cpp
index 3d1a4014e339f8ab7e5252162143aded979543e5..392f87b18fb3a3f96c74934c8ab3efca207d0bb1 100644 (file)
--- a/folly/io/test/IOBufCursorTest.cpp
+++ b/folly/io/test/IOBufCursorTest.cpp
@@ -415,6 +415,47 @@ TEST(IOBuf, cloneAndInsert) {
   }
 }
 
+TEST(IOBuf, cloneWithEmptyBufAtStart) {
+  folly::IOBufEqual eq;
+  auto empty = IOBuf::create(0);
+  auto hel = IOBuf::create(3);
+  append(hel, "hel");
+  auto lo = IOBuf::create(2);
+  append(lo, "lo");
+
+  auto iobuf = empty->clone();
+  iobuf->prependChain(hel->clone());
+  iobuf->prependChain(lo->clone());
+  iobuf->prependChain(empty->clone());
+  iobuf->prependChain(hel->clone());
+  iobuf->prependChain(lo->clone());
+  iobuf->prependChain(empty->clone());
+  iobuf->prependChain(lo->clone());
+  iobuf->prependChain(hel->clone());
+  iobuf->prependChain(lo->clone());
+  iobuf->prependChain(lo->clone());
+
+  Cursor cursor(iobuf.get());
+  std::unique_ptr<IOBuf> cloned;
+  char data[3];
+  cursor.pull(&data, 3);
+  cursor.clone(cloned, 2);
+  EXPECT_EQ(1, cloned->countChainElements());
+  EXPECT_EQ(2, cloned->length());
+  EXPECT_TRUE(eq(lo, cloned));
+
+  cursor.pull(&data, 3);
+  EXPECT_EQ("hel", std::string(data, sizeof(data)));
+
+  cursor.skip(2);
+  cursor.clone(cloned, 2);
+  EXPECT_TRUE(eq(lo, cloned));
+
+  std::string hello = cursor.readFixedString(5);
+  cursor.clone(cloned, 2);
+  EXPECT_TRUE(eq(lo, cloned));
+}
+
 TEST(IOBuf, Appender) {
   std::unique_ptr<IOBuf> head(IOBuf::create(10));
   append(head, "hello");