Commit 780643b3 authored by Dan Melnic's avatar Dan Melnic Committed by Facebook GitHub Bot

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
parent 07112a71
...@@ -320,15 +320,7 @@ unique_ptr<IOBuf> IOBuf::create(std::size_t capacity) { ...@@ -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 we do not have space for the overhead, allocate the mem separateley
if (mallocSize < minSize) { if (mallocSize < minSize) {
auto* buf = checkedMalloc(mallocSize); auto* buf = checkedMalloc(mallocSize);
return takeOwnership( return takeOwnership(SIZED_FREE, buf, mallocSize, 0, 0);
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);
} }
} }
...@@ -426,8 +418,31 @@ IOBuf::IOBuf( ...@@ -426,8 +418,31 @@ IOBuf::IOBuf(
}); });
setSharedInfo(new SharedInfo(freeFn, userData)); setSharedInfo(new SharedInfo(freeFn, userData));
rollback.dismiss(); 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); io_buf_alloc_cb(buf, capacity);
} }
} }
...@@ -477,7 +492,7 @@ unique_ptr<IOBuf> IOBuf::takeOwnership( ...@@ -477,7 +492,7 @@ unique_ptr<IOBuf> IOBuf::takeOwnership(
if (io_buf_alloc_cb) { if (io_buf_alloc_cb) {
io_buf_alloc_cb(storage, mallocSize); 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() // 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 // since we will call io_buf_free_cb() on destruction, and we want these
// calls to be 1:1. // calls to be 1:1.
......
...@@ -228,9 +228,9 @@ class IOBuf { ...@@ -228,9 +228,9 @@ class IOBuf {
enum WrapBufferOp { WRAP_BUFFER }; enum WrapBufferOp { WRAP_BUFFER };
enum TakeOwnershipOp { TAKE_OWNERSHIP }; enum TakeOwnershipOp { TAKE_OWNERSHIP };
enum CopyBufferOp { COPY_BUFFER }; enum CopyBufferOp { COPY_BUFFER };
enum SizedFree { SIZED_FREE };
enum class CombinedOption { DEFAULT, COMBINED, SEPARATE }; enum class CombinedOption { DEFAULT, COMBINED, SEPARATE };
enum class TakeOwnershipOption { DEFAULT, STORE_SIZE };
typedef ByteRange value_type; typedef ByteRange value_type;
typedef Iterator iterator; typedef Iterator iterator;
...@@ -328,7 +328,14 @@ class IOBuf { ...@@ -328,7 +328,14 @@ class IOBuf {
void* userData = nullptr, void* userData = nullptr,
bool freeOnError = true) { bool freeOnError = true) {
return takeOwnership( return takeOwnership(
buf, capacity, 0, capacity, freeFn, userData, freeOnError); buf,
capacity,
0,
capacity,
freeFn,
userData,
freeOnError,
TakeOwnershipOption::DEFAULT);
} }
IOBuf( IOBuf(
TakeOwnershipOp op, TakeOwnershipOp op,
...@@ -347,7 +354,14 @@ class IOBuf { ...@@ -347,7 +354,14 @@ class IOBuf {
void* userData = nullptr, void* userData = nullptr,
bool freeOnError = true) { bool freeOnError = true) {
return takeOwnership( return takeOwnership(
buf, capacity, 0, length, freeFn, userData, freeOnError); buf,
capacity,
0,
length,
freeFn,
userData,
freeOnError,
TakeOwnershipOption::DEFAULT);
} }
IOBuf( IOBuf(
TakeOwnershipOp op, TakeOwnershipOp op,
...@@ -366,8 +380,36 @@ class IOBuf { ...@@ -366,8 +380,36 @@ class IOBuf {
std::size_t length, std::size_t length,
FreeFunction freeFn = nullptr, FreeFunction freeFn = nullptr,
void* userData = nullptr, void* userData = nullptr,
bool freeOnError = true, bool freeOnError = true) {
TakeOwnershipOption option = TakeOwnershipOption::DEFAULT); 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( IOBuf(
TakeOwnershipOp, TakeOwnershipOp,
void* buf, void* buf,
...@@ -378,6 +420,15 @@ class IOBuf { ...@@ -378,6 +420,15 @@ class IOBuf {
void* userData = nullptr, void* userData = nullptr,
bool freeOnError = true); 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 * Create a new IOBuf pointing to an existing data buffer made up of
* count objects of a given standard-layout type. * count objects of a given standard-layout type.
...@@ -1391,6 +1442,17 @@ class IOBuf { ...@@ -1391,6 +1442,17 @@ class IOBuf {
IOBuf& operator=(const IOBuf& other); IOBuf& operator=(const IOBuf& other);
private: 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 { enum FlagsEnum : uintptr_t {
// Adding any more flags would not work on 32-bit architectures, // Adding any more flags would not work on 32-bit architectures,
// as these flags are stashed in the least significant 2 bits of a // as these flags are stashed in the least significant 2 bits of a
......
...@@ -81,7 +81,8 @@ TEST(IOBufCB, TakeOwnership) { ...@@ -81,7 +81,8 @@ TEST(IOBufCB, TakeOwnership) {
static constexpr size_t kAllocSize = 1024; 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(); size_t newSize = gIOBufAlloc.load();
CHECK_GE(newSize, initialSize + kAllocSize); CHECK_GE(newSize, initialSize + kAllocSize);
} }
...@@ -90,7 +91,12 @@ TEST(IOBufCB, TakeOwnership) { ...@@ -90,7 +91,12 @@ TEST(IOBufCB, TakeOwnership) {
{ {
folly::IOBuf buf( 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(); size_t newSize = gIOBufAlloc.load();
CHECK_GE(newSize, initialSize + kAllocSize); CHECK_GE(newSize, initialSize + kAllocSize);
} }
...@@ -105,7 +111,8 @@ TEST(IOBufCB, MoveToFbString) { ...@@ -105,7 +111,8 @@ TEST(IOBufCB, MoveToFbString) {
LOG(INFO) << gIOBufAlloc.load(); 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(); LOG(INFO) << gIOBufAlloc.load();
size_t newSize = gIOBufAlloc.load(); size_t newSize = gIOBufAlloc.load();
CHECK_GE(newSize, initialSize + kAllocSize); CHECK_GE(newSize, initialSize + kAllocSize);
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment