From 559830c8a51872c237b01021a71998fcadcec665 Mon Sep 17 00:00:00 2001 From: Dave Watson Date: Thu, 15 Nov 2012 10:41:35 -0800 Subject: [PATCH] add clone / insert methods Summary: cursor.clone() will clone a length of the chain. insert(std::move(buf)) will splice in a length of chain I want this to change thrift2 binary type to IOBuf: we will clone() the network data for zero-copy userspace data, and insert() it if the return value is a IOBuf. Test Plan: added unittest Reviewed By: afrind@fb.com FB internal diff: D632073 --- folly/experimental/io/Cursor.h | 69 ++++++++++++++++++ .../experimental/io/test/IOBufCursorTest.cpp | 73 +++++++++++++++++++ 2 files changed, 142 insertions(+) diff --git a/folly/experimental/io/Cursor.h b/folly/experimental/io/Cursor.h index fbd77e5a..5b6e8b26 100644 --- a/folly/experimental/io/Cursor.h +++ b/folly/experimental/io/Cursor.h @@ -127,6 +127,12 @@ class CursorBase { } } + void clone(std::unique_ptr& buf, size_t length) { + if (UNLIKELY(cloneAtMost(buf, length) != length)) { + throw std::out_of_range("underflow"); + } + } + void skip(size_t length) { if (UNLIKELY(skipAtMost(length) != length)) { throw std::out_of_range("underflow"); @@ -155,6 +161,43 @@ class CursorBase { } } + size_t cloneAtMost(std::unique_ptr& buf, size_t len) { + buf.reset(nullptr); + + std::unique_ptr tmp; + size_t copied = 0; + for (;;) { + // Fast path: it all fits in one buffer. + size_t available = length(); + if (LIKELY(available >= len)) { + tmp = crtBuf_->cloneOne(); + tmp->trimStart(offset_); + tmp->trimEnd(tmp->length() - len); + offset_ += len; + if (!buf) { + buf = std::move(tmp); + } else { + buf->prependChain(std::move(tmp)); + } + return copied + len; + } + + tmp = crtBuf_->cloneOne(); + tmp->trimStart(offset_); + if (!buf) { + buf = std::move(tmp); + } else { + buf->prependChain(std::move(tmp)); + } + + copied += available; + if (UNLIKELY(!tryAdvanceBuffer())) { + return copied; + } + len -= available; + } + } + size_t skipAtMost(size_t len) { size_t skipped = 0; for (;;) { @@ -293,6 +336,32 @@ class RWCursor } } + void insert(std::unique_ptr buf) { + folly::IOBuf* nextBuf; + if (this->offset_ == 0) { + // Can just prepend + nextBuf = buf.get(); + this->crtBuf_->prependChain(std::move(buf)); + } else { + std::unique_ptr remaining; + if (this->crtBuf_->length() - this->offset_ > 0) { + // Need to split current IOBuf in two. + remaining = this->crtBuf_->cloneOne(); + remaining->trimStart(this->offset_); + nextBuf = remaining.get(); + buf->prependChain(std::move(remaining)); + } else { + // Can just append + nextBuf = this->crtBuf_->next(); + } + this->crtBuf_->trimEnd(this->length()); + this->crtBuf_->appendChain(std::move(buf)); + } + // Jump past the new links + this->offset_ = 0; + this->crtBuf_ = nextBuf; + } + uint8_t* writableData() { return this->crtBuf_->writableData() + this->offset_; } diff --git a/folly/experimental/io/test/IOBufCursorTest.cpp b/folly/experimental/io/test/IOBufCursorTest.cpp index 3ed91f04..6ef4d0a5 100644 --- a/folly/experimental/io/test/IOBufCursorTest.cpp +++ b/folly/experimental/io/test/IOBufCursorTest.cpp @@ -236,6 +236,79 @@ TEST(IOBuf, PullAndPeek) { } } +TEST(IOBuf, cloneAndInsert) { + std::unique_ptr iobuf1(IOBuf::create(10)); + append(iobuf1, "he"); + std::unique_ptr iobuf2(IOBuf::create(10)); + append(iobuf2, "llo "); + std::unique_ptr iobuf3(IOBuf::create(10)); + append(iobuf3, "world"); + iobuf1->prependChain(std::move(iobuf2)); + iobuf1->prependChain(std::move(iobuf3)); + EXPECT_EQ(3, iobuf1->countChainElements()); + EXPECT_EQ(11, iobuf1->computeChainDataLength()); + + std::unique_ptr cloned; + + Cursor(iobuf1.get()).clone(cloned, 3); + EXPECT_EQ(2, cloned->countChainElements()); + EXPECT_EQ(3, cloned->computeChainDataLength()); + + + EXPECT_EQ(11, Cursor(iobuf1.get()).cloneAtMost(cloned, 20)); + EXPECT_EQ(3, cloned->countChainElements()); + EXPECT_EQ(11, cloned->computeChainDataLength()); + + + EXPECT_THROW({Cursor(iobuf1.get()).clone(cloned, 20);}, + std::out_of_range); + + { + // Check that inserting in the middle of an iobuf splits + RWPrivateCursor cursor(iobuf1.get()); + Cursor(iobuf1.get()).clone(cloned, 3); + EXPECT_EQ(2, cloned->countChainElements()); + EXPECT_EQ(3, cloned->computeChainDataLength()); + + cursor.skip(1); + + cursor.insert(std::move(cloned)); + EXPECT_EQ(6, iobuf1->countChainElements()); + EXPECT_EQ(14, iobuf1->computeChainDataLength()); + // Check that nextBuf got set correctly + cursor.read(); + } + + { + // Check that inserting at the end doesn't create empty buf + RWPrivateCursor cursor(iobuf1.get()); + Cursor(iobuf1.get()).clone(cloned, 1); + EXPECT_EQ(1, cloned->countChainElements()); + EXPECT_EQ(1, cloned->computeChainDataLength()); + + cursor.skip(1); + + cursor.insert(std::move(cloned)); + EXPECT_EQ(7, iobuf1->countChainElements()); + EXPECT_EQ(15, iobuf1->computeChainDataLength()); + // Check that nextBuf got set correctly + cursor.read(); + } + { + // Check that inserting at the beginning doesn't create empty buf + RWPrivateCursor cursor(iobuf1.get()); + Cursor(iobuf1.get()).clone(cloned, 1); + EXPECT_EQ(1, cloned->countChainElements()); + EXPECT_EQ(1, cloned->computeChainDataLength()); + + cursor.insert(std::move(cloned)); + EXPECT_EQ(8, iobuf1->countChainElements()); + EXPECT_EQ(16, iobuf1->computeChainDataLength()); + // Check that nextBuf got set correctly + cursor.read(); + } +} + TEST(IOBuf, Appender) { std::unique_ptr head(IOBuf::create(10)); append(head, "hello"); -- 2.34.1