add IOBuf::moveToFbString

Summary:
Convert a IOBuf chain into a fbstring, with zero copies in the common case
(unshared, unchained, no headroom) and one copy in all other cases.

Test Plan: tests added

Reviewed By: simpkins@fb.com

FB internal diff: D621210

folly/experimental/io/IOBuf.cpp

@@ -353,6 +353,15 @@ void IOBuf::coalesceSlow(size_t maxLength) {
 
   uint64_t newHeadroom = headroom();
   uint64_t newTailroom = end->prev_->tailroom();
+  coalesceAndReallocate(newHeadroom, newLength, end, newTailroom);
+  // We should be only element left in the chain now
+  assert(length_ >= maxLength || !isChained());
+}
+
+void IOBuf::coalesceAndReallocate(size_t newHeadroom,
+                                  size_t newLength,
+                                  IOBuf* end,
+                                  size_t newTailroom) {
   uint64_t newCapacity = newLength + newHeadroom + newTailroom;
   if (newCapacity > UINT32_MAX) {
     throw std::overflow_error("IOBuf chain too large to coalesce");
@@ -370,11 +379,15 @@ void IOBuf::coalesceSlow(size_t maxLength) {
   uint8_t* newData = newBuf + newHeadroom;
   uint8_t* p = newData;
   IOBuf* current = this;
+  size_t remaining = newLength;
   do {
+    assert(current->length_ <= remaining);
+    remaining -= current->length_;
     memcpy(p, current->data_, current->length_);
     p += current->length_;
     current = current->next_;
   } while (current != end);
+  assert(remaining == 0);
 
   // Point at the new buffer
   if (flags_ & kFlagExt) {
@@ -395,10 +408,9 @@ void IOBuf::coalesceSlow(size_t maxLength) {
   // Separate from the rest of our chain.
   // Since we don't store the unique_ptr returned by separateChain(),
   // this will immediately delete the returned subchain.
-  (void)separateChain(next_, end->prev_);
-
-  // We should be only element left in the chain now
-  assert(length_ >= maxLength || !isChained());
+  if (isChained()) {
+    (void)separateChain(next_, current->prev_);
+  }
 }
 
 void IOBuf::decrementRefcount() {
@@ -598,4 +610,29 @@ void IOBuf::initExtBuffer(uint8_t* buf, size_t mallocSize,
   *infoReturn = sharedInfo;
 }
 
+fbstring IOBuf::moveToFbString() {
+  // Externally allocated buffers (malloc) are just fine, everything else needs
+  // to be turned into one.
+  if (flags_ != kFlagExt ||  // not malloc()-ed
+      headroom() != 0 ||     // malloc()-ed block doesn't start at beginning
+      tailroom() == 0 ||     // no room for NUL terminator
+      isShared() ||          // shared
+      isChained()) {         // chained
+    // We might as well get rid of all head and tailroom if we're going
+    // to reallocate; we need 1 byte for NUL terminator.
+    coalesceAndReallocate(0, computeChainDataLength(), this, 1);
+  }
+
+  // Ensure NUL terminated
+  *writableTail() = 0;
+  fbstring str(reinterpret_cast<char*>(writableData()),
+               length(),  capacity(),
+               AcquireMallocatedString());
+
+  // Reset to internal buffer.
+  flags_ = 0;
+  clear();
+  return str;
+}
+
 } // folly

folly/experimental/io/IOBuf.h

@@ -27,6 +27,8 @@
 #include <limits>
 #include <type_traits>
 
+#include "folly/FBString.h"
+
 namespace folly {
 
 /**
@@ -896,6 +898,13 @@ class IOBuf {
   void* operator new(size_t size, void* ptr);
   void operator delete(void* ptr);
 
+  /**
+   * Destructively convert this IOBuf to a fbstring efficiently.
+   * We rely on fbstring's AcquireMallocatedString constructor to
+   * transfer memory.
+   */
+  fbstring moveToFbString();
+
  private:
   enum FlagsEnum {
     kFlagExt = 0x1,
@@ -966,6 +975,13 @@ class IOBuf {
   void unshareOneSlow();
   void unshareChained();
   void coalesceSlow(size_t maxLength=std::numeric_limits<size_t>::max());
+  // newLength must be the entire length of the buffers between this and
+  // end (no truncation)
+  void coalesceAndReallocate(
+      size_t newHeadroom,
+      size_t newLength,
+      IOBuf* end,
+      size_t newTailroom);
   void decrementRefcount();
   void reserveSlow(uint32_t minHeadroom, uint32_t minTailroom);
 

folly/experimental/io/test/IOBufTest.cpp

@@ -17,6 +17,9 @@
 #include "folly/experimental/io/IOBuf.h"
 #include "folly/experimental/io/TypedIOBuf.h"
 
+// googletest requires std::tr1::tuple, not std::tuple
+#include <tr1/tuple>
+
 #include <gflags/gflags.h>
 #include <boost/random.hpp>
 #include <gtest/gtest.h>
@@ -24,6 +27,7 @@
 #include "folly/Malloc.h"
 #include "folly/Range.h"
 
+using folly::fbstring;
 using folly::IOBuf;
 using folly::TypedIOBuf;
 using folly::StringPiece;
@@ -661,6 +665,64 @@ TEST(TypedIOBuf, Simple) {
   }
 }
 
+// chain element size, number of elements in chain, shared
+class MoveToFbStringTest
+  : public ::testing::TestWithParam<std::tr1::tuple<int, int, bool>> {
+ protected:
+  void SetUp() {
+    std::tr1::tie(elementSize_, elementCount_, shared_) = GetParam();
+    buf_ = makeBuf();
+    for (int i = 0; i < elementCount_ - 1; ++i) {
+      buf_->prependChain(makeBuf());
+    }
+    EXPECT_EQ(elementCount_, buf_->countChainElements());
+    EXPECT_EQ(elementCount_ * elementSize_, buf_->computeChainDataLength());
+    if (shared_) {
+      buf2_ = buf_->clone();
+      EXPECT_EQ(elementCount_, buf2_->countChainElements());
+      EXPECT_EQ(elementCount_ * elementSize_, buf2_->computeChainDataLength());
+    }
+  }
+
+  std::unique_ptr<IOBuf> makeBuf() {
+    auto buf = IOBuf::create(elementSize_);
+    memset(buf->writableTail(), 'x', elementSize_);
+    buf->append(elementSize_);
+    return buf;
+  }
+
+  void check(std::unique_ptr<IOBuf>& buf) {
+    fbstring str = buf->moveToFbString();
+    EXPECT_EQ(elementCount_ * elementSize_, str.size());
+    EXPECT_EQ(elementCount_ * elementSize_, strspn(str.c_str(), "x"));
+    EXPECT_EQ(0, buf->length());
+    EXPECT_EQ(1, buf->countChainElements());
+    EXPECT_EQ(0, buf->computeChainDataLength());
+    EXPECT_FALSE(buf->isChained());
+  }
+
+  int elementSize_;
+  int elementCount_;
+  bool shared_;
+  std::unique_ptr<IOBuf> buf_;
+  std::unique_ptr<IOBuf> buf2_;
+};
+
+TEST_P(MoveToFbStringTest, Simple) {
+  check(buf_);
+  if (shared_) {
+    check(buf2_);
+  }
+}
+
+INSTANTIATE_TEST_CASE_P(
+    MoveToFbString,
+    MoveToFbStringTest,
+    ::testing::Combine(
+        ::testing::Values(0, 1, 24, 256, 1 << 10, 1 << 20),  // element size
+        ::testing::Values(1, 2, 10),                         // element count
+        ::testing::Bool()));                                 // shared
+
 int main(int argc, char** argv) {
   testing::InitGoogleTest(&argc, argv);
   google::ParseCommandLineFlags(&argc, &argv, true);