Add TakeOwnershipOption::STORE_SIZE

Summary: Add TakeOwnershipOption::STORE_SIZE

Reviewed By: simpkins

Differential Revision: D28663081

fbshipit-source-id: 5f2b0cd5c9ebdb9a8475b35f8741bc0b17717441

folly/io/IOBuf.cpp

@@ -320,7 +320,15 @@ 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));
+      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);
     }
   }
 
@@ -413,11 +421,6 @@ IOBuf::IOBuf(
   // since we use that for folly::sizedFree
   DCHECK(!userData || (userData && freeFn));
 
-  // add support for sized dealloc if freeFn is null
-  if (!userData && !freeFn) {
-    userData = reinterpret_cast<void*>(capacity);
-  }
-
   auto rollback = makeGuard([&] { //
     takeOwnershipError(freeOnError, buf, freeFn, userData);
   });
@@ -436,15 +439,14 @@ unique_ptr<IOBuf> IOBuf::takeOwnership(
     std::size_t length,
     FreeFunction freeFn,
     void* userData,
-    bool freeOnError) {
+    bool freeOnError,
+    TakeOwnershipOption option) {
   // do not allow only user data without a freeFn
   // since we use that for folly::sizedFree
-  DCHECK(!userData || (userData && freeFn));
 
-  // add support for sized dealloc if freeFn is null
-  if (!userData && !freeFn) {
-    userData = reinterpret_cast<void*>(capacity);
-  }
+  DCHECK(
+      !userData || (userData && freeFn) ||
+      (userData && !freeFn && (option == TakeOwnershipOption::STORE_SIZE)));
 
   HeapFullStorage* storage = nullptr;
   auto rollback = makeGuard([&] {

folly/io/IOBuf.h

@@ -230,6 +230,7 @@ class IOBuf {
   enum CopyBufferOp { COPY_BUFFER };
 
   enum class CombinedOption { DEFAULT, COMBINED, SEPARATE };
+  enum class TakeOwnershipOption { DEFAULT, STORE_SIZE };
 
   typedef ByteRange value_type;
   typedef Iterator iterator;
@@ -365,7 +366,8 @@ class IOBuf {
       std::size_t length,
       FreeFunction freeFn = nullptr,
       void* userData = nullptr,
-      bool freeOnError = true);
+      bool freeOnError = true,
+      TakeOwnershipOption option = TakeOwnershipOption::DEFAULT);
   IOBuf(
       TakeOwnershipOp,
       void* buf,