From aef2fb0936a5e8b8eec583e5b9a4b246d6d1b52e Mon Sep 17 00:00:00 2001 From: Tudor Bosman Date: Mon, 5 Nov 2012 15:35:16 -0800 Subject: [PATCH] add IOBuf::moveToFbString Summary: Convert a IOBuf chain into a fbstring, with zero copies in the common case (unshared, unchained, no headroom) and one copy in all other cases. Test Plan: tests added Reviewed By: simpkins@fb.com FB internal diff: D621210 --- folly/experimental/io/IOBuf.cpp | 45 +++++++++++++++-- folly/experimental/io/IOBuf.h | 16 ++++++ folly/experimental/io/test/IOBufTest.cpp | 62 ++++++++++++++++++++++++ 3 files changed, 119 insertions(+), 4 deletions(-) diff --git a/folly/experimental/io/IOBuf.cpp b/folly/experimental/io/IOBuf.cpp index 42abf330..931af87f 100644 --- a/folly/experimental/io/IOBuf.cpp +++ b/folly/experimental/io/IOBuf.cpp @@ -353,6 +353,15 @@ void IOBuf::coalesceSlow(size_t maxLength) { uint64_t newHeadroom = headroom(); uint64_t newTailroom = end->prev_->tailroom(); + coalesceAndReallocate(newHeadroom, newLength, end, newTailroom); + // We should be only element left in the chain now + assert(length_ >= maxLength || !isChained()); +} + +void IOBuf::coalesceAndReallocate(size_t newHeadroom, + size_t newLength, + IOBuf* end, + size_t newTailroom) { uint64_t newCapacity = newLength + newHeadroom + newTailroom; if (newCapacity > UINT32_MAX) { throw std::overflow_error("IOBuf chain too large to coalesce"); @@ -370,11 +379,15 @@ void IOBuf::coalesceSlow(size_t maxLength) { uint8_t* newData = newBuf + newHeadroom; uint8_t* p = newData; IOBuf* current = this; + size_t remaining = newLength; do { + assert(current->length_ <= remaining); + remaining -= current->length_; memcpy(p, current->data_, current->length_); p += current->length_; current = current->next_; } while (current != end); + assert(remaining == 0); // Point at the new buffer if (flags_ & kFlagExt) { @@ -395,10 +408,9 @@ void IOBuf::coalesceSlow(size_t maxLength) { // Separate from the rest of our chain. // Since we don't store the unique_ptr returned by separateChain(), // this will immediately delete the returned subchain. - (void)separateChain(next_, end->prev_); - - // We should be only element left in the chain now - assert(length_ >= maxLength || !isChained()); + if (isChained()) { + (void)separateChain(next_, current->prev_); + } } void IOBuf::decrementRefcount() { @@ -598,4 +610,29 @@ void IOBuf::initExtBuffer(uint8_t* buf, size_t mallocSize, *infoReturn = sharedInfo; } +fbstring IOBuf::moveToFbString() { + // Externally allocated buffers (malloc) are just fine, everything else needs + // to be turned into one. + if (flags_ != kFlagExt || // not malloc()-ed + headroom() != 0 || // malloc()-ed block doesn't start at beginning + tailroom() == 0 || // no room for NUL terminator + isShared() || // shared + isChained()) { // chained + // We might as well get rid of all head and tailroom if we're going + // to reallocate; we need 1 byte for NUL terminator. + coalesceAndReallocate(0, computeChainDataLength(), this, 1); + } + + // Ensure NUL terminated + *writableTail() = 0; + fbstring str(reinterpret_cast(writableData()), + length(), capacity(), + AcquireMallocatedString()); + + // Reset to internal buffer. + flags_ = 0; + clear(); + return str; +} + } // folly diff --git a/folly/experimental/io/IOBuf.h b/folly/experimental/io/IOBuf.h index a9c5fcf3..e48650ac 100644 --- a/folly/experimental/io/IOBuf.h +++ b/folly/experimental/io/IOBuf.h @@ -27,6 +27,8 @@ #include #include +#include "folly/FBString.h" + namespace folly { /** @@ -896,6 +898,13 @@ class IOBuf { void* operator new(size_t size, void* ptr); void operator delete(void* ptr); + /** + * Destructively convert this IOBuf to a fbstring efficiently. + * We rely on fbstring's AcquireMallocatedString constructor to + * transfer memory. + */ + fbstring moveToFbString(); + private: enum FlagsEnum { kFlagExt = 0x1, @@ -966,6 +975,13 @@ class IOBuf { void unshareOneSlow(); void unshareChained(); void coalesceSlow(size_t maxLength=std::numeric_limits::max()); + // newLength must be the entire length of the buffers between this and + // end (no truncation) + void coalesceAndReallocate( + size_t newHeadroom, + size_t newLength, + IOBuf* end, + size_t newTailroom); void decrementRefcount(); void reserveSlow(uint32_t minHeadroom, uint32_t minTailroom); diff --git a/folly/experimental/io/test/IOBufTest.cpp b/folly/experimental/io/test/IOBufTest.cpp index 2e123297..b16511a4 100644 --- a/folly/experimental/io/test/IOBufTest.cpp +++ b/folly/experimental/io/test/IOBufTest.cpp @@ -17,6 +17,9 @@ #include "folly/experimental/io/IOBuf.h" #include "folly/experimental/io/TypedIOBuf.h" +// googletest requires std::tr1::tuple, not std::tuple +#include + #include #include #include @@ -24,6 +27,7 @@ #include "folly/Malloc.h" #include "folly/Range.h" +using folly::fbstring; using folly::IOBuf; using folly::TypedIOBuf; using folly::StringPiece; @@ -661,6 +665,64 @@ TEST(TypedIOBuf, Simple) { } } +// chain element size, number of elements in chain, shared +class MoveToFbStringTest + : public ::testing::TestWithParam> { + protected: + void SetUp() { + std::tr1::tie(elementSize_, elementCount_, shared_) = GetParam(); + buf_ = makeBuf(); + for (int i = 0; i < elementCount_ - 1; ++i) { + buf_->prependChain(makeBuf()); + } + EXPECT_EQ(elementCount_, buf_->countChainElements()); + EXPECT_EQ(elementCount_ * elementSize_, buf_->computeChainDataLength()); + if (shared_) { + buf2_ = buf_->clone(); + EXPECT_EQ(elementCount_, buf2_->countChainElements()); + EXPECT_EQ(elementCount_ * elementSize_, buf2_->computeChainDataLength()); + } + } + + std::unique_ptr makeBuf() { + auto buf = IOBuf::create(elementSize_); + memset(buf->writableTail(), 'x', elementSize_); + buf->append(elementSize_); + return buf; + } + + void check(std::unique_ptr& buf) { + fbstring str = buf->moveToFbString(); + EXPECT_EQ(elementCount_ * elementSize_, str.size()); + EXPECT_EQ(elementCount_ * elementSize_, strspn(str.c_str(), "x")); + EXPECT_EQ(0, buf->length()); + EXPECT_EQ(1, buf->countChainElements()); + EXPECT_EQ(0, buf->computeChainDataLength()); + EXPECT_FALSE(buf->isChained()); + } + + int elementSize_; + int elementCount_; + bool shared_; + std::unique_ptr buf_; + std::unique_ptr buf2_; +}; + +TEST_P(MoveToFbStringTest, Simple) { + check(buf_); + if (shared_) { + check(buf2_); + } +} + +INSTANTIATE_TEST_CASE_P( + MoveToFbString, + MoveToFbStringTest, + ::testing::Combine( + ::testing::Values(0, 1, 24, 256, 1 << 10, 1 << 20), // element size + ::testing::Values(1, 2, 10), // element count + ::testing::Bool())); // shared + int main(int argc, char** argv) { testing::InitGoogleTest(&argc, argv); google::ParseCommandLineFlags(&argc, &argv, true); -- 2.34.1