AsyncWork.h

Go to the documentation of this file.

// Copyright Epic Games, Inc. All Rights Reserved.
 
/*=============================================================================
    AsyncWork.h: Definition of queued work classes
=============================================================================*/
 
#pragma once
 
#include "CoreTypes.h"
#include "Misc/AssertionMacros.h"
#include "Misc/Compression.h"
#include "HAL/ThreadSafeCounter.h"
#include "Stats/Stats.h"
#include "HAL/Event.h"
#include "HAL/PlatformProcess.h"
#include "Async/InheritedContext.h"
#include "Misc/IQueuedWork.h"
#include "Misc/QueuedThreadPool.h"
#include "Async/Fundamental/Scheduler.h"
 
template<typename TTask>
class FAutoDeleteAsyncTask : 
    private UE::FInheritedContextBase, 
    private IQueuedWork
{
    TTask Task;
 
    /* Generic start function, not called directly
     * @param bForceSynchronous if true, this job will be started synchronously, now, on this thread
     **/
    void Start(bool bForceSynchronous, FQueuedThreadPool* InQueuedPool, EQueuedWorkPriority InPriority = EQueuedWorkPriority::Normal)
    {
        CaptureInheritedContext();
        FPlatformMisc::MemoryBarrier();
        FQueuedThreadPool* QueuedPool = InQueuedPool;
        if (bForceSynchronous)
        {
            QueuedPool = 0;
        }
        if (QueuedPool)
        {
            QueuedPool->AddQueuedWork(this, InPriority);
        }
        else
        {
            // we aren't doing async stuff
            DoWork();
        }
    }
 
    void DoWork()
    {
        UE::FInheritedContextScope InheritedContextScope = RestoreInheritedContext();
        FScopeCycleCounter Scope(Task.GetStatId(), true);
 
        Task.DoWork();
        delete this;
    }
 
    virtual void DoThreadedWork()
    {
        DoWork();
    }
 
    virtual void Abandon(void)
    {
        if (Task.CanAbandon())
        {
            Task.Abandon();
            delete this;
        }
        else
        {
            DoWork();
        }
    }
 
public:
    template<typename...T>
    explicit FAutoDeleteAsyncTask(T&&... Args) : Task(Forward<T>(Args)...)
    {
    }
 
    void StartSynchronousTask()
    {
        Start(true, nullptr);
    }
 
    void StartBackgroundTask(FQueuedThreadPool* InQueuedPool = GThreadPool, EQueuedWorkPriority InPriority = EQueuedWorkPriority::Normal)
    {
        Start(false, InQueuedPool, InPriority);
    }
};
 
 
class FAsyncTaskBase
    : private UE::FInheritedContextBase
    , private IQueuedWork
{
    FThreadSafeCounter  WorkNotFinishedCounter;
    FEvent*             DoneEvent = nullptr;
    FQueuedThreadPool*  QueuedPool = nullptr;
    EQueuedWorkPriority Priority = EQueuedWorkPriority::Normal;
    EQueuedWorkFlags Flags = EQueuedWorkFlags::None;
    int64 RequiredMemory = -1;
    const TCHAR * DebugName = nullptr;
    TStatId StatId;
 
    /* Internal function to destroy the completion event
    **/
    void DestroyEvent()
    {
        FPlatformProcess::ReturnSynchEventToPool(DoneEvent);
        DoneEvent = nullptr;
    }
 
    EQueuedWorkFlags GetQueuedWorkFlags() const final
    {
        return Flags;
    }
 
    /* Generic start function, not called directly
        * @param bForceSynchronous if true, this job will be started synchronously, now, on this thread
    **/
    void Start(bool bForceSynchronous, FQueuedThreadPool* InQueuedPool, EQueuedWorkPriority InQueuedWorkPriority, EQueuedWorkFlags InQueuedWorkFlags, int64 InRequiredMemory, const TCHAR * InDebugName)
    {
        CaptureInheritedContext();
 
        FScopeCycleCounter Scope(StatId, true);
        DECLARE_SCOPE_CYCLE_COUNTER( TEXT( "FAsyncTask::Start" ), STAT_FAsyncTask_Start, STATGROUP_ThreadPoolAsyncTasks );
        // default arg has InRequiredMemory == -1
        RequiredMemory = InRequiredMemory;
        DebugName = InDebugName;
 
        FPlatformMisc::MemoryBarrier();
        CheckIdle();  // can't start a job twice without it being completed first
        WorkNotFinishedCounter.Increment();
        QueuedPool = InQueuedPool;
        Priority = InQueuedWorkPriority;
        Flags = InQueuedWorkFlags;
        if (bForceSynchronous)
        {
            QueuedPool = 0;
        }
        if (QueuedPool)
        {
            if (!DoneEvent)
            {
                DoneEvent = FPlatformProcess::GetSynchEventFromPool(true);
            }
            DoneEvent->Reset();
            QueuedPool->AddQueuedWork(this, InQueuedWorkPriority);
        }
        else 
        {
            // we aren't doing async stuff
            DestroyEvent();
            DoWork();
        }
    }
 
    void DoWork()
    {
        UE::FInheritedContextScope InheritedContextScope = RestoreInheritedContext();
        FScopeCycleCounter Scope(StatId, true);
 
        DoTaskWork();
        check(WorkNotFinishedCounter.GetValue() == 1);
        WorkNotFinishedCounter.Decrement();
    }
 
    void FinishThreadedWork()
    {
        check(QueuedPool);
        if (DoneEvent)
        {
            FScopeCycleCounter Scope(StatId, true);
            DECLARE_SCOPE_CYCLE_COUNTER( TEXT( "FAsyncTask::FinishThreadedWork" ), STAT_FAsyncTask_FinishThreadedWork, STATGROUP_ThreadPoolAsyncTasks );        
            DoneEvent->Trigger();
        }
    }
 
    void DoThreadedWork() final
    {
        DoWork();
        FinishThreadedWork();
    }
 
    void Abandon() final
    {
        if (TryAbandonTask())
        {
            check(WorkNotFinishedCounter.GetValue() == 1);
            WorkNotFinishedCounter.Decrement();
        }
        else
        {
            DoWork();
        }
        FinishThreadedWork();
    }
 
    void SyncCompletion(bool bIsLatencySensitive)
    {
        FPlatformMisc::MemoryBarrier();
        if (QueuedPool)
        {
            TRACE_CPUPROFILER_EVENT_SCOPE(FAsyncTask::SyncCompletion);
            FScopeCycleCounter Scope(StatId);
            DECLARE_SCOPE_CYCLE_COUNTER(TEXT("FAsyncTask::SyncCompletion"), STAT_FAsyncTask_SyncCompletion, STATGROUP_ThreadPoolAsyncTasks);
 
            check(DoneEvent); // if it is not done yet, we must have an event
            DoneEvent->Wait();
            QueuedPool = 0;
        }
        CheckIdle();
    }
 
protected:
    void Init(TStatId InStatId)
    {
        StatId = InStatId;
    }
 
    void CheckIdle() const
    {
        check(WorkNotFinishedCounter.GetValue() == 0);
        if (QueuedPool) FDebug::DumpStackTraceToLog(ELogVerbosity::Warning);
        check(!QueuedPool);
    }
 
    virtual void DoTaskWork() = 0;
 
    virtual bool TryAbandonTask() = 0;
 
    virtual void MarkAsCanceled()
    {
    }
 
public:
    virtual ~FAsyncTaskBase()
    {
        // destroying an unfinished task is a bug
        CheckIdle();
        DestroyEvent();
    }
 
    int64 GetRequiredMemory() const final
    {
        return RequiredMemory;
    }
    
    const TCHAR * GetDebugName() const final
    {
        return DebugName;
    }
 
    void StartSynchronousTask(EQueuedWorkPriority InQueuedWorkPriority = EQueuedWorkPriority::Normal, EQueuedWorkFlags InQueuedWorkFlags = EQueuedWorkFlags::None, int64 InRequiredMemory = -1, const TCHAR * InDebugName = nullptr)
    {
        Start(true, GThreadPool, InQueuedWorkPriority, InQueuedWorkFlags, InRequiredMemory, InDebugName);
    }
 
    void StartBackgroundTask(FQueuedThreadPool* InQueuedPool = GThreadPool, EQueuedWorkPriority InQueuedWorkPriority = EQueuedWorkPriority::Normal, EQueuedWorkFlags InQueuedWorkFlags = EQueuedWorkFlags::None, int64 InRequiredMemory = -1, const TCHAR * InDebugName = nullptr)
    {
        Start(false, InQueuedPool, InQueuedWorkPriority, InQueuedWorkFlags, InRequiredMemory, InDebugName);
    }
 
    void EnsureCompletion(bool bDoWorkOnThisThreadIfNotStarted = true, bool bIsLatencySensitive = false)
    {
        bool DoSyncCompletion = true;
        if (bDoWorkOnThisThreadIfNotStarted)
        {
            if (QueuedPool)
            {
                if (QueuedPool->RetractQueuedWork(this))
                {
                    // we got the job back, so do the work now and no need to synchronize
                    DoSyncCompletion = false;
                    DoWork(); 
                    FinishThreadedWork();
                    QueuedPool = 0;
                }
            }
            else if (WorkNotFinishedCounter.GetValue())  // in the synchronous case, if we haven't done it yet, do it now
            {
                DoWork(); 
            }
        }
        if (DoSyncCompletion)
        {
            SyncCompletion(bIsLatencySensitive);
        }
        CheckIdle(); // Must have had bDoWorkOnThisThreadIfNotStarted == false and needed it to be true for a synchronous job
    }
    
    bool Reschedule(FQueuedThreadPool* InQueuedPool = GThreadPool, EQueuedWorkPriority InQueuedWorkPriority = EQueuedWorkPriority::Normal)
    {
        if (QueuedPool)
        {
            if (QueuedPool->RetractQueuedWork(this))
            {
                QueuedPool = InQueuedPool;
                Priority = InQueuedWorkPriority;
                QueuedPool->AddQueuedWork(this, InQueuedWorkPriority);
                
                return true;
            }
        }
        return false;
    }
 
    bool Cancel()
    {
        if (QueuedPool)
        {
            if (QueuedPool->RetractQueuedWork(this))
            {
                check(WorkNotFinishedCounter.GetValue() == 1);
                WorkNotFinishedCounter.Decrement();
                FinishThreadedWork();
                QueuedPool = 0;
                return true;
            }
 
            if (!IsWorkDone())
            {
                MarkAsCanceled();
            }
        }
        return false;
    }
 
    bool WaitCompletionWithTimeout(float TimeLimitSeconds)
    {
        if (TimeLimitSeconds <= 0.0f)
        {
            return IsDone();
        }
 
        FPlatformMisc::MemoryBarrier();
        if (QueuedPool)
        {
            FScopeCycleCounter Scope(StatId);
            DECLARE_SCOPE_CYCLE_COUNTER(TEXT("FAsyncTask::SyncCompletion"), STAT_FAsyncTask_SyncCompletion, STATGROUP_ThreadPoolAsyncTasks);
 
            uint32 Ms = uint32(TimeLimitSeconds * 1000.0f) + 1;
            check(Ms);
 
            check(DoneEvent); // if it is not done yet, we must have an event
            if (DoneEvent->Wait(Ms))
            {
                QueuedPool = 0;
                CheckIdle();
                return true;
            }
            return false;
        }
        CheckIdle();
        return true;
    }
 
    bool IsDone()
    {
        if (!IsWorkDone())
        {
            return false;
        }
        SyncCompletion(/*bIsLatencySensitive = */false);
        return true;
    }
 
    bool IsWorkDone() const
    {
        if (WorkNotFinishedCounter.GetValue())
        {
            return false;
        }
        return true;
    }
 
    bool IsIdle() const
    {
        return WorkNotFinishedCounter.GetValue() == 0 && QueuedPool == 0;
    }
 
    bool SetPriority(EQueuedWorkPriority QueuedWorkPriority)
    {
        return Reschedule(QueuedPool, QueuedWorkPriority);
    }
 
    EQueuedWorkPriority GetPriority() const
    {
        return Priority;
    }
};
 
