main.cc

Go to the documentation of this file.

#include <stdint.h>
#include <string.h>
 
#include <algorithm>
#include <condition_variable>
#include <cstdint>
#include <limits>
#include <mutex>
#include <numeric>
#include <print>
#include <queue>
#include <random>
#include <stack>
#include <string>
#include <thread>
 
#include "include/broom-vector.h"
#include "include/broom.h"
 
constexpr const int KB = 1024;
 
#if defined(__clang__) || defined(__GNUC__)
#define NOINLINE __attribute__((noinline))
#elif defined(_MSC_VER)
#define NOINLINE __declspec(noinline)
#endif
 
class Job;
 
// Global to make things simpler.
std::mutex job_queue_mtx;
std::mutex stdout_mtx;
std::queue<broom::sharing_ptr<Job>> job_queue;
std::condition_variable cv;
bool done = false;
std::atomic<bool> exit_producer = false;
 
class BigThing {
 public:
  explicit BigThing() : bytes_(nullptr), size_(0) {}
  // Simply allocates bytes_ as a managed region.
  explicit BigThing(size_t size)
      : bytes_(broom::allocate<uint8_t>(size, 0)), size_(size) {
    assert(bytes_ != nullptr);
  }
  // Don't need to manage bytes_, it's GC'd.
  ~BigThing() = default;
 
  void FillWithRandomBytes() {
    std::random_device rd;
    std::mt19937 mt(rd());
    std::uniform_int_distribution<uint16_t> dist(0, 255);
    for (size_t i = 0; i < size_; ++i)
      bytes_[i] = static_cast<uint8_t>(dist(mt));
  }
 
  BigThing(const BigThing& other) : bytes_(other.bytes_), size_(other.size_) {}
  BigThing& operator=(const BigThing& other) = delete;
  BigThing& operator=(BigThing&& other) = delete;
 
  uint8_t At(size_t index) const { return bytes_[index]; }
 
 private:
  uint8_t* bytes_;
  size_t size_;
};
 
enum class Kind {
  kPrint,
  kAddNumbers,
  kProcessBigThing,
};
 
class Job {
 public:
  // Pointers inside objects are fine.
  explicit Job(BigThing big_thing)
      : kind_(Kind::kProcessBigThing), big_thing_(big_thing) {}
  // Will be ignored since it's not managed memory.
  explicit Job(std::vector<int>&& numbers)
      : kind_(Kind::kAddNumbers), numbers_(numbers) {}
  // Regular pointers are fine too.
  explicit Job(const std::string string_to_print) : kind_(Kind::kPrint) {
    size_t s = string_to_print.size();
    string_to_print_ = broom::allocate<char>(s + 1);
    assert(string_to_print_ != nullptr);
    memcpy(string_to_print_, &string_to_print[0], s);
    string_to_print_[s] = '\0';
  }
  // We don't need to care about what happens to any of the inner pointers.
  ~Job() = default;
 
  BigThing GetBigThing() const {
    assert(kind_ == Kind::kProcessBigThing);
    return big_thing_;
  }
 
  const std::vector<int>& GetNumbers() const {
    assert(kind_ == Kind::kAddNumbers);
    return numbers_;
  }
 
  const char* GetStringToPrint() const {
    assert(kind_ == Kind::kPrint);
    return string_to_print_;
  }
 
  const Kind GetKind() const { return kind_; }
 
 private:
  Kind kind_;
  // They all exist for simplicity.
  BigThing big_thing_;
  std::vector<int> numbers_;
  char* string_to_print_;
};
 
