Stress-test and fix thread pool when work queues are full (#4690)

While investigating an unrelated issue, I noticed that the thread pool may drop tasks when a burst of 1024+ tasks is submitted by a thread from inside the pool. Today, in general, we execute work synchronously in this case. However, there is a bug where work submitted by a thread already inside the pool will be discarded instead of executed. Currently the only scenario where I can see this occurring is when the parallel executor is used with a model in which such a large number of nodes become eligible to run all at once. This PR fixes the underlying issue and adds a test case for burst-submission of work.

include/onnxruntime/core/platform/EigenNonBlockingThreadPool.h

@@ -455,6 +455,11 @@ class ThreadPoolTempl : public onnxruntime::concurrency::ExtendedThreadPoolInter
     for (size_t i = 0; i < worker_data_.size(); ++i) worker_data_[i].thread.reset();
   }
 
+  // Run fn().  Ordinarily, the function will be added to the thread pool and executed
+  // by a worker thread.  If the thread pool rejects the work then fn() will instead
+  // execute synchronously during Schedule(fn).  Currently the thread pool will only
+  // reject work if the queue of pending work is full.
+
   void Schedule(std::function<void()> fn) override {
     Task t = env_.CreateTask(std::move(fn));
     PerThread* pt = GetPerThread();
@@ -469,14 +474,16 @@ class ThreadPoolTempl : public onnxruntime::concurrency::ExtendedThreadPoolInter
       WorkerData &td = worker_data_[q_idx];
       Queue& q = td.queue;
       t = q.PushBack(std::move(t));
-      if (t.f) {
-        // The queue rejected the work; run it directly
-        env_.ExecuteTask(t);
-      } else {
+      if (!t.f) {
         // The queue accepted the work; ensure that the thread will pick it up
         td.EnsureAwake();
       }
     }
+
+    // Run the work directly if the queue rejected the work
+    if (t.f) {
+      env_.ExecuteTask(t);
+    }
   }
 
 // The thread pool maintains a set of hints for which threads will be good to distribute

onnxruntime/test/platform/threadpool_test.cc

@@ -104,6 +104,39 @@ void TestMultipleParallelFor(const std::string& name, int num_threads, int num_c
   }
 }
 
+void TestBurstScheduling(const std::string& name, int num_tasks) {
+  // Test submitting a burst of functions for executing.  The aim is to provoke cases such
+  // as the thread pool's work queues being full.
+  for (int rep = 0; rep < 5; rep++) {
+    std::atomic<int> ctr{0};
+    // Schedule a burst of num_tasks back-to-back, and then cleanly shut down the thread
+    // pool.  The synchronization barrier during shut down should ensure that all of the
+    // tasks are complete.  Note that if the thread pool's work queues are full, then a
+    // call to tp->Schedule() may run its argument synchronously.  In any case, we expect
+    // ctr==num_tasks.
+    CreateThreadPoolAndTest(name, 2, [&](ThreadPool* tp) {
+      // First variant : schedule from outside the pool
+      for (int tasks = 0; tasks < num_tasks; tasks++) {
+        tp->Schedule([&]() {
+          ctr++;
+        });
+      }
+    });
+    ASSERT_TRUE(ctr == num_tasks);
+    CreateThreadPoolAndTest(name, 2, [&](ThreadPool* tp) {
+      // Second variant : schedule from inside the pool
+      tp->Schedule([&, tp]() {
+        for (int tasks = 0; tasks < num_tasks; tasks++) {
+          tp->Schedule([&]() {
+            ctr++;
+          });
+        }
+      });
+    });
+    ASSERT_TRUE(ctr == num_tasks*2);
+  }
+}
+
 }  // namespace
 
 namespace onnxruntime {
@@ -202,6 +235,24 @@ TEST(ThreadPoolTest, TestMultipleParallelFor_4Thread_4Conc_8Tasks) {
 TEST(ThreadPoolTest, TestMultipleParallelFor_4Thread_4Conc_1MTasks) {
   TestMultipleParallelFor("TestMultipleParallelFor_4Thread_4Conc_1MTasks", 4, 4, 1000000);
 }
+
+TEST(ThreadPoolTest, TestBurstScheduling_0Tasks) {
+  TestBurstScheduling("TestBurstScheduling_0Tasks", 0);
+}
+
+TEST(ThreadPoolTest, TestBurstScheduling_1Task) {
+  TestBurstScheduling("TestBurstScheduling_1Task", 1);
+}
+
+TEST(ThreadPoolTest, TestBurstScheduling_16Tasks) {
+  TestBurstScheduling("TestBurstScheduling_16Tasks", 16);
+}
+
+TEST(ThreadPoolTest, TestBurstScheduling_65536Task) {
+  // Attempt to exhaust the size of the queues used in the thread pool to
+  // buffer tasks.
+  TestBurstScheduling("TestBurstScheduling_65536Tasks", 65536);
+}
 #ifdef _WIN32
 TEST(ThreadPoolTest, TestStackSize) {
   ThreadOptions to;