AuroraRuntime/Source/Async/WorkItem.cpp

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

    File: WorkItem.cpp
    Date: 2021-6-26
    Author: Reece
***/
#include <Source/RuntimeInternal.hpp>
#include "Async.hpp"
#include "WorkItem.hpp"
#include "AsyncApp.hpp"
#include "Schedular.hpp"

namespace Aurora::Async
{
    WorkItem::WorkItem(IThreadPoolInternal *owner, const WorkerId_t &worker, const AuSPtr<IWorkItemHandler> &task, bool supportsBlocking) : worker_(worker), task_(task), owner_(owner)
    {
        if (supportsBlocking)
        {
            this->finishedEvent_ = AuThreadPrimitives::EventUnique(false, true, true);
            SysAssert(this->finishedEvent_);
        }
    }

    WorkItem::~WorkItem()
    {
        //Fail();
    }

    AuSPtr<IWorkItem> WorkItem::WaitFor(const AuSPtr<IWorkItem> &workItem)
    {
        bool status {};

        {
            auto dependency = AuReinterpretCast<WorkItem>(workItem);

            AU_LOCK_GUARD(this->lock);
            AU_LOCK_GUARD(dependency->lock);

            if (dependency->HasFailed())
            {
                Fail();
                return AU_SHARED_FROM_THIS;
            }

            if (!AuTryInsert(dependency->waiters_, AuSharedFromThis()))
            {
                Fail();
                return AU_SHARED_FROM_THIS;
            }

            if (!AuTryInsert(this->waitOn_, workItem))
            {
                AuTryRemove(dependency->waiters_, AuSharedFromThis());
                Fail();
                return AU_SHARED_FROM_THIS;
            }
        }
        
        return AU_SHARED_FROM_THIS;
    }

    bool WorkItem::WaitForLocked(const AuList<AuSPtr<IWorkItem>> &workItems)
    {
        for (auto &workItem : workItems)
        {
            auto dependency = AuReinterpretCast<WorkItem>(workItem);
            AU_LOCK_GUARD(dependency->lock);

            if (dependency->HasFailed())
            {
                return false;
            }

            if (!AuTryInsert(dependency->waiters_, AuSharedFromThis()))
            {
                return false;
            }

            if (!AuTryInsert(this->waitOn_, workItem))
            {
                AuTryRemove(dependency->waiters_, AuSharedFromThis());
                return false;
            }
        }

        return true;
    }

    AuSPtr<IWorkItem> WorkItem::WaitFor(const AuList<AuSPtr<IWorkItem>> &workItems)
    {
        bool status {};

        {
            AU_LOCK_GUARD(this->lock);
            status = WaitForLocked(workItems);
        }

        if (!status)
        {
            Fail();
        }

        return AU_SHARED_FROM_THIS;
    }
    
    AuSPtr<IWorkItem> WorkItem::Then(const AuSPtr<IWorkItem> &next)
    {
        auto that = AU_SHARED_FROM_THIS;
        next->WaitFor(that);
        next->Dispatch();
        return that;
    }
    
    AuSPtr<IWorkItem> WorkItem::SetSchedTimeNs(AuUInt64 ns)
    {
        this->dispatchTimeNs_  = Time::CurrentClockNS() + ns;
        return AU_SHARED_FROM_THIS;
    }
    
    AuSPtr<IWorkItem> WorkItem::SetSchedTimeAbs(AuUInt32 ms)
    {
        this->dispatchTimeNs_  = AuUInt64(ms) * AuMSToNS<AuUInt64>(ms);
        return AU_SHARED_FROM_THIS;
    }

    AuSPtr<IWorkItem> WorkItem::SetSchedTimeNsAbs(AuUInt64 ns)
    {
        this->dispatchTimeNs_  = ns;
        return AU_SHARED_FROM_THIS;
    }

    AuSPtr<IWorkItem> WorkItem::SetSchedTime(AuUInt32 ms)
    {
        this->dispatchTimeNs_  = Time::CurrentClockNS() + AuMSToNS<AuUInt64>(ms);
        return AU_SHARED_FROM_THIS;
    }
    
    AuSPtr<IWorkItem> WorkItem::AddDelayTime(AuUInt32 ms)
    {
        this->delayTimeNs_ += AuMSToNS<AuUInt64>(ms);
        return AU_SHARED_FROM_THIS;
    }

    AuSPtr<IWorkItem> WorkItem::AddDelayTimeNs(AuUInt64 ns)
    {
        this->delayTimeNs_ += ns;
        return AU_SHARED_FROM_THIS;
    }

    AuSPtr<IWorkItem> WorkItem::Dispatch()
    {
        DispatchEx(false);
        return AU_SHARED_FROM_THIS;
    }

    void WorkItem::DispatchEx(bool check)
    {
        AU_LOCK_GUARD(lock);

        DispatchExLocked(check);
    }

    void WorkItem::DispatchExLocked(bool check)
    {
        if (check)
        {
            if (this->dispatchPending_)
            {
                return;
            }
        }

        if (HasFailed())
        {
            return;
        }

        for (auto itr = waitOn_.begin(); itr != waitOn_.end(); )
        {
            auto &waitable = *itr;

            if (!waitable->HasFinished())
            {
                return;
            }
            
            itr = waitOn_.erase(itr);
        }

        this->dispatchPending_ = true;

        if (Time::CurrentClockNS() < this->dispatchTimeNs_)
        {
            Schedule();
            return;
        }
        
        if (auto delay = AuExchange(delayTimeNs_, {}))
        {
            this->dispatchTimeNs_ = delay + Time::CurrentClockNS();
            Schedule();
            return;
        }

        SendOff();
    }

    float WorkItem::GetPrio()
    {
        return prio_;
    }

