allocator.h

Go to the documentation of this file.

#pragma once
 
#include <absl/container/btree_set.h>
 
#include <cstdint>
#include <limits>
#include <type_traits>
#include <vector>
 
#include "src/atomics.h"
#include "src/globals.h"
#include "src/macros.h"
 
namespace broom {
// MemoryRegion is used to encapsulate a pair of {start, size} with helper
// methods. MemoryRegions should never own the underlying pointer and change its
// state.
class MemoryRegion {
 private:
  static bool SetCompareFn(const MemoryRegion& s1, const MemoryRegion& s2) {
    return s1.Size() < s2.Size();
  }
 
 public:
  MemoryRegion(const void* start, size_t size)
      : start_(SafeCast<UintPtr>(start)), size_(size) {}
  MemoryRegion(UintPtr start, UintPtr end) : start_(start), size_(end - start) {
    BASSERT_GE(end, start, "end must be greater than start: {:#x} >= {:#x}",
               start, end);
  }
  MemoryRegion(const void* start, const void* end)
      : MemoryRegion(SafeCast<UintPtr>(start), SafeCast<UintPtr>(end)) {}
  ~MemoryRegion() = default;
 
  size_t Size() const { return size_; }
  UintPtr Start() const { return start_; }
  UintPtr End() const { return start_ + size_; }
  void* StartPtr() const { return SafeCast<void*>(start_); }
  void* EndPtr() const { return SafeCast<void*>(start_ + size_); }
  bool AttemptGrow(size_t howmuch) {
    return !AddOverflowChecked(size_, howmuch, size_);
  }
  void IncrementStart(size_t how_much) {
    size_ -= how_much;
    start_ += how_much;
    BENSURE_GT(size_, 0, "Size must not be 0");
  }
 
  bool operator==(const MemoryRegion& other) const {
    return start_ == other.start_ && size_ == other.size_;
  }
  bool operator!=(const MemoryRegion& other) const { return !(*this == other); }
 
  using SetCompare =
      std::integral_constant<decltype(&SetCompareFn), &SetCompareFn>;
 
 protected:
  UintPtr start_;
  size_t size_;
};
 
static_assert(sizeof(MemoryRegion) == kPointerSize + sizeof(size_t));
static_assert(std::is_trivially_destructible_v<MemoryRegion>);
 
// MemorySpace is meant as a light-weight abstraction for an allocated space and
// a top pointer to keep track how much of the space was allocated. More complex
// allocation strategies should be implemented in a class derived from
// Allocator, not MemorySpace.
//
// MemorySpace is also meant to be disjoint between threads, and no two threads
// should have overlapping MemorySpaces.
class MemorySpace : public MemoryRegion {
 public:
  MemorySpace(const void* start, size_t size)
      : MemoryRegion(start, size), top_(SafeCast<UintPtr>(start)) {}
  ~MemorySpace() = default;
 
  UintPtr Top() const { return top_; }
  // Simple bump allocation in the memory space. Memory spaces don't support
  // freeing and the top_ pointer only ever grows rather than shrinks.
  UintPtr Allocate(size_t size) {
    // Sanity check.
    BASSERT(IsAligned<kUintPtrSize>(size), "size ({}) must be pointer aligned",
            size);
    BASSERT_GE(top_, Start(),
               "top address must be greater than or equal to the start "
               "address: {:#x} < {:#x}",
               top_, start_);
    if (top_ >= Start() && (top_ + size) <= End()) {
      UintPtr old_top = top_;
      top_ += size;
      return old_top;
    }
    return kNullUintPtr;
  }
 
  // Sets the top_ to the end of the space.
  void Fill() { top_ = End(); }
  size_t RemainingBytes() const { return End() - top_; }
 
  bool IsValid() const { return top_ != kNullUintPtr; }
 
 private:
  UintPtr top_;
};
 
static_assert(sizeof(MemorySpace) == sizeof(MemoryRegion) + kPointerSize);
static_assert(std::is_trivially_destructible_v<MemorySpace>);
 
// TODO(gc): Currently not exposed to consumers in any meaningful way.
enum class MemoryPermissions {
  kReadOnly,
  kWriteOnly,
  kReadWrite,
  kNone,
};
 
class AllocatorCommonMetadata {
 public:
  // None of these are used in a concurrent setting, and therefore are safe to
  // access with non-atomic operations.
  bool IsCurrentlyAllocated() const { return chunk_bits_ & kAllocatedBit; }
  bool AlreadyScanned() const { return chunk_bits_ & kMarkBit; }
  size_t Size() const { return chunk_bits_ >> kSizeShift; }
  bool IsValid() const { return chunk_bits_ != 0; }
  void ZeroBytes() { chunk_bits_ = 0; }
 
