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  protected:
 220   BufType* crtBuf_;
 221   size_t offset_;
 222
 223   ~CursorBase(){}
 224
 225   bool tryAdvanceBuffer() {
 226     BufType* nextBuf = crtBuf_->next();
 227     if (UNLIKELY(nextBuf == buffer_)) {
 228       offset_ = crtBuf_->length();
 229       return false;
 230     }
 231
 232     offset_ = 0;
 233     crtBuf_ = nextBuf;
 234     static_cast<Derived*>(this)->advanceDone();
 235     return true;
 236   }
 237
 238  private:
 239   void advanceDone() {
 240   }
 241
 242   BufType* buffer_;
 243 };
 244
 245 template <class Derived>
 246 class Writable {
 247  public:
 248   template <class T>
 249   typename std::enable_if<std::is_integral<T>::value>::type
 250   write(T value) {
 251     const uint8_t* u8 = reinterpret_cast<const uint8_t*>(&value);
 252     push(u8, sizeof(T));
 253   }
 254
 255   template <class T>
 256   void writeBE(T value) {
 257     write(Endian::big(value));
 258   }
 259
 260   template <class T>
 261   void writeLE(T value) {
 262     write(Endian::little(value));
 263   }
 264
 265   void push(const uint8_t* buf, size_t len) {
 266     Derived* d = static_cast<Derived*>(this);
 267     if (d->pushAtMost(buf, len) != len) {
 268       throw std::out_of_range("overflow");
 269     }
 270   }
 271 };
 272
 273 } // namespace detail
 274
 275 class Cursor : public detail::CursorBase<Cursor, const IOBuf> {
 276  public:
 277   explicit Cursor(const IOBuf* buf)
 278     : detail::CursorBase<Cursor, const IOBuf>(buf) {}
 279
 280   template <class CursorType>
 281   explicit Cursor(CursorType& cursor)
 282     : detail::CursorBase<Cursor, const IOBuf>(cursor) {}
 283 };
 284
 285 enum class CursorAccess {
 286   PRIVATE,
 287   UNSHARE
 288 };
 289
 290 template <CursorAccess access>
 291 class RWCursor
 292   : public detail::CursorBase<RWCursor<access>, IOBuf>,
 293     public detail::Writable<RWCursor<access>> {
 294   friend class detail::CursorBase<RWCursor<access>, IOBuf>;
 295  public:
 296   explicit RWCursor(IOBuf* buf)
 297     : detail::CursorBase<RWCursor<access>, IOBuf>(buf),
 298       maybeShared_(true) {}
 299
 300   template <class CursorType>
 301   explicit RWCursor(CursorType& cursor)
 302     : detail::CursorBase<RWCursor<access>, IOBuf>(cursor),
 303       maybeShared_(true) {}
 304   /**
 305    * Gather at least n bytes contiguously into the current buffer,
 306    * by coalescing subsequent buffers from the chain as necessary.
 307    */
 308   void gather(size_t n) {
 309     this->crtBuf_->gather(this->offset_ + n);
 310   }
 311
 312   size_t pushAtMost(const uint8_t* buf, size_t len) {
 313     size_t copied = 0;
 314     for (;;) {
 315       // Fast path: the current buffer is big enough.
 316       size_t available = this->length();
 317       if (LIKELY(available >= len)) {
 318         if (access == CursorAccess::UNSHARE) {
 319           maybeUnshare();
 320         }
 321         memcpy(writableData(), buf, len);
 322         this->offset_ += len;
 323         return copied + len;
 324       }
 325
 326       if (access == CursorAccess::UNSHARE) {
 327         maybeUnshare();
 328       }
 329       memcpy(writableData(), buf, available);
 330       copied += available;
 331       if (UNLIKELY(!this->tryAdvanceBuffer())) {
 332         return copied;
 333       }
 334       buf += available;
 335       len -= available;
 336     }
 337   }
 338
 339   void insert(std::unique_ptr<folly::IOBuf> buf) {
 340     folly::IOBuf* nextBuf;
 341     if (this->offset_ == 0) {
 342       // Can just prepend
 343       nextBuf = buf.get();
 344       this->crtBuf_->prependChain(std::move(buf));
 345     } else {
 346       std::unique_ptr<folly::IOBuf> remaining;
 347       if (this->crtBuf_->length() - this->offset_ > 0) {
 348         // Need to split current IOBuf in two.
 349         remaining = this->crtBuf_->cloneOne();
 350         remaining->trimStart(this->offset_);
 351         nextBuf = remaining.get();
 352         buf->prependChain(std::move(remaining));
 353       } else {
 354         // Can just append
 355         nextBuf = this->crtBuf_->next();
 356       }
 357       this->crtBuf_->trimEnd(this->length());
 358       this->crtBuf_->appendChain(std::move(buf));
 359     }
 360     // Jump past the new links
 361     this->offset_ = 0;
 362     this->crtBuf_ = nextBuf;
 363   }
 364
 365   uint8_t* writableData() {
 366     return this->crtBuf_->writableData() + this->offset_;
 367   }
 368
 369  private:
 370   void maybeUnshare() {
 371     if (UNLIKELY(maybeShared_)) {
 372       this->crtBuf_->unshareOne();
 373       maybeShared_ = false;
 374     }
 375   }
 376
 377   void advanceDone() {
 378     maybeShared_ = true;
 379   }
 380
 381   bool maybeShared_;
 382 };
 383
 384 typedef RWCursor<CursorAccess::PRIVATE> RWPrivateCursor;
 385 typedef RWCursor<CursorAccess::UNSHARE> RWUnshareCursor;
 386
 387 /**
 388  * Append to the end of a buffer chain, growing the chain (by allocating new
 389  * buffers) in increments of at least growth bytes every time.  Won't grow
 390  * (and push() and ensure() will throw) if growth == 0.
 391  *
 392  * TODO(tudorb): add a flavor of Appender that reallocates one IOBuf instead
 393  * of chaining.
 394  */
 395 class Appender : public detail::Writable<Appender> {
 396  public:
 397   Appender(IOBuf* buf, uint32_t growth)
 398     : buffer_(buf),
 399       crtBuf_(buf->prev()),
 400       growth_(growth) {
 401   }
 402
 403   uint8_t* writableData() {
 404     return crtBuf_->writableTail();
 405   }
 406
 407   size_t length() const {
 408     return crtBuf_->tailroom();
 409   }
 410
 411   /**
 412    * Mark n bytes (must be <= length()) as appended, as per the
 413    * IOBuf::append() method.
 414    */
 415   void append(size_t n) {
 416     crtBuf_->append(n);
 417   }
 418
 419   /**
 420    * Ensure at least n contiguous bytes available to write.
 421    * Postcondition: length() >= n.
 422    */
 423   void ensure(uint32_t n) {
 424     if (LIKELY(length() >= n)) {
 425       return;
 426     }
 427
 428     // Waste the rest of the current buffer and allocate a new one.
 429     // Don't make it too small, either.
 430     if (growth_ == 0) {
 431       throw std::out_of_range("can't grow buffer chain");
 432     }
 433
 434     n = std::max(n, growth_);
 435     buffer_->prependChain(IOBuf::create(n));
 436     crtBuf_ = buffer_->prev();
 437   }
 438
 439   size_t pushAtMost(const uint8_t* buf, size_t len) {
 440     size_t copied = 0;
 441     for (;;) {
 442       // Fast path: it all fits in one buffer.
 443       size_t available = length();
 444       if (LIKELY(available >= len)) {
 445         memcpy(writableData(), buf, len);
 446         append(len);
 447         return copied + len;
 448       }
 449
 450       memcpy(writableData(), buf, available);
 451       append(available);
 452       copied += available;
 453       if (UNLIKELY(!tryGrowChain())) {
 454         return copied;
 455       }
 456       buf += available;
 457       len -= available;
 458     }
 459   }
 460
 461  private:
 462   bool tryGrowChain() {
 463     assert(crtBuf_->next() == buffer_);
 464     if (growth_ == 0) {
 465       return false;
 466     }
 467
 468     buffer_->prependChain(IOBuf::create(growth_));
 469     crtBuf_ = buffer_->prev();
 470     return true;
 471   }
 472
 473   IOBuf* buffer_;
 474   IOBuf* crtBuf_;
 475   uint32_t growth_;
 476 };
 477
 478 }}  // folly::io
 479
 480 #endif // FOLLY_CURSOR_H