Permalink
Browse files

Add IOBufQueue::prepend, fix bug in IOBuf::prepend.

Summary:
IOBuf::prepend needs to increment length_.

Added IOBufQueue::prepend, which uses the headroom in the first buffer
instead of growing the queue at the head.

Test Plan: tests added

Reviewed By: simpkins@fb.com

FB internal diff: D513676
  • Loading branch information...
1 parent b384b6b commit 53d4209fdaa204747f5bb894b616d9770cb767a9 @tudor tudor committed Jul 7, 2012
@@ -497,6 +497,7 @@ class IOBuf {
void prepend(uint32_t amount) {
CHECK(amount <= headroom());
data_ -= amount;
+ length_ += amount;
}
/**
@@ -69,6 +69,37 @@ IOBufQueue& IOBufQueue::operator=(IOBufQueue&& other) {
return *this;
}
+std::pair<void*, uint32_t>
+IOBufQueue::headroom() {
+ if (head_) {
+ return std::make_pair(head_->writableBuffer(), head_->headroom());
+ } else {
+ return std::make_pair(nullptr, 0);
+ }
+}
+
+void
+IOBufQueue::markPrepended(uint32_t n) {
+ if (n == 0) {
+ return;
+ }
+ assert(head_);
+ head_->prepend(n);
+ if (options_.cacheChainLength) {
+ chainLength_ += n;
+ }
+}
+
+void
+IOBufQueue::prepend(const void* buf, uint32_t n) {
+ auto p = headroom();
+ if (n > p.second) {
+ throw std::overflow_error("Not enough room to prepend");
+ }
+ memcpy(static_cast<char*>(p.first) + p.second - n, buf, n);
+ markPrepended(n);
+}
+
void
IOBufQueue::append(unique_ptr<IOBuf>&& buf) {
if (!buf) {
@@ -28,6 +28,9 @@ namespace folly {
* An IOBufQueue encapsulates a chain of IOBufs and provides
* convenience functions to append data to the back of the chain
* and remove data from the front.
+ *
+ * You may also prepend data into the headroom of the first buffer in the
+ * chain, if any.
*/
class IOBufQueue {
public:
@@ -49,6 +52,22 @@ class IOBufQueue {
explicit IOBufQueue(const Options& options = Options());
/**
+ * Return a space to prepend bytes and the amount of headroom available.
+ */
+ std::pair<void*, uint32_t> headroom();
+
+ /**
+ * Indicate that n bytes from the headroom have been used.
+ */
+ void markPrepended(uint32_t n);
+
+ /**
+ * Prepend an existing range; throws std::overflow_error if not enough
+ * room.
+ */
+ void prepend(const void* buf, uint32_t n);
+
+ /**
* Add a buffer or buffer chain to the end of this queue. The
* queue takes ownership of buf.
*/
@@ -239,6 +239,25 @@ TEST(IOBufQueue, trim) {
checkConsistency(queue);
}
+TEST(IOBufQueue, Prepend) {
+ folly::IOBufQueue queue;
+
+ auto buf = folly::IOBuf::create(10);
+ buf->advance(5);
+ queue.append(std::move(buf));
+
+ queue.append(SCL(" World"));
+ queue.prepend(SCL("Hello"));
+
+ EXPECT_THROW(queue.prepend(SCL("x")), std::overflow_error);
+
+ auto out = queue.move();
+ out->coalesce();
+ EXPECT_EQ("Hello World",
+ StringPiece(reinterpret_cast<const char*>(out->data()),
+ out->length()));
+}
+
int main(int argc, char** argv) {
testing::InitGoogleTest(&argc, argv);
google::ParseCommandLineFlags(&argc, &argv, true);
@@ -35,6 +35,12 @@ void append(std::unique_ptr<IOBuf>& buf, StringPiece str) {
buf->append(str.size());
}
+void prepend(std::unique_ptr<IOBuf>& buf, StringPiece str) {
+ EXPECT_LE(str.size(), buf->headroom());
+ memcpy(buf->writableData() - str.size(), str.data(), str.size());
+ buf->prepend(str.size());
+}
+
TEST(IOBuf, Simple) {
unique_ptr<IOBuf> buf(IOBuf::create(100));
uint32_t cap = buf->capacity();
@@ -43,13 +49,19 @@ TEST(IOBuf, Simple) {
EXPECT_EQ(0, buf->length());
EXPECT_EQ(cap, buf->tailroom());
- append(buf, "hello");
+ append(buf, "world");
buf->advance(10);
EXPECT_EQ(10, buf->headroom());
EXPECT_EQ(5, buf->length());
EXPECT_EQ(cap - 15, buf->tailroom());
+
+ prepend(buf, "hello ");
+ EXPECT_EQ(4, buf->headroom());
+ EXPECT_EQ(11, buf->length());
+ EXPECT_EQ(cap - 15, buf->tailroom());
+
const char* p = reinterpret_cast<const char*>(buf->data());
- EXPECT_EQ("hello", std::string(p, buf->length()));
+ EXPECT_EQ("hello world", std::string(p, buf->length()));
buf->clear();
EXPECT_EQ(0, buf->headroom());

0 comments on commit 53d4209

Please sign in to comment.