  void SetSize(size_t size) {
    BENSURE_LE(size, kMaxSize, "{} is too large of an allocation size.", size);
    chunk_bits_ = size << kSizeShift;
  }
  void MarkAllocated(size_t size) {
    BENSURE_LE(size, kMaxSize, "{} is too large of an allocation size.", size);
    chunk_bits_ = kAllocatedBit | size << kSizeShift;
  }
  void MarkDisposed() { chunk_bits_ &= ~kAllocatedBit; }
  void MarkScanned() { chunk_bits_ |= kMarkBit; }
  void UnmarkScanned() { chunk_bits_ &= ~kMarkBit; }
  bool AttemptGrow(uint64_t howmuch) {
    // Make sure that the top bits are clear.
    BENSURE_EQ((howmuch << kSizeShift) >> kSizeShift, howmuch,
               "Expected top bits of {} to be clear", howmuch);
    return !AddOverflowChecked(chunk_bits_, howmuch << kSizeShift, chunk_bits_);
  }
 
  static constexpr size_t MaxSize() { return kMaxSize; }
 
 private:
  size_t SizeWithoutAllocatedBit() const {
    return chunk_bits_ & ~kAllocatedBit;
  }
 
  static constexpr const size_t kAllocatedBit = 1;
  static constexpr const size_t kMarkBit = 2;
  static constexpr const size_t kSizeShift = 5;
  // 5 bits free, largest allocation is 0.5 EiB.
  static constexpr const size_t kMaxSize = 0x7FF'FFFF'FFFF'FFFFull;
  // The format of the metadata is as follows, 0-indexed:
  //   - Bit 0:       Marks if the chunk is allocated or not.
  //   - Bit 1:       Used during GC marking.
  //   - Bits 63 - 5: Store the size of the allocation.
  uint64_t chunk_bits_;
};
 
// We don't need to keep track of pin information at the back.
using AllocatorBackMetadata = AllocatorCommonMetadata;
 
// XXX(gc): Consider inheritance instead.
class AllocatorFrontMetadata {
 public:
  // Forward to common metadata.
  bool IsCurrentlyAllocated() const { return common_.IsCurrentlyAllocated(); }
  bool AlreadyScanned() const { return common_.AlreadyScanned(); }
  size_t Size() const { return common_.Size(); }
  bool IsValid() const { return common_.IsValid(); }
  void ZeroBytes() {
    common_.ZeroBytes();
    pin_count_ = 0;
  }
 
  void SetSize(size_t size) { common_.SetSize(size); }
  void MarkAllocated(size_t size) {
    common_.MarkAllocated(size);
    pin_count_ = 0;
  }
  void MarkDisposed() {
    common_.MarkDisposed();
    BASSERT_EQ(pin_count_, 0, "Trying to mark pinned allocation as disposed");
  }
  void MarkScanned() { common_.MarkScanned(); }
  void UnmarkScanned() { common_.UnmarkScanned(); }
  bool AttemptGrow(size_t howmuch) { return common_.AttemptGrow(howmuch); }
 
  // Front metadata specific methods. These are used to implement sharing across
  // threads, and therefore need to be operated on using atomics. Relaxed memory
  // semantics are sufficient here, since we don't care about the ordering.
  void Pin(uint32_t how_many) {
    BROOM_UNUSED uint32_t previous =
        Atomic::RelaxedFetchAdd(&pin_count_, uint32_t{how_many});
    BASSERT_LT(previous, std::numeric_limits<uint32_t>::max(),
               "Overflow of pin_count_");
  }
  void Unpin() {
    BROOM_UNUSED uint32_t previous =
        Atomic::RelaxedFetchSub(&pin_count_, uint32_t{1});
    BASSERT_GT(previous, 0, "Expected pin_count_ > 0, {}", previous);
  }
  bool IsPinned() { return Atomic::RelaxedLoad(&pin_count_) > 0; }
 
