From 3d5106cd605ed237a5625dd8d45223075a57bcae Mon Sep 17 00:00:00 2001 From: Adam Simpkins Date: Sun, 24 Nov 2013 15:46:46 -0800 Subject: [PATCH] fix IOBuf::reserve() when operating on a user-supplied buffer Summary: The IOBuf::reserveSlow() code assumed that external buffers were always allocated with malloc. This would cause problems when if reserve() was ever used on a buffer created with IOBuf::takeOwnership(). This changes the code to now check if a user-specified free function has been supplied. If so, then it does not try to use realloc()/rallocm(), and it now frees the old buffer using the specified free function rather than free(). User-supplied buffers also have a separately allocated SharedInfo object, which must be freed when we no longer need it. Test Plan: Added additional unit tests to check calling reserve() on IOBufs created with IOBuf::takeOwnership(). Reviewed By: davejwatson@fb.com FB internal diff: D1072292 --- folly/io/IOBuf.cpp | 52 +++++++++++++++++++++++++------------ folly/io/IOBuf.h | 1 + folly/io/test/IOBufTest.cpp | 26 +++++++++++++++++++ 3 files changed, 62 insertions(+), 17 deletions(-) diff --git a/folly/io/IOBuf.cpp b/folly/io/IOBuf.cpp index 13ebdac3..5bab8716 100644 --- a/folly/io/IOBuf.cpp +++ b/folly/io/IOBuf.cpp @@ -453,18 +453,7 @@ void IOBuf::decrementRefcount() { } // We were the last user. Free the buffer - if (ext_.sharedInfo->freeFn != NULL) { - try { - ext_.sharedInfo->freeFn(ext_.buf, ext_.sharedInfo->userData); - } catch (...) { - // The user's free function should never throw. Otherwise we might - // throw from the IOBuf destructor. Other code paths like coalesce() - // also assume that decrementRefcount() cannot throw. - abort(); - } - } else { - free(ext_.buf); - } + freeExtBuffer(); // Free the SharedInfo if it was allocated separately. // @@ -485,6 +474,11 @@ void IOBuf::reserveSlow(uint32_t minHeadroom, uint32_t minTailroom) { size_t newCapacity = (size_t)length_ + minHeadroom + minTailroom; DCHECK_LT(newCapacity, UINT32_MAX); + // reserveSlow() is dangerous if anyone else is sharing the buffer, as we may + // reallocate and free the original buffer. It should only ever be called if + // we are the only user of the buffer. + DCHECK(!isSharedOne()); + // We'll need to reallocate the buffer. // There are a few options. // - If we have enough total room, move the data around in the buffer @@ -511,11 +505,15 @@ void IOBuf::reserveSlow(uint32_t minHeadroom, uint32_t minTailroom) { uint32_t newHeadroom = 0; uint32_t oldHeadroom = headroom(); - if ((flags_ & kFlagExt) && length_ != 0 && oldHeadroom >= minHeadroom) { + // If we have a buffer allocated with malloc and we just need more tailroom, + // try to use realloc()/rallocm() to grow the buffer in place. + if ((flags_ & (kFlagExt | kFlagUserOwned)) == kFlagExt && + (ext_.sharedInfo->freeFn == nullptr) && + length_ != 0 && oldHeadroom >= minHeadroom) { if (usingJEMalloc()) { size_t headSlack = oldHeadroom - minHeadroom; // We assume that tailroom is more useful and more important than - // tailroom (not least because realloc / rallocm allow us to grow the + // headroom (not least because realloc / rallocm allow us to grow the // buffer at the tail, but not at the head) So, if we have more headroom // than we need, we consider that "wasted". We arbitrarily define "too // much" headroom to be 25% of the capacity. @@ -559,8 +557,8 @@ void IOBuf::reserveSlow(uint32_t minHeadroom, uint32_t minTailroom) { } newBuffer = static_cast(p); memcpy(newBuffer + minHeadroom, data_, length_); - if (flags_ & kFlagExt) { - free(ext_.buf); + if ((flags_ & (kFlagExt | kFlagUserOwned)) == kFlagExt) { + freeExtBuffer(); } newHeadroom = minHeadroom; } @@ -569,8 +567,11 @@ void IOBuf::reserveSlow(uint32_t minHeadroom, uint32_t minTailroom) { uint32_t cap; initExtBuffer(newBuffer, newAllocatedCapacity, &info, &cap); - flags_ = kFlagExt; + if (flags_ & kFlagFreeSharedInfo) { + delete ext_.sharedInfo; + } + flags_ = kFlagExt; ext_.capacity = cap; ext_.type = kExtAllocated; ext_.buf = newBuffer; @@ -579,6 +580,23 @@ void IOBuf::reserveSlow(uint32_t minHeadroom, uint32_t minTailroom) { // length_ is unchanged } +void IOBuf::freeExtBuffer() { + DCHECK((flags_ & (kFlagExt | kFlagUserOwned)) == kFlagExt); + + if (ext_.sharedInfo->freeFn) { + try { + ext_.sharedInfo->freeFn(ext_.buf, ext_.sharedInfo->userData); + } catch (...) { + // The user's free function should never throw. Otherwise we might + // throw from the IOBuf destructor. Other code paths like coalesce() + // also assume that decrementRefcount() cannot throw. + abort(); + } + } else { + free(ext_.buf); + } +} + void IOBuf::allocExtBuffer(uint32_t minCapacity, uint8_t** bufReturn, SharedInfo** infoReturn, diff --git a/folly/io/IOBuf.h b/folly/io/IOBuf.h index 03f3e657..e9048626 100644 --- a/folly/io/IOBuf.h +++ b/folly/io/IOBuf.h @@ -1078,6 +1078,7 @@ class IOBuf { size_t newTailroom); void decrementRefcount(); void reserveSlow(uint32_t minHeadroom, uint32_t minTailroom); + void freeExtBuffer(); static size_t goodExtBufferSize(uint32_t minCapacity); static void initExtBuffer(uint8_t* buf, size_t mallocSize, diff --git a/folly/io/test/IOBufTest.cpp b/folly/io/test/IOBufTest.cpp index b17be708..a0ca939f 100644 --- a/folly/io/test/IOBufTest.cpp +++ b/folly/io/test/IOBufTest.cpp @@ -494,6 +494,12 @@ TEST(IOBuf, Chaining) { EXPECT_EQ(3, iob2->computeChainDataLength()); } +void reserveTestFreeFn(void* buffer, void* ptr) { + uint32_t* freeCount = static_cast(ptr);; + delete[] static_cast(buffer); + ++(*freeCount); +}; + TEST(IOBuf, Reserve) { uint32_t fillSeed = 0x23456789; boost::mt19937 gen(fillSeed); @@ -555,6 +561,26 @@ TEST(IOBuf, Reserve) { EXPECT_EQ(0, iob->headroom()); EXPECT_LE(2000, iob->tailroom()); } + + // Test reserving from a user-allocated buffer. + { + uint8_t* buf = static_cast(malloc(100)); + auto iob = IOBuf::takeOwnership(buf, 100); + iob->reserve(0, 2000); + EXPECT_EQ(0, iob->headroom()); + EXPECT_LE(2000, iob->tailroom()); + } + + // Test reserving from a user-allocated with a custom free function. + { + uint32_t freeCount{0}; + uint8_t* buf = new uint8_t[100]; + auto iob = IOBuf::takeOwnership(buf, 100, reserveTestFreeFn, &freeCount); + iob->reserve(0, 2000); + EXPECT_EQ(0, iob->headroom()); + EXPECT_LE(2000, iob->tailroom()); + EXPECT_EQ(1, freeCount); + } } TEST(IOBuf, copyBuffer) { -- 2.34.1