support stack-allocated IOBufs

Summary:
Previously, all IOBuf APIs required that IOBufs always be allocated on
the heap.  The only methods provided to create IOBufs returned
unique_ptr<IOBuf>.

This adds new methods to support creating IOBufs on the stack.  This is
useful in cases where the IOBuf will be short-lived, and the overhead of
the heap allocation is undesirable.  (One use case is to wrap an
existing buffer in a short-lived IOBuf so that it can be used with the
Cursor API.)

I have currently made IOBufs movable but not copyable.  (Move operations
clearly should move only a single IOBuf, but it is not clear if the copy
operators should copy only a single IOBuf or the entire chain.)

Test Plan:
Updated the unit tests to test the new CREATE, WRAP_BUFFER,
TAKE_OWNERSHIP, and COPY_BUFFER constructors, as well as the move
constructor and assignment operator.

Reviewed By: davejwatson@fb.com

FB internal diff: D1067341

folly/io/IOBuf.cpp

@@ -18,8 +18,11 @@
 
 #include "folly/io/IOBuf.h"
 
-#include "folly/Malloc.h"
+#include "folly/Conv.h"
 #include "folly/Likely.h"
+#include "folly/Malloc.h"
+#include "folly/Memory.h"
+#include "folly/ScopeGuard.h"
 
 #include <stdexcept>
 #include <assert.h>
@@ -51,8 +54,29 @@ enum : uint32_t {
   kDefaultCombinedBufSize = 1024
 };
 
+// Helper function for IOBuf::takeOwnership()
+void takeOwnershipError(bool freeOnError, void* buf,
+                        folly::IOBuf::FreeFunction freeFn,
+                        void* userData) {
+  if (!freeOnError) {
+    return;
+  }
+  if (!freeFn) {
+    free(buf);
+    return;
+  }
+  try {
+    freeFn(buf, userData);
+  } catch (...) {
+    // The user's free function is not allowed to throw.
+    // (We are already in the middle of throwing an exception, so
+    // we cannot let this exception go unhandled.)
+    abort();
+  }
 }
 
