AThreadPool.h

/*
 * AUI Framework - Declarative UI toolkit for modern C++20
 * Copyright (C) 2020-2024 Alex2772 and Contributors
 *
 * SPDX-License-Identifier: MPL-2.0
 *
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/.
 */
 
#pragma once
 
#include <AUI/Core.h>
#include <cassert>
#include <atomic>
 
#include <AUI/Common/AVector.h>
#include <AUI/Common/AQueue.h>
#include <AUI/Common/AException.h>
#include <AUI/Thread/AThread.h>
#include <glm/glm.hpp>
#include <utility>
#include "AUI/Traits/concepts.h"
 
template<typename T>
class AFuture;
 
class API_AUI_CORE AThreadPool {
   public:
    class API_AUI_CORE Worker : public AThread {
       private:
        bool mEnabled = true;
        bool processQueue(std::unique_lock<std::mutex>& mutex, AQueue<std::function<void()>>& queue);
        AThreadPool& mTP;
 
        void iteration(std::unique_lock<std::mutex>& tpLock);
        void wait(std::unique_lock<std::mutex>& tpLock);
 
       public:
        Worker(AThreadPool& tp, size_t index);
        ~Worker();
        void aboutToDelete();
 
        template<aui::predicate ShouldContinue>
        void loop(ShouldContinue&& shouldContinue) {
            std::unique_lock lock(mTP.mQueueLock);
            while (shouldContinue()) {
                iteration(lock);
                if (!shouldContinue()) {
                    return;
                }
                wait(lock);
            }
        }
 
        AThreadPool& threadPool() noexcept { return mTP; }
    };
 
    enum Priority {
        PRIORITY_HIGHEST,
        PRIORITY_MEDIUM,
        PRIORITY_LOWEST,
    };
 
   protected:
    typedef std::function<void()> task;
    AVector<_<Worker>> mWorkers;
    AQueue<task> mQueueHighest;
    AQueue<task> mQueueMedium;
    AQueue<task> mQueueLowest;
    AQueue<task> mQueueTryLater;
    std::mutex mQueueLock;
    std::condition_variable mCV;
    size_t mIdleWorkers = 0;
 
   public:
    AThreadPool(size_t size);
 
    AThreadPool();
    ~AThreadPool();
    size_t getPendingTaskCount();
    void run(const std::function<void()>& fun, Priority priority = PRIORITY_MEDIUM);
    void clear();
    void runLaterTasks();
    static void enqueue(const std::function<void()>& fun, Priority priority = PRIORITY_MEDIUM);
 
    void setWorkersCount(std::size_t workersCount);
 
    void wakeUpAll() {
        std::unique_lock lck(mQueueLock);
        mCV.notify_all();
    }
 
    static AThreadPool& global();
 
    size_t getTotalWorkerCount() const { return mWorkers.size(); }
    size_t getIdleWorkerCount() const { return mIdleWorkers; }
 
    template <typename Iterator, typename Functor>
    auto parallel(Iterator begin, Iterator end, Functor&& functor);
 
    template <aui::invocable Callable>
    [[nodiscard]] inline auto operator*(Callable fun) {
        using Value = std::invoke_result_t<Callable>;
        AFuture<Value> future(std::move(fun));
        run(
            [innerWeak = future.inner().weak()]() {
                /*
                 * Avoid holding a strong reference - we need to keep future cancellation on reference count exceeding
                 * even while actual future execution.
                 */
                if (auto lock = innerWeak.lock()) {
                    auto innerUnsafePointer = lock->ptr().get();   // using .get() here in order to bypass
                                                                   // null check in operator->
 
                    lock = nullptr;   // destroy strong ref
 
                    innerUnsafePointer->tryExecute(
                        innerWeak);   // there's a check inside tryExecute to check its validity
                }
            },
            AThreadPool::PRIORITY_LOWEST);
        return future;
    }
 
    class TryLaterException {};
};
 
#include <AUI/Thread/AFuture.h>
 
template <typename T = void>
class AFutureSet : public AVector<AFuture<T>> {
   private:
    using super = AVector<AFuture<T>>;
 
   public:
    using AVector<AFuture<T>>::AVector;
 
    void waitForAll() {
        // wait from the end to avoid idling (see AFuture::wait for details)
        for (const AFuture<T>& v : aui::reverse_iterator_wrap(*this)) {
            v.operator*();
        }
    }
 
    void checkForExceptions() const {
        for (const AFuture<T>& v : *this) {
            if (v.hasResult()) {
                v.operator*();   // TODO bad design
            }
        }
    }
 
    template <aui::invocable OnComplete>
    void onAllComplete(OnComplete&& onComplete) {
        // check if all futures is already complete.
        for (const AFuture<T>& v : *this) {
            if (!v.hasResult()) {
                goto setupTheHell;
            }
        }
        onComplete();
        return;
 
    setupTheHell:
        struct Temporary {
            OnComplete onComplete;
            AFutureSet myCopy;
            std::atomic_bool canBeCalled = true;
        };
        auto temporary = _new<Temporary>(std::forward<OnComplete>(onComplete), *this);
        for (const AFuture<T>& v : *this) {
            v.onSuccess([temporary](const auto& v) {
                for (const AFuture<T>& v : temporary->myCopy) {
                    if (!v.hasResult()) {
                        return;
                    }
                }
                // yay! all tasks are completed. the last thing to check if the callback is already called
                if (temporary->canBeCalled.exchange(false)) {
                    temporary->onComplete();
                }
            });
        }
    }
};
 
template <typename Iterator, typename Functor>
auto AThreadPool::parallel(Iterator begin, Iterator end, Functor&& functor) {
    using ResultType = decltype(std::declval<Functor>()(std::declval<Iterator>(), std::declval<Iterator>()));
    AFutureSet<ResultType> futureSet;
 
    size_t itemCount = end - begin;
    size_t affinity = (glm::min)(AThreadPool::global().getTotalWorkerCount(), itemCount);
    if (affinity == 0)
        return futureSet;
    size_t itemsPerThread = itemCount / affinity;
 
    for (size_t threadIndex = 0; threadIndex < affinity; ++threadIndex) {
        auto forThreadBegin = begin;
        begin += itemsPerThread;
        auto forThreadEnd = threadIndex + 1 == affinity ? end : begin;
        futureSet.push_back(*this *
                                [functor = std::forward<Functor>(functor), forThreadBegin,
                                 forThreadEnd]() -> decltype(auto) { return functor(forThreadBegin, forThreadEnd); });
    }
 
    return futureSet;
}
 
#include <AUI/Reflect/AReflect.h>