folly/io/Cursor.h

   1 /*
   2  * Copyright 2013 Facebook, Inc.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *   http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 #ifndef FOLLY_CURSOR_H
  18 #define FOLLY_CURSOR_H
  19
  20 #include <assert.h>
  21 #include <stdexcept>
  22 #include <string.h>
  23 #include <type_traits>
  24 #include <memory>
  25
  26 #include "folly/Bits.h"
  27 #include "folly/io/IOBuf.h"
  28 #include "folly/Likely.h"
  29
  30 /**
  31  * Cursor class for fast iteration over IOBuf chains.
  32  *
  33  * Cursor - Read-only access
  34  *
  35  * RWPrivateCursor - Read-write access, assumes private access to IOBuf chain
  36  * RWUnshareCursor - Read-write access, calls unshare on write (COW)
  37  * Appender        - Write access, assumes private access to IOBuf chian
  38  *
  39  * Note that RW cursors write in the preallocated part of buffers (that is,
  40  * between the buffer's data() and tail()), while Appenders append to the end
  41  * of the buffer (between the buffer's tail() and bufferEnd()).  Appenders
  42  * automatically adjust the buffer pointers, so you may only use one
  43  * Appender with a buffer chain; for this reason, Appenders assume private
  44  * access to the buffer (you need to call unshare() yourself if necessary).
  45  **/
  46 namespace folly { namespace io {
  47 namespace detail {
  48
  49 template <class Derived, typename BufType>
  50 class CursorBase {
  51  public:
  52   const uint8_t* data() const {
  53     return crtBuf_->data() + offset_;
  54   }
  55
  56   // Space available in the current IOBuf.  May be 0; use peek() instead which
  57   // will always point to a non-empty chunk of data or at the end of the
  58   // chain.
  59   size_t length() const {
  60     return crtBuf_->length() - offset_;
  61   }
  62
  63   Derived& operator+=(size_t offset) {
  64     Derived* p = static_cast<Derived*>(this);
  65     p->skip(offset);
  66     return *p;
  67   }
  68
  69   template <class T>
  70   typename std::enable_if<std::is_integral<T>::value, T>::type
  71   read() {
  72     T val;
  73     pull(&val, sizeof(T));
  74     return val;
  75   }
  76
  77   template <class T>
  78   T readBE() {
  79     return Endian::big(read<T>());
  80   }
  81
  82   template <class T>
  83   T readLE() {
  84     return Endian::little(read<T>());
  85   }
  86
  87   explicit CursorBase(BufType* buf)
  88     : crtBuf_(buf)
  89     , offset_(0)
  90     , buffer_(buf) {}
  91
  92   // Make all the templated classes friends for copy constructor.
  93   template <class D, typename B> friend class CursorBase;
  94
  95   template <class T>
  96   explicit CursorBase(const T& cursor) {
  97     crtBuf_ = cursor.crtBuf_;
  98     offset_ = cursor.offset_;
  99     buffer_ = cursor.buffer_;
 100   }
 101
 102   // reset cursor to point to a new buffer.
 103   void reset(BufType* buf) {
 104     crtBuf_ = buf;
 105     buffer_ = buf;
 106     offset_ = 0;
 107   }
 108
 109   /**
 110    * Return the available data in the current buffer.
 111    * If you want to gather more data from the chain into a contiguous region
 112    * (for hopefully zero-copy access), use gather() before peek().
 113    */
 114   std::pair<const uint8_t*, size_t> peek() {
 115     // Ensure that we're pointing to valid data
 116     size_t available = length();
 117     while (UNLIKELY(available == 0 && tryAdvanceBuffer())) {
 118       available = length();
 119     }
 120
 121     return std::make_pair(data(), available);
 122   }
 123
 124   void pull(void* buf, size_t length) {
 125     if (UNLIKELY(pullAtMost(buf, length) != length)) {
 126       throw std::out_of_range("underflow");
 127     }
 128   }
 129
 130   void clone(std::unique_ptr<folly::IOBuf>& buf, size_t length) {
 131     if (UNLIKELY(cloneAtMost(buf, length) != length)) {
 132       throw std::out_of_range("underflow");
 133     }
 134   }
 135
 136   void skip(size_t length) {
 137     if (UNLIKELY(skipAtMost(length) != length)) {
 138       throw std::out_of_range("underflow");
 139     }
 140   }
 141
 142   size_t pullAtMost(void* buf, size_t len) {
 143     uint8_t* p = reinterpret_cast<uint8_t*>(buf);
 144     size_t copied = 0;
 145     for (;;) {
 146       // Fast path: it all fits in one buffer.
 147       size_t available = length();
 148       if (LIKELY(available >= len)) {
 149         memcpy(p, data(), len);
 150         offset_ += len;
 151         return copied + len;
 152       }
 153
 154       memcpy(p, data(), available);
 155       copied += available;
 156       if (UNLIKELY(!tryAdvanceBuffer())) {
 157         return copied;
 158       }
 159       p += available;
 160       len -= available;
 161     }
 162   }
 163
 164   size_t cloneAtMost(std::unique_ptr<folly::IOBuf>& buf, size_t len) {
 165     buf.reset(nullptr);
 166
 167     std::unique_ptr<folly::IOBuf> tmp;
 168     size_t copied = 0;
 169     for (;;) {
 170       // Fast path: it all fits in one buffer.
 171       size_t available = length();
 172       if (LIKELY(available >= len)) {
 173         tmp = crtBuf_->cloneOne();
 174         tmp->trimStart(offset_);
 175         tmp->trimEnd(tmp->length() - len);
 176         offset_ += len;
 177         if (!buf) {
 178           buf = std::move(tmp);
 179         } else {
 180           buf->prependChain(std::move(tmp));
 181         }
 182         return copied + len;
 183       }
 184
 185       tmp = crtBuf_->cloneOne();
 186       tmp->trimStart(offset_);
 187       if (!buf) {
 188         buf = std::move(tmp);
 189       } else {
 190         buf->prependChain(std::move(tmp));
 191       }
 192
 193       copied += available;
 194       if (UNLIKELY(!tryAdvanceBuffer())) {
 195         return copied;
 196       }
 197       len -= available;
 198     }
 199   }
 200
 201   size_t skipAtMost(size_t len) {
 202     size_t skipped = 0;
 203     for (;;) {
 204       // Fast path: it all fits in one buffer.
 205       size_t available = length();
 206       if (LIKELY(available >= len)) {
 207         offset_ += len;
 208         return skipped + len;
 209       }
 210
 211       skipped += available;
 212       if (UNLIKELY(!tryAdvanceBuffer())) {
 213         return skipped;
 214       }
 215       len -= available;
 216     }
 217   }
 218
 219   /**
 220    * Return the distance between two cursors.
 221    */
 222   size_t operator-(const CursorBase& other) const {
 223     BufType *otherBuf = other.crtBuf_;
 224     size_t len = 0;
 225
 226     if (otherBuf != crtBuf_) {
 227       len += otherBuf->length() - other.offset_;
 228
 229       for (otherBuf = otherBuf->next();
 230            otherBuf != crtBuf_ && otherBuf != other.buffer_;
 231            otherBuf = otherBuf->next()) {
 232         len += otherBuf->length();
 233       }
 234
 235       if (otherBuf == other.buffer_) {
 236         throw std::out_of_range("wrap-around");
 237       }
 238
 239       len += offset_;
 240     } else {
 241       if (offset_ < other.offset_) {
 242         throw std::out_of_range("underflow");
 243       }
 244
 245       len += offset_ - other.offset_;
 246     }
 247
 248     return len;
 249   }
 250
 251   /**
 252    * Return the distance from the given IOBuf to the this cursor.
 253    */
 254   size_t operator-(const BufType* buf) const {
 255     size_t len = 0;
 256
 257     BufType *curBuf = buf;
 258     while (curBuf != crtBuf_) {
 259       len += curBuf->length();
 260       curBuf = curBuf->next();
 261       if (curBuf == buf || curBuf == buffer_) {
 262         throw std::out_of_range("wrap-around");
 263       }
 264     }
 265
 266     len += offset_;
 267     return len;
 268   }
 269
 270  protected:
 271   BufType* crtBuf_;
 272   size_t offset_;
 273
 274   ~CursorBase(){}
 275
 276   bool tryAdvanceBuffer() {
 277     BufType* nextBuf = crtBuf_->next();
 278     if (UNLIKELY(nextBuf == buffer_)) {
 279       offset_ = crtBuf_->length();
 280       return false;
 281     }
 282
 283     offset_ = 0;
 284     crtBuf_ = nextBuf;
 285     static_cast<Derived*>(this)->advanceDone();
 286     return true;
 287   }
 288
 289  private:
 290   void advanceDone() {
 291   }
 292
 293   BufType* buffer_;
 294 };
 295
 296 template <class Derived>
 297 class Writable {
 298  public:
 299   template <class T>
 300   typename std::enable_if<std::is_integral<T>::value>::type
 301   write(T value) {
 302     const uint8_t* u8 = reinterpret_cast<const uint8_t*>(&value);
 303     push(u8, sizeof(T));
 304   }
 305
 306   template <class T>
 307   void writeBE(T value) {
 308     write(Endian::big(value));
 309   }
 310
 311   template <class T>
 312   void writeLE(T value) {
 313     write(Endian::little(value));
 314   }
 315
 316   void push(const uint8_t* buf, size_t len) {
 317     Derived* d = static_cast<Derived*>(this);
 318     if (d->pushAtMost(buf, len) != len) {
 319       throw std::out_of_range("overflow");
 320     }
 321   }
 322 };
 323
 324 } // namespace detail
 325
 326 class Cursor : public detail::CursorBase<Cursor, const IOBuf> {
 327  public:
 328   explicit Cursor(const IOBuf* buf)
 329     : detail::CursorBase<Cursor, const IOBuf>(buf) {}
 330
 331   template <class CursorType>
 332   explicit Cursor(CursorType& cursor)
 333     : detail::CursorBase<Cursor, const IOBuf>(cursor) {}
 334 };
 335
 336 enum class CursorAccess {
 337   PRIVATE,
 338   UNSHARE
 339 };
 340
 341 template <CursorAccess access>
 342 class RWCursor
 343   : public detail::CursorBase<RWCursor<access>, IOBuf>,
 344     public detail::Writable<RWCursor<access>> {
 345   friend class detail::CursorBase<RWCursor<access>, IOBuf>;
 346  public:
 347   explicit RWCursor(IOBuf* buf)
 348     : detail::CursorBase<RWCursor<access>, IOBuf>(buf),
 349       maybeShared_(true) {}
 350
 351   template <class CursorType>
 352   explicit RWCursor(CursorType& cursor)
 353     : detail::CursorBase<RWCursor<access>, IOBuf>(cursor),
 354       maybeShared_(true) {}
 355   /**
 356    * Gather at least n bytes contiguously into the current buffer,
 357    * by coalescing subsequent buffers from the chain as necessary.
 358    */
 359   void gather(size_t n) {
 360     this->crtBuf_->gather(this->offset_ + n);
 361   }
 362
 363   size_t pushAtMost(const uint8_t* buf, size_t len) {
 364     size_t copied = 0;
 365     for (;;) {
 366       // Fast path: the current buffer is big enough.
 367       size_t available = this->length();
 368       if (LIKELY(available >= len)) {
 369         if (access == CursorAccess::UNSHARE) {
 370           maybeUnshare();
 371         }
 372         memcpy(writableData(), buf, len);
 373         this->offset_ += len;
 374         return copied + len;
 375       }
 376
 377       if (access == CursorAccess::UNSHARE) {
 378         maybeUnshare();
 379       }
 380       memcpy(writableData(), buf, available);
 381       copied += available;
 382       if (UNLIKELY(!this->tryAdvanceBuffer())) {
 383         return copied;
 384       }
 385       buf += available;
 386       len -= available;
 387     }
 388   }
 389
 390   void insert(std::unique_ptr<folly::IOBuf> buf) {
 391     folly::IOBuf* nextBuf;
 392     if (this->offset_ == 0) {
 393       // Can just prepend
 394       nextBuf = buf.get();
 395       this->crtBuf_->prependChain(std::move(buf));
 396     } else {
 397       std::unique_ptr<folly::IOBuf> remaining;
 398       if (this->crtBuf_->length() - this->offset_ > 0) {
 399         // Need to split current IOBuf in two.
 400         remaining = this->crtBuf_->cloneOne();
 401         remaining->trimStart(this->offset_);
 402         nextBuf = remaining.get();
 403         buf->prependChain(std::move(remaining));
 404       } else {
 405         // Can just append
 406         nextBuf = this->crtBuf_->next();
 407       }
 408       this->crtBuf_->trimEnd(this->length());
 409       this->crtBuf_->appendChain(std::move(buf));
 410     }
 411     // Jump past the new links
 412     this->offset_ = 0;
 413     this->crtBuf_ = nextBuf;
 414   }
 415
 416   uint8_t* writableData() {
 417     return this->crtBuf_->writableData() + this->offset_;
 418   }
 419
 420  private:
 421   void maybeUnshare() {
 422     if (UNLIKELY(maybeShared_)) {
 423       this->crtBuf_->unshareOne();
 424       maybeShared_ = false;
 425     }
 426   }
 427
 428   void advanceDone() {
 429     maybeShared_ = true;
 430   }
 431
 432   bool maybeShared_;
 433 };
 434
 435 typedef RWCursor<CursorAccess::PRIVATE> RWPrivateCursor;
 436 typedef RWCursor<CursorAccess::UNSHARE> RWUnshareCursor;
 437
 438 /**
 439  * Append to the end of a buffer chain, growing the chain (by allocating new
 440  * buffers) in increments of at least growth bytes every time.  Won't grow
 441  * (and push() and ensure() will throw) if growth == 0.
 442  *
 443  * TODO(tudorb): add a flavor of Appender that reallocates one IOBuf instead
 444  * of chaining.
 445  */
 446 class Appender : public detail::Writable<Appender> {
 447  public:
 448   Appender(IOBuf* buf, uint32_t growth)
 449     : buffer_(buf),
 450       crtBuf_(buf->prev()),
 451       growth_(growth) {
 452   }
 453
 454   uint8_t* writableData() {
 455     return crtBuf_->writableTail();
 456   }
 457
 458   size_t length() const {
 459     return crtBuf_->tailroom();
 460   }
 461
 462   /**
 463    * Mark n bytes (must be <= length()) as appended, as per the
 464    * IOBuf::append() method.
 465    */
 466   void append(size_t n) {
 467     crtBuf_->append(n);
 468   }
 469
 470   /**
 471    * Ensure at least n contiguous bytes available to write.
 472    * Postcondition: length() >= n.
 473    */
 474   void ensure(uint32_t n) {
 475     if (LIKELY(length() >= n)) {
 476       return;
 477     }
 478
 479     // Waste the rest of the current buffer and allocate a new one.
 480     // Don't make it too small, either.
 481     if (growth_ == 0) {
 482       throw std::out_of_range("can't grow buffer chain");
 483     }
 484
 485     n = std::max(n, growth_);
 486     buffer_->prependChain(IOBuf::create(n));
 487     crtBuf_ = buffer_->prev();
 488   }
 489
 490   size_t pushAtMost(const uint8_t* buf, size_t len) {
 491     size_t copied = 0;
 492     for (;;) {
 493       // Fast path: it all fits in one buffer.
 494       size_t available = length();
 495       if (LIKELY(available >= len)) {
 496         memcpy(writableData(), buf, len);
 497         append(len);
 498         return copied + len;
 499       }
 500
 501       memcpy(writableData(), buf, available);
 502       append(available);
 503       copied += available;
 504       if (UNLIKELY(!tryGrowChain())) {
 505         return copied;
 506       }
 507       buf += available;
 508       len -= available;
 509     }
 510   }
 511
 512  private:
 513   bool tryGrowChain() {
 514     assert(crtBuf_->next() == buffer_);
 515     if (growth_ == 0) {
 516       return false;
 517     }
 518
 519     buffer_->prependChain(IOBuf::create(growth_));
 520     crtBuf_ = buffer_->prev();
 521     return true;
 522   }
 523
 524   IOBuf* buffer_;
 525   IOBuf* crtBuf_;
 526   uint32_t growth_;
 527 };
 528
 529 }}  // folly::io
 530
 531 #endif // FOLLY_CURSOR_H