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 2aa4227..5d12b6f 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 3d1a401..392f87b 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");