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/io/IOBufQueue.h"
  29 #include "folly/Likely.h"
  30
  31 /**
  32  * Cursor class for fast iteration over IOBuf chains.
  33  *
  34  * Cursor - Read-only access
  35  *
  36  * RWPrivateCursor - Read-write access, assumes private access to IOBuf chain
  37  * RWUnshareCursor - Read-write access, calls unshare on write (COW)
  38  * Appender        - Write access, assumes private access to IOBuf chian
  39  *
  40  * Note that RW cursors write in the preallocated part of buffers (that is,
  41  * between the buffer's data() and tail()), while Appenders append to the end
  42  * of the buffer (between the buffer's tail() and bufferEnd()).  Appenders
  43  * automatically adjust the buffer pointers, so you may only use one
  44  * Appender with a buffer chain; for this reason, Appenders assume private
  45  * access to the buffer (you need to call unshare() yourself if necessary).
  46  **/
  47 namespace folly { namespace io {
  48 namespace detail {
  49
  50 template <class Derived, typename BufType>
  51 class CursorBase {
  52  public:
  53   const uint8_t* data() const {
  54     return crtBuf_->data() + offset_;
  55   }
  56
  57   /*
  58    * Return the remaining space available in the current IOBuf.
  59    *
  60    * May return 0 if the cursor is at the end of an IOBuf.  Use peek() instead
  61    * if you want to avoid this.  peek() will advance to the next non-empty
  62    * IOBuf (up to the end of the chain) if the cursor is currently pointing at
  63    * the end of a buffer.
  64    */
  65   size_t length() const {
  66     return crtBuf_->length() - offset_;
  67   }
  68
  69   /*
  70    * Return the space available until the end of the entire IOBuf chain.
  71    */
  72   size_t totalLength() const {
  73     if (crtBuf_ == buffer_) {
  74       return crtBuf_->computeChainDataLength() - offset_;
  75     }
  76     CursorBase end(buffer_->prev());
  77     end.offset_ = end.buffer_->length();
  78     return end - *this;
  79   }
  80
  81   Derived& operator+=(size_t offset) {
  82     Derived* p = static_cast<Derived*>(this);
  83     p->skip(offset);
  84     return *p;
  85   }
  86
  87   template <class T>
  88   typename std::enable_if<std::is_integral<T>::value, T>::type
  89   read() {
  90     T val;
  91     pull(&val, sizeof(T));
  92     return val;
  93   }
  94
  95   template <class T>
  96   T readBE() {
  97     return Endian::big(read<T>());
  98   }
  99
 100   template <class T>
 101   T readLE() {
 102     return Endian::little(read<T>());
 103   }
 104
 105   /**
 106    * Read a fixed-length string.
 107    *
 108    * The std::string-based APIs should probably be avoided unless you
 109    * ultimately want the data to live in an std::string. You're better off
 110    * using the pull() APIs to copy into a raw buffer otherwise.
 111    */
 112   std::string readFixedString(size_t len) {
 113     std::string str;
 114
 115     str.reserve(len);
 116     for (;;) {
 117       // Fast path: it all fits in one buffer.
 118       size_t available = length();
 119       if (LIKELY(available >= len)) {
 120         str.append(reinterpret_cast<const char*>(data()), len);
 121         offset_ += len;
 122         return str;
 123       }
 124
 125       str.append(reinterpret_cast<const char*>(data()), available);
 126       if (UNLIKELY(!tryAdvanceBuffer())) {
 127         throw std::out_of_range("string underflow");
 128       }
 129       len -= available;
 130     }
 131   }
 132
 133   /**
 134    * Read a string consisting of bytes until the given terminator character is
 135    * seen. Raises an std::length_error if maxLength bytes have been processed
 136    * before the terminator is seen.
 137    *
 138    * See comments in readFixedString() about when it's appropriate to use this
 139    * vs. using pull().
 140    */
 141   std::string readTerminatedString(
 142     char termChar = '\0',
 143     size_t maxLength = std::numeric_limits<size_t>::max()) {
 144     std::string str;
 145
 146     for (;;) {
 147       const uint8_t* buf = data();
 148       size_t buflen = length();
 149
 150       size_t i = 0;
 151       while (i < buflen && buf[i] != termChar) {
 152         ++i;
 153
 154         // Do this check after incrementing 'i', as even though we start at the
 155         // 0 byte, it still represents a single character
 156         if (str.length() + i >= maxLength) {
 157           throw std::length_error("string overflow");
 158         }
 159       }
 160
 161       str.append(reinterpret_cast<const char*>(buf), i);
 162       if (i < buflen) {
 163         skip(i + 1);
 164         return str;
 165       }
 166
 167       skip(i);
 168
 169       if (UNLIKELY(!tryAdvanceBuffer())) {
 170         throw std::out_of_range("string underflow");
 171       }
 172     }
 173   }
 174
 175   explicit CursorBase(BufType* buf)
 176     : crtBuf_(buf)
 177     , offset_(0)
 178     , buffer_(buf) {}
 179
 180   // Make all the templated classes friends for copy constructor.
 181   template <class D, typename B> friend class CursorBase;
 182
 183   template <class T>
 184   explicit CursorBase(const T& cursor) {
 185     crtBuf_ = cursor.crtBuf_;
 186     offset_ = cursor.offset_;
 187     buffer_ = cursor.buffer_;
 188   }
 189
 190   // reset cursor to point to a new buffer.
 191   void reset(BufType* buf) {
 192     crtBuf_ = buf;
 193     buffer_ = buf;
 194     offset_ = 0;
 195   }
 196
 197   /**
 198    * Return the available data in the current buffer.
 199    * If you want to gather more data from the chain into a contiguous region
 200    * (for hopefully zero-copy access), use gather() before peek().
 201    */
 202   std::pair<const uint8_t*, size_t> peek() {
 203     // Ensure that we're pointing to valid data
 204     size_t available = length();
 205     while (UNLIKELY(available == 0 && tryAdvanceBuffer())) {
 206       available = length();
 207     }
 208
 209     return std::make_pair(data(), available);
 210   }
 211
 212   void pull(void* buf, size_t len) {
 213     if (UNLIKELY(pullAtMost(buf, len) != len)) {
 214       throw std::out_of_range("underflow");
 215     }
 216   }
 217
 218   void clone(std::unique_ptr<folly::IOBuf>& buf, size_t len) {
 219     if (UNLIKELY(cloneAtMost(buf, len) != len)) {
 220       throw std::out_of_range("underflow");
 221     }
 222   }
 223
 224   void skip(size_t len) {
 225     if (UNLIKELY(skipAtMost(len) != len)) {
 226       throw std::out_of_range("underflow");
 227     }
 228   }
 229
 230   size_t pullAtMost(void* buf, size_t len) {
 231     uint8_t* p = reinterpret_cast<uint8_t*>(buf);
 232     size_t copied = 0;
 233     for (;;) {
 234       // Fast path: it all fits in one buffer.
 235       size_t available = length();
 236       if (LIKELY(available >= len)) {
 237         memcpy(p, data(), len);
 238         offset_ += len;
 239         return copied + len;
 240       }
 241
 242       memcpy(p, data(), available);
 243       copied += available;
 244       if (UNLIKELY(!tryAdvanceBuffer())) {
 245         return copied;
 246       }
 247       p += available;
 248       len -= available;
 249     }
 250   }
 251
 252   size_t cloneAtMost(std::unique_ptr<folly::IOBuf>& buf, size_t len) {
 253     buf.reset(nullptr);
 254
 255     std::unique_ptr<folly::IOBuf> tmp;
 256     size_t copied = 0;
 257     for (;;) {
 258       // Fast path: it all fits in one buffer.
 259       size_t available = length();
 260       if (LIKELY(available >= len)) {
 261         tmp = crtBuf_->cloneOne();
 262         tmp->trimStart(offset_);
 263         tmp->trimEnd(tmp->length() - len);
 264         offset_ += len;
 265         if (!buf) {
 266           buf = std::move(tmp);
 267         } else {
 268           buf->prependChain(std::move(tmp));
 269         }
 270         return copied + len;
 271       }
 272
 273       tmp = crtBuf_->cloneOne();
 274       tmp->trimStart(offset_);
 275       if (!buf) {
 276         buf = std::move(tmp);
 277       } else {
 278         buf->prependChain(std::move(tmp));
 279       }
 280
 281       copied += available;
 282       if (UNLIKELY(!tryAdvanceBuffer())) {
 283         return copied;
 284       }
 285       len -= available;
 286     }
 287   }
 288
 289   size_t skipAtMost(size_t len) {
 290     size_t skipped = 0;
 291     for (;;) {
 292       // Fast path: it all fits in one buffer.
 293       size_t available = length();
 294       if (LIKELY(available >= len)) {
 295         offset_ += len;
 296         return skipped + len;
 297       }
 298
 299       skipped += available;
 300       if (UNLIKELY(!tryAdvanceBuffer())) {
 301         return skipped;
 302       }
 303       len -= available;
 304     }
 305   }
 306
 307   /**
 308    * Return the distance between two cursors.
 309    */
 310   size_t operator-(const CursorBase& other) const {
 311     BufType *otherBuf = other.crtBuf_;
 312     size_t len = 0;
 313
 314     if (otherBuf != crtBuf_) {
 315       len += otherBuf->length() - other.offset_;
 316
 317       for (otherBuf = otherBuf->next();
 318            otherBuf != crtBuf_ && otherBuf != other.buffer_;
 319            otherBuf = otherBuf->next()) {
 320         len += otherBuf->length();
 321       }
 322
 323       if (otherBuf == other.buffer_) {
 324         throw std::out_of_range("wrap-around");
 325       }
 326
 327       len += offset_;
 328     } else {
 329       if (offset_ < other.offset_) {
 330         throw std::out_of_range("underflow");
 331       }
 332
 333       len += offset_ - other.offset_;
 334     }
 335
 336     return len;
 337   }
 338
 339   /**
 340    * Return the distance from the given IOBuf to the this cursor.
 341    */
 342   size_t operator-(const BufType* buf) const {
 343     size_t len = 0;
 344
 345     BufType *curBuf = buf;
 346     while (curBuf != crtBuf_) {
 347       len += curBuf->length();
 348       curBuf = curBuf->next();
 349       if (curBuf == buf || curBuf == buffer_) {
 350         throw std::out_of_range("wrap-around");
 351       }
 352     }
 353
 354     len += offset_;
 355     return len;
 356   }
 357
 358  protected:
 359   BufType* crtBuf_;
 360   size_t offset_;
 361
 362   ~CursorBase(){}
 363
 364   BufType* head() {
 365     return buffer_;
 366   }
 367
 368   bool tryAdvanceBuffer() {
 369     BufType* nextBuf = crtBuf_->next();
 370     if (UNLIKELY(nextBuf == buffer_)) {
 371       offset_ = crtBuf_->length();
 372       return false;
 373     }
 374
 375     offset_ = 0;
 376     crtBuf_ = nextBuf;
 377     static_cast<Derived*>(this)->advanceDone();
 378     return true;
 379   }
 380
 381  private:
 382   void advanceDone() {
 383   }
 384
 385   BufType* buffer_;
 386 };
 387
 388 template <class Derived>
 389 class Writable {
 390  public:
 391   template <class T>
 392   typename std::enable_if<std::is_integral<T>::value>::type
 393   write(T value) {
 394     const uint8_t* u8 = reinterpret_cast<const uint8_t*>(&value);
 395     Derived* d = static_cast<Derived*>(this);
 396     d->push(u8, sizeof(T));
 397   }
 398
 399   template <class T>
 400   void writeBE(T value) {
 401     Derived* d = static_cast<Derived*>(this);
 402     d->write(Endian::big(value));
 403   }
 404
 405   template <class T>
 406   void writeLE(T value) {
 407     Derived* d = static_cast<Derived*>(this);
 408     d->write(Endian::little(value));
 409   }
 410
 411   void push(const uint8_t* buf, size_t len) {
 412     Derived* d = static_cast<Derived*>(this);
 413     if (d->pushAtMost(buf, len) != len) {
 414       throw std::out_of_range("overflow");
 415     }
 416   }
 417 };
 418
 419 } // namespace detail
 420
 421 class Cursor : public detail::CursorBase<Cursor, const IOBuf> {
 422  public:
 423   explicit Cursor(const IOBuf* buf)
 424     : detail::CursorBase<Cursor, const IOBuf>(buf) {}
 425
 426   template <class CursorType>
 427   explicit Cursor(CursorType& cursor)
 428     : detail::CursorBase<Cursor, const IOBuf>(cursor) {}
 429 };
 430
 431 enum class CursorAccess {
 432   PRIVATE,
 433   UNSHARE
 434 };
 435
 436 template <CursorAccess access>
 437 class RWCursor
 438   : public detail::CursorBase<RWCursor<access>, IOBuf>,
 439     public detail::Writable<RWCursor<access>> {
 440   friend class detail::CursorBase<RWCursor<access>, IOBuf>;
 441  public:
 442   explicit RWCursor(IOBuf* buf)
 443     : detail::CursorBase<RWCursor<access>, IOBuf>(buf),
 444       maybeShared_(true) {}
 445
 446   template <class CursorType>
 447   explicit RWCursor(CursorType& cursor)
 448     : detail::CursorBase<RWCursor<access>, IOBuf>(cursor),
 449       maybeShared_(true) {}
 450   /**
 451    * Gather at least n bytes contiguously into the current buffer,
 452    * by coalescing subsequent buffers from the chain as necessary.
 453    */
 454   void gather(size_t n) {
 455     // Forbid attempts to gather beyond the end of this IOBuf chain.
 456     // Otherwise we could try to coalesce the head of the chain and end up
 457     // accidentally freeing it, invalidating the pointer owned by external
 458     // code.
 459     //
 460     // If crtBuf_ == head() then IOBuf::gather() will perform all necessary
 461     // checking.  We only have to perform an explicit check here when calling
 462     // gather() on a non-head element.
 463     if (this->crtBuf_ != this->head() && this->totalLength() < n) {
 464       throw std::overflow_error("cannot gather() past the end of the chain");
 465     }
 466     this->crtBuf_->gather(this->offset_ + n);
 467   }
 468   void gatherAtMost(size_t n) {
 469     size_t size = std::min(n, this->totalLength());
 470     return this->crtBuf_->gather(this->offset_ + size);
 471   }
 472
 473   size_t pushAtMost(const uint8_t* buf, size_t len) {
 474     size_t copied = 0;
 475     for (;;) {
 476       // Fast path: the current buffer is big enough.
 477       size_t available = this->length();
 478       if (LIKELY(available >= len)) {
 479         if (access == CursorAccess::UNSHARE) {
 480           maybeUnshare();
 481         }
 482         memcpy(writableData(), buf, len);
 483         this->offset_ += len;
 484         return copied + len;
 485       }
 486
 487       if (access == CursorAccess::UNSHARE) {
 488         maybeUnshare();
 489       }
 490       memcpy(writableData(), buf, available);
 491       copied += available;
 492       if (UNLIKELY(!this->tryAdvanceBuffer())) {
 493         return copied;
 494       }
 495       buf += available;
 496       len -= available;
 497     }
 498   }
 499
 500   void insert(std::unique_ptr<folly::IOBuf> buf) {
 501     folly::IOBuf* nextBuf;
 502     if (this->offset_ == 0) {
 503       // Can just prepend
 504       nextBuf = buf.get();
 505       this->crtBuf_->prependChain(std::move(buf));
 506     } else {
 507       std::unique_ptr<folly::IOBuf> remaining;
 508       if (this->crtBuf_->length() - this->offset_ > 0) {
 509         // Need to split current IOBuf in two.
 510         remaining = this->crtBuf_->cloneOne();
 511         remaining->trimStart(this->offset_);
 512         nextBuf = remaining.get();
 513         buf->prependChain(std::move(remaining));
 514       } else {
 515         // Can just append
 516         nextBuf = this->crtBuf_->next();
 517       }
 518       this->crtBuf_->trimEnd(this->length());
 519       this->crtBuf_->appendChain(std::move(buf));
 520     }
 521     // Jump past the new links
 522     this->offset_ = 0;
 523     this->crtBuf_ = nextBuf;
 524   }
 525
 526   uint8_t* writableData() {
 527     return this->crtBuf_->writableData() + this->offset_;
 528   }
 529
 530  private:
 531   void maybeUnshare() {
 532     if (UNLIKELY(maybeShared_)) {
 533       this->crtBuf_->unshareOne();
 534       maybeShared_ = false;
 535     }
 536   }
 537
 538   void advanceDone() {
 539     maybeShared_ = true;
 540   }
 541
 542   bool maybeShared_;
 543 };
 544
 545 typedef RWCursor<CursorAccess::PRIVATE> RWPrivateCursor;
 546 typedef RWCursor<CursorAccess::UNSHARE> RWUnshareCursor;
 547
 548 /**
 549  * Append to the end of a buffer chain, growing the chain (by allocating new
 550  * buffers) in increments of at least growth bytes every time.  Won't grow
 551  * (and push() and ensure() will throw) if growth == 0.
 552  *
 553  * TODO(tudorb): add a flavor of Appender that reallocates one IOBuf instead
 554  * of chaining.
 555  */
 556 class Appender : public detail::Writable<Appender> {
 557  public:
 558   Appender(IOBuf* buf, uint32_t growth)
 559     : buffer_(buf),
 560       crtBuf_(buf->prev()),
 561       growth_(growth) {
 562   }
 563
 564   uint8_t* writableData() {
 565     return crtBuf_->writableTail();
 566   }
 567
 568   size_t length() const {
 569     return crtBuf_->tailroom();
 570   }
 571
 572   /**
 573    * Mark n bytes (must be <= length()) as appended, as per the
 574    * IOBuf::append() method.
 575    */
 576   void append(size_t n) {
 577     crtBuf_->append(n);
 578   }
 579
 580   /**
 581    * Ensure at least n contiguous bytes available to write.
 582    * Postcondition: length() >= n.
 583    */
 584   void ensure(uint32_t n) {
 585     if (LIKELY(length() >= n)) {
 586       return;
 587     }
 588
 589     // Waste the rest of the current buffer and allocate a new one.
 590     // Don't make it too small, either.
 591     if (growth_ == 0) {
 592       throw std::out_of_range("can't grow buffer chain");
 593     }
 594
 595     n = std::max(n, growth_);
 596     buffer_->prependChain(IOBuf::create(n));
 597     crtBuf_ = buffer_->prev();
 598   }
 599
 600   size_t pushAtMost(const uint8_t* buf, size_t len) {
 601     size_t copied = 0;
 602     for (;;) {
 603       // Fast path: it all fits in one buffer.
 604       size_t available = length();
 605       if (LIKELY(available >= len)) {
 606         memcpy(writableData(), buf, len);
 607         append(len);
 608         return copied + len;
 609       }
 610
 611       memcpy(writableData(), buf, available);
 612       append(available);
 613       copied += available;
 614       if (UNLIKELY(!tryGrowChain())) {
 615         return copied;
 616       }
 617       buf += available;
 618       len -= available;
 619     }
 620   }
 621
 622  private:
 623   bool tryGrowChain() {
 624     assert(crtBuf_->next() == buffer_);
 625     if (growth_ == 0) {
 626       return false;
 627     }
 628
 629     buffer_->prependChain(IOBuf::create(growth_));
 630     crtBuf_ = buffer_->prev();
 631     return true;
 632   }
 633
 634   IOBuf* buffer_;
 635   IOBuf* crtBuf_;
 636   uint32_t growth_;
 637 };
 638
 639 class QueueAppender : public detail::Writable<QueueAppender> {
 640  public:
 641   /**
 642    * Create an Appender that writes to a IOBufQueue.  When we allocate
 643    * space in the queue, we grow no more than growth bytes at once
 644    * (unless you call ensure() with a bigger value yourself).
 645    */
 646   QueueAppender(IOBufQueue* queue, uint32_t growth) {
 647     reset(queue, growth);
 648   }
 649
 650   void reset(IOBufQueue* queue, uint32_t growth) {
 651     queue_ = queue;
 652     growth_ = growth;
 653   }
 654
 655   uint8_t* writableData() {
 656     return static_cast<uint8_t*>(queue_->writableTail());
 657   }
 658
 659   size_t length() const { return queue_->tailroom(); }
 660
 661   void append(size_t n) { queue_->postallocate(n); }
 662
 663   // Ensure at least n contiguous; can go above growth_, throws if
 664   // not enough room.
 665   void ensure(uint32_t n) { queue_->preallocate(n, growth_); }
 666
 667   template <class T>
 668   typename std::enable_if<std::is_integral<T>::value>::type
 669   write(T value) {
 670     // We can't fail.
 671     auto p = queue_->preallocate(sizeof(T), growth_);
 672     storeUnaligned(p.first, value);
 673     queue_->postallocate(sizeof(T));
 674   }
 675
 676
 677   size_t pushAtMost(const uint8_t* buf, size_t len) {
 678     size_t remaining = len;
 679     while (remaining != 0) {
 680       auto p = queue_->preallocate(std::min(remaining, growth_),
 681                                    growth_,
 682                                    remaining);
 683       memcpy(p.first, buf, p.second);
 684       queue_->postallocate(p.second);
 685       buf += p.second;
 686       remaining -= p.second;
 687     }
 688
 689     return len;
 690   }
 691
 692   void insert(std::unique_ptr<folly::IOBuf> buf) {
 693     if (buf) {
 694       queue_->append(std::move(buf), true);
 695     }
 696   }
 697
 698  private:
 699   folly::IOBufQueue* queue_;
 700   size_t growth_;
 701 };
 702
 703 }}  // folly::io
 704
 705 #endif // FOLLY_CURSOR_H