 private:
  // The front metadata must mirror the metadata at the back of an allocation.
  AllocatorCommonMetadata common_;
  // Atomically counts how many times an allocation was pinned. This is used
  // during sharing in a multi-threaded environment.
  uint32_t pin_count_;
  // Reserved bits, with the dual purpose of alignment.
  BROOM_UNUSED uint32_t reserved_;
};
 
static_assert(alignof(AllocatorFrontMetadata) == kPointerSize);
static_assert(alignof(AllocatorBackMetadata) == kPointerSize);
static_assert(std::is_trivially_destructible_v<AllocatorBackMetadata>);
static_assert(std::is_trivially_destructible_v<AllocatorFrontMetadata>);
 
enum class MemoryPressure {
  kLight,
  kMedium,
  kCritical,
};
 
// The Broom allocator interface. The allocator state is assumed to be
// thread_local. This allocator interface is private to Broom and shouldn't be
// exposed outside of this library. However, in order for the garbage collector
// to work, all of this functionality must be provided.
class Allocator {
 public:
  // XXX(gc): Decouple the metadata format from the allocator interface?
  static constexpr const int kFrontMetadataSize =
      sizeof(AllocatorFrontMetadata);
  static constexpr const int kBackMetadataSize = sizeof(AllocatorBackMetadata);
  static constexpr const int kTotalMetadataSize =
      kFrontMetadataSize + kBackMetadataSize;
  // XXX(gc): Decouple space size from the allocator interface?
  static constexpr const size_t kDefaultSpaceSize = 2 * MB;
  static constexpr const size_t kDefaultLargeSpaceSize = 64 * MB;
  static constexpr const size_t kDefaultLargeSpaceThreshold = 1 * MB;
  static constexpr const MemoryPermissions kDefaultPermissions =
      MemoryPermissions::kReadWrite;
  static constexpr const size_t kGiantMappingSize = kDefaultLargeSpaceSize;
 
  virtual ~Allocator() = default;
 
  virtual bool Grow(size_t requested_size, MemorySpace** space_to_update) = 0;
  virtual size_t LargeSpaceThreshold() const = 0;
 
  virtual void* Allocate(size_t requested_size) = 0;
  virtual void Dispose(void* definitely_uintptr) = 0;
  virtual const void* GetBasePointerOfMaybeInnerPointer(
      const void* maybe_inner) const = 0;
 
  MemoryPressure CalculateMemoryPressure() const {
    if (stats_.total_mapped_bytes == 0) return MemoryPressure::kLight;
    double alloc_to_mapped_ratio =
        static_cast<double>(stats_.allocated_bytes) /
        static_cast<double>(stats_.total_mapped_bytes);
    if (100 * alloc_to_mapped_ratio >= 75.) return MemoryPressure::kCritical;
    return MemoryPressure::kLight;
  }
 
 protected:
  struct AllocatorStatistics {
    size_t total_mapped_bytes = 0;
    size_t allocated_bytes = 0;
    size_t free_bytes = 0;
  };
  Allocator() = default;
 
  // XXX(gc): Provide different implementations for different allocators? This
  // is especially relevant if the metadata is decoupled from the allocator
  // interface.
  static void PinBasePointer(const void* pointer, uint32_t how_many) {
    AllocatorFrontMetadata* metadata = UnsafeCast<AllocatorFrontMetadata*>(
        SafeCast<UintPtr>(pointer) - Allocator::kFrontMetadataSize);
    metadata->Pin(how_many);
  }
  static void UnpinBasePointer(const void* pointer) {
    AllocatorFrontMetadata* metadata = UnsafeCast<AllocatorFrontMetadata*>(
        SafeCast<UintPtr>(pointer) - Allocator::kFrontMetadataSize);
    metadata->Unpin();
  }
  static bool IsPinnedBasePointer(const void* pointer) {
    AllocatorFrontMetadata* metadata = UnsafeCast<AllocatorFrontMetadata*>(
        SafeCast<UintPtr>(pointer) - Allocator::kFrontMetadataSize);
    return metadata->IsPinned();
  }
  static size_t SpaceMask(size_t space_size) {
    BASSERT(IsAligned<2>(space_size), "space_size ({}) must be a power of 2",
            space_size);
    return ~(space_size - 1);
  }
 
  AllocatorStatistics stats_;
  friend class GarbageCollector;
};
 
// The FreeListAllocator class implements an allocator which holds free lists of
// free'd allocations in the form of a B-tree ordered by the size of the free'd
// allocation.
class FreeListAllocator : public Allocator {
 public:
  explicit FreeListAllocator(
      size_t space_size = kDefaultSpaceSize,
      size_t large_space_size = kDefaultLargeSpaceSize,
      size_t large_space_threshold = kDefaultLargeSpaceThreshold)
      : space_size_(space_size),
        large_space_size_(large_space_size),
        large_space_threshold_(large_space_threshold) {
    BENSURE_LT(large_space_threshold, space_size,
               "large_space_threshold must be less than space_size: {} >= {}",
               large_space_threshold, space_size);
  }
  virtual ~FreeListAllocator();
 