void ProduceJobs() {
  // We need to share our pointers, so use a broom::sharing_scope.
  broom::sharing_scope broom_scope;
  // Needs to be in a broom::vector because anything inside a regular
  // std::vector is not visible to Broom.
  broom::vector<BigThing> big_things;
  // 10k iterations, 3 jobs per iteration.
  // Any call to broom::allocate will potentially cause a garbage collection
  // pass to trigger depending on the memory pressure.
  for (size_t i = 0; i < 10000; ++i) {
    {
      Job* new_job = broom::allocate<Job>(1, BigThing(1 * KB));
      // Simulate some being collected, while others are still live.
      if (i % 1000 == 0) {
        // Keep one every so often. All others should be GC'd.
        big_things.push_back(new_job->GetBigThing());
      }
      assert(new_job != nullptr);
      {
        std::unique_lock lock(job_queue_mtx);
        job_queue.push(broom::share<Job>(new_job));
      }
    }
    cv.notify_all();
 
    {
      Job* new_job =
          broom::allocate<Job>(1, "hello world from: " + std::to_string(i));
      assert(new_job != nullptr);
      {
        std::unique_lock lock(job_queue_mtx);
        job_queue.push(broom::share<Job>(new_job));
      }
    }
    cv.notify_all();
 
    {
      std::vector<int> v(i / 2 > 10 ? i / 2 : 10);
      std::iota(v.begin(), v.end(), 0);
      Job* new_job = broom::allocate<Job>(1, std::move(v));
      assert(new_job != nullptr);
      {
        std::unique_lock lock(job_queue_mtx);
        job_queue.push(broom::share<Job>(new_job));
      }
    }
    cv.notify_all();
    // Simulate sometimes going to sleep because there are no jobs to enqueue.
    if (i % 500 == 0) {
#ifdef FORCE_COLLECT
      // Since we will go to sleep, we can force a full sweep here.
      broom::force_slow_collection();
#endif
      std::this_thread::sleep_for(std::chrono::milliseconds(200));
    }
  }
 
  {
    std::unique_lock lock(job_queue_mtx);
    done = true;
    cv.notify_all();
  }
 
  while (!exit_producer.load(std::memory_order_acquire));
 
  // Should be alive throughout
  for (const auto& big_thing : big_things) {
    if (big_thing.At(1024) == 255) {
      std::unique_lock stdout_lock(stdout_mtx);
      std::println("Found 255 @ 1024!");
    }
  }
  // Broom will run destructors here.
}
 
void ProcessJob(Job* job) {
  switch (job->GetKind()) {
    case Kind::kPrint: {
      std::unique_lock stdout_lock(stdout_mtx);
      std::println("{}", job->GetStringToPrint());
    } break;
    case Kind::kAddNumbers: {
      const auto& numbers = job->GetNumbers();
      int64_t result =
          std::accumulate(std::begin(numbers), std::end(numbers), 0);
      std::unique_lock stdout_lock(stdout_mtx);
      std::println("result = {}", result);
    } break;
    case Kind::kProcessBigThing:
      job->GetBigThing().FillWithRandomBytes();
      std::unique_lock stdout_lock(stdout_mtx);
      std::println("Filled with random bytes");
      break;
  }
}
 
void Worker() {
  // No broom scope needed here. We aren't GCing anything on this thread.
  std::unique_lock<std::mutex> lock(job_queue_mtx);
  while (!(done && job_queue.empty())) {
    cv.wait(lock, [&]() { return done || !job_queue.empty(); });
    if (!job_queue.empty()) {
      // Anchor the lifetime of the job in the current thread inside this scope.
      // This won't cause the pointer to become garbage collected. Instead, it
      // will simply let the garbage collector know that this thread is done
      // using the pointer and it no longer needs to be pinned for this
      // particular use.
      broom::lifetime_anchor<Job> job = job_queue.front();
      job_queue.pop();
      // We can safely use this pointer up until the end of the scope.
      ProcessJob(job.to_unsafe_ptr());
    }
  }
}
 
int main() {
  std::thread job_producer(ProduceJobs);
  const auto core_count = std::thread::hardware_concurrency();
  std::vector<std::thread> workers(core_count / 2);
  for (size_t i = 0; i < core_count / 2; ++i) {
    workers[i] = std::thread(Worker);
  }
 
  for (size_t i = 0; i < core_count / 2; ++i) {
    workers[i].join();
  }
  exit_producer.store(true, std::memory_order_release);
  job_producer.join();
  return 0;
}