Commit 559830c8 authored by Dave Watson's avatar Dave Watson Committed by Jordan DeLong

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
parent cb771e84
......@@ -127,6 +127,12 @@ class CursorBase {
}
}
void clone(std::unique_ptr<folly::IOBuf>& 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<folly::IOBuf>& buf, size_t len) {
buf.reset(nullptr);
std::unique_ptr<folly::IOBuf> 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<folly::IOBuf> buf) {
folly::IOBuf* nextBuf;
if (this->offset_ == 0) {
// Can just prepend
nextBuf = buf.get();
this->crtBuf_->prependChain(std::move(buf));
} else {
std::unique_ptr<folly::IOBuf> 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_;
}
......
......@@ -236,6 +236,79 @@ TEST(IOBuf, PullAndPeek) {
}
}
TEST(IOBuf, cloneAndInsert) {
std::unique_ptr<IOBuf> iobuf1(IOBuf::create(10));
append(iobuf1, "he");
std::unique_ptr<IOBuf> iobuf2(IOBuf::create(10));
append(iobuf2, "llo ");
std::unique_ptr<IOBuf> 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<IOBuf> 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<uint8_t>();
}
{
// 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<uint8_t>();
}
{
// 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<uint8_t>();
}
}
TEST(IOBuf, Appender) {
std::unique_ptr<IOBuf> head(IOBuf::create(10));
append(head, "hello");
......
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