ThreadPool: allow task rescheduling by returning the NOT_STARTED status from the Run() method

Author: TheMostDiligent

Common/interface/ThreadPool.h

@@ -78,7 +78,9 @@ DILIGENT_BEGIN_INTERFACE(IAsyncTask, IObject)
     ///
     ///             Before returning from the function, the task implementation must
     ///             set the task status to either ASYNC_TASK_STATUS_CANCELLED or
-    ///             ASYNC_TASK_STATUS_COMPLETE.
+    ///             ASYNC_TASK_STATUS_COMPLETE to indicate that the task is finished,
+    ///             or to ASYNC_TASK_STATUS_NOT_STARTED to request the task to be
+    ///             rescheduled.
     VIRTUAL void METHOD(Run)(THIS_
                              Uint32 ThreadId) PURE;
 

Common/src/ThreadPool.cpp

@@ -125,21 +125,22 @@ class ThreadPoolImpl final : public ObjectBase<IThreadPool>
                 }
             }
 
+            bool TaskFinished = false;
             if (PrerequisitesMet)
             {
                 TaskInfo.pTask->SetStatus(ASYNC_TASK_STATUS_RUNNING);
                 TaskInfo.pTask->Run(ThreadId);
-                DEV_CHECK_ERR((TaskInfo.pTask->GetStatus() == ASYNC_TASK_STATUS_COMPLETE ||
-                               TaskInfo.pTask->GetStatus() == ASYNC_TASK_STATUS_CANCELLED),
-                              "Finished tasks must be in COMPLETE or CANCELLED state");
+                TaskFinished = TaskInfo.pTask->IsFinished();
+                DEV_CHECK_ERR((TaskFinished || TaskInfo.pTask->GetStatus() == ASYNC_TASK_STATUS_NOT_STARTED),
+                              "Finished tasks must be in COMPLETE, CANCELLED or NOT_STARTED state");
             }
 
             {
                 std::unique_lock<std::mutex> lock{m_TasksQueueMtx};
 
                 const auto NumRunningTasks = m_NumRunningTasks.fetch_add(-1) - 1;
 
-                if (PrerequisitesMet)
+                if (TaskFinished)
                 {
                     if (m_TasksQueue.empty() && NumRunningTasks == 0)
                     {
@@ -148,14 +149,15 @@ class ThreadPoolImpl final : public ObjectBase<IThreadPool>
                 }
                 else
                 {
-                    // If prerequisites are not met, re-enqueue the task with the minimum prerequisite priority
+                    // If prerequisites are not met or the task requested to be re-run,
+                    // re-enqueue the task with the minimum prerequisite priority
                     if (TaskInfo.pTask->GetPriority() > MinPrereqPriority)
                         TaskInfo.pTask->SetPriority(MinPrereqPriority);
                     m_TasksQueue.emplace(TaskInfo.pTask->GetPriority(), std::move(TaskInfo));
                 }
             }
 
-            if (!PrerequisitesMet)
+            if (!TaskFinished)
             {
                 m_NextTaskCond.notify_one();
             }