Skip to content

Commit

Permalink
Browse files Browse the repository at this point in the history
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
tudor committed Jul 13, 2012
1 parent b384b6b commit 53d4209
Show file tree
Hide file tree
Showing 5 changed files with 84 additions and 2 deletions.
1 change: 1 addition & 0 deletions folly/experimental/io/IOBuf.h
Expand Up @@ -497,6 +497,7 @@ class IOBuf {
void prepend(uint32_t amount) {
CHECK(amount <= headroom());
data_ -= amount;
length_ += amount;
}

/**
Expand Down
31 changes: 31 additions & 0 deletions folly/experimental/io/IOBufQueue.cpp
Expand Up @@ -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) {
Expand Down
19 changes: 19 additions & 0 deletions folly/experimental/io/IOBufQueue.h
Expand Up @@ -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:
Expand All @@ -48,6 +51,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.
Expand Down
19 changes: 19 additions & 0 deletions folly/experimental/io/test/IOBufQueueTest.cpp
Expand Up @@ -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);
Expand Down
16 changes: 14 additions & 2 deletions folly/experimental/io/test/IOBufTest.cpp
Expand Up @@ -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();
Expand All @@ -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());
Expand Down

0 comments on commit 53d4209

Please sign in to comment.