+} // unnamed namespace
+
 namespace folly {
 
 struct IOBuf::HeapPrefix {
@@ -160,6 +184,30 @@ void IOBuf::freeInternalBuf(void* buf, void* userData) {
   releaseStorage(storage, kDataInUse);
 }
 
+IOBuf::IOBuf(CreateOp, uint32_t capacity)
+  : next_(this),
+    prev_(this),
+    data_(nullptr),
+    length_(0),
+    flags_(0),
+    type_(kExtAllocated) {
+  allocExtBuffer(capacity, &buf_, &sharedInfo_, &capacity_);
+  data_ = buf_;
+}
+
+IOBuf::IOBuf(CopyBufferOp op, const void* buf, uint32_t size,
+             uint32_t headroom, uint32_t minTailroom)
+  : IOBuf(CREATE, headroom + size + minTailroom) {
+  advance(headroom);
+  memcpy(writableData(), buf, size);
+  append(size);
+}
+
+IOBuf::IOBuf(CopyBufferOp op, ByteRange br,
+             uint32_t headroom, uint32_t minTailroom)
+  : IOBuf(op, br.data(), br.size(), headroom, minTailroom) {
+}
+
 unique_ptr<IOBuf> IOBuf::create(uint32_t capacity) {
   // For smaller-sized buffers, allocate the IOBuf, SharedInfo, and the buffer
   // all with a single allocation.
@@ -196,22 +244,7 @@ unique_ptr<IOBuf> IOBuf::createCombined(uint32_t capacity) {
 }
 
 unique_ptr<IOBuf> IOBuf::createSeparate(uint32_t capacity) {
-  // Allocate an external buffer
-  uint8_t* buf;
-  SharedInfo* sharedInfo;
-  uint32_t actualCapacity;
-  allocExtBuffer(capacity, &buf, &sharedInfo, &actualCapacity);
-
-  // Allocate the IOBuf header
-  try {
-    return unique_ptr<IOBuf>(new IOBuf(kExtAllocated, 0,
-                                       buf, actualCapacity,
-                                       buf, 0,
-                                       sharedInfo));
-  } catch (...) {
-    free(buf);
-    throw;
-  }
+  return make_unique<IOBuf>(CREATE, capacity);
 }
 
 unique_ptr<IOBuf> IOBuf::createChain(
@@ -230,52 +263,71 @@ unique_ptr<IOBuf> IOBuf::createChain(
   return out;
 }
 
+IOBuf::IOBuf(TakeOwnershipOp, void* buf, uint32_t capacity, uint32_t length,
+             FreeFunction freeFn, void* userData,
+             bool freeOnError)
+  : next_(this),
+    prev_(this),
+    data_(static_cast<uint8_t*>(buf)),
+    buf_(static_cast<uint8_t*>(buf)),
+    length_(length),
+    capacity_(capacity),
+    flags_(kFlagFreeSharedInfo),
+    type_(kExtUserSupplied) {
+  try {
+    sharedInfo_ = new SharedInfo(freeFn, userData);
+  } catch (...) {
+    takeOwnershipError(freeOnError, buf, freeFn, userData);
+    throw;
+  }
+}
+
 unique_ptr<IOBuf> IOBuf::takeOwnership(void* buf, uint32_t capacity,
                                        uint32_t length,
                                        FreeFunction freeFn,
                                        void* userData,
                                        bool freeOnError) {
-  SharedInfo* sharedInfo = NULL;
   try {
     // TODO: We could allocate the IOBuf object and SharedInfo all in a single
     // memory allocation.  We could use the existing HeapStorage class, and
     // define a new kSharedInfoInUse flag.  We could change our code to call
     // releaseStorage(kFlagFreeSharedInfo) when this kFlagFreeSharedInfo,
     // rather than directly calling delete.
-    sharedInfo = new SharedInfo(freeFn, userData);
-
-    uint8_t* bufPtr = static_cast<uint8_t*>(buf);
-    return unique_ptr<IOBuf>(new IOBuf(kExtUserSupplied, kFlagFreeSharedInfo,
-                                       bufPtr, capacity,
-                                       bufPtr, length,
-                                       sharedInfo));
+    //
+    // Note that we always pass freeOnError as false to the constructor.
+    // If the constructor throws we'll handle it below.  (We have to handle
+    // allocation failures from make_unique too.)
+    return make_unique<IOBuf>(TAKE_OWNERSHIP, buf, capacity, length,
+                              freeFn, userData, false);
   } catch (...) {
-    delete sharedInfo;
-    if (freeOnError) {
-      if (freeFn) {
-        try {
-          freeFn(buf, userData);
-        } catch (...) {
-          // The user's free function is not allowed to throw.
-          abort();
-        }
-      } else {
-        free(buf);
-      }
-    }
+    takeOwnershipError(freeOnError, buf, freeFn, userData);
     throw;
   }
 }
 
+IOBuf::IOBuf(WrapBufferOp, const void* buf, uint32_t capacity)
+  : IOBuf(kExtUserOwned, kFlagUserOwned,
+          // We cast away the const-ness of the buffer here.
+          // This is okay since IOBuf users must use unshare() to create a copy
+          // of this buffer before writing to the buffer.
+          static_cast<uint8_t*>(const_cast<void*>(buf)), capacity,
+          static_cast<uint8_t*>(const_cast<void*>(buf)), capacity,
+          nullptr) {
+}
+
+IOBuf::IOBuf(WrapBufferOp op, ByteRange br)
+  : IOBuf(op, br.data(), folly::to<uint32_t>(br.size())) {
+}
+
 unique_ptr<IOBuf> IOBuf::wrapBuffer(const void* buf, uint32_t capacity) {
-  // We cast away the const-ness of the buffer here.
-  // This is okay since IOBuf users must use unshare() to create a copy of
-  // this buffer before writing to the buffer.
-  uint8_t* bufPtr = static_cast<uint8_t*>(const_cast<void*>(buf));
-  return unique_ptr<IOBuf>(new IOBuf(kExtUserSupplied, kFlagUserOwned,
-                                     bufPtr, capacity,
-                                     bufPtr, capacity,
-                                     NULL));
+  return make_unique<IOBuf>(WRAP_BUFFER, buf, capacity);
+}
+
+IOBuf::IOBuf() noexcept {
+}
+
+IOBuf::IOBuf(IOBuf&& other) noexcept {
+  *this = std::move(other);
 }
 
 IOBuf::IOBuf(ExtBufTypeEnum type,
@@ -313,6 +365,54 @@ IOBuf::~IOBuf() {
   decrementRefcount();
 }
 
+IOBuf& IOBuf::operator=(IOBuf&& other) noexcept {
+  // If we are part of a chain, delete the rest of the chain.
+  while (next_ != this) {
+    // Since unlink() returns unique_ptr() and we don't store it,
+    // it will automatically delete the unlinked element.
+    (void)next_->unlink();
+  }
+
+  // Decrement our refcount on the current buffer
+  decrementRefcount();
+
+  // Take ownership of the other buffer's data
+  data_ = other.data_;
+  buf_ = other.buf_;
+  length_ = other.length_;
+  capacity_ = other.capacity_;
+  flags_ = other.flags_;
+  type_ = other.type_;
+  sharedInfo_ = other.sharedInfo_;
+  // Reset other so it is a clean state to be destroyed.
+  other.data_ = nullptr;
+  other.buf_ = nullptr;
+  other.length_ = 0;
+  other.capacity_ = 0;
+  other.flags_ = kFlagUserOwned;
+  other.type_ = kExtUserOwned;
+  other.sharedInfo_ = nullptr;
+
+  // If other was part of the chain, assume ownership of the rest of its chain.
+  // (It's only valid to perform move assignment on the head of a chain.)
+  if (other.next_ != &other) {
+    next_ = other.next_;
+    next_->prev_ = this;
+    other.next_ = &other;
+
+    prev_ = other.prev_;
+    prev_->next_ = this;
+    other.prev_ = &other;
+  }
+
+  // Sanity check to make sure that other is in a valid state to be destroyed.
+  DCHECK_EQ(other.prev_, &other);
+  DCHECK_EQ(other.next_, &other);
+  DCHECK(other.flags_ & kFlagUserOwned);
+
+  return *this;
+}
+
 bool IOBuf::empty() const {
   const IOBuf* current = this;
   do {

folly/io/test/IOBufTest.cpp

@@ -142,7 +142,26 @@ TEST(IOBuf, TakeOwnership) {
   iobuf3.reset();
   EXPECT_EQ(1, deleteCount);
 
-
+  deleteCount = 0;
+  {
+    uint32_t size4 = 1234;
+    uint8_t *buf4 = new uint8_t[size4];
+    uint32_t length4 = 48;
+    IOBuf iobuf4(IOBuf::TAKE_OWNERSHIP, buf4, size4, length4,
+                 deleteArrayBuffer, &deleteCount);
+    EXPECT_EQ(buf4, iobuf4.data());
+    EXPECT_EQ(length4, iobuf4.length());
+    EXPECT_EQ(buf4, iobuf4.buffer());
+    EXPECT_EQ(size4, iobuf4.capacity());
+
+    IOBuf iobuf5 = std::move(iobuf4);
+    EXPECT_EQ(buf4, iobuf5.data());
+    EXPECT_EQ(length4, iobuf5.length());
+    EXPECT_EQ(buf4, iobuf5.buffer());
+    EXPECT_EQ(size4, iobuf5.capacity());
+    EXPECT_EQ(0, deleteCount);
+  }
+  EXPECT_EQ(1, deleteCount);
 }
 
 TEST(IOBuf, WrapBuffer) {
@@ -161,6 +180,14 @@ TEST(IOBuf, WrapBuffer) {
   EXPECT_EQ(size2, iobuf2->length());
   EXPECT_EQ(buf2.get(), iobuf2->buffer());
   EXPECT_EQ(size2, iobuf2->capacity());
+
+  uint32_t size3 = 4321;
+  unique_ptr<uint8_t[]> buf3(new uint8_t[size3]);
+  IOBuf iobuf3(IOBuf::WRAP_BUFFER, buf3.get(), size3);
+  EXPECT_EQ(buf3.get(), iobuf3.data());
+  EXPECT_EQ(size3, iobuf3.length());
+  EXPECT_EQ(buf3.get(), iobuf3.buffer());
+  EXPECT_EQ(size3, iobuf3.capacity());
 }
 
 TEST(IOBuf, CreateCombined) {
@@ -660,6 +687,13 @@ TEST(IOBuf, copyBuffer) {
   EXPECT_EQ(3, buf->headroom());
   EXPECT_EQ(0, buf->length());
   EXPECT_LE(6, buf->tailroom());
+
+  // A stack-allocated version
+  IOBuf stackBuf(IOBuf::COPY_BUFFER, s, 1, 2);
+  EXPECT_EQ(1, stackBuf.headroom());
+  EXPECT_EQ(s, std::string(reinterpret_cast<const char*>(stackBuf.data()),
+                           stackBuf.length()));
+  EXPECT_LE(2, stackBuf.tailroom());
 }
 
 TEST(IOBuf, maybeCopyBuffer) {
@@ -931,6 +965,47 @@ TEST(IOBuf, getIov) {
   EXPECT_EQ(count - 3, iov.size());
 }
 
+TEST(IOBuf, move) {
+  // Default allocate an IOBuf on the stack
+  IOBuf outerBuf;
+  char data[] = "foobar";
+  uint32_t length = sizeof(data);
+  uint32_t actualCapacity{0};
+  const void* ptr{nullptr};
+
+  {
+    // Create a small IOBuf on the stack.
+    // Note that IOBufs created on the stack always use an external buffer.
+    IOBuf b1(IOBuf::CREATE, 10);
+    actualCapacity = b1.capacity();
+    EXPECT_GE(actualCapacity, 10);
+    EXPECT_EQ(0, b1.length());
+    EXPECT_FALSE(b1.isShared());
+    ptr = b1.data();
+    ASSERT_TRUE(ptr != nullptr);
+    memcpy(b1.writableTail(), data, length);
+    b1.append(length);
+    EXPECT_EQ(length, b1.length());
+
+    // Use the move constructor
+    IOBuf b2(std::move(b1));
+    EXPECT_EQ(ptr, b2.data());
+    EXPECT_EQ(length, b2.length());
+    EXPECT_EQ(actualCapacity, b2.capacity());
+    EXPECT_FALSE(b2.isShared());
+
+    // Use the move assignment operator
+    outerBuf = std::move(b2);
+    // Close scope, destroying b1 and b2
+    // (which are both be invalid now anyway after moving out of them)
+  }
+
+  EXPECT_EQ(ptr, outerBuf.data());
+  EXPECT_EQ(length, outerBuf.length());
+  EXPECT_EQ(actualCapacity, outerBuf.capacity());
+  EXPECT_FALSE(outerBuf.isShared());
+}
+
 int main(int argc, char** argv) {
   testing::InitGoogleTest(&argc, argv);
   google::ParseCommandLineFlags(&argc, &argv, true);