Add IOBuf SizedFree API

Summary:
Add IOBuf SizedFree API

(Note: this ignores all push blocking failures!)

Reviewed By: simpkins

Differential Revision: D28667639

fbshipit-source-id: 6264d829948110077bad9dead8dad785795c211b

folly/io/IOBuf.cpp

@@ -320,15 +320,7 @@ unique_ptr<IOBuf> IOBuf::create(std::size_t capacity) {
     // if we do not have space for the overhead, allocate the mem separateley
     if (mallocSize < minSize) {
       auto* buf = checkedMalloc(mallocSize);
-      return takeOwnership(
-          buf,
-          mallocSize,
-          static_cast<size_t>(0),
-          static_cast<size_t>(0),
-          nullptr, /*freeFn*/
-          reinterpret_cast<void*>(mallocSize), /*userData - used for sizedFree*/
-          true, /*freeOnError*/
-          TakeOwnershipOption::STORE_SIZE);
+      return takeOwnership(SIZED_FREE, buf, mallocSize, 0, 0);
     }
   }
 
@@ -426,8 +418,31 @@ IOBuf::IOBuf(
   });
   setSharedInfo(new SharedInfo(freeFn, userData));
   rollback.dismiss();
+}
+
+IOBuf::IOBuf(
+    TakeOwnershipOp,
+    SizedFree,
+    void* buf,
+    std::size_t capacity,
+    std::size_t offset,
+    std::size_t length,
+    bool freeOnError)
+    : next_(this),
+      prev_(this),
+      data_(static_cast<uint8_t*>(buf) + offset),
+      buf_(static_cast<uint8_t*>(buf)),
+      length_(length),
+      capacity_(capacity),
+      flagsAndSharedInfo_(
+          packFlagsAndSharedInfo(kFlagFreeSharedInfo, nullptr)) {
+  auto rollback = makeGuard([&] { //
+    takeOwnershipError(freeOnError, buf, nullptr, nullptr);
+  });
+  setSharedInfo(new SharedInfo(nullptr, reinterpret_cast<void*>(capacity)));
+  rollback.dismiss();
 
-  if (io_buf_alloc_cb && !freeFn && capacity) {
+  if (io_buf_alloc_cb && capacity) {
     io_buf_alloc_cb(buf, capacity);
   }
 }
@@ -477,7 +492,7 @@ unique_ptr<IOBuf> IOBuf::takeOwnership(
 
   if (io_buf_alloc_cb) {
     io_buf_alloc_cb(storage, mallocSize);
-    if (userData && !freeFn) {
+    if (userData && !freeFn && (option == TakeOwnershipOption::STORE_SIZE)) {
       // Even though we did not allocate the buffer, call io_buf_alloc_cb()
       // since we will call io_buf_free_cb() on destruction, and we want these
       // calls to be 1:1.

folly/io/IOBuf.h

@@ -228,9 +228,9 @@ class IOBuf {
   enum WrapBufferOp { WRAP_BUFFER };
   enum TakeOwnershipOp { TAKE_OWNERSHIP };
   enum CopyBufferOp { COPY_BUFFER };
+  enum SizedFree { SIZED_FREE };
 
   enum class CombinedOption { DEFAULT, COMBINED, SEPARATE };
-  enum class TakeOwnershipOption { DEFAULT, STORE_SIZE };
 
   typedef ByteRange value_type;
   typedef Iterator iterator;
@@ -328,7 +328,14 @@ class IOBuf {
       void* userData = nullptr,
       bool freeOnError = true) {
     return takeOwnership(
-        buf, capacity, 0, capacity, freeFn, userData, freeOnError);
+        buf,
+        capacity,
+        0,
+        capacity,
+        freeFn,
+        userData,
+        freeOnError,
+        TakeOwnershipOption::DEFAULT);
   }
   IOBuf(
       TakeOwnershipOp op,
@@ -347,7 +354,14 @@ class IOBuf {
       void* userData = nullptr,
       bool freeOnError = true) {
     return takeOwnership(
-        buf, capacity, 0, length, freeFn, userData, freeOnError);
+        buf,
+        capacity,
+        0,
+        length,
+        freeFn,
+        userData,
+        freeOnError,
+        TakeOwnershipOption::DEFAULT);
   }
   IOBuf(
       TakeOwnershipOp op,
@@ -366,8 +380,36 @@ class IOBuf {
       std::size_t length,
       FreeFunction freeFn = nullptr,
       void* userData = nullptr,
-      bool freeOnError = true,
-      TakeOwnershipOption option = TakeOwnershipOption::DEFAULT);
+      bool freeOnError = true) {
+    return takeOwnership(
+        buf,
+        capacity,
+        offset,
+        length,
+        freeFn,
+        userData,
+        freeOnError,
+        TakeOwnershipOption::DEFAULT);
+  }
+
+  static std::unique_ptr<IOBuf> takeOwnership(
+      SizedFree,
+      void* buf,
+      std::size_t capacity,
+      std::size_t offset,
+      std::size_t length,
+      bool freeOnError = true) {
+    return takeOwnership(
+        buf,
+        capacity,
+        offset,
+        length,
+        nullptr,
+        reinterpret_cast<void*>(capacity),
+        freeOnError,
+        TakeOwnershipOption::STORE_SIZE);
+  }
+
   IOBuf(
       TakeOwnershipOp,
       void* buf,
@@ -378,6 +420,15 @@ class IOBuf {
       void* userData = nullptr,
       bool freeOnError = true);
 
+  IOBuf(
+      TakeOwnershipOp,
+      SizedFree,
+      void* buf,
+      std::size_t capacity,
+      std::size_t offset,
+      std::size_t length,
+      bool freeOnError = true);
+
   /**
    * Create a new IOBuf pointing to an existing data buffer made up of
    * count objects of a given standard-layout type.
@@ -1391,6 +1442,17 @@ class IOBuf {
   IOBuf& operator=(const IOBuf& other);
 
  private:
+  enum class TakeOwnershipOption { DEFAULT, STORE_SIZE };
+  static std::unique_ptr<IOBuf> takeOwnership(
+      void* buf,
+      std::size_t capacity,
+      std::size_t offset,
+      std::size_t length,
+      FreeFunction freeFn,
+      void* userData,
+      bool freeOnError,
+      TakeOwnershipOption option);
+
   enum FlagsEnum : uintptr_t {
     // Adding any more flags would not work on 32-bit architectures,
     // as these flags are stashed in the least significant 2 bits of a

folly/io/test/IOBufCBTest.cpp

@@ -81,7 +81,8 @@ TEST(IOBufCB, TakeOwnership) {
   static constexpr size_t kAllocSize = 1024;
 
   {
-    auto buf = folly::IOBuf::takeOwnership(malloc(kAllocSize), kAllocSize);
+    auto buf = folly::IOBuf::takeOwnership(
+        folly::IOBuf::SIZED_FREE, malloc(kAllocSize), kAllocSize, 0, 0);
     size_t newSize = gIOBufAlloc.load();
     CHECK_GE(newSize, initialSize + kAllocSize);
   }
@@ -90,7 +91,12 @@ TEST(IOBufCB, TakeOwnership) {
 
   {
     folly::IOBuf buf(
-        folly::IOBuf::TAKE_OWNERSHIP, malloc(kAllocSize), kAllocSize);
+        folly::IOBuf::TAKE_OWNERSHIP,
+        folly::IOBuf::SIZED_FREE,
+        malloc(kAllocSize),
+        kAllocSize,
+        0,
+        0);
     size_t newSize = gIOBufAlloc.load();
     CHECK_GE(newSize, initialSize + kAllocSize);
   }
@@ -105,7 +111,8 @@ TEST(IOBufCB, MoveToFbString) {
   LOG(INFO) << gIOBufAlloc.load();
 
   {
-    auto buf = folly::IOBuf::takeOwnership(malloc(kAllocSize), kAllocSize);
+    auto buf = folly::IOBuf::takeOwnership(
+        folly::IOBuf::SIZED_FREE, malloc(kAllocSize), kAllocSize, 0, 0);
     LOG(INFO) << gIOBufAlloc.load();
     size_t newSize = gIOBufAlloc.load();
     CHECK_GE(newSize, initialSize + kAllocSize);