membuf.h

/*
 * 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.
 */
 
/*
 * The code in this file if adapated from the IOBuf of folly:
 * https://github.com/facebook/folly/blob/master/folly/io/IOBuf.h
 */
 
#pragma once
 
#include <hicn/transport/portability/portability.h>
#include <hicn/transport/utils/branch_prediction.h>
#include <stdlib.h>
 
#include <atomic>
#include <cassert>
#include <cinttypes>
#include <cstddef>
#include <cstring>
#include <iterator>
#include <limits>
#include <memory>
#include <type_traits>
#include <vector>
 
#ifndef _WIN32
TRANSPORT_GNU_DISABLE_WARNING("-Wshadow")
#endif
 
namespace utils {
 
class MemBuf {
 public:
  enum CreateOp { CREATE };
  enum WrapBufferOp { WRAP_BUFFER };
  enum TakeOwnershipOp { TAKE_OWNERSHIP };
  enum CopyBufferOp { COPY_BUFFER };
 
  using Ptr = std::shared_ptr<MemBuf>;
 
  typedef void (*FreeFunction)(void* buf, void* userData);
 
  static std::unique_ptr<MemBuf> create(std::size_t capacity);
  MemBuf(CreateOp, std::size_t capacity);
 
  static std::unique_ptr<MemBuf> createCombined(std::size_t capacity);
 
  static std::unique_ptr<MemBuf> createSeparate(std::size_t capacity);
 
  static std::unique_ptr<MemBuf> createChain(size_t totalCapacity,
                                             std::size_t maxBufCapacity);
 
  static std::unique_ptr<MemBuf> takeOwnership(void* buf, std::size_t capacity,
                                               FreeFunction freeFn = nullptr,
                                               void* userData = nullptr,
                                               bool freeOnError = true) {
    return takeOwnership(buf, capacity, capacity, freeFn, userData,
                         freeOnError);
  }
 
  MemBuf(TakeOwnershipOp op, void* buf, std::size_t capacity,
         FreeFunction freeFn = nullptr, void* userData = nullptr,
         bool freeOnError = true)
      : MemBuf(op, buf, capacity, capacity, freeFn, userData, freeOnError) {}
 
  static std::unique_ptr<MemBuf> takeOwnership(void* buf, std::size_t capacity,
                                               std::size_t length,
                                               FreeFunction freeFn = nullptr,
                                               void* userData = nullptr,
                                               bool freeOnError = true);
 
  MemBuf(TakeOwnershipOp, void* buf, std::size_t capacity, std::size_t length,
         FreeFunction freeFn = nullptr, void* userData = nullptr,
         bool freeOnError = true);
 
  static std::unique_ptr<MemBuf> wrapBuffer(const void* buf, std::size_t length,
                                            std::size_t capacity);
 
  static MemBuf wrapBufferAsValue(const void* buf, std::size_t length,
                                  std::size_t capacity) noexcept;
 
  MemBuf(WrapBufferOp op, const void* buf, std::size_t length,
         std::size_t capacity) noexcept;
 
  static std::unique_ptr<MemBuf> copyBuffer(const void* buf, std::size_t size,
                                            std::size_t headroom = 0,
                                            std::size_t minTailroom = 0);
 
  MemBuf(CopyBufferOp op, const void* buf, std::size_t size,
         std::size_t headroom = 0, std::size_t minTailroom = 0);
 
  static void destroy(std::unique_ptr<MemBuf>&& data) {
    auto destroyer = std::move(data);
  }
 
  ~MemBuf();
 
  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_; }
 
  void setLength(std::size_t length) { length_ = 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_; }
 
  MemBuf* next() { return next_; }
 
  const MemBuf* next() const { return next_; }
 
  MemBuf* prev() { return prev_; }
 