    void WorkItem::SetPrio(float val)
    {
        prio_ = val;
    }

    void WorkItem::CancelAsync()
    {
        AU_TRY_LOCK_GUARD_NAMED(this->lock, asd);
        Fail();
    }


    void WorkItem::RunAsync()
    {
        AU_LOCK_GUARD(this->lock);
        RunAsyncLocked();
    }

    void WorkItem::RunAsyncLocked()
    {

        IWorkItemHandler::ProcessInfo info(true);
        info.pool = this->owner_->ToThreadPool();

        if (this->task_)
        {
            this->task_->DispatchFrame(info);
        }

        switch (info.type)
        {
        case ETickType::eFinished:
        {

            // do nothing
            break;
        }
        case ETickType::eEnumInvalid:
        {
            SysPanic("Handle Invalid");
            break;
        }
        case ETickType::eSchedule:
        {
            if (info.reschedMs)
            {
                SetSchedTime(info.reschedMs);
            } 
            else if (info.reschedNs)
            {
                SetSchedTimeNs(info.reschedNs);
            }
            else if (info.reschedClockAbsMs)
            {
                SetSchedTimeAbs(info.reschedMs);
            }
            else if (info.reschedClockAbsNs)
            {
                SetSchedTimeNsAbs(info.reschedNs);
            }

            if (!WaitForLocked(info.waitFor))
            {
                Fail();
            }

        }
        [[fallthrough]];
        case ETickType::eRerun:
        {
            DispatchExLocked(false);
            return;
        }
        case ETickType::eFailed:
        {
            Fail();
            return;
        }
        }

        this->finished = true;
        if (this->finishedEvent_)
        {
            this->finishedEvent_->Set();
        }

        for (auto &waiter : this->waiters_)
        {
            AuReinterpretCast<WorkItem>(waiter)->DispatchEx(true);
        }
    }

    void WorkItem::Fail()
    {
        failed = true;

        if (auto task_ = AuExchange(this->task_, {}))
        {
            task_->OnFailure();
        }

        for (auto &waiter : this->waiters_)
        {
            AuReinterpretCast<WorkItem>(waiter)->Fail();
        }

        this->waiters_.clear();
        this->waitOn_.clear();
     
        if (this->finishedEvent_)
        {
            this->finishedEvent_->Set();
        }
    }
    
    bool WorkItem::BlockUntilComplete()
    {
        if (!this->finishedEvent_) return false;
        return this->owner_->WaitFor(this->worker_, AuUnsafeRaiiToShared(this->finishedEvent_), 0);
    }
    
    bool WorkItem::HasFinished()
    {
        return this->finished;
    }

    void WorkItem::Cancel()
    {
        AU_LOCK_GUARD(this->lock);
        Fail();
    }

    bool WorkItem::HasFailed()
    {
        return this->failed;
    }
    
    void WorkItem::Schedule()
    {
        Async::Schedule(this->dispatchTimeNs_, this->owner_, this->worker_, AuSharedFromThis());
    }

    void WorkItem::SendOff()
    {
        if (!this->task_)
        {
            // If we aren't actually calling a task interface, we may as well just dispatch objects waiting on us from here
            RunAsyncLocked();
        }
        else
        {
            this->owner_->Run(this->worker_, AuSharedFromThis());
        }
    }
    static auto GetWorkerInternal()
    {
        return static_cast<AsyncApp *>(GetAsyncApp());
    }

    inline auto GetWorkerInternal(const AuSPtr<IThreadPool> &pool)
    {
        if (pool.get() == AuStaticCast<IAsyncApp>(gAsyncApp))
        {
            return AuUnsafeRaiiToShared(AuStaticCast<ThreadPool>(gAsyncApp));
        }

        return AuStaticPointerCast<ThreadPool>(pool);
    }

    AUKN_SYM AuSPtr<IWorkItem> NewWorkItem(const WorkerId_t &worker, const AuSPtr<IWorkItemHandler> &task, bool supportsBlocking)
    {
        if (!task)
        {
            SysPushErrorArg("WorkItem has null task. Running out of memory?");
            return {};
        }

        return AuMakeShared<WorkItem>(GetWorkerInternal(), worker, task, supportsBlocking);
    }

    AUKN_SYM AuSPtr<IWorkItem> NewWorkItem(const WorkerPId_t &worker, const AuSPtr<IWorkItemHandler> &task, bool supportsBlocking)
    {
        if (!task)
        {
            SysPushErrorArg("WorkItem has null task. Running out of memory?");
            return {};
        }
        
        if (!worker.pool)
        {
            SysPushErrorArg("WorkItem has null pool");
            return {};
        }

        return AuMakeShared<WorkItem>(GetWorkerInternal(worker.pool).get(), worker, task, supportsBlocking);
    }

    AUKN_SYM AuSPtr<IWorkItem> NewFence()
    {
        return AuMakeShared<WorkItem>(GetWorkerInternal(), WorkerId_t{}, AuSPtr<IWorkItemHandler>{}, true);
    }
    
    void *WorkItem::GetPrivateData()
    {
        if (!this->task_)
        {
            return nullptr;
        }
 
        return this->task_->GetPrivateData();
    }

    AuOptional<void *> WorkItem::ToWorkResultT()
    {
        if (!this->task_)
        {
            return nullptr;
        }
        
        auto priv = reinterpret_cast<Async::WorkPriv *>(this->task_->GetPrivateData());
        if (!priv)
        {
            return nullptr;
        }
        
        if (priv->magic == AuConvertMagicTag32("BWOT"))
        {
            
            return reinterpret_cast<Async::BasicWorkCtx *>(priv)->opt;
        }

        return {};
    }
}