garbage-collector.h

Go to the documentation of this file.

#pragma once
 
#include <absl/container/btree_set.h>
 
#include <list>
#include <memory>
#include <type_traits>
 
#include "include/broom.h"
#include "src/allocator.h"
#include "src/globals.h"
#include "src/macros.h"
#include "src/roots.h"
 
namespace broom {
// A node representing each managed pointer in the GC. It comprises of the base
// pointer to a managed allocation and, optionally, the destructor to run upon
// its collection.
class GcNode {
 private:
  static bool SetCompareFn(const GcNode& n1, const GcNode& n2) {
    return SafeCast<MemoryAddress>(n1.Pointer()) <
           SafeCast<MemoryAddress>(n2.Pointer());
  }
 
 public:
  explicit GcNode(const void* pointer, Destructor destructor)
      : pointer_(pointer), destructor_(destructor) {}
  explicit GcNode(const void* pointer)
      : pointer_(pointer), destructor_(nullptr) {}
  const void* Pointer() const { return pointer_; }
  bool HasDestructor() const { return destructor_ != nullptr; }
  void RunDestructor() const { destructor_(pointer_); }
 
  using SetCompare =
      std::integral_constant<decltype(&SetCompareFn), &SetCompareFn>;
 
 private:
  const void* pointer_;
  Destructor destructor_;
  friend class GarbageCollector;
};
 
// The thread-local Broom garbage collector state.
class GarbageCollector {
 private:
  static constexpr auto cmp_helper = []<typename T>(const T& n1, const T& n2) {
    return SafeCast<MemoryAddress>(n1.Pointer()) <
           SafeCast<MemoryAddress>(n2.Pointer());
  };
 
 public:
  // TODO(gc): Consider a swiss table instead.
  using GcSet = absl::btree_set<GcNode, GcNode::SetCompare>;
  using ExternalRootSet =
      absl::btree_set<ExternalRoot, ExternalRoot::SetCompare>;
 
  explicit GarbageCollector(broom_configuration config, Allocator* allocator)
      : configuration_(config),
        allocator_(allocator),
        root_visitor_(RootVisitor(allocator, this)) {
    BENSURE_NE(allocator_, nullptr, "Failed to allocate the allocator");
  }
  explicit GarbageCollector() : GarbageCollector({}, new FreeListAllocator) {}
  explicit GarbageCollector(broom_configuration config)
      : GarbageCollector(config, new FreeListAllocator) {}
  explicit GarbageCollector(Allocator* allocator)
      : GarbageCollector({}, allocator) {}
  ~GarbageCollector() {
    BASSERT_EQ(external_roots_.size(), 0,
               "Must deallocate all external roots, have: {}",
               external_roots_.size());
    BASSERT_EQ(precise_roots_.size(), 0,
               "Must deallocate all precise roots, have {}",
               precise_roots_.size());
    BASSERT_EQ(to_collect_.size(), 0, "Exited during collection.");
 
    // If we don't need to run the destructors, we don't need to iterate over
    // the objects that are alive since deleting the allocator will unmap all
    // the mappings it owns.
    if (configuration_.run_destructors_on_shutdown) {
      for (auto it = alive_.begin(); it != alive_.end(); ++it) {
        if (it->HasDestructor()) {
          it->RunDestructor();
        }
      }
    }
    delete allocator_;
  }
 
  void ForceFastCollection() {
    AtomicMark();
    AtomicSweep();
  }
  void ForceSlowCollection() {
    AtomicMark();
    AtomicSweep();
  }
 
  void MarkAlive(const void* pointer) {
    auto it = std::ranges::find(to_collect_, pointer, &GcNode::pointer_);
    // If we can't find it, that means it must already be alive.
    if (it == to_collect_.end()) return;
    alive_.insert(*it);
    to_collect_.erase(it);
  }
  BROOM_INLINE void RegisterExternalRoot(ExternalRoot root) {
    external_roots_.emplace(std::move(root));
  }
  BROOM_INLINE void UnregisterExternalRoot(ExternalRoot root) {
    external_roots_.erase(std::move(root));
  }
  BROOM_INLINE void* RegisterPreciseRoot(PreciseRoot root) {
    PreciseRoot& root_ref = precise_roots_.emplace_front(root);
    auto it = precise_roots_.begin();
    root_ref.UpdateReference(it);
    return SafeCast<void*>(&*it);
  }
  BROOM_INLINE void UnregisterPreciseRoot(void* type_erased_root) {
    PreciseRoot* root = SafeCast<PreciseRoot*>(type_erased_root);
    precise_roots_.erase(root->GetReference());
  }
  void AtomicMark();
  void AtomicSweep();
 
  BROOM_INLINE void RegisterPointer(const void* pointer,
                                    Destructor destructor = nullptr) {
    alive_.emplace(GcNode{pointer, destructor});
  }
 
  BROOM_INLINE void Pin(const void* pointer, uint32_t how_many = 1) {
    PinBasePointer(pointer, how_many);
  }
  BROOM_INLINE void Unpin(const void* pointer) { UnpinBasePointer(pointer); }
 
  BROOM_INLINE void DisableAutomatedCollection() { collect_ = false; }
  BROOM_INLINE void EnableAutomatedCollection() { collect_ = true; }
 
  BROOM_INLINE void CollectIfNecessary() {
    MemoryPressure pressure = allocator_->CalculateMemoryPressure();
    // clang-format off
    // In the current version of Broom, this will never happen. We currently
    // don't distinguish between slow and fast collection.
    if (collect_ && pressure == MemoryPressure::kMedium) BROOM_UNLIKELY {
      ForceFastCollection();
    } else if (collect_ && pressure == MemoryPressure::kCritical) BROOM_LIKELY {
      ForceSlowCollection();
    }
    // clang-format on
  }
#ifdef BROOM_FOR_TESTING
  // See include/broom.h for the implementation. We currently only use this for
  // testing, however we might want to expose the GarbageCollector interface
  // publicly at some point.
  template <typename T, typename... Args>
  T* Allocate(size_t count, Args&&... args) {
    CollectIfNecessary();
    __BROOM_IMPL_ALLOCATE_T(T, RegisterPointer, allocator_->Allocate, args)
  }
#endif
 
  // Checks for memory pressure and calls into the allocator. This function does
  // **not** register the new allocation with the GC, it is intended to be used
  // by the API functions.
  void* AllocateRawDontRegister(size_t size) {
    CollectIfNecessary();
    return allocator_->Allocate(size);
  }
 
#ifdef BROOM_FOR_TESTING
  // Only used during testing.
  const GcSet& Alive() const { return alive_; }
#endif
 
 private:
  BROOM_INLINE void PinBasePointer(const void* pointer, uint32_t how_many) {
    Allocator::PinBasePointer(pointer, how_many);
  }
  BROOM_INLINE void UnpinBasePointer(const void* pointer) {
    Allocator::UnpinBasePointer(pointer);
  }
 
  BROOM_UNUSED broom_configuration configuration_;
  Allocator* allocator_;
  ExternalRootSet external_roots_;
  PreciseRoot::List precise_roots_;
  RootVisitor root_visitor_;
  GcSet to_collect_;
  GcSet alive_;
  bool collect_ = true;
};
}  // namespace broom