Merge pull request #48582 from Dr15Jones/improveConcurrencyTest

Fix race condition in SerialTaskQueue

Author: cmsbuild

FWCore/Concurrency/README.md

@@ -11,7 +11,7 @@ Exactly how to _dispose_ of the class can be customized via the virtual function
 ## `edm::WaitingTask`
 This class inherits from `edm::TaskBase` and represents a set of code to be run once other activities have completed. This includes the ability to hold a `std::exception_ptr` which can hold an exception which was generated in a dependent task.
 
-A raw pointer to a  `edm::WaitingTask` is not supposed to be handled directly. Instead, one should use the helpers `edm::WaitingTaskHolder`, `edm::WaitingTaskWithArenaHolder` or `edm::WaitingTaskList` to properly manage the internal reference count such that when the count drops to 0 the `execute()` method will be run followed by `recycle()`. In addition, these helper classes will handled passing along any `edm::exception_ptr` generated from a dependent task.
+A raw pointer to a  `edm::WaitingTask` is not supposed to be handled directly. Instead, one should use the helpers `edm::WaitingTaskHolder`, `edm::WaitingTaskWithArenaHolder` or `edm::WaitingTaskList` to properly manage the internal reference count such that when the count drops to 0 the `execute()` method will be run followed by `recycle()`. In addition, these helper classes will handle passing along any `edm::exception_ptr` generated from a dependent task.
 
 The easiest way to create an `edm::WaitingTask` is to call `edm::make_waiting_task` and pass in a lambda of the form `void(std::exception_ptr const*)`.
 ```C++
@@ -34,7 +34,7 @@ In the case where one is doing a synchronous wait on a series of asynchronous ta
 
 Note that the function `wait` will rethrow any exception stored in `finalTask`. There is an alternative function named `waitNoThrow` which will return the `std::exception_ptr`.
 
-WARNING: It important that the finalTask not execute before completion of the construction of all `WaitingTaskHolders` that will be constructed directly from finalTask. The following would be a bug:
+WARNING: It is important that the finalTask not be executed before completion of the construction of all `WaitingTaskHolders` that will be constructed directly from finalTask. The following would be a bug:
 
 ```C++
    oneapi::tbb::task_group group;

FWCore/Concurrency/interface/SerialTaskQueue.h

