scheduling_tests.cpp

#include <catch.hpp>

#include <pls/pls.h>

using namespace pls;

class once_sub_task : public task {
  std::atomic<int> *counter_;
  int children_;

 protected:
  void execute_internal() override {
    (*counter_)++;
    for (int i = 0; i < children_; i++) {
      spawn_child(once_sub_task(counter_, children_ - 1));
    }
  }

 public:
  explicit once_sub_task(std::atomic<int> *counter, int children) :
      task{},
      counter_{counter},
      children_{children} {}
};

class force_steal_sub_task : public task {
  std::atomic<int> *parent_counter_;
  std::atomic<int> *overall_counter_;

 protected:
  void execute_internal() override {
    (*overall_counter_)--;
    if (overall_counter_->load() > 0) {
      std::atomic<int> counter{1};
      spawn_child(force_steal_sub_task(&counter, overall_counter_));
      while (counter.load() > 0); // Spin...
    }

    (*parent_counter_)--;
  }

 public:
  explicit force_steal_sub_task(std::atomic<int> *parent_counter, std::atomic<int> *overall_counter) :
      task{},
      parent_counter_{parent_counter},
      overall_counter_{overall_counter} {}
};

TEST_CASE("tbb task are scheduled correctly", "[internal/scheduling/fork_join_task.h]") {
  malloc_scheduler_memory my_scheduler_memory{8, 2 << 12};

  SECTION("tasks are executed exactly once") {
    scheduler my_scheduler{&my_scheduler_memory, 2};
    int start_counter = 4;
    int total_tasks = 1 + 4 + 4 * 3 + 4 * 3 * 2 + 4 * 3 * 2 * 1;
    std::atomic<int> counter{0};

    my_scheduler.perform_work([&]() {
      once_sub_task sub_task{&counter, start_counter};
      scheduler::spawn_child(sub_task);
    });

    REQUIRE(counter.load() == total_tasks);
    my_scheduler.terminate(true);
  }

  SECTION("tasks can be stolen") {
    scheduler my_scheduler{&my_scheduler_memory, 8};
    my_scheduler.perform_work([&]() {
      std::atomic<int> dummy_parent{1}, overall_counter{8};
      force_steal_sub_task sub_task{&dummy_parent, &overall_counter};
      scheduler::spawn_child(sub_task);

      // Required, as child operates on our stack's memory!!!
      scheduler::wait_for_all();
    });
    my_scheduler.terminate(true);
  }
}