template<typename TTask>
class FAsyncTask
    : public FAsyncTaskBase
{
    TTask Task;
 
    struct CMarkAsCanceled
    {
        template <typename TaskType>
        auto Requires(TaskType& Task) -> decltype(
            Task.MarkAsCanceled()
        );
    };
 
    void MarkAsCanceled() final
    {
        if constexpr (TModels_V<CMarkAsCanceled, TTask>)
        {
            Task.MarkAsCanceled();
        }
    }
public:
    FAsyncTask()
        : Task()
    {
        // Cache the StatId to remain backward compatible with TTask that declare GetStatId as non-const.
        Init(Task.GetStatId());
    }
 
    template <typename Arg0Type, typename... ArgTypes>
    FAsyncTask(Arg0Type&& Arg0, ArgTypes&&... Args)
        : Task(Forward<Arg0Type>(Arg0), Forward<ArgTypes>(Args)...)
    {
        // Cache the StatId to remain backward compatible with TTask that declare GetStatId as non-const.
        Init(Task.GetStatId());
    }
 
    /* Retrieve embedded user job, not legal to call while a job is in process
    * @return reference to embedded user job
    **/
    TTask& GetTask()
    {
        CheckIdle();  // can't modify a job without it being completed first
        return Task;
    }
 
    /* Retrieve embedded user job, not legal to call while a job is in process
    * @return reference to embedded user job
    **/
    const TTask& GetTask() const
    {
        CheckIdle();  // could be safe, but I won't allow it anyway because the data could be changed while it is being read
        return Task;
    }
 
    bool TryAbandonTask() final
    {
        if (Task.CanAbandon())
        {
            Task.Abandon();
            return true;
        }
 
        return false;
    }
 
    void DoTaskWork() final
    {
        Task.DoWork();
    }
};
 
class FNonAbandonableTask
{
public:
    bool CanAbandon()
    {
        return false;
    }
    void Abandon()
    {
    }
};