@@ -66,13 +66,14 @@
 namespace edm {
   class SerialTaskQueue {
   public:
-    SerialTaskQueue() : m_taskChosen(false), m_pauseCount{0} {}
+    SerialTaskQueue() : m_pauseCount{0}, m_taskChosen{false}, m_pickingNextTask{false} {}
 
     SerialTaskQueue(SerialTaskQueue&& iOther)
         : m_tasks(std::move(iOther.m_tasks)),
+          m_pauseCount(iOther.m_pauseCount.exchange(0)),
           m_taskChosen(iOther.m_taskChosen.exchange(false)),
-          m_pauseCount(iOther.m_pauseCount.exchange(0)) {
-      assert(m_tasks.empty() and m_taskChosen == false);
+          m_pickingNextTask(false) {
+      assert(m_tasks.empty() and m_taskChosen == false and iOther.m_pickingNextTask == false);
     }
     SerialTaskQueue(const SerialTaskQueue&) = delete;
     const SerialTaskQueue& operator=(const SerialTaskQueue&) = delete;
@@ -159,8 +160,9 @@ namespace edm {
 
     // ---------- member data --------------------------------
     oneapi::tbb::concurrent_queue<TaskBase*> m_tasks;
-    std::atomic<bool> m_taskChosen;
     std::atomic<unsigned long> m_pauseCount;
+    std::atomic<bool> m_taskChosen;
+    std::atomic<bool> m_pickingNextTask;
   };
 
   template <typename T>

FWCore/Concurrency/src/SerialTaskQueue.cc

@@ -18,6 +18,7 @@
 #include "FWCore/Concurrency/interface/SerialTaskQueue.h"
 
 #include "FWCore/Utilities/interface/Likely.h"
+#include "FWCore/Utilities/interface/make_sentry.h"
 
 using namespace edm;
 
@@ -86,25 +87,24 @@ SerialTaskQueue::TaskBase* SerialTaskQueue::finishedTask() {
 }
 
 SerialTaskQueue::TaskBase* SerialTaskQueue::pickNextTask() {
+  if UNLIKELY (0 != m_pauseCount)
+    return nullptr;
   bool expect = false;
-  if LIKELY (0 == m_pauseCount and m_taskChosen.compare_exchange_strong(expect, true)) {
+  //need pop task and setting m_taskChosen to be atomic to avoid
+  // case where thread pauses just after try_pop failed but then
+  // a task is added and that call fails the check on m_taskChosen
+  while (not m_pickingNextTask.compare_exchange_strong(expect, true)) {
+    expect = false;
+  }
+  auto sentry = edm::make_sentry(&m_pickingNextTask, [](auto* v) { v->store(false); });
+
+  if LIKELY (m_taskChosen.compare_exchange_strong(expect, true)) {
     TaskBase* t = nullptr;
     if LIKELY (m_tasks.try_pop(t)) {
       return t;
     }
     //no task was actually pulled
     m_taskChosen.store(false);
-
-    //was a new entry added after we called 'try_pop' but before we did the clear?
-    expect = false;
-    if (not m_tasks.empty() and m_taskChosen.compare_exchange_strong(expect, true)) {
-      t = nullptr;
-      if (m_tasks.try_pop(t)) {
-        return t;
-      }
-      //no task was still pulled since a different thread beat us to it
-      m_taskChosen.store(false);
-    }
   }
   return nullptr;
 }

FWCore/Concurrency/test/test2_catch2_limitedtaskqueue.cc

@@ -11,39 +11,49 @@
 #include <memory>
 #include <atomic>
 #include <thread>
+#include <mutex>
 #include "oneapi/tbb/task_arena.h"
 #include "FWCore/Concurrency/interface/WaitingTask.h"
+#include "FWCore/Concurrency/interface/FinalWaitingTask.h"
+#include "FWCore/Concurrency/interface/WaitingTaskHolder.h"
 #include "FWCore/Concurrency/interface/LimitedTaskQueue.h"
 #include "FWCore/Concurrency/interface/FunctorTask.h"
 
 using namespace std::chrono_literals;
 
+namespace {
+  std::mutex g_requiresMutex;
+
+}
+//catch2 REQUIRE is not thread safe
+#define SAFE_REQUIRE(__var__)           \
+  {                                     \
+    std::lock_guard g{g_requiresMutex}; \
+    REQUIRE(__var__);                   \
+  }
+
 TEST_CASE("LimitedTaskQueue", "[LimitedTaskQueue]") {
   SECTION("push") {
     {
       std::atomic<unsigned int> count{0};
       edm::LimitedTaskQueue queue{1};
       {
-        std::atomic<int> waitingTasks{3};
         oneapi::tbb::task_group group;
-        queue.push(group, [&count, &waitingTasks] {
+        edm::FinalWaitingTask lastTask(group);
+        edm::WaitingTaskHolder waitingTask(group, &lastTask);
+        queue.push(group, [&count, waitingTask] {
           REQUIRE(count++ == 0u);
           std::this_thread::sleep_for(10us);
-          --waitingTasks;
         });
-        queue.push(group, [&count, &waitingTasks] {
+        queue.push(group, [&count, waitingTask] {
           REQUIRE(count++ == 1u);
           std::this_thread::sleep_for(10us);
-          --waitingTasks;
         });
-        queue.push(group, [&count, &waitingTasks] {
+        queue.push(group, [&count, lastTask = std::move(waitingTask)] {
           REQUIRE(count++ == 2u);
           std::this_thread::sleep_for(10us);
-          --waitingTasks;
         });
-        do {
-          group.wait();
-        } while (0 != waitingTasks.load());
+        lastTask.wait();
         REQUIRE(count == 3u);
       }
     }
@@ -52,29 +62,25 @@ TEST_CASE("LimitedTaskQueue", "[LimitedTaskQueue]") {
       constexpr unsigned int kMax = 2;
       edm::LimitedTaskQueue queue{kMax};
       {
-        std::atomic<int> waitingTasks{3};
         oneapi::tbb::task_group group;
-        queue.push(group, [&count, &waitingTasks, kMax] {
-          REQUIRE(count++ < kMax);
+        edm::FinalWaitingTask lastTask(group);
+        edm::WaitingTaskHolder waitingTask(group, &lastTask);
+        queue.push(group, [&count, waitingTask, kMax] {
+          SAFE_REQUIRE(count++ < kMax);
           std::this_thread::sleep_for(10us);
           --count;
-          --waitingTasks;
         });
-        queue.push(group, [&count, &waitingTasks, kMax] {
-          REQUIRE(count++ < kMax);
+        queue.push(group, [&count, waitingTask, kMax] {
+          SAFE_REQUIRE(count++ < kMax);
           std::this_thread::sleep_for(10us);
           --count;
-          --waitingTasks;
         });
-        queue.push(group, [&count, &waitingTasks, kMax] {
-          REQUIRE(count++ < kMax);
+        queue.push(group, [&count, lastTask = std::move(waitingTask), kMax] {
+          SAFE_REQUIRE(count++ < kMax);
           std::this_thread::sleep_for(10us);
           --count;
-          --waitingTasks;
         });
-        do {
-          group.wait();
-        } while (0 != waitingTasks);
+        lastTask.wait();
         REQUIRE(count == 0u);
       }
     }
@@ -85,48 +91,45 @@ TEST_CASE("LimitedTaskQueue", "[LimitedTaskQueue]") {
     edm::LimitedTaskQueue queue{1};
     {
       {
-        std::atomic<int> waitingTasks{3};
         oneapi::tbb::task_group group;
+        edm::FinalWaitingTask lastTask(group);
+        edm::WaitingTaskHolder waitingTask(group, &lastTask);
+
         edm::LimitedTaskQueue::Resumer resumer;
         std::atomic<bool> resumerSet{false};
         std::exception_ptr e1;
-        queue.pushAndPause(
-            group, [&resumer, &resumerSet, &count, &waitingTasks, &e1](edm::LimitedTaskQueue::Resumer iResumer) {
-              resumer = std::move(iResumer);
-              resumerSet = true;
-              try {
-                REQUIRE(++count == 1u);
-              } catch (...) {
-                e1 = std::current_exception();
-              }
-              --waitingTasks;
-            });
+        queue.pushAndPause(group,
+                           [&resumer, &resumerSet, &count, waitingTask, &e1](edm::LimitedTaskQueue::Resumer iResumer) {
+                             resumer = std::move(iResumer);
+                             resumerSet = true;
+                             try {
+                               SAFE_REQUIRE(++count == 1u);
+                             } catch (...) {
+                               e1 = std::current_exception();
+                             }
+                           });
         std::exception_ptr e2;
-        queue.push(group, [&count, &waitingTasks, &e2] {
+        queue.push(group, [&count, waitingTask, &e2] {
           try {
-            REQUIRE(++count == 2u);
+            SAFE_REQUIRE(++count == 2u);
           } catch (...) {
             e2 = std::current_exception();
           }
-          --waitingTasks;
         });
         std::exception_ptr e3;
-        queue.push(group, [&count, &waitingTasks, &e3] {
+        queue.push(group, [&count, lastTask = std::move(waitingTask), &e3] {
           try {
-            REQUIRE(++count == 3u);
+            SAFE_REQUIRE(++count == 3u);
           } catch (...) {
             e3 = std::current_exception();
           }
-          --waitingTasks;
         });
         std::this_thread::sleep_for(100us);
         REQUIRE(2u >= count);
         while (not resumerSet) {
         }
-        REQUIRE(resumer.resume());
-        do {
-          group.wait();
-        } while (0 != waitingTasks.load());
+        SAFE_REQUIRE(resumer.resume());
+        lastTask.wait();
         REQUIRE(count == 3u);
         if (e1) {
           std::rethrow_exception(e1);
@@ -147,54 +150,46 @@ TEST_CASE("LimitedTaskQueue", "[LimitedTaskQueue]") {
     edm::LimitedTaskQueue queue{kMax};
     unsigned int index = 100;
     const unsigned int nTasks = 1000;
-    //catch2 REQUIRE is not thread safe
-    std::mutex mutex;
     while (0 != --index) {
-      std::atomic<int> waiting{1};
+      edm::FinalWaitingTask lastTask(group);
+
       std::atomic<unsigned int> count{0};
       std::atomic<unsigned int> nRunningTasks{0};
       std::atomic<bool> waitToStart{true};
       {
-        group.run([&queue, &waitToStart, &group, &waiting, &count, &nRunningTasks, &mutex, kMax] {
+        edm::WaitingTaskHolder waitingTask(group, &lastTask);
+
+        group.run([&queue, &waitToStart, &group, waitingTask, &count, &nRunningTasks, kMax] {
           while (waitToStart) {
           }
           for (unsigned int i = 0; i < nTasks; ++i) {
-            ++waiting;
-            queue.push(group, [&count, &waiting, &nRunningTasks, &mutex, kMax] {
-              std::shared_ptr<std::atomic<int>> guardAgain{&waiting, [](auto* v) { --(*v); }};
+            queue.push(group, [&count, waitingTask, &nRunningTasks, kMax] {
               auto nrt = nRunningTasks++;
               if (nrt >= kMax) {
                 std::cout << "ERROR " << nRunningTasks << " >= " << kMax << std::endl;
-                std::lock_guard lock{mutex};
-                REQUIRE(nrt < kMax);
+                SAFE_REQUIRE(nrt < kMax);
               }
               ++count;
               --nRunningTasks;
             });
           }
         });
-        group.run([&queue, &waitToStart, &group, &waiting, &count, &nRunningTasks, &mutex, kMax] {
+        group.run([&queue, &waitToStart, &group, waitingTask, &count, &nRunningTasks, kMax] {
           waitToStart = false;
           for (unsigned int i = 0; i < nTasks; ++i) {
-            ++waiting;
-            queue.push(group, [&count, &waiting, &nRunningTasks, &mutex, kMax] {
-              std::shared_ptr<std::atomic<int>> guardAgain{&waiting, [](auto* v) { --(*v); }};
+            queue.push(group, [&count, waitingTask, &nRunningTasks, kMax] {
               auto nrt = nRunningTasks++;
               if (nrt >= kMax) {
                 std::cout << "ERROR " << nRunningTasks << " >= " << kMax << std::endl;
-                std::lock_guard lock{mutex};
-                REQUIRE(nrt < kMax);
+                SAFE_REQUIRE(nrt < kMax);
               }
               ++count;
               --nRunningTasks;
             });
           }
-          --waiting;
         });
       }
-      do {
-        group.wait();
-      } while (0 != waiting.load());
+      lastTask.wait();
       REQUIRE(nRunningTasks == 0u);
       REQUIRE(2 * nTasks == count);
     }