add function to push data from another cursor

Summary:
Add a function to Writable to push data from another cursor. Data in the input
cursor could be in a buffer chain.

Added test in IOBufCurosTest.cpp to test the new function.

Test Plan: fbmake runtests;

Reviewed By: simpkins@fb.com

Subscribers: njormrod, net-systems@

FB internal diff: D1563209

Tasks: 4886058

Blame Revision:

diff --git a/folly/io/Cursor.h b/folly/io/Cursor.h
index 59fa044016c390f2eed0e809f20e63e64b7667e0..7963f193fe825d00750bd9b33c10a910615a572f 100644 (file)
--- a/folly/io/Cursor.h
+++ b/folly/io/Cursor.h
@@ -428,6 +428,20 @@ class CursorBase {
   BufType* buffer_;
 };
 
+} //namespace detail
+
+class Cursor : public detail::CursorBase<Cursor, const IOBuf> {
+ public:
+  explicit Cursor(const IOBuf* buf)
+    : detail::CursorBase<Cursor, const IOBuf>(buf) {}
+
+  template <class OtherDerived, class OtherBuf>
+  explicit Cursor(const detail::CursorBase<OtherDerived, OtherBuf>& cursor)
+    : detail::CursorBase<Cursor, const IOBuf>(cursor) {}
+};
+
+namespace detail {
+
 template <class Derived>
 class Writable {
  public:
@@ -457,20 +471,47 @@ class Writable {
       throw std::out_of_range("overflow");
     }
   }
-};
 
-} // namespace detail
+  /**
+   * push len bytes of data from input cursor, data could be in an IOBuf chain.
+   * If input cursor contains less than len bytes, or this cursor has less than
+   * len bytes writable space, an out_of_range exception will be thrown.
+   */
+  void push(Cursor cursor, size_t len) {
+    if (this->pushAtMost(cursor, len) != len) {
+      throw std::out_of_range("overflow");
+    }
+  }
 
-class Cursor : public detail::CursorBase<Cursor, const IOBuf> {
- public:
-  explicit Cursor(const IOBuf* buf)
-    : detail::CursorBase<Cursor, const IOBuf>(buf) {}
+  size_t pushAtMost(Cursor cursor, size_t len) {
+    size_t written = 0;
+    for(;;) {
+      auto currentBuffer = cursor.peek();
+      const uint8_t* crtData = currentBuffer.first;
+      size_t available = currentBuffer.second;
+      if (available == 0) {
+        // end of buffer chain
+        return written;
+      }
+      // all data is in current buffer
+      if (available >= len) {
+        this->push(crtData, len);
+        cursor.skip(len);
+        return written + len;
+      }
+
+      // write the whole current IOBuf
+      this->push(crtData, available);
+      cursor.skip(available);
+      written += available;
+      len -= available;
+    }
+  }
 
-  template <class OtherDerived, class OtherBuf>
-  explicit Cursor(const detail::CursorBase<OtherDerived, OtherBuf>& cursor)
-    : detail::CursorBase<Cursor, const IOBuf>(cursor) {}
 };
 
+} // namespace detail
+
 enum class CursorAccess {
   PRIVATE,
   UNSHARE

diff --git a/folly/io/test/IOBufCursorTest.cpp b/folly/io/test/IOBufCursorTest.cpp
index 3ec3055eacab4be8444408fff2b3c4ce5da9cbfb..58ad41d7e6eaec01b2c05ded1a9ff2587b9fd4f9 100644 (file)
--- a/folly/io/test/IOBufCursorTest.cpp
+++ b/folly/io/test/IOBufCursorTest.cpp
@@ -244,6 +244,52 @@ TEST(IOBuf, PullAndPeek) {
   }
 }
 
+TEST(IOBuf, pushCursorData) {
+  unique_ptr<IOBuf> iobuf1(IOBuf::create(20));
+  iobuf1->append(15);
+  iobuf1->trimStart(5);
+  unique_ptr<IOBuf> iobuf2(IOBuf::create(10));
+  unique_ptr<IOBuf> iobuf3(IOBuf::create(10));
+  iobuf3->append(10);
+
+  iobuf1->prependChain(std::move(iobuf2));
+  iobuf1->prependChain(std::move(iobuf3));
+  EXPECT_TRUE(iobuf1->isChained());
+
+  //write 20 bytes to the buffer chain
+  RWPrivateCursor wcursor(iobuf1.get());
+  wcursor.writeBE<uint64_t>(1);
+  wcursor.writeBE<uint64_t>(10);
+  wcursor.writeBE<uint32_t>(20);
+
+  // create a read buffer for the buffer chain
+  Cursor rcursor(iobuf1.get());
+  EXPECT_EQ(1, rcursor.readBE<uint64_t>());
+  EXPECT_EQ(10, rcursor.readBE<uint64_t>());
+  EXPECT_EQ(20, rcursor.readBE<uint32_t>());
+  EXPECT_EQ(0, rcursor.totalLength());
+  rcursor.reset(iobuf1.get());
+  EXPECT_EQ(20, rcursor.totalLength());
+
+  // create another write buffer
+  unique_ptr<IOBuf> iobuf4(IOBuf::create(30));
+  iobuf4->append(30);
+  RWPrivateCursor wcursor2(iobuf4.get());
+  // write buffer chain data into it, now wcursor2 should only
+  // have 10 bytes writable space
+  wcursor2.push(rcursor, 20);
+  EXPECT_EQ(wcursor2.totalLength(), 10);
+  // write again with not enough space in rcursor
+  EXPECT_THROW(wcursor2.push(rcursor, 20), std::out_of_range);
+
+  // create a read cursor to check iobuf3 data back
+  Cursor rcursor2(iobuf4.get());
+  EXPECT_EQ(1, rcursor2.readBE<uint64_t>());
+  EXPECT_EQ(10, rcursor2.readBE<uint64_t>());
+  EXPECT_EQ(20, rcursor2.readBE<uint32_t>());
+
+}
+
 TEST(IOBuf, Gather) {
   std::unique_ptr<IOBuf> iobuf1(IOBuf::create(10));
   append(iobuf1, "he");