Commit ce4efaad authored by Bo Liu's avatar Bo Liu Committed by Dave Watson

add clone() of stack allocated memory for folly::Cursor and folly::IOBuf

Summary: as title

Test Plan: fbconfig folly/io/test/ && fbmake runtests

Reviewed By: simpkins@fb.com

FB internal diff: D1176922
parent 23bddf52
...@@ -27,6 +27,7 @@ ...@@ -27,6 +27,7 @@
#include "folly/io/IOBuf.h" #include "folly/io/IOBuf.h"
#include "folly/io/IOBufQueue.h" #include "folly/io/IOBufQueue.h"
#include "folly/Likely.h" #include "folly/Likely.h"
#include "folly/Memory.h"
/** /**
* Cursor class for fast iteration over IOBuf chains. * Cursor class for fast iteration over IOBuf chains.
...@@ -221,6 +222,12 @@ class CursorBase { ...@@ -221,6 +222,12 @@ class CursorBase {
} }
} }
void clone(folly::IOBuf& buf, size_t len) {
if (UNLIKELY(cloneAtMost(buf, len) != len)) {
throw std::out_of_range("underflow");
}
}
void skip(size_t len) { void skip(size_t len) {
if (UNLIKELY(skipAtMost(len) != len)) { if (UNLIKELY(skipAtMost(len) != len)) {
throw std::out_of_range("underflow"); throw std::out_of_range("underflow");
...@@ -249,33 +256,38 @@ class CursorBase { ...@@ -249,33 +256,38 @@ class CursorBase {
} }
} }
size_t cloneAtMost(std::unique_ptr<folly::IOBuf>& buf, size_t len) { size_t cloneAtMost(folly::IOBuf& buf, size_t len) {
buf.reset(nullptr); buf = folly::IOBuf();
std::unique_ptr<folly::IOBuf> tmp; std::unique_ptr<folly::IOBuf> tmp;
size_t copied = 0; size_t copied = 0;
for (;;) { for (int loopCount = 0; true; ++loopCount) {
// Fast path: it all fits in one buffer. // Fast path: it all fits in one buffer.
size_t available = length(); size_t available = length();
if (LIKELY(available >= len)) { if (LIKELY(available >= len)) {
if (loopCount == 0) {
crtBuf_->cloneOneInto(buf);
buf.trimStart(offset_);
buf.trimEnd(buf.length() - len);
} else {
tmp = crtBuf_->cloneOne(); tmp = crtBuf_->cloneOne();
tmp->trimStart(offset_); tmp->trimStart(offset_);
tmp->trimEnd(tmp->length() - len); tmp->trimEnd(tmp->length() - len);
offset_ += len; buf.prependChain(std::move(tmp));
if (!buf) {
buf = std::move(tmp);
} else {
buf->prependChain(std::move(tmp));
} }
offset_ += len;
return copied + len; return copied + len;
} }
if (loopCount == 0) {
crtBuf_->cloneOneInto(buf);
buf.trimStart(offset_);
} else {
tmp = crtBuf_->cloneOne(); tmp = crtBuf_->cloneOne();
tmp->trimStart(offset_); tmp->trimStart(offset_);
if (!buf) { buf.prependChain(std::move(tmp));
buf = std::move(tmp);
} else {
buf->prependChain(std::move(tmp));
} }
copied += available; copied += available;
...@@ -286,6 +298,14 @@ class CursorBase { ...@@ -286,6 +298,14 @@ class CursorBase {
} }
} }
size_t cloneAtMost(std::unique_ptr<folly::IOBuf>& buf, size_t len) {
if (!buf) {
buf = make_unique<folly::IOBuf>();
}
return cloneAtMost(*buf, len);
}
size_t skipAtMost(size_t len) { size_t skipAtMost(size_t len) {
size_t skipped = 0; size_t skipped = 0;
for (;;) { for (;;) {
......
...@@ -459,27 +459,39 @@ void IOBuf::prependChain(unique_ptr<IOBuf>&& iobuf) { ...@@ -459,27 +459,39 @@ void IOBuf::prependChain(unique_ptr<IOBuf>&& iobuf) {
} }
unique_ptr<IOBuf> IOBuf::clone() const { unique_ptr<IOBuf> IOBuf::clone() const {
unique_ptr<IOBuf> newHead(cloneOne()); unique_ptr<IOBuf> ret = make_unique<IOBuf>();
cloneInto(*ret);
return ret;
}
unique_ptr<IOBuf> IOBuf::cloneOne() const {
unique_ptr<IOBuf> ret = make_unique<IOBuf>();
cloneOneInto(*ret);
return ret;
}
void IOBuf::cloneInto(IOBuf& other) const {
IOBuf tmp;
cloneOneInto(tmp);
for (IOBuf* current = next_; current != this; current = current->next_) { for (IOBuf* current = next_; current != this; current = current->next_) {
newHead->prependChain(current->cloneOne()); tmp.prependChain(current->cloneOne());
} }
return newHead; other = std::move(tmp);
} }
unique_ptr<IOBuf> IOBuf::cloneOne() const { void IOBuf::cloneOneInto(IOBuf& other) const {
if (sharedInfo_) { if (sharedInfo_) {
flags_ |= kFlagMaybeShared; flags_ |= kFlagMaybeShared;
} }
unique_ptr<IOBuf> iobuf(new IOBuf(static_cast<ExtBufTypeEnum>(type_), other = IOBuf(static_cast<ExtBufTypeEnum>(type_),
flags_, buf_, capacity_, flags_, buf_, capacity_,
data_, length_, data_, length_,
sharedInfo_)); sharedInfo_);
if (sharedInfo_) { if (sharedInfo_) {
sharedInfo_->refcount.fetch_add(1, std::memory_order_acq_rel); sharedInfo_->refcount.fetch_add(1, std::memory_order_acq_rel);
} }
return iobuf;
} }
void IOBuf::unshareOneSlow() { void IOBuf::unshareOneSlow() {
......
...@@ -1026,6 +1026,18 @@ class IOBuf { ...@@ -1026,6 +1026,18 @@ class IOBuf {
*/ */
std::unique_ptr<IOBuf> cloneOne() const; std::unique_ptr<IOBuf> cloneOne() const;
/**
* Similar to Clone(). But use other as the head node. Other nodes in the
* chain (if any) will be allocted on heap.
*/
void cloneInto(IOBuf& other) const;
/**
* Similar to CloneOne(). But to fill an existing IOBuf instead of a new
* IOBuf.
*/
void cloneOneInto(IOBuf& other) const;
/** /**
* Return an iovector suitable for e.g. writev() * Return an iovector suitable for e.g. writev()
* *
......
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