  virtual bool Grow(size_t requested_size,
                    MemorySpace** space_to_update) override;
  virtual void* Allocate(size_t requested_size) override;
  virtual void Dispose(void* definitely_uintptr) override;
  virtual size_t LargeSpaceThreshold() const override {
    return kLargeSpaceThreshold;
  }
  virtual const void* GetBasePointerOfMaybeInnerPointer(
      const void* maybe_inner) const override;
 
 protected:
  // Ensure out B-trees are compared by size.
  using FreeList = absl::btree_multiset<MemoryRegion, MemoryRegion::SetCompare>;
#ifdef BROOM_DEBUG_SLOW
  // Debugging helpers. Both of these traverse the trees and should be used
  // sparingly in strategic locations.
  bool VerifyLists() const;
  void PrintList(std::string prefix, const FreeList& free_list) const;
#endif
  // Find the exact MemoryRegion we are looking for. Simply calling find is not
  // sufficient as it will find the first matching size. Instead, we need to
  // keep iterating through the matching sizes until we find the correct region.
  FreeList::iterator FindExactRegion(FreeList& free_list,
                                     const MemoryRegion region) {
    FreeList::iterator it = free_list.find(region);
    while (it != free_list.end() && region != *it) {
      ++it;
    }
    return it;
  }
  // Allocation/disposal helpers.
  void* SplitFreeListEntry(FreeList& free_list, FreeList::iterator& pos,
                           size_t requested_size);
  void* TryAllocateFromFreeListImpl(FreeList& free_list, size_t requested_size);
  void* TryAllocateFromLargeFreeList(size_t requested_size);
  void* TryAllocateFromFreeList(size_t requested_size);
  void* TryAllocateExactOrFindNext(FreeList& free_list, FreeList::iterator& pos,
                                   size_t requested_size);
  void AddOrCoalesceLargeFreeList(const MemoryRegion current,
                                  const MemoryRegion previous,
                                  const MemoryRegion next,
                                  const MemoryRegion coalesced);
  void AddOrCoalesceFreeList(const MemoryRegion current,
                             const MemoryRegion previous,
                             const MemoryRegion next,
                             const MemoryRegion coalesced);
  void AddOrCoalesceFreeListImpl(FreeList& free_list,
                                 const MemoryRegion current,
                                 const MemoryRegion previous,
                                 const MemoryRegion next,
                                 const MemoryRegion coalesced);
  void FillSpaceAndAddToFreeList(FreeList* list, MemorySpace* space);
 
#ifdef BROOM_FOR_TESTING
  // These shouldn't be used in Broom and only exist for testing.
  const FreeList GetFreeList() { return free_list_; }
  const FreeList GetLargeFreeList() { return large_free_list_; }
#endif
 
  // XXX(gc): Experiment with this.
  static constexpr const size_t kLargeSpaceThreshold = 1 * MB;
  static constexpr size_t MaximumAllocationSize() {
    return AllocatorCommonMetadata::MaxSize();
  }
  // Returns the requested allocation size aligned to pointer size, along with
  // the added size needed to store the allocator metdata.
  static constexpr size_t GetRealAllocationSize(size_t requested_size) {
    return RoundUp<kPointerSize>(requested_size) +
           Allocator::kFrontMetadataSize + Allocator::kBackMetadataSize;
  }
  void* AllocateRawPointer(size_t requested_size);
  void DisposeRawPointer(void* definitely_pointer);
 
  // XXX(gc): Split into multiple spaces?
  std::vector<MemorySpace> spaces_;
  std::vector<MemorySpace> large_spaces_;
  size_t space_size_;
  size_t large_space_size_;
  size_t large_space_threshold_;
  MemorySpace* current_space_ = nullptr;
  MemorySpace* current_large_space_ = nullptr;
  FreeList free_list_;
  FreeList large_free_list_;
 
#ifdef BROOM_FOR_TESTING
  // Needed to expose internal information of the allocator for testing
  // purposes.
  friend class TestFreeListAllocator;
#endif
};
}  // namespace broom