Frostbite2D/Frostbite2D/include/frostbite2D/core/task_system.h

#pragma once

#include <atomic>
#include <condition_variable>
#include <exception>
#include <future>
#include <memory>
#include <mutex>
#include <queue>
#include <thread>
#include <type_traits>
#include <utility>
#include <vector>

#include <frostbite2D/types/type_alias.h>

namespace frostbite2D {

struct TaskSystemConfig {
  uint32 workerCount = 0;
  uint32 maxMainThreadCallbacksPerFrame = 0;
  uint32 threadStackSize = 0;
};

class TaskSystem {
public:
  static TaskSystem& get();

  TaskSystem(const TaskSystem&) = delete;
  TaskSystem& operator=(const TaskSystem&) = delete;

  bool init(const TaskSystemConfig& config = {});
  void shutdown();

  bool isInitialized() const { return initialized_; }
  bool isMainThread() const;
  size_t workerCount() const;
  uint32 maxMainThreadCallbacksPerFrame() const;

  template <typename F>
  void post(F&& task) {
    enqueueWorkerTask(Function<void()>(std::forward<F>(task)));
  }

  template <typename F>
  auto submit(F&& task) -> std::future<std::invoke_result_t<F>> {
    using Result = std::invoke_result_t<F>;

    auto packagedTask =
        std::make_shared<std::packaged_task<Result()>>(std::forward<F>(task));
    std::future<Result> future = packagedTask->get_future();

    enqueueWorkerTask([packagedTask]() mutable { (*packagedTask)(); });
    return future;
  }

  template <typename F>
  void postToMainThread(F&& callback) {
    enqueueMainThreadTask(Function<void()>(std::forward<F>(callback)));
  }

  template <typename Task, typename OnComplete>
  void submitThen(Task&& task, OnComplete&& onComplete) {
    using Result = std::invoke_result_t<Task>;
    using Storage = std::decay_t<Result>;

    auto taskFn = std::make_shared<std::decay_t<Task>>(std::forward<Task>(task));
    auto completeFn =
        std::make_shared<std::decay_t<OnComplete>>(std::forward<OnComplete>(onComplete));

    enqueueWorkerTask([this, taskFn, completeFn]() mutable {
      try {
        if constexpr (std::is_void_v<Result>) {
          std::invoke(*taskFn);
          enqueueMainThreadTask([completeFn]() mutable { std::invoke(*completeFn); });
        } else {
          auto result = std::make_shared<Storage>(std::invoke(*taskFn));
          enqueueMainThreadTask([completeFn, result]() mutable {
            std::invoke(*completeFn, std::move(*result));
          });
        }
      } catch (...) {
        logUnhandledTaskException(std::current_exception());
      }
    });
  }

  template <typename Task, typename OnComplete, typename OnError>
  void submitThen(Task&& task, OnComplete&& onComplete, OnError&& onError) {
    using Result = std::invoke_result_t<Task>;
    using Storage = std::decay_t<Result>;

    auto taskFn = std::make_shared<std::decay_t<Task>>(std::forward<Task>(task));
    auto completeFn =
        std::make_shared<std::decay_t<OnComplete>>(std::forward<OnComplete>(onComplete));
    auto errorFn =
        std::make_shared<std::decay_t<OnError>>(std::forward<OnError>(onError));

    enqueueWorkerTask([this, taskFn, completeFn, errorFn]() mutable {
      try {
        if constexpr (std::is_void_v<Result>) {
          std::invoke(*taskFn);
          enqueueMainThreadTask([completeFn]() mutable { std::invoke(*completeFn); });
        } else {
          auto result = std::make_shared<Storage>(std::invoke(*taskFn));
          enqueueMainThreadTask([completeFn, result]() mutable {
            std::invoke(*completeFn, std::move(*result));
          });
        }
      } catch (...) {
        auto error = std::make_shared<std::exception_ptr>(std::current_exception());
        enqueueMainThreadTask([errorFn, error]() mutable { std::invoke(*errorFn, *error); });
      }
    });
  }

  void drainMainThreadTasks(uint32 maxTasks = 0);

private:
  TaskSystem() = default;
  ~TaskSystem();

  void enqueueWorkerTask(Function<void()> task);
  void enqueueMainThreadTask(Function<void()> task);
  size_t resolveWorkerCount(uint32 requestedCount) const;
  static void logUnhandledTaskException(std::exception_ptr error);

  TaskSystemConfig config_;
  std::vector<std::thread> workers_;
  std::queue<Function<void()>> workerTasks_;
  std::queue<Function<void()>> mainThreadTasks_;

  mutable std::mutex workerMutex_;
  mutable std::mutex mainThreadMutex_;
  std::condition_variable workerCondition_;

  std::atomic<bool> initialized_ = false;
  std::atomic<bool> acceptingTasks_ = false;
  std::thread::id mainThreadId_;
};

} // namespace frostbite2D