Use semaphore. Numberous fixes.

Signed-off-by: Michał Zientkiewicz <mzient@gmail.com>

Author: mzient

dali/core/exec/thread_pool_base.cc

@@ -21,6 +21,8 @@ Job::~Job() noexcept(false) {
   if (!tasks_.empty() && !waited_for_) {
     throw std::logic_error("The job is not empty, but hasn't been scrapped or waited for.");
   }
+  while (running_)
+    std::this_thread::yield();
 }
 
 void Job::Wait() {
@@ -30,8 +32,8 @@ void Job::Wait() {
   if (waited_for_)
     throw std::logic_error("This job has already been waited for.");
 
+  auto ready = [&]() { return num_pending_tasks_ == 0; };
   if (ThreadPoolBase::this_thread_pool() != nullptr) {
-    auto ready = [&]() { return num_pending_tasks_ == 0; };
     bool result = ThreadPoolBase::this_thread_pool()->WaitOrRunTasks(cv_, ready);
     waited_for_ = true;
     if (!result)
@@ -62,6 +64,7 @@ void Job::Run(ThreadPoolBase &tp, bool wait) {
   if (started_)
     throw std::logic_error("This job has already been started.");
   started_ = true;
+  running_ = !tasks_.empty();
   {
     auto batch = tp.BulkAdd();
     num_pending_tasks_ += tasks_.size();
@@ -86,6 +89,8 @@ IncrementalJob::~IncrementalJob() noexcept(false) {
   if (!tasks_.empty() && !waited_for_) {
     throw std::logic_error("The job is not empty, but hasn't been scrapped or waited for.");
   }
+  while (running_)
+    std::this_thread::yield();
 }
 
 void IncrementalJob::Wait() {
@@ -95,8 +100,8 @@ void IncrementalJob::Wait() {
   if (waited_for_)
     throw std::logic_error("This job has already been waited for.");
 
+  auto ready = [&]() { return num_pending_tasks_ == 0; };
   if (ThreadPoolBase::this_thread_pool() != nullptr) {
-    auto ready = [&]() { return num_pending_tasks_ == 0; };
     bool result = ThreadPoolBase::this_thread_pool()->WaitOrRunTasks(cv_, ready);
     waited_for_ = true;
     if (!result)
@@ -134,6 +139,8 @@ void IncrementalJob::Run(ThreadPoolBase &tp, bool wait) {
       batch.Add(std::move(it->func));
       last_task_run_ = it;
     }
+    running_ = batch.Size() > 0;
+    batch.Submit();
   }
   if (wait && !tasks_.empty())
     Wait();
@@ -148,7 +155,7 @@ void IncrementalJob::Scrap() {
 ///////////////////////////////////////////////////////////////////////////
 
 thread_local ThreadPoolBase *ThreadPoolBase::this_thread_pool_ = nullptr;
-thread_local int ThreadPoolBase::this_thread_idx_ = -1;;
+thread_local int ThreadPoolBase::this_thread_idx_ = -1;
 
 void ThreadPoolBase::Init(int num_threads, const std::function<OnThreadStartFn> &on_thread_start) {
   if (shutdown_pending_)
@@ -162,20 +169,19 @@ void ThreadPoolBase::Init(int num_threads, const std::function<OnThreadStartFn>
 }
 
 void ThreadPoolBase::Shutdown(bool join) {
-  if (shutdown_pending_ && !join)
+  if ((shutdown_pending_ && !join) || threads_.empty())
     return;
   {
     std::lock_guard<std::mutex> g(mtx_);
     if (shutdown_pending_ && !join)
       return;
     shutdown_pending_ = true;
-    cv_.notify_all();
+    sem_.release(threads_.size());
   }
 
   for (auto &t : threads_)
     t.join();
-
-  assert(tasks_.empty());
+  threads_.clear();
 }
 
 void ThreadPoolBase::AddTaskNoLock(TaskFunc &&f) {
@@ -189,7 +195,6 @@ void ThreadPoolBase::AddTask(TaskFunc &&f) {
     std::lock_guard<std::mutex> g(mtx_);
     AddTaskNoLock(std::move(f));
   }
-  cv_.notify_one();
 }
 
 void ThreadPoolBase::Run(
@@ -200,11 +205,12 @@ void ThreadPoolBase::Run(
   std::any scope;
   if (on_thread_start)
     scope = on_thread_start(index);
-  std::unique_lock lock(mtx_);
   while (!shutdown_pending_ || !tasks_.empty()) {
-    cv_.wait(lock, [&]() { return shutdown_pending_ || !tasks_.empty(); });
-    if (tasks_.empty())
+    sem_.acquire();
+    std::unique_lock lock(mtx_);
+    if (shutdown_pending_)
       break;
+    assert(!tasks_.empty() && "Semaphore acquired but no tasks present.");
     PopAndRunTask(lock);
   }
 }
@@ -228,11 +234,12 @@ bool ThreadPoolBase::WaitOrRunTasks(std::condition_variable &cv, Condition &&con
 
     if (ret || condition())  // re-evaluate the condition, just in case
       return true;
-    if (tasks_.empty()) {
-      assert(shutdown_pending_);
+    if (shutdown_pending_)
       return condition();
-    }
+    if (!sem_.try_acquire())
+      continue;
 
+    assert(!tasks_.empty() && "Semaphore acquired but no tasks present.");
     PopAndRunTask(lock);
   }
   return condition();

include/dali/core/exec/thread_pool_base.h

@@ -29,6 +29,7 @@
 #include <vector>
 #include <mutex>
 #include <condition_variable>
+#include "dali/core/semaphore.h"
 #include "dali/core/api_helper.h"
 #include "dali/core/multi_error.h"
 #include "dali/core/mm/detail/aux_alloc.h"
@@ -66,10 +67,14 @@ class DLL_PUBLIC Job {
 
         if (--num_pending_tasks_ == 0) {
           num_pending_tasks_.notify_all();
-          std::cerr << make_string((void *)this, " notified.") << std::endl;
+          (void)std::lock_guard(mtx_);
           cv_.notify_all();
+          // We need this second flag to avoid a race condition where the
+          // desctructor is called between decrementing num_pending_tasks_ and notification_
+          // without excessive use of mutexes.  This must be the very last operation in the task
+          // function that touches `this`.
+          running_ = false;
         }
-        assert(num_pending_tasks_ >= 0);
       };
     } catch (...) {  // if, for whatever reason, we cannot initialize the task, we should erase it
       tasks_.erase(it);
@@ -88,8 +93,10 @@ class DLL_PUBLIC Job {
 
  private:
   // atomic wait has no timeout, so we're stuck with condvar for reentrance
+  std::mutex mtx_;
   std::condition_variable cv_;
   std::atomic_int num_pending_tasks_{0};
+  std::atomic_bool running_{false};
   bool started_ = false;
   bool waited_for_ = false;
 
@@ -139,8 +146,10 @@ class DLL_PUBLIC IncrementalJob {
   const void *executor_ = nullptr;
   bool waited_for_ = false;
   // atomic wait has no timeout, so we're stuck with condvar for reentrance
+  std::mutex mtx_;
   std::condition_variable cv_;
   std::atomic_int num_pending_tasks_{0};
+  std::atomic_bool running_{false};
   using task_list_t = std::list<Task, mm::detail::object_pool_allocator<Task>>;
   task_list_t tasks_;
   std::optional<task_list_t::iterator> last_task_run_;
@@ -192,18 +201,20 @@ class DLL_PUBLIC ThreadPoolBase {
     void Submit() {
       if (lock.owns_lock()) {
         lock.unlock();
-        if (tasks_added > 1)
-          owner->cv_.notify_all();
-        else
-          owner->cv_.notify_one();
+        owner->sem_.release(tasks_added);
       }
     }
+
+    int Size() const {
+      return tasks_added;
+    }
+
    private:
     friend class ThreadPoolBase;
     explicit TaskBulkAdd(ThreadPoolBase *o) : owner(o), lock(o->mtx_, std::defer_lock) {}
-    ThreadPoolBase *owner;
+    ThreadPoolBase *owner = nullptr;
     std::unique_lock<std::mutex> lock;
-    int tasks_added;
+    int tasks_added = 0;
   };
   friend class TaskBulkAdd;
 
@@ -247,7 +258,7 @@ class DLL_PUBLIC ThreadPoolBase {
   void Run(int index, const std::function<OnThreadStartFn> &on_thread_start) noexcept;
 
   std::mutex mtx_;
-  std::condition_variable cv_;
+  counting_semaphore sem_{0};
   bool shutdown_pending_ = false;
   std::queue<TaskFunc> tasks_;
   std::vector<std::thread> threads_;
@@ -289,13 +300,15 @@ void Job::Run(Executor &executor, bool wait) {
     if (started_)
       throw std::logic_error("This job has already been started.");
     started_ = true;
+    running_ = !tasks_.empty();
     for (auto &x : tasks_) {
       num_pending_tasks_++;
       try {
         executor.AddTask(std::move(x.second.func));
       } catch (...) {
         if (--num_pending_tasks_ == 0) {
           num_pending_tasks_.notify_all();
+          (void)std::lock_guard(mtx_);
           cv_.notify_all();
         }
         throw;
@@ -325,9 +338,14 @@ IncrementalJob::AddTask(Runnable &&runnable) {
 
       if (--num_pending_tasks_ == 0) {
         num_pending_tasks_.notify_all();
+        (void)std::lock_guard(mtx_);
         cv_.notify_all();
+        // We need this second flag to avoid a race condition where the
+        // desctructor is called between decrementing num_pending_tasks_ and notification_
+        // without excessive use of mutexes. This must be the very last operation in the task
+        // function that touches `this`.
+        running_ = false;
       }
-      assert(num_pending_tasks_ >= 0);
     };
   } catch (...) {  // if, for whatever reason, we cannot initialize the task, we should erase it
     tasks_.erase(it);
@@ -346,6 +364,7 @@ void IncrementalJob::Run(Executor &executor, bool wait) {
     executor_ = &executor;
     auto it = last_task_run_.has_value() ? std::next(*last_task_run_) : tasks_.begin();
     for (; it != tasks_.end(); ++it) {
+      running_ = true;
       executor.AddTask(std::move(it->func));
       last_task_run_ = it;
     }