  const MemBuf* 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) {
    data_ -= amount;
    length_ += amount;
  }
 
  void append(std::size_t amount) { length_ += amount; }
 
  void trimStart(std::size_t amount) {
    data_ += amount;
    length_ -= amount;
  }
 
  void trimEnd(std::size_t amount) { length_ -= amount; }
 
  // Never call clear on cloned membuf sharing different
  // portions of the same underlying buffer.
  // Use the trim functions instead.
  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<MemBuf>&& iobuf);
 
  void appendChain(std::unique_ptr<MemBuf>&& iobuf) {
    // Just use prependChain() on the next element in our chain
    next_->prependChain(std::move(iobuf));
  }
 
  std::unique_ptr<MemBuf> unlink() {
    next_->prev_ = prev_;
    prev_->next_ = next_;
    prev_ = this;
    next_ = this;
    return std::unique_ptr<MemBuf>(this);
  }
 
  std::unique_ptr<MemBuf> pop() {
    MemBuf* next = next_;
    next_->prev_ = prev_;
    prev_->next_ = next_;
    prev_ = this;
    next_ = this;
    return std::unique_ptr<MemBuf>((next == this) ? nullptr : next);
  }
 
  std::unique_ptr<MemBuf> separateChain(MemBuf* head, MemBuf* 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<MemBuf>(head);
  }
 
  bool isShared() const {
    const MemBuf* current = this;
    while (true) {
      if (current->isSharedOne()) {
        return true;
      }
      current = current->next_;
      if (current == this) {
        return false;
      }
    }
  }
 
  bool isManaged() const {
    const MemBuf* 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 ((TRANSPORT_EXPECT_FALSE(!sharedInfo()))) {
      return true;
    }
 
    if ((TRANSPORT_EXPECT_FALSE(sharedInfo()->externallyShared))) {
      return true;
    }
 
    if ((TRANSPORT_EXPECT_TRUE(!(flags() & flag_maybe_shared)))) {
      return false;
    }
 
    // flag_maybe_shared 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 flag_maybe_shared if possible.)
    bool shared = sharedInfo()->refcount.load(std::memory_order_acquire) > 1;
    if (!shared) {
      // we're the last one left
      clearFlags(flag_maybe_shared);
    }
    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();
    }
  }
 
  // /**
  //  * Coalesce this MemBuf chain into a single buffer.
  //  *
  //  * This method moves all of the data in this MemBuf chain into a single
  //  * contiguous buffer, if it is not already in one buffer.  After coalesce()
  //  * returns, this MemBuf will be a chain of length one.  Other MemBufs in
  //  the
  //  * chain will be automatically deleted.
  //  *
  //  * After coalescing, the MemBuf will have at least as much headroom as the
  //  * first MemBuf in the chain, and at least as much tailroom as the last
  //  MemBuf
  //  * in the chain.
  //  *
  //  * Throws std::bad_alloc on error.  On error the MemBuf chain will be
  //  * unmodified.
  //  *
  //  * Returns ByteRange that points to the data MemBuf stores.
  //  */
  // ByteRange coalesce() {
  //   const std::size_t newHeadroom = headroom();
  //   const std::size_t newTailroom = prev()->tailroom();
  //   return coalesceWithHeadroomTailroom(newHeadroom, newTailroom);
  // }
 
  // /**
  //  * This is similar to the coalesce() method, except this allows to set a
  //  * headroom and tailroom after coalescing.
  //  *
  //  * Returns ByteRange that points to the data MemBuf stores.
  //  */
  // 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<MemBuf> clone() const;
 
  MemBuf cloneAsValue() const;
 
  std::unique_ptr<MemBuf> cloneOne() const;
 
  MemBuf cloneOneAsValue() const;
 
  std::unique_ptr<MemBuf> cloneCoalesced() const;
 
  std::unique_ptr<MemBuf> cloneCoalescedWithHeadroomTailroom(
      std::size_t newHeadroom, std::size_t newTailroom) const;
 
  MemBuf cloneCoalescedAsValue() const;
 
  MemBuf cloneCoalescedAsValueWithHeadroomTailroom(
      std::size_t newHeadroom, std::size_t newTailroom) const;
 
  void cloneInto(MemBuf& other) const { other = cloneAsValue(); }
 
  void cloneOneInto(MemBuf& other) const { other = cloneOneAsValue(); }
 
  std::vector<struct iovec> getIov() const;
 
  void appendToIov(std::vector<struct iovec>* iov) const;
 
  size_t fillIov(struct iovec* iov, size_t len) const;
 
  static std::unique_ptr<MemBuf> wrapIov(const iovec* vec, size_t count);
 
  static std::unique_ptr<MemBuf> takeOwnershipIov(const iovec* vec,
                                                  size_t count,
                                                  FreeFunction freeFn = nullptr,
                                                  void* userData = nullptr,
                                                  bool freeOnError = true);
 
  bool ensureCapacity(std::size_t capacity);
 
  bool ensureCapacityAndFillUnused(std::size_t capacity, uint8_t placeholder);
 
  /*
   * Overridden operator new and delete.
   * These perform specialized memory management to help support
   * createCombined(), which allocates MemBuf 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);
 
  bool operator==(const MemBuf& other);
  bool operator!=(const MemBuf& other);
 
  // /**
  //  * Iteration support: a chain of MemBufs may be iterated through using
  //  * STL-style iterators over const ByteRanges.  Iterators are only
  //  invalidated
  //  * if the MemBuf that they currently point to is removed.
  //  */
  // Iterator cbegin() const;
  // Iterator cend() const;
  // Iterator begin() const;
  // Iterator end() const;
 
  MemBuf() noexcept;
 
  MemBuf(MemBuf&& other) noexcept;
  MemBuf& operator=(MemBuf&& other) noexcept;
 
  MemBuf(const MemBuf& other);
  MemBuf& operator=(const MemBuf& 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.
    flag_free_shared_info = 0x1,
    flag_maybe_shared = 0x2,
    flag_mask = flag_free_shared_info | flag_maybe_shared
  };
 
  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
  MemBuf(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, MemBuf* end,
                             size_t newTailroom);
  void coalesceAndReallocate(size_t newLength, MemBuf* 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 MemBuf 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 MemBuf is the
   * only element in the chain, next_ and prev_ will both point to this.
   */
  MemBuf* next_{this};
  MemBuf* prev_{this};
 
  /*
   * A pointer to the start of the data referenced by this MemBuf, 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 flags_and_shared_info_{0};
 
  static inline uintptr_t packFlagsAndSharedInfo(uintptr_t flags,
                                                 SharedInfo* info) {
    uintptr_t uinfo = reinterpret_cast<uintptr_t>(info);
    return flags | uinfo;
  }
 
  inline SharedInfo* sharedInfo() const {
    return reinterpret_cast<SharedInfo*>(flags_and_shared_info_ & ~flag_mask);
  }
 
  inline void setSharedInfo(SharedInfo* info) {
    uintptr_t uinfo = reinterpret_cast<uintptr_t>(info);
    flags_and_shared_info_ = (flags_and_shared_info_ & flag_mask) | uinfo;
  }
 
  inline uintptr_t flags() const { return flags_and_shared_info_ & flag_mask; }
 
  // flags_ are changed from const methods
  inline void setFlags(uintptr_t flags) const {
    flags_and_shared_info_ |= flags;
  }
 
  inline void clearFlags(uintptr_t flags) const {
    flags_and_shared_info_ &= ~flags;
  }
 
  inline void setFlagsAndSharedInfo(uintptr_t flags, SharedInfo* info) {
    flags_and_shared_info_ = 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);
  }
};
 
// template <class UniquePtr>
// typename std::enable_if<
//     detail::IsUniquePtrToSL<UniquePtr>::value,
//     std::unique_ptr<MemBuf>>::type
// MemBuf::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, &MemBuf::freeUniquePtrBuffer, deleter);
// }
 
inline std::unique_ptr<MemBuf> MemBuf::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<MemBuf> buf = MemBuf::create(capacity);
  buf->advance(headroom);
  if (size != 0) {
    memcpy(buf->writableData(), data, size);
  }
  buf->append(size);
  return buf;
}
 
}  // namespace utils