Scheduler.mpp

export module CppUtils.Thread.Scheduler;

import std;
import CppUtils.Chrono.Concept;
import CppUtils.Execution.ScopeGuard;
import CppUtils.Thread.UniqueLocker;
import CppUtils.Thread.SharedLocker;
import CppUtils.Thread.ThreadLoop;
import CppUtils.Thread.ThreadPool;

export namespace CppUtils::Thread
{
	class Scheduler final
	{
	public:
		using Clock = std::chrono::steady_clock;
		using TimePoint = Clock::time_point;

	private:
		using Task = std::function<void()>;

		struct Item final
		{
			std::size_t id;
			TimePoint time;
			Task task;
			std::atomic_bool cancelled = false;
			const void* owner = nullptr;
		};

		struct Compare final
		{
			inline auto operator()(const std::shared_ptr<Item>& a,
				const std::shared_ptr<Item>& b) const noexcept -> bool
			{
				if (a->time != b->time)
					return a->time < b->time;
				return a->id < b->id;
			}
		};

		struct Items final
		{
			using Set = std::set<std::shared_ptr<Item>, Compare>;
			using Map = std::unordered_map<std::size_t, std::shared_ptr<Item>>;

			Set set;
			Map map;
			std::vector<std::shared_ptr<Item>> activeItems;
			std::size_t processingTasks = 0;
		};

	public:
		inline Scheduler(
			Clock::duration step = Clock::duration::zero(),
			std::size_t numberThreads = std::thread::hardware_concurrency(),
			std::function<void(std::exception_ptr)> onError = nullptr,
			std::function<void()> finally = nullptr):
			m_step{step},
			m_finally{std::move(finally)},
			m_threadPool{
				numberThreads,
				onError,
				[this] {
					auto _ = Execution::ScopeGuard{[&] {
						auto accessor = m_items.access();
						if (--accessor.value().processingTasks == 0)
							m_finishedCondition.notify_all();
					}};
					if (m_finally)
						m_finally();
				}},
			m_threadLoop{
				[this] { runLoop(); },
				[this] {
					m_workCondition.notify_all();
					m_finishedCondition.notify_all();
				},
				onError}
		{
			m_threadLoop.start();
		}

		inline ~Scheduler() noexcept
		{
			m_threadLoop.requestStop();
			cancelAll();
		}

		inline auto schedule(Task task, TimePoint when, const void* owner = nullptr) -> std::size_t
		{
			auto item = std::make_shared<Item>(m_nextId.fetch_add(1, std::memory_order_relaxed), when, std::move(task), false, owner);
			auto accessor = m_items.access();
			auto& [set, map, activeItems, processingTasks] = accessor.value();
			set.insert(item);
			map.emplace(item->id, item);
			m_workCondition.notify_all();
			return item->id;
		}

		inline auto schedule(Task task, Chrono::Duration auto delay, const void* owner = nullptr) -> std::size_t
		{
			return schedule(std::move(task), Clock::now() + delay, owner);
		}

		inline auto waitUntilFinished() -> void
		{
			auto accessor = m_items.access();
			m_finishedCondition.wait(accessor.getLockGuard(), [&] {
				return std::ranges::empty(accessor.value().set) and accessor.value().processingTasks == 0;
			});
		}

		inline auto getScheduledTasksCount() const -> std::size_t
		{
			auto accessor = m_items.access();
			return std::ranges::size(accessor.value().set) + accessor.value().processingTasks;
		}

		inline auto flush() -> void
		{
			while (true)
			{
				auto tasksToRun = std::vector<std::shared_ptr<Item>>{};
				{
					auto accessor = m_items.access();
					auto& [set, map, activeItems, processingTasks] = accessor.value();

					if (std::ranges::empty(set))
					{
						if (processingTasks == 0)
							break;

						m_finishedCondition.wait(accessor.getLockGuard());
						continue;
					}

					tasksToRun.assign(std::ranges::begin(set), std::ranges::end(set));
					processingTasks += std::ranges::size(set);
					map.clear();
					set.clear();
				}

				dispatchTasks(std::move(tasksToRun));
			}

			m_workCondition.notify_all();
			m_finishedCondition.notify_all();
			waitUntilFinished();
		}

		inline auto cancel(std::size_t id) -> void
		{
			auto accessor = m_items.access();
			auto& [set, map, activeItems, processingTasks] = accessor.value();
			if (auto it = map.find(id); it != std::ranges::end(map))
			{
				it->second->cancelled = true;
				set.erase(it->second);
				map.erase(it);
				m_workCondition.notify_all();
				m_finishedCondition.notify_all();
			}
		}

		inline auto cancelOwner(const void* owner) -> void
		{
			if (owner == nullptr)
				return;

			auto accessor = m_items.access();
			auto& [set, map, activeItems, processingTasks] = accessor.value();

			for (auto it = std::ranges::begin(map); it != std::ranges::end(map);)
				if (it->second->owner == owner)
				{
					it->second->cancelled = true;
					set.erase(it->second);
					it = map.erase(it);
				}
				else
					++it;

			for (auto& item : activeItems)
				if (item->owner == owner)
					item->cancelled = true;

			m_finishedCondition.wait(accessor.getLockGuard(), [&] {
				return std::ranges::none_of(activeItems, [&](const auto& item) {
					return item->owner == owner;
				});
			});

			m_workCondition.notify_all();
			m_finishedCondition.notify_all();
		}

		inline auto cancelAll() -> void
		{
			auto accessor = m_items.access();
			auto& [set, map, activeItems, processingTasks] = accessor.value();
			for (auto& [_, item] : map)
				item->cancelled = true;
			for (auto& item : activeItems)
				item->cancelled = true;
			map.clear();
			set.clear();
			m_workCondition.notify_all();
			m_finishedCondition.notify_all();
		}

	private:
		inline auto runLoop() -> void
		{
			auto tasksToRun = std::vector<std::shared_ptr<Item>>{};
			{
				auto accessor = m_items.access();
				auto& [set, map, activeItems, processingTasks] = accessor.value();

				if (std::ranges::empty(set))
					m_workCondition.wait(accessor.getLockGuard(), [this, &set] {
						return not std::ranges::empty(set) or m_threadLoop.isStopRequested();
					});
				else if (auto first = *std::ranges::begin(set); first->time > Clock::now())
					m_workCondition.wait_until(accessor.getLockGuard(), first->time, [this, &set, first] {
						return m_threadLoop.isStopRequested() or std::ranges::empty(set) or (*std::ranges::begin(set) != first) or ((*std::ranges::begin(set))->time <= Clock::now());
					});

				if (m_threadLoop.isStopRequested() or std::ranges::empty(set) or (*std::ranges::begin(set))->time > Clock::now())
					return;

				auto first = *std::ranges::begin(set);
				const auto endTime = (m_step == Clock::duration::zero()) ? first->time : first->time + m_step;

				for (auto it = std::ranges::begin(set); it != std::ranges::end(set);)
					if ((*it)->time <= endTime)
					{
						tasksToRun.push_back(*it);
						map.erase((*it)->id);
						it = set.erase(it);
						++accessor.value().processingTasks;
					}
					else
						break;

				if (std::ranges::empty(set))
					m_finishedCondition.notify_all();
			}

			dispatchTasks(std::move(tasksToRun));
		}

	private:
		inline auto dispatchTasks(std::vector<std::shared_ptr<Item>> tasks) -> void
		{
			for (auto&& item : tasks)
			{
				{
					auto accessor = m_items.access();
					accessor.value().activeItems.push_back(item);
				}

				try
				{
					m_threadPool.call([this, item] {
						auto scope = Execution::ScopeGuard{[this, &item] {
							auto accessor = m_items.access();
							std::erase(accessor.value().activeItems, item);
							m_finishedCondition.notify_all();
						}};

						if (not item->cancelled)
							std::invoke(item->task);
					});
				}
				catch (...)
				{
					auto accessor = m_items.access();
					std::erase(accessor.value().activeItems, item);
					if (--accessor.value().processingTasks == 0)
						m_finishedCondition.notify_all();
				}
			}
		}

		std::atomic_size_t m_nextId = 0;
		Clock::duration m_step;
		UniqueLocker<Items> m_items;

		std::function<void()> m_finally;

		std::condition_variable m_workCondition;
		std::condition_variable m_finishedCondition;
		ThreadPool m_threadPool;
		ThreadLoop m_threadLoop;
	};
}