diff --git a/folly/io/Cursor.h b/folly/io/Cursor.h index 5f65fb0e01a39b73507526eca628b76632767f4c..67a8624315ee52a7b809d7b7f589060437a7a6da 100644 --- a/folly/io/Cursor.h +++ b/folly/io/Cursor.h @@ -27,6 +27,7 @@ #include "folly/io/IOBuf.h" #include "folly/io/IOBufQueue.h" #include "folly/Likely.h" +#include "folly/Memory.h" /** * Cursor class for fast iteration over IOBuf chains. @@ -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) { if (UNLIKELY(skipAtMost(len) != len)) { throw std::out_of_range("underflow"); @@ -249,33 +256,38 @@ class CursorBase { } } - size_t cloneAtMost(std::unique_ptr<folly::IOBuf>& buf, size_t len) { - buf.reset(nullptr); + size_t cloneAtMost(folly::IOBuf& buf, size_t len) { + buf = folly::IOBuf(); std::unique_ptr<folly::IOBuf> tmp; size_t copied = 0; - for (;;) { + for (int loopCount = 0; true; ++loopCount) { // 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); + if (loopCount == 0) { + crtBuf_->cloneOneInto(buf); + buf.trimStart(offset_); + buf.trimEnd(buf.length() - len); } else { - buf->prependChain(std::move(tmp)); + tmp = crtBuf_->cloneOne(); + tmp->trimStart(offset_); + tmp->trimEnd(tmp->length() - len); + buf.prependChain(std::move(tmp)); } + + offset_ += len; return copied + len; } - tmp = crtBuf_->cloneOne(); - tmp->trimStart(offset_); - if (!buf) { - buf = std::move(tmp); + + if (loopCount == 0) { + crtBuf_->cloneOneInto(buf); + buf.trimStart(offset_); } else { - buf->prependChain(std::move(tmp)); + tmp = crtBuf_->cloneOne(); + tmp->trimStart(offset_); + buf.prependChain(std::move(tmp)); } copied += available; @@ -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 skipped = 0; for (;;) { diff --git a/folly/io/IOBuf.cpp b/folly/io/IOBuf.cpp index b3cd4125999edd2c1696dcaa2df329dd29ccf986..3bcef5bd027448d140e6d1884635ef4622032cf8 100644 --- a/folly/io/IOBuf.cpp +++ b/folly/io/IOBuf.cpp @@ -459,27 +459,39 @@ void IOBuf::prependChain(unique_ptr<IOBuf>&& iobuf) { } 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_) { - 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_) { flags_ |= kFlagMaybeShared; } - unique_ptr<IOBuf> iobuf(new IOBuf(static_cast<ExtBufTypeEnum>(type_), - flags_, buf_, capacity_, - data_, length_, - sharedInfo_)); + other = IOBuf(static_cast<ExtBufTypeEnum>(type_), + flags_, buf_, capacity_, + data_, length_, + sharedInfo_); if (sharedInfo_) { sharedInfo_->refcount.fetch_add(1, std::memory_order_acq_rel); } - return iobuf; } void IOBuf::unshareOneSlow() { diff --git a/folly/io/IOBuf.h b/folly/io/IOBuf.h index 8e063822ae8a552edc1d5befbeddfc77dccab07c..ba216061ff89ae4e57351e0e46cc4dd2e6f6e466 100644 --- a/folly/io/IOBuf.h +++ b/folly/io/IOBuf.h @@ -1026,6 +1026,18 @@ class IOBuf { */ 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() *