Track absolute position of the cursor

Summary:
Start tracking the position of the cursor from the head of IOBuf chain. This comes at almost no cost (one arithmetic operation on IOBuf advance).

The main use case for this cursor is Thrift deserialization code. It allows us to stop accumulating `xfer` on every single byte/field/element write and instead get it from Cursor in the end (when we're exiting Thrift code).
This allows achieving ~10% better performance of deserialization.

Reviewed By: yfeldblum

Differential Revision: D6646813

fbshipit-source-id: 8f796854a24a411698e96afe037695e816813022

diff --git a/folly/io/Cursor.h b/folly/io/Cursor.h
index bdaeabe3d292bf0a9d67f2b3e9be1a9fd34b3783..838d366bb52257222d89116bf468c472bde2a2ed 100644 (file)
--- a/folly/io/Cursor.h
+++ b/folly/io/Cursor.h
@@ -73,10 +73,11 @@ class CursorBase {
   template <class OtherDerived, class OtherBuf>
   explicit CursorBase(const CursorBase<OtherDerived, OtherBuf>& cursor)
       : crtBuf_(cursor.crtBuf_),
+        buffer_(cursor.buffer_),
         crtBegin_(cursor.crtBegin_),
         crtEnd_(cursor.crtEnd_),
         crtPos_(cursor.crtPos_),
-        buffer_(cursor.buffer_) {}
+        absolutePos_(cursor.absolutePos_) {}
 
   /**
    * Reset cursor to point to a new buffer.
@@ -84,12 +85,20 @@ class CursorBase {
   void reset(BufType* buf) {
     crtBuf_ = buf;
     buffer_ = buf;
+    absolutePos_ = 0;
     if (crtBuf_) {
       crtPos_ = crtBegin_ = crtBuf_->data();
       crtEnd_ = crtBuf_->tail();
     }
   }
 
+  /**
+   * Get the current Cursor position relative to the head of IOBuf chain.
+   */
+  size_t getCurrentPosition() const {
+    return (crtPos_ - crtBegin_) + absolutePos_;
+  }
+
   const uint8_t* data() const {
     return crtPos_;
   }
@@ -167,10 +176,21 @@ class CursorBase {
    * Advances the cursor to the end of the entire IOBuf chain.
    */
   void advanceToEnd() {
-    crtBegin_ = buffer_->prev()->data();
-    crtPos_ = crtEnd_ = buffer_->prev()->tail();
-    if (crtBuf_ != buffer_->prev()) {
-      crtBuf_ = buffer_->prev();
+    // Simple case, we're already in the last IOBuf.
+    if (crtBuf_ == buffer_->prev()) {
+      crtPos_ = crtEnd_;
+      return;
+    }
+
+    auto* nextBuf = crtBuf_->next();
+    while (nextBuf != buffer_) {
+      absolutePos_ += crtEnd_ - crtBegin_;
+
+      crtBuf_ = nextBuf;
+      nextBuf = crtBuf_->next();
+      crtBegin_ = crtBuf_->data();
+      crtPos_ = crtEnd_ = crtBuf_->tail();
+
       static_cast<Derived*>(this)->advanceDone();
     }
   }
@@ -544,6 +564,7 @@ class CursorBase {
       return false;
     }
 
+    absolutePos_ += crtEnd_ - crtBegin_;
     crtBuf_ = nextBuf;
     crtPos_ = crtBegin_ = crtBuf_->data();
     crtEnd_ = crtBuf_->tail();
@@ -559,6 +580,7 @@ class CursorBase {
     crtBuf_ = crtBuf_->prev();
     crtBegin_ = crtBuf_->data();
     crtPos_ = crtEnd_ = crtBuf_->tail();
+    absolutePos_ -= crtEnd_ - crtBegin_;
     static_cast<Derived*>(this)->advanceDone();
     return true;
   }
@@ -570,9 +592,11 @@ class CursorBase {
   }
 
   BufType* crtBuf_;
+  BufType* buffer_;
   const uint8_t* crtBegin_{nullptr};
   const uint8_t* crtEnd_{nullptr};
   const uint8_t* crtPos_{nullptr};
+  size_t absolutePos_{0};
 
  private:
   template <class T>
@@ -660,8 +684,6 @@ class CursorBase {
 
   void advanceDone() {
   }
-
-  BufType* buffer_;
 };
 
 } // namespace detail
@@ -845,12 +867,12 @@ class RWCursor
   }
 
   void insert(std::unique_ptr<folly::IOBuf> buf) {
-    folly::IOBuf* nextBuf;
-    if (this->crtPos_ == this->crtBegin_) {
+    this->absolutePos_ += buf->computeChainDataLength();
+    if (this->crtPos_ == this->crtBegin_ && this->crtBuf_ != this->buffer_) {
       // Can just prepend
-      nextBuf = this->crtBuf_;
       this->crtBuf_->prependChain(std::move(buf));
     } else {
+      IOBuf* nextBuf;
       std::unique_ptr<folly::IOBuf> remaining;
       if (this->crtPos_ != this->crtEnd_) {
         // Need to split current IOBuf in two.
@@ -863,6 +885,7 @@ class RWCursor
         nextBuf = this->crtBuf_->next();
       }
       this->crtBuf_->trimEnd(this->length());
+      this->absolutePos_ += this->crtPos_ - this->crtBegin_;
       this->crtBuf_->appendChain(std::move(buf));
 
       // Jump past the new links

diff --git a/folly/io/test/IOBufCursorTest.cpp b/folly/io/test/IOBufCursorTest.cpp
index d9756a91c71e0f3c18a95028a0b5cc021d839bec..1f013da621bf1db55d117d0e45a792280a9d3089 100644 (file)
--- a/folly/io/test/IOBufCursorTest.cpp
+++ b/folly/io/test/IOBufCursorTest.cpp
@@ -13,7 +13,6 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-
 #include <folly/io/IOBuf.h>
 
 #include <folly/Format.h>
@@ -369,6 +368,7 @@ TEST(IOBuf, cloneAndInsert) {
 
     cursor.insert(std::move(cloned));
     cursor.insert(folly::IOBuf::create(0));
+    EXPECT_EQ(4, cursor.getCurrentPosition());
     EXPECT_EQ(7, iobuf1->countChainElements());
     EXPECT_EQ(14, iobuf1->computeChainDataLength());
     // Check that nextBuf got set correctly to the buffer with 1 byte left
@@ -386,24 +386,46 @@ TEST(IOBuf, cloneAndInsert) {
     cursor.skip(1);
 
     cursor.insert(std::move(cloned));
+    EXPECT_EQ(2, cursor.getCurrentPosition());
     EXPECT_EQ(8, iobuf1->countChainElements());
     EXPECT_EQ(15, iobuf1->computeChainDataLength());
     // Check that nextBuf got set correctly
     cursor.read<uint8_t>();
   }
   {
-    // Check that inserting at the beginning doesn't create empty buf
+    // Check that inserting at the beginning of a chunk (except first one)
+    // doesn't create empty buf
     RWPrivateCursor cursor(iobuf1.get());
     Cursor(iobuf1.get()).clone(cloned, 1);
     EXPECT_EQ(1, cloned->countChainElements());
     EXPECT_EQ(1, cloned->computeChainDataLength());
 
+    cursor.skip(1);
+
     cursor.insert(std::move(cloned));
+    EXPECT_EQ(2, cursor.getCurrentPosition());
+    EXPECT_EQ(14, cursor.totalLength());
     EXPECT_EQ(9, iobuf1->countChainElements());
     EXPECT_EQ(16, iobuf1->computeChainDataLength());
     // Check that nextBuf got set correctly
     cursor.read<uint8_t>();
   }
+  {
+    // Check that inserting at the beginning of a chain DOES keep an empty
+    // buffer.
+    RWPrivateCursor cursor(iobuf1.get());
+    Cursor(iobuf1.get()).clone(cloned, 1);
+    EXPECT_EQ(1, cloned->countChainElements());
+    EXPECT_EQ(1, cloned->computeChainDataLength());
+
+    cursor.insert(std::move(cloned));
+    EXPECT_EQ(1, cursor.getCurrentPosition());
+    EXPECT_EQ(16, cursor.totalLength());
+    EXPECT_EQ(11, iobuf1->countChainElements());
+    EXPECT_EQ(17, iobuf1->computeChainDataLength());
+    // Check that nextBuf got set correctly
+    cursor.read<uint8_t>();
+  }
 }
 
 TEST(IOBuf, cloneWithEmptyBufAtStart) {
@@ -1132,3 +1154,40 @@ TEST(IOBuf, pushEmptyByteRange) {
   app.push(emptyBytes);
   EXPECT_EQ(0, buf.computeChainDataLength());
 }
+
+TEST(IOBuf, positionTracking) {
+  unique_ptr<IOBuf> iobuf1(IOBuf::create(6));
+  iobuf1->append(6);
+  unique_ptr<IOBuf> iobuf2(IOBuf::create(24));
+  iobuf2->append(24);
+  iobuf1->prependChain(std::move(iobuf2));
+
+  Cursor cursor(iobuf1.get());
+
+  EXPECT_EQ(0, cursor.getCurrentPosition());
+  EXPECT_EQ(6, cursor.length());
+
+  cursor.skip(3);
+  EXPECT_EQ(3, cursor.getCurrentPosition());
+  EXPECT_EQ(3, cursor.length());
+
+  // Test that we properly handle advancing to the next chunk.
+  cursor.skip(4);
+  EXPECT_EQ(7, cursor.getCurrentPosition());
+  EXPECT_EQ(23, cursor.length());
+
+  // Test that we properly handle doing to the previous chunk.
+  cursor.retreat(2);
+  EXPECT_EQ(5, cursor.getCurrentPosition());
+  EXPECT_EQ(1, cursor.length());
+
+  // Test that we properly handle advanceToEnd
+  cursor.advanceToEnd();
+  EXPECT_EQ(30, cursor.getCurrentPosition());
+  EXPECT_EQ(0, cursor.totalLength());
+
+  // Reset to 0.
+  cursor.reset(iobuf1.get());
+  EXPECT_EQ(0, cursor.getCurrentPosition());
+  EXPECT_EQ(30, cursor.totalLength());
+}