AuroraRuntime/Include/Aurora/Async/WorkPairImpl.hpp

/***
    Copyright (C) 2021 J Reece Wilson (a/k/a "Reece"). All rights reserved.

    File: WorkPairImpl.hpp
    Date: 2021-11-1
    Author: Reece
***/
 #pragma once

namespace Aurora::Threading::Primitives
{
    class SpinLoop;
}

namespace Aurora::Async
{
    struct BasicWorkStdFunc;

    /// @hideinitializer
    struct BasicWorkCtx : WorkPriv
    {
        BasicWorkCtx()
        {
            magic = AuConvertMagicTag32("BWOT");
            opt = nullptr;
        }
        void *opt;
    };
    
    /// @hideinitializer
    template<typename Info_t = AVoid, typename Result_t = AVoid, typename Task_t = FTask<Info_t, Result_t>, typename Job_t = FJob<Info_t, Result_t>>
    struct WorkPairImpl : IWorkItemHandler, std::enable_shared_from_this<IWorkItemHandler>
    {
        WorkPairImpl() : caller_(Async::GetCurrentWorkerPId())
        {}
        
        WorkPairImpl(Task_t &&task) : task(std::move(task)), caller_(Async::GetCurrentWorkerPId())
        {}
        
        WorkPairImpl(Task_t &&task, Job_t &&callback) : task(std::move(task)), callback(std::move(callback)), caller_(Async::GetCurrentWorkerPId())
        {}

        WorkPairImpl(const Task_t &task) : task(task), caller_(Async::GetCurrentWorkerPId())
        {}
        
        WorkPairImpl(const Task_t &task, const Job_t &callback) : task(task), callback(callback), caller_(Async::GetCurrentWorkerPId())
        {}
        
        WorkPairImpl(const Task_t &task, const Job_t &callback, const Info_t &info) : task(task), callback(callback), input(info), caller_(Async::GetCurrentWorkerPId())
        {}
        
        WorkPairImpl(Task_t &&task, const Job_t &callback, const Info_t &info) : task(std::move(task)), callback(callback), input(info), caller_(Async::GetCurrentWorkerPId())
        {}
        
        WorkPairImpl(Task_t &&task, Job_t &&callback, const Info_t &info) : task(std::move(task)), callback(std::move(callback)), input(info), caller_(Async::GetCurrentWorkerPId())
        {}
        
        WorkPairImpl(Task_t &&task, Job_t &&callback, Info_t &&info) : task(std::move(task)), callback(std::move(callback)), input(std::move(info)), caller_(Async::GetCurrentWorkerPId())
        {}
        
        WorkPairImpl(const Task_t &task, const Job_t &callback, Info_t &&info) : task(task), callback(callback), input(info), caller_(Async::GetCurrentWorkerPId())
        {}

        Info_t input;
        Task_t task;
        Job_t callback;

        
        WorkPairImpl<Info_t, Result_t, Task_t, Job_t> &SetTask(const Task_t &task)
        {
            this->task = task;
            return *this;
        }
        
        WorkPairImpl<Info_t, Result_t, Task_t, Job_t> &SetTask(const Job_t &callback)
        {
            this->callback = callback;
            return *this;
        }

    private:
        
        static constexpr bool IsCallbackPtr = std::is_pointer_v<Job_t>  || AuIsBaseOfTemplate<AURORA_RUNTIME_AU_SHARED_PTR, Job_t>::value  || AuIsBaseOfTemplate<AURORA_RUNTIME_AU_UNIQUE_PTR, Job_t>::value;
        static constexpr bool IsTaskPtr     = std::is_pointer_v<Task_t> || AuIsBaseOfTemplate<AURORA_RUNTIME_AU_SHARED_PTR, Task_t>::value || AuIsBaseOfTemplate<AURORA_RUNTIME_AU_UNIQUE_PTR, Task_t>::value;

        //WorkerId_t caller;
        WorkerPId_t caller_;
        Threading::Primitives::SpinLock lock_;
        
        BasicWorkCtx secretContext_;
        Result_t resultValue_;

        virtual void *GetPrivateData() override { return &secretContext_; }

        void DispatchFrame(ProcessInfo &info) override
        {
            AU_LOCK_GUARD(this->lock_);

            try
            {
                if constexpr (IsTaskPtr)
                {
                    if (task)
                    {
                        resultValue_ = task->OnFrame(input);
                    }
                }
                else
                {
                    resultValue_ = task.OnFrame(input);
                }

                task = {};
            }
            catch (...)
            {
                Debug::PrintError();
                Shutdown();
                return;
            }

            auto pin = AuSharedFromThis();

            std::function<void()> func = [pin]()
            {
                try
                {
                    pin->secretContext_.opt = &pin->resultValue_;
                    if constexpr (IsCallbackPtr)
                    {
                        pin->callback->onSuccess(pin->input, pin->resultValue_);
                    }
                    else
                    {
                        if (pin->callback.onSuccess)
                        {
                            pin->callback.onSuccess(pin->input,pin->resultValue_);
                        }
                    }
                }
                catch (...)
                {
                    Debug::PrintError();
                }
                
                pin->callback = {};
                pin->caller_ = {};
            };

            try
            {
                if (caller_ == Async::GetCurrentWorkerPId())
                {
                    func();
                }
                else
                {
                    AuFunction<void()> err = [pin]()
                    {
                        pin->CallOnFailure();
                    };

                    // TODO: this is somewhat evil. double alloc when we could reuse this
                    if (!caller_.pool->NewWorkItem(caller_, AuMakeShared<BasicWorkStdFunc>(func, err))->Dispatch())
                    {
                        pin->CallOnFailure();
                    }
                }

                caller_ = {};
            } 
            catch (...)
            {
                Debug::PrintError();
                Shutdown();
            }
        }

        void Shutdown() override
        {
            AU_LOCK_GUARD(this->lock_);
            ShutdownNoLock();
        }

        inline void ShutdownNoLock()
        {
            caller_ = {};
            try
            {
                CallOnFailure();
            }
            catch (...)
            {
                Debug::PrintError();
            }
            callback = {};
            task = {};
        }

        inline void CallOnFailure()
        {
            if constexpr (IsCallbackPtr)
            {
                if constexpr (AuIsBaseOfTemplate<std::function, decltype(callback->onFailure)>::value)
                {
                    if (!callback->onFailure)
                    {
                        return;
                    }
                }
                
                callback->onFailure(input);
            }
            else
            {
                if constexpr (AuIsBaseOfTemplate<std::function, decltype(callback.onFailure)>::value)
                {
                    if (!callback.onFailure)
                    {
                        return;
                    }
                }
                callback.onFailure(input);
            }
        }
    };
}