folly/experimental/io/Cursor.h

   1 /*
   2  * Copyright 2012 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/experimental/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 skip(size_t length) {
 131     if (UNLIKELY(skipAtMost(length) != length)) {
 132       throw std::out_of_range("underflow");
 133     }
 134   }
 135
 136   size_t pullAtMost(void* buf, size_t len) {
 137     uint8_t* p = reinterpret_cast<uint8_t*>(buf);
 138     size_t copied = 0;
 139     for (;;) {
 140       // Fast path: it all fits in one buffer.
 141       size_t available = length();
 142       if (LIKELY(available >= len)) {
 143         memcpy(p, data(), len);
 144         offset_ += len;
 145         return copied + len;
 146       }
 147
 148       memcpy(p, data(), available);
 149       copied += available;
 150       if (UNLIKELY(!tryAdvanceBuffer())) {
 151         return copied;
 152       }
 153       p += available;
 154       len -= available;
 155     }
 156   }
 157
 158   size_t skipAtMost(size_t len) {
 159     size_t skipped = 0;
 160     for (;;) {
 161       // Fast path: it all fits in one buffer.
 162       size_t available = length();
 163       if (LIKELY(available >= len)) {
 164         offset_ += len;
 165         return skipped + len;
 166       }
 167
 168       skipped += available;
 169       if (UNLIKELY(!tryAdvanceBuffer())) {
 170         return skipped;
 171       }
 172       len -= available;
 173     }
 174   }
 175
 176  protected:
 177   BufType* crtBuf_;
 178   size_t offset_;
 179
 180   ~CursorBase(){}
 181
 182   bool tryAdvanceBuffer() {
 183     BufType* nextBuf = crtBuf_->next();
 184     if (UNLIKELY(nextBuf == buffer_)) {
 185       offset_ = crtBuf_->length();
 186       return false;
 187     }
 188
 189     offset_ = 0;
 190     crtBuf_ = nextBuf;
 191     static_cast<Derived*>(this)->advanceDone();
 192     return true;
 193   }
 194
 195  private:
 196   void advanceDone() {
 197   }
 198
 199   BufType* buffer_;
 200 };
 201
 202 template <class Derived>
 203 class Writable {
 204  public:
 205   template <class T>
 206   typename std::enable_if<std::is_integral<T>::value>::type
 207   write(T value) {
 208     const uint8_t* u8 = reinterpret_cast<const uint8_t*>(&value);
 209     push(u8, sizeof(T));
 210   }
 211
 212   template <class T>
 213   void writeBE(T value) {
 214     write(Endian::big(value));
 215   }
 216
 217   template <class T>
 218   void writeLE(T value) {
 219     write(Endian::little(value));
 220   }
 221
 222   void push(const uint8_t* buf, size_t len) {
 223     Derived* d = static_cast<Derived*>(this);
 224     if (d->pushAtMost(buf, len) != len) {
 225       throw std::out_of_range("overflow");
 226     }
 227   }
 228 };
 229
 230 } // namespace detail
 231
 232 class Cursor : public detail::CursorBase<Cursor, const IOBuf> {
 233  public:
 234   explicit Cursor(const IOBuf* buf)
 235     : detail::CursorBase<Cursor, const IOBuf>(buf) {}
 236
 237   template <class CursorType>
 238   explicit Cursor(CursorType& cursor)
 239     : detail::CursorBase<Cursor, const IOBuf>(cursor) {}
 240 };
 241
 242 enum class CursorAccess {
 243   PRIVATE,
 244   UNSHARE
 245 };
 246
 247 template <CursorAccess access>
 248 class RWCursor
 249   : public detail::CursorBase<RWCursor<access>, IOBuf>,
 250     public detail::Writable<RWCursor<access>> {
 251   friend class detail::CursorBase<RWCursor<access>, IOBuf>;
 252  public:
 253   explicit RWCursor(IOBuf* buf)
 254     : detail::CursorBase<RWCursor<access>, IOBuf>(buf),
 255       maybeShared_(true) {}
 256
 257   template <class CursorType>
 258   explicit RWCursor(CursorType& cursor)
 259     : detail::CursorBase<RWCursor<access>, IOBuf>(cursor),
 260       maybeShared_(true) {}
 261   /**
 262    * Gather at least n bytes contiguously into the current buffer,
 263    * by coalescing subsequent buffers from the chain as necessary.
 264    */
 265   void gather(size_t n) {
 266     this->crtBuf_->gather(this->offset_ + n);
 267   }
 268
 269   size_t pushAtMost(const uint8_t* buf, size_t len) {
 270     size_t copied = 0;
 271     for (;;) {
 272       // Fast path: the current buffer is big enough.
 273       size_t available = this->length();
 274       if (LIKELY(available >= len)) {
 275         if (access == CursorAccess::UNSHARE) {
 276           maybeUnshare();
 277         }
 278         memcpy(writableData(), buf, len);
 279         this->offset_ += len;
 280         return copied + len;
 281       }
 282
 283       if (access == CursorAccess::UNSHARE) {
 284         maybeUnshare();
 285       }
 286       memcpy(writableData(), buf, available);
 287       copied += available;
 288       if (UNLIKELY(!this->tryAdvanceBuffer())) {
 289         return copied;
 290       }
 291       buf += available;
 292       len -= available;
 293     }
 294   }
 295
 296   uint8_t* writableData() {
 297     return this->crtBuf_->writableData() + this->offset_;
 298   }
 299
 300  private:
 301   void maybeUnshare() {
 302     if (UNLIKELY(maybeShared_)) {
 303       this->crtBuf_->unshareOne();
 304       maybeShared_ = false;
 305     }
 306   }
 307
 308   void advanceDone() {
 309     maybeShared_ = true;
 310   }
 311
 312   bool maybeShared_;
 313 };
 314
 315 typedef RWCursor<CursorAccess::PRIVATE> RWPrivateCursor;
 316 typedef RWCursor<CursorAccess::UNSHARE> RWUnshareCursor;
 317
 318 /**
 319  * Append to the end of a buffer chain, growing the chain (by allocating new
 320  * buffers) in increments of at least growth bytes every time.  Won't grow
 321  * (and push() and ensure() will throw) if growth == 0.
 322  *
 323  * TODO(tudorb): add a flavor of Appender that reallocates one IOBuf instead
 324  * of chaining.
 325  */
 326 class Appender : public detail::Writable<Appender> {
 327  public:
 328   Appender(IOBuf* buf, uint32_t growth)
 329     : buffer_(buf),
 330       crtBuf_(buf->prev()),
 331       growth_(growth) {
 332   }
 333
 334   uint8_t* writableData() {
 335     return crtBuf_->writableTail();
 336   }
 337
 338   size_t length() const {
 339     return crtBuf_->tailroom();
 340   }
 341
 342   /**
 343    * Mark n bytes (must be <= length()) as appended, as per the
 344    * IOBuf::append() method.
 345    */
 346   void append(size_t n) {
 347     crtBuf_->append(n);
 348   }
 349
 350   /**
 351    * Ensure at least n contiguous bytes available to write.
 352    * Postcondition: length() >= n.
 353    */
 354   void ensure(uint32_t n) {
 355     if (LIKELY(length() >= n)) {
 356       return;
 357     }
 358
 359     // Waste the rest of the current buffer and allocate a new one.
 360     // Don't make it too small, either.
 361     if (growth_ == 0) {
 362       throw std::out_of_range("can't grow buffer chain");
 363     }
 364
 365     n = std::max(n, growth_);
 366     buffer_->prependChain(IOBuf::create(n));
 367     crtBuf_ = buffer_->prev();
 368   }
 369
 370   size_t pushAtMost(const uint8_t* buf, size_t len) {
 371     size_t copied = 0;
 372     for (;;) {
 373       // Fast path: it all fits in one buffer.
 374       size_t available = length();
 375       if (LIKELY(available >= len)) {
 376         memcpy(writableData(), buf, len);
 377         append(len);
 378         return copied + len;
 379       }
 380
 381       memcpy(writableData(), buf, available);
 382       append(available);
 383       copied += available;
 384       if (UNLIKELY(!tryGrowChain())) {
 385         return copied;
 386       }
 387       buf += available;
 388       len -= available;
 389     }
 390   }
 391
 392  private:
 393   bool tryGrowChain() {
 394     assert(crtBuf_->next() == buffer_);
 395     if (growth_ == 0) {
 396       return false;
 397     }
 398
 399     buffer_->prependChain(IOBuf::create(growth_));
 400     crtBuf_ = buffer_->prev();
 401     return true;
 402   }
 403
 404   IOBuf* buffer_;
 405   IOBuf* crtBuf_;
 406   uint32_t growth_;
 407 };
 408
 409 }}  // folly::io
 410
 411 #endif // FOLLY_CURSOR_H