ThreadPool: fixed potential issue with the IAsyncTask::IsFinished() returning true while the task is still running

Author: TheMostDiligent

Common/interface/ThreadPool.h

@@ -72,17 +72,20 @@ DILIGENT_BEGIN_INTERFACE(IAsyncTask, IObject)
 
     /// \param [in] ThreadId - Id of the thread that is running this task.
     ///
-    /// \remarks    This method is only called once by the thread pool.
-    ///             Before starting the task, the thread pool sets its
+    /// \remarks    Before starting the task, the thread pool sets its
     ///             status to ASYNC_TASK_STATUS_RUNNING.
     ///
-    ///             Before returning from the function, the task implementation must
-    ///             set the task status to either ASYNC_TASK_STATUS_CANCELLED or
-    ///             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;
+    ///             The method must return one of the following values:
+    ///               - ASYNC_TASK_STATUS_CANCELLED to indicate that the task was cancelled.
+    ///               - ASYNC_TASK_STATUS_COMPLETE to indicate that the task is finished successfully.
+    ///               - ASYNC_TASK_STATUS_NOT_STARTED to request the task to be rescheduled.
+    ///
+    ///             The thread pool will set the task status to the returned value after
+    ///             the Run() method completes. This way if the GetStatus() method returns
+    ///             any value other than ASYNC_TASK_STATUS_RUNNING, it is guaranteed that the task
+    ///             is not executed by any thread.
+    VIRTUAL ASYNC_TASK_STATUS METHOD(Run)(THIS_
+                                     Uint32 ThreadId) PURE;
 
     /// Cancel the task, if possible.
     ///

Common/interface/ThreadPool.hpp

@@ -166,6 +166,9 @@ class AsyncTaskBase : public ObjectBase<IAsyncTask>
 };
 
 
+/// Enqueues a function to be executed asynchronously by the thread pool.
+/// For the list of parameters, see Diligent::IThreadPool::EnqueueTask() method.
+/// The handler function must return the task status, see Diligent::IAsyncTask::Run() method.
 template <typename HanlderType>
 RefCntAutoPtr<IAsyncTask> EnqueueAsyncWork(IThreadPool* pThreadPool,
                                            IAsyncTask** ppPrerequisites,
@@ -183,12 +186,11 @@ RefCntAutoPtr<IAsyncTask> EnqueueAsyncWork(IThreadPool* pThreadPool,
             m_Handler{std::move(Handler)}
         {}
 
-        virtual void DILIGENT_CALL_TYPE Run(Uint32 ThreadId) override final
+        virtual ASYNC_TASK_STATUS DILIGENT_CALL_TYPE Run(Uint32 ThreadId) override final
         {
-            ASYNC_TASK_STATUS TaskStatus = !m_bSafelyCancel.load() ?
+            return !m_bSafelyCancel.load() ?
                 m_Handler(ThreadId) :
                 ASYNC_TASK_STATUS_CANCELLED;
-            SetStatus(TaskStatus);
         }
 
     private:

Common/src/ThreadPool.cpp

@@ -129,7 +129,11 @@ class ThreadPoolImpl final : public ObjectBase<IThreadPool>
             if (PrerequisitesMet)
             {
                 TaskInfo.pTask->SetStatus(ASYNC_TASK_STATUS_RUNNING);
-                TaskInfo.pTask->Run(ThreadId);
+                ASYNC_TASK_STATUS ReturnStatus = TaskInfo.pTask->Run(ThreadId);
+                // NB: It is essential to set the task status after the Run() method returns.
+                //     This way if the GetStatus() method returns any value other than ASYNC_TASK_STATUS_RUNNING,
+                //     it is guaranteed that the task is not executed by any thread.
+                TaskInfo.pTask->SetStatus(ReturnStatus);
                 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");

Tests/DiligentCoreTest/src/Common/ThreadPoolTest.cpp

@@ -173,10 +173,10 @@ class WaitTask : public AsyncTaskBase
         m_WaitSignal{WaitSignal}
     {}
 
-    virtual void DILIGENT_CALL_TYPE Run(Uint32 ThreadId) override final
+    virtual ASYNC_TASK_STATUS DILIGENT_CALL_TYPE Run(Uint32 ThreadId) override final
     {
         m_WaitSignal.Wait();
-        SetStatus(ASYNC_TASK_STATUS_COMPLETE);
+        return ASYNC_TASK_STATUS_COMPLETE;
     }
 
 private:
@@ -191,9 +191,9 @@ class DummyTask : public AsyncTaskBase
         AsyncTaskBase{pRefCounters, fPriority}
     {}
 
-    virtual void DILIGENT_CALL_TYPE Run(Uint32 ThreadId) override final
+    virtual ASYNC_TASK_STATUS DILIGENT_CALL_TYPE Run(Uint32 ThreadId) override final
     {
-        SetStatus(ASYNC_TASK_STATUS_COMPLETE);
+        return ASYNC_TASK_STATUS_COMPLETE;
     }
 };