IOBuf.h

Go to the documentation of this file.

/*
 * Copyright 2013-present Facebook, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#pragma once

#include <glog/logging.h>
#include <atomic>
#include <cassert>
#include <cinttypes>
#include <cstddef>
#include <cstring>
#include <iterator>
#include <limits>
#include <memory>
#include <type_traits>

#include <folly/FBString.h>
#include <folly/FBVector.h>
#include <folly/Portability.h>
#include <folly/Range.h>
#include <folly/detail/Iterators.h>
#include <folly/lang/Ordering.h>
#include <folly/portability/SysUio.h>

// Ignore shadowing warnings within this file, so includers can use -Wshadow.
FOLLY_PUSH_WARNING
FOLLY_GNU_DISABLE_WARNING("-Wshadow")

namespace folly {

namespace detail {
// Is T a unique_ptr<> to a standard-layout type?
template <typename T>
struct IsUniquePtrToSL : std::false_type {};
template <typename T, typename D>
struct IsUniquePtrToSL<std::unique_ptr<T, D>> : std::is_standard_layout<T> {};
} // namespace detail

class IOBuf {
 public:
  class Iterator;

  enum CreateOp { CREATE };
  enum WrapBufferOp { WRAP_BUFFER };
  enum TakeOwnershipOp { TAKE_OWNERSHIP };
  enum CopyBufferOp { COPY_BUFFER };

  typedef ByteRange value_type;
  typedef Iterator iterator;
  typedef Iterator const_iterator;

  typedef void (*FreeFunction)(void* buf, void* userData);

  static std::unique_ptr<IOBuf> create(std::size_t capacity);
  IOBuf(CreateOp, std::size_t capacity);

  static std::unique_ptr<IOBuf> createCombined(std::size_t capacity);

  static std::unique_ptr<IOBuf> createSeparate(std::size_t capacity);

  static std::unique_ptr<IOBuf> createChain(
      size_t totalCapacity,
      std::size_t maxBufCapacity);

  static std::unique_ptr<IOBuf> takeOwnership(
      void* buf,
      std::size_t capacity,
      FreeFunction freeFn = nullptr,
      void* userData = nullptr,
      bool freeOnError = true) {
    return takeOwnership(
        buf, capacity, capacity, freeFn, userData, freeOnError);
  }
  IOBuf(
      TakeOwnershipOp op,
      void* buf,
      std::size_t capacity,
      FreeFunction freeFn = nullptr,
      void* userData = nullptr,
      bool freeOnError = true)
      : IOBuf(op, buf, capacity, capacity, freeFn, userData, freeOnError) {}

  static std::unique_ptr<IOBuf> takeOwnership(
      void* buf,
      std::size_t capacity,
      std::size_t length,
      FreeFunction freeFn = nullptr,
      void* userData = nullptr,
      bool freeOnError = true);
  IOBuf(
      TakeOwnershipOp,
      void* buf,
      std::size_t capacity,
      std::size_t length,
      FreeFunction freeFn = nullptr,
      void* userData = nullptr,
      bool freeOnError = true);

  template <class UniquePtr>
  static typename std::enable_if<
      detail::IsUniquePtrToSL<UniquePtr>::value,
      std::unique_ptr<IOBuf>>::type
  takeOwnership(UniquePtr&& buf, size_t count = 1);

  static std::unique_ptr<IOBuf> wrapBuffer(
      const void* buf,
      std::size_t capacity);
  static std::unique_ptr<IOBuf> wrapBuffer(ByteRange br) {
    return wrapBuffer(br.data(), br.size());
  }

  static IOBuf wrapBufferAsValue(
      const void* buf,
      std::size_t capacity) noexcept;
  static IOBuf wrapBufferAsValue(ByteRange br) noexcept {
    return wrapBufferAsValue(br.data(), br.size());
  }

  IOBuf(WrapBufferOp op, const void* buf, std::size_t capacity) noexcept;
  IOBuf(WrapBufferOp op, ByteRange br) noexcept;

  static std::unique_ptr<IOBuf> copyBuffer(
      const void* buf,
      std::size_t size,
      std::size_t headroom = 0,
      std::size_t minTailroom = 0);
  static std::unique_ptr<IOBuf> copyBuffer(
      ByteRange br,
      std::size_t headroom = 0,
      std::size_t minTailroom = 0) {
    return copyBuffer(br.data(), br.size(), headroom, minTailroom);
  }
  IOBuf(
      CopyBufferOp op,
      const void* buf,
      std::size_t size,
      std::size_t headroom = 0,
      std::size_t minTailroom = 0);
  IOBuf(
      CopyBufferOp op,
      ByteRange br,
      std::size_t headroom = 0,
      std::size_t minTailroom = 0);

  static std::unique_ptr<IOBuf> copyBuffer(
      const std::string& buf,
      std::size_t headroom = 0,
      std::size_t minTailroom = 0);
  IOBuf(
      CopyBufferOp op,
      const std::string& buf,
      std::size_t headroom = 0,
      std::size_t minTailroom = 0)
      : IOBuf(op, buf.data(), buf.size(), headroom, minTailroom) {}

  static std::unique_ptr<IOBuf> maybeCopyBuffer(
      const std::string& buf,
      std::size_t headroom = 0,
      std::size_t minTailroom = 0);

  static void destroy(std::unique_ptr<IOBuf>&& data) {
    auto destroyer = std::move(data);
  }

  ~IOBuf();

  bool empty() const;

  const uint8_t* data() const {
    return data_;
  }

  uint8_t* writableData() {
    return data_;
  }

  const uint8_t* tail() const {
    return data_ + length_;
  }

  uint8_t* writableTail() {
    return data_ + length_;
  }

  std::size_t length() const {
    return length_;
  }

  std::size_t headroom() const {
    return std::size_t(data_ - buffer());
  }

  std::size_t tailroom() const {
    return std::size_t(bufferEnd() - tail());
  }

  const uint8_t* buffer() const {
    return buf_;
  }

  uint8_t* writableBuffer() {
    return buf_;
  }

  const uint8_t* bufferEnd() const {
    return buf_ + capacity_;
  }

  std::size_t capacity() const {
    return capacity_;
  }

  IOBuf* next() {
    return next_;
  }
  const IOBuf* next() const {
    return next_;
  }

  IOBuf* prev() {
    return prev_;
  }
  const IOBuf* prev() const {
    return prev_;
  }

  void advance(std::size_t amount) {
    // In debug builds, assert if there is a problem.
    assert(amount <= tailroom());

    if (length_ > 0) {
      memmove(data_ + amount, data_, length_);
    }
    data_ += amount;
  }

  void retreat(std::size_t amount) {
    // In debug builds, assert if there is a problem.
    assert(amount <= headroom());

    if (length_ > 0) {
      memmove(data_ - amount, data_, length_);
    }
    data_ -= amount;
  }

  void prepend(std::size_t amount) {
    DCHECK_LE(amount, headroom());
    data_ -= amount;
    length_ += amount;
  }

  void append(std::size_t amount) {
    DCHECK_LE(amount, tailroom());
    length_ += amount;
  }

  void trimStart(std::size_t amount) {
    DCHECK_LE(amount, length_);
    data_ += amount;
    length_ -= amount;
  }

  void trimEnd(std::size_t amount) {
    DCHECK_LE(amount, length_);
    length_ -= amount;
  }

  void clear() {
    data_ = writableBuffer();
    length_ = 0;
  }

  void reserve(std::size_t minHeadroom, std::size_t minTailroom) {
    // Maybe we don't need to do anything.
    if (headroom() >= minHeadroom && tailroom() >= minTailroom) {
      return;
    }
    // If the buffer is empty but we have enough total room (head + tail),
    // move the data_ pointer around.
    if (length() == 0 && headroom() + tailroom() >= minHeadroom + minTailroom) {
      data_ = writableBuffer() + minHeadroom;
      return;
    }
    // Bah, we have to do actual work.
    reserveSlow(minHeadroom, minTailroom);
  }

  bool isChained() const {
    assert((next_ == this) == (prev_ == this));
    return next_ != this;
  }

  size_t countChainElements() const;

  std::size_t computeChainDataLength() const;

  void prependChain(std::unique_ptr<IOBuf>&& iobuf);

  void appendChain(std::unique_ptr<IOBuf>&& iobuf) {
    // Just use prependChain() on the next element in our chain
    next_->prependChain(std::move(iobuf));
  }

  std::unique_ptr<IOBuf> unlink() {
    next_->prev_ = prev_;
    prev_->next_ = next_;
    prev_ = this;
    next_ = this;
    return std::unique_ptr<IOBuf>(this);
  }

  std::unique_ptr<IOBuf> pop() {
    IOBuf* next = next_;
    next_->prev_ = prev_;
    prev_->next_ = next_;
    prev_ = this;
    next_ = this;
    return std::unique_ptr<IOBuf>((next == this) ? nullptr : next);
  }

  std::unique_ptr<IOBuf> separateChain(IOBuf* head, IOBuf* tail) {
    assert(head != this);
    assert(tail != this);

    head->prev_->next_ = tail->next_;
    tail->next_->prev_ = head->prev_;

    head->prev_ = tail;
    tail->next_ = head;

    return std::unique_ptr<IOBuf>(head);
  }

  bool isShared() const {
    const IOBuf* current = this;
    while (true) {
      if (current->isSharedOne()) {
        return true;
      }
      current = current->next_;
      if (current == this) {
        return false;
      }
    }
  }

  bool isManaged() const {
    const IOBuf* current = this;
    while (true) {
      if (!current->isManagedOne()) {
        return false;
      }
      current = current->next_;
      if (current == this) {
        return true;
      }
    }
  }

  bool isManagedOne() const {
    return sharedInfo();
  }

  bool isSharedOne() const {
    // If this is a user-owned buffer, it is always considered shared
    if (UNLIKELY(!sharedInfo())) {
      return true;
    }

    if (UNLIKELY(sharedInfo()->externallyShared)) {
      return true;
    }

    if (LIKELY(!(flags() & kFlagMaybeShared))) {
      return false;
    }

    // kFlagMaybeShared is set, so we need to check the reference count.
    // (Checking the reference count requires an atomic operation, which is why
    // we prefer to only check kFlagMaybeShared if possible.)
    bool shared = sharedInfo()->refcount.load(std::memory_order_acquire) > 1;
    if (!shared) {
      // we're the last one left
      clearFlags(kFlagMaybeShared);
    }
    return shared;
  }

  void unshare() {
    if (isChained()) {
      unshareChained();
    } else {
      unshareOne();
    }
  }

  void unshareOne() {
    if (isSharedOne()) {
      unshareOneSlow();
    }
  }

  void markExternallyShared();

  void markExternallySharedOne() {
    SharedInfo* info = sharedInfo();
    if (info) {
      info->externallyShared = true;
    }
  }

  void makeManaged() {
    if (isChained()) {
      makeManagedChained();
    } else {
      makeManagedOne();
    }
  }

  void makeManagedOne() {
    if (!isManagedOne()) {
      // We can call the internal function directly; unmanaged implies shared.
      unshareOneSlow();
    }
  }

  ByteRange coalesce() {
    const std::size_t newHeadroom = headroom();
    const std::size_t newTailroom = prev()->tailroom();
    return coalesceWithHeadroomTailroom(newHeadroom, newTailroom);
  }

  ByteRange coalesceWithHeadroomTailroom(
      std::size_t newHeadroom,
      std::size_t newTailroom) {
    if (isChained()) {
      coalesceAndReallocate(
          newHeadroom, computeChainDataLength(), this, newTailroom);
    }
    return ByteRange(data_, length_);
  }

  void gather(std::size_t maxLength) {
    if (!isChained() || length_ >= maxLength) {
      return;
    }
    coalesceSlow(maxLength);
  }

  std::unique_ptr<IOBuf> clone() const;

  IOBuf cloneAsValue() const;

  std::unique_ptr<IOBuf> cloneOne() const;

  IOBuf cloneOneAsValue() const;

  std::unique_ptr<IOBuf> cloneCoalesced() const;

  std::unique_ptr<IOBuf> cloneCoalescedWithHeadroomTailroom(
      std::size_t newHeadroom,
      std::size_t newTailroom) const;

  IOBuf cloneCoalescedAsValue() const;

  IOBuf cloneCoalescedAsValueWithHeadroomTailroom(
      std::size_t newHeadroom,
      std::size_t newTailroom) const;

  void cloneInto(IOBuf& other) const {
    other = cloneAsValue();
  }

  void cloneOneInto(IOBuf& other) const {
    other = cloneOneAsValue();
  }

  folly::fbvector<struct iovec> getIov() const;

  void appendToIov(folly::fbvector<struct iovec>* iov) const;

  size_t fillIov(struct iovec* iov, size_t len) const;

  static std::unique_ptr<IOBuf> wrapIov(const iovec* vec, size_t count);

  static std::unique_ptr<IOBuf> takeOwnershipIov(
      const iovec* vec,
      size_t count,
      FreeFunction freeFn = nullptr,
      void* userData = nullptr,
      bool freeOnError = true);

  /*
   * Overridden operator new and delete.
   * These perform specialized memory management to help support
   * createCombined(), which allocates IOBuf objects together with the buffer
   * data.
   */
  void* operator new(size_t size);
  void* operator new(size_t size, void* ptr);
  void operator delete(void* ptr);
  void operator delete(void* ptr, void* placement);

  fbstring moveToFbString();

  Iterator cbegin() const;
  Iterator cend() const;
  Iterator begin() const;
  Iterator end() const;

  IOBuf() noexcept;

  IOBuf(IOBuf&& other) noexcept;
  IOBuf& operator=(IOBuf&& other) noexcept;

  IOBuf(const IOBuf& other);
  IOBuf& operator=(const IOBuf& other);

 private:
  enum FlagsEnum : uintptr_t {
    // Adding any more flags would not work on 32-bit architectures,
    // as these flags are stashed in the least significant 2 bits of a
    // max-align-aligned pointer.
    kFlagFreeSharedInfo = 0x1,
    kFlagMaybeShared = 0x2,
    kFlagMask = kFlagFreeSharedInfo | kFlagMaybeShared
  };

  struct SharedInfo {
    SharedInfo();
    SharedInfo(FreeFunction fn, void* arg);

    // A pointer to a function to call to free the buffer when the refcount
    // hits 0.  If this is null, free() will be used instead.
    FreeFunction freeFn;
    void* userData;
    std::atomic<uint32_t> refcount;
    bool externallyShared{false};
  };
  // Helper structs for use by operator new and delete
  struct HeapPrefix;
  struct HeapStorage;
  struct HeapFullStorage;

  struct InternalConstructor {}; // avoid conflicts
  IOBuf(
      InternalConstructor,
      uintptr_t flagsAndSharedInfo,
      uint8_t* buf,
      std::size_t capacity,
      uint8_t* data,
      std::size_t length) noexcept;

  void unshareOneSlow();
  void unshareChained();
  void makeManagedChained();
  void coalesceSlow();
  void coalesceSlow(size_t maxLength);
  // newLength must be the entire length of the buffers between this and
  // end (no truncation)
  void coalesceAndReallocate(
      size_t newHeadroom,
      size_t newLength,
      IOBuf* end,
      size_t newTailroom);
  void coalesceAndReallocate(size_t newLength, IOBuf* end) {
    coalesceAndReallocate(headroom(), newLength, end, end->prev_->tailroom());
  }
  void decrementRefcount();
  void reserveSlow(std::size_t minHeadroom, std::size_t minTailroom);
  void freeExtBuffer();

  static size_t goodExtBufferSize(std::size_t minCapacity);
  static void initExtBuffer(
      uint8_t* buf,
      size_t mallocSize,
      SharedInfo** infoReturn,
      std::size_t* capacityReturn);
  static void allocExtBuffer(
      std::size_t minCapacity,
      uint8_t** bufReturn,
      SharedInfo** infoReturn,
      std::size_t* capacityReturn);
  static void releaseStorage(HeapStorage* storage, uint16_t freeFlags);
  static void freeInternalBuf(void* buf, void* userData);

  /*
   * Member variables
   */

  /*
   * Links to the next and the previous IOBuf in this chain.
   *
   * The chain is circularly linked (the last element in the chain points back
   * at the head), and next_ and prev_ can never be null.  If this IOBuf is the
   * only element in the chain, next_ and prev_ will both point to this.
   */
  IOBuf* next_{this};
  IOBuf* prev_{this};

  /*
   * A pointer to the start of the data referenced by this IOBuf, and the
   * length of the data.
   *
   * This may refer to any subsection of the actual buffer capacity.
   */
  uint8_t* data_{nullptr};
  uint8_t* buf_{nullptr};
  std::size_t length_{0};
  std::size_t capacity_{0};

  // Pack flags in least significant 2 bits, sharedInfo in the rest
  mutable uintptr_t flagsAndSharedInfo_{0};

  static inline uintptr_t packFlagsAndSharedInfo(
      uintptr_t flags,
      SharedInfo* info) {
    uintptr_t uinfo = reinterpret_cast<uintptr_t>(info);
    DCHECK_EQ(flags & ~kFlagMask, 0u);
    DCHECK_EQ(uinfo & kFlagMask, 0u);
    return flags | uinfo;
  }

  inline SharedInfo* sharedInfo() const {
    return reinterpret_cast<SharedInfo*>(flagsAndSharedInfo_ & ~kFlagMask);
  }

  inline void setSharedInfo(SharedInfo* info) {
    uintptr_t uinfo = reinterpret_cast<uintptr_t>(info);
    DCHECK_EQ(uinfo & kFlagMask, 0u);
    flagsAndSharedInfo_ = (flagsAndSharedInfo_ & kFlagMask) | uinfo;
  }

  inline uintptr_t flags() const {
    return flagsAndSharedInfo_ & kFlagMask;
  }

  // flags_ are changed from const methods
  inline void setFlags(uintptr_t flags) const {
    DCHECK_EQ(flags & ~kFlagMask, 0u);
    flagsAndSharedInfo_ |= flags;
  }

  inline void clearFlags(uintptr_t flags) const {
    DCHECK_EQ(flags & ~kFlagMask, 0u);
    flagsAndSharedInfo_ &= ~flags;
  }

  inline void setFlagsAndSharedInfo(uintptr_t flags, SharedInfo* info) {
    flagsAndSharedInfo_ = packFlagsAndSharedInfo(flags, info);
  }

  struct DeleterBase {
    virtual ~DeleterBase() {}
    virtual void dispose(void* p) = 0;
  };

  template <class UniquePtr>
  struct UniquePtrDeleter : public DeleterBase {
    typedef typename UniquePtr::pointer Pointer;
    typedef typename UniquePtr::deleter_type Deleter;

    explicit UniquePtrDeleter(Deleter deleter) : deleter_(std::move(deleter)) {}
    void dispose(void* p) override {
      try {
        deleter_(static_cast<Pointer>(p));
        delete this;
      } catch (...) {
        abort();
      }
    }

   private:
    Deleter deleter_;
  };

  static void freeUniquePtrBuffer(void* ptr, void* userData) {
    static_cast<DeleterBase*>(userData)->dispose(ptr);
  }
};

struct IOBufHash {
  size_t operator()(const IOBuf& buf) const noexcept;
  size_t operator()(const std::unique_ptr<IOBuf>& buf) const noexcept {
    return operator()(buf.get());
  }
  size_t operator()(const IOBuf* buf) const noexcept {
    return buf ? (*this)(*buf) : 0;
  }
};

struct IOBufCompare {
  ordering operator()(const IOBuf& a, const IOBuf& b) const {
    return &a == &b ? ordering::eq : impl(a, b);
  }
  ordering operator()(
      const std::unique_ptr<IOBuf>& a,
      const std::unique_ptr<IOBuf>& b) const {
    return operator()(a.get(), b.get());
  }
  ordering operator()(const IOBuf* a, const IOBuf* b) const {
    // clang-format off
    return
        !a && !b ? ordering::eq :
        !a && b ? ordering::lt :
        a && !b ? ordering::gt :
        operator()(*a, *b);
    // clang-format on
  }

 private:
  ordering impl(IOBuf const& a, IOBuf const& b) const noexcept;
};

struct IOBufEqualTo : compare_equal_to<IOBufCompare> {};

struct IOBufNotEqualTo : compare_not_equal_to<IOBufCompare> {};

struct IOBufLess : compare_less<IOBufCompare> {};

struct IOBufLessEqual : compare_less_equal<IOBufCompare> {};

struct IOBufGreater : compare_greater<IOBufCompare> {};

struct IOBufGreaterEqual : compare_greater_equal<IOBufCompare> {};

template <class UniquePtr>
typename std::enable_if<
    detail::IsUniquePtrToSL<UniquePtr>::value,
    std::unique_ptr<IOBuf>>::type
IOBuf::takeOwnership(UniquePtr&& buf, size_t count) {
  size_t size = count * sizeof(typename UniquePtr::element_type);
  auto deleter = new UniquePtrDeleter<UniquePtr>(buf.get_deleter());
  return takeOwnership(
      buf.release(), size, &IOBuf::freeUniquePtrBuffer, deleter);
}

inline std::unique_ptr<IOBuf> IOBuf::copyBuffer(
    const void* data,
    std::size_t size,
    std::size_t headroom,
    std::size_t minTailroom) {
  std::size_t capacity = headroom + size + minTailroom;
  std::unique_ptr<IOBuf> buf = create(capacity);
  buf->advance(headroom);
  if (size != 0) {
    memcpy(buf->writableData(), data, size);
  }
  buf->append(size);
  return buf;
}

inline std::unique_ptr<IOBuf> IOBuf::copyBuffer(
    const std::string& buf,
    std::size_t headroom,
    std::size_t minTailroom) {
  return copyBuffer(buf.data(), buf.size(), headroom, minTailroom);
}

inline std::unique_ptr<IOBuf> IOBuf::maybeCopyBuffer(
    const std::string& buf,
    std::size_t headroom,
    std::size_t minTailroom) {
  if (buf.empty()) {
    return nullptr;
  }
  return copyBuffer(buf.data(), buf.size(), headroom, minTailroom);
}

class IOBuf::Iterator : public detail::IteratorFacade<
                            IOBuf::Iterator,
                            ByteRange const,
                            std::forward_iterator_tag> {
 public:
  // Note that IOBufs are stored as a circular list without a guard node,
  // so pos == end is ambiguous (it may mean "begin" or "end").  To solve
  // the ambiguity (at the cost of one extra comparison in the "increment"
  // code path), we define end iterators as having pos_ == end_ == nullptr
  // and we only allow forward iteration.
  explicit Iterator(const IOBuf* pos, const IOBuf* end) : pos_(pos), end_(end) {
    // Sadly, we must return by const reference, not by value.
    if (pos_) {
      setVal();
    }
  }

  Iterator() {}

  Iterator(Iterator const& rhs) : Iterator(rhs.pos_, rhs.end_) {}

  Iterator& operator=(Iterator const& rhs) {
    pos_ = rhs.pos_;
    end_ = rhs.end_;
    if (pos_) {
      setVal();
    }
    return *this;
  }

  const ByteRange& dereference() const {
    return val_;
  }

  bool equal(const Iterator& other) const {
    // We must compare end_ in addition to pos_, because forward traversal
    // requires that if two iterators are equal (a == b) and dereferenceable,
    // then ++a == ++b.
    return pos_ == other.pos_ && end_ == other.end_;
  }

  void increment() {
    pos_ = pos_->next();
    adjustForEnd();
  }

 private:
  void setVal() {
    val_ = ByteRange(pos_->data(), pos_->tail());
  }

  void adjustForEnd() {
    if (pos_ == end_) {
      pos_ = end_ = nullptr;
      val_ = ByteRange();
    } else {
      setVal();
    }
  }

  const IOBuf* pos_{nullptr};
  const IOBuf* end_{nullptr};
  ByteRange val_;
};

inline IOBuf::Iterator IOBuf::begin() const {
  return cbegin();
}
inline IOBuf::Iterator IOBuf::end() const {
  return cend();
}

} // namespace folly

FOLLY_POP_WARNING