[llvm] [orc-rt] Refactor QueueingTaskDispatcher to use an external TaskQueue. (PR #190920)

Lang Hames via llvm-commits llvm-commits at lists.llvm.org
Wed Apr 8 00:05:44 PDT 2026


https://github.com/lhames created https://github.com/llvm/llvm-project/pull/190920

QueueingTaskDispatcher now takes a TaskQueue by reference rather than maintaining an internal queue. This lets API clients retain direct access to the queue after transferring dispatcher ownership to the Session.

TaskQueue operations (takeFirstIn, takeLastIn) are blocking: callers wait until a task arrives or the queue is shut down. This enables a simple client idiom:

```
  QueueingTaskDispatcher::TaskQueue TQ;
  Session S(std::make_unique<QueueingTaskDispatcher>(TQ), ...);
  S.attach(<controller access>);

  while (auto T = TQ.takeFirstIn())
    T->run();
```

>From b8c1cf81f24e7fdacd2bfcfbfd13fc466aa491ab Mon Sep 17 00:00:00 2001
From: Lang Hames <lhames at gmail.com>
Date: Wed, 8 Apr 2026 16:16:25 +1000
Subject: [PATCH] [orc-rt] Refactor QueueingTaskDispatcher to use an external
 TaskQueue.

QueueingTaskDispatcher now takes a TaskQueue by reference rather than
maintaining an internal queue. This lets API clients retain direct access
to the queue after transferring dispatcher ownership to the Session.

TaskQueue operations (takeFirstIn, takeLastIn) are blocking: callers
wait until a task arrives or the queue is shut down. This enables a
simple client idiom:

```
  QueueingTaskDispatcher::TaskQueue TQ;
  Session S(std::make_unique<QueueingTaskDispatcher>(TQ), ...);
  S.attach(<controller access>);

  while (auto T = TQ.takeFirstIn())
    T->run();
```
---
 .../include/orc-rt/QueueingTaskDispatcher.h   |  46 ++-
 .../lib/executor/QueueingTaskDispatcher.cpp   |  41 ++-
 .../unittests/QueueingTaskDispatcherTest.cpp  | 343 ++++++++----------
 3 files changed, 199 insertions(+), 231 deletions(-)

diff --git a/orc-rt/include/orc-rt/QueueingTaskDispatcher.h b/orc-rt/include/orc-rt/QueueingTaskDispatcher.h
index e1bbc8ea05765..d11b7fdf94248 100644
--- a/orc-rt/include/orc-rt/QueueingTaskDispatcher.h
+++ b/orc-rt/include/orc-rt/QueueingTaskDispatcher.h
@@ -15,6 +15,7 @@
 
 #include "orc-rt/TaskDispatcher.h"
 
+#include <condition_variable>
 #include <deque>
 #include <memory>
 #include <mutex>
@@ -29,27 +30,42 @@ namespace orc_rt {
 /// discouraged, and alternatives like ThreadPoolTaskDispatcher are preferred.
 class QueueingTaskDispatcher : public TaskDispatcher {
 public:
-  void dispatch(std::unique_ptr<Task> T) override;
-  void shutdown() override;
+  class TaskQueue {
+  public:
+    /// Append a task to the queue.
+    void addTask(std::unique_ptr<Task> T);
+
+    /// Shut down the queue. Further calls to addTask will be ignored (the task
+    /// arguments will be discarded).
+    void shutdown();
+
+    /// Take the task most recently added to the queue. Blocks until a task is
+    /// available or the dispatcher shuts down.
+    std::unique_ptr<Task> takeLastIn();
 
-  /// Take a task from the back of the queue. If there are no tasks, returns
-  /// nullptr.
-  std::unique_ptr<Task> pop_back();
+    /// Take the earliest task from the queue. Blocks until a task is available
+    /// or the dispatcher shuts down.
+    std::unique_ptr<Task> takeFirstIn();
 
-  /// Take a task from the front of the queue. If there are no tasks, returns
-  /// nullptr.
-  std::unique_ptr<Task> pop_front();
+    /// Run tasks in last-in-first-out order until the queue is empty.
+    void runLIFOUntilEmpty();
 
-  /// Run tasks in last-in-first-out order until the queue is empty.
-  void runLIFOUntilEmpty();
+    /// Run tasks in first-in-first-out order until the queue is empty.
+    void runFIFOUntilEmpty();
 
-  /// Run tasks in first-in-first-out order until the queue is empty.
-  void runFIFOUntilEmpty();
+  private:
+    std::mutex M;
+    std::condition_variable CV;
+    enum { Running, Shutdown } State = Running;
+    std::deque<std::unique_ptr<Task>> Tasks;
+  };
+
+  QueueingTaskDispatcher(TaskQueue &Q) : Q(Q) {}
+  void dispatch(std::unique_ptr<Task> T) override;
+  void shutdown() override;
 
 private:
-  std::mutex M;
-  enum { Running, Shutdown } State = Running;
-  std::deque<std::unique_ptr<Task>> Tasks;
+  TaskQueue &Q;
 };
 
 } // namespace orc_rt
diff --git a/orc-rt/lib/executor/QueueingTaskDispatcher.cpp b/orc-rt/lib/executor/QueueingTaskDispatcher.cpp
index 506010f9c4b41..65d862bfccfda 100644
--- a/orc-rt/lib/executor/QueueingTaskDispatcher.cpp
+++ b/orc-rt/lib/executor/QueueingTaskDispatcher.cpp
@@ -17,24 +17,26 @@
 
 namespace orc_rt {
 
-void QueueingTaskDispatcher::dispatch(std::unique_ptr<Task> T) {
-  std::scoped_lock<std::mutex> Lock(M);
-  if (State == Running)
-    Tasks.push_back(std::move(T));
+void QueueingTaskDispatcher::TaskQueue::addTask(std::unique_ptr<Task> T) {
+  {
+    std::scoped_lock<std::mutex> Lock(M);
+    if (State == Running)
+      Tasks.push_back(std::move(T));
+  }
+  CV.notify_one();
 }
 
-void QueueingTaskDispatcher::shutdown() {
-  std::deque<std::unique_ptr<Task>> ResidualTasks;
+void QueueingTaskDispatcher::TaskQueue::shutdown() {
   {
     std::scoped_lock<std::mutex> Lock(M);
     State = Shutdown;
-    ResidualTasks = std::move(Tasks);
   }
-  // ResidualTask destruction can run Task destructors outside the lock.
+  CV.notify_all();
 }
 
-std::unique_ptr<Task> QueueingTaskDispatcher::pop_back() {
-  std::scoped_lock<std::mutex> Lock(M);
+std::unique_ptr<Task> QueueingTaskDispatcher::TaskQueue::takeLastIn() {
+  std::unique_lock<std::mutex> Lock(M);
+  CV.wait(Lock, [&]() { return !Tasks.empty() || State == Shutdown; });
   if (Tasks.empty())
     return nullptr;
   auto T = std::move(Tasks.back());
@@ -42,8 +44,9 @@ std::unique_ptr<Task> QueueingTaskDispatcher::pop_back() {
   return T;
 }
 
-std::unique_ptr<Task> QueueingTaskDispatcher::pop_front() {
-  std::scoped_lock<std::mutex> Lock(M);
+std::unique_ptr<Task> QueueingTaskDispatcher::TaskQueue::takeFirstIn() {
+  std::unique_lock<std::mutex> Lock(M);
+  CV.wait(Lock, [&]() { return !Tasks.empty() || State == Shutdown; });
   if (Tasks.empty())
     return nullptr;
   auto T = std::move(Tasks.front());
@@ -51,14 +54,20 @@ std::unique_ptr<Task> QueueingTaskDispatcher::pop_front() {
   return T;
 }
 
-void QueueingTaskDispatcher::runLIFOUntilEmpty() {
-  while (auto T = pop_back())
+void QueueingTaskDispatcher::TaskQueue::runLIFOUntilEmpty() {
+  while (auto T = takeLastIn())
     T->run();
 }
 
-void QueueingTaskDispatcher::runFIFOUntilEmpty() {
-  while (auto T = pop_front())
+void QueueingTaskDispatcher::TaskQueue::runFIFOUntilEmpty() {
+  while (auto T = takeFirstIn())
     T->run();
 }
 
+void QueueingTaskDispatcher::dispatch(std::unique_ptr<Task> T) {
+  Q.addTask(std::move(T));
+}
+
+void QueueingTaskDispatcher::shutdown() { Q.shutdown(); }
+
 } // namespace orc_rt
diff --git a/orc-rt/unittests/QueueingTaskDispatcherTest.cpp b/orc-rt/unittests/QueueingTaskDispatcherTest.cpp
index 1528f177e47be..06d0347b444cd 100644
--- a/orc-rt/unittests/QueueingTaskDispatcherTest.cpp
+++ b/orc-rt/unittests/QueueingTaskDispatcherTest.cpp
@@ -18,331 +18,274 @@ using namespace orc_rt;
 
 namespace {
 
-TEST(QueueingTaskDispatcherTest, EmptyDispatcher) {
-  // Test that a newly created dispatcher has no tasks
-  QueueingTaskDispatcher Dispatcher;
-
-  EXPECT_EQ(Dispatcher.pop_back(), nullptr);
-  EXPECT_EQ(Dispatcher.pop_front(), nullptr);
-
-  Dispatcher.shutdown();
-}
-
 TEST(QueueingTaskDispatcherTest, BasicTaskDispatch) {
-  // Test basic task dispatching and retrieval
-  QueueingTaskDispatcher Dispatcher;
+  // Test basic task dispatching and retrieval.
+  QueueingTaskDispatcher::TaskQueue Q;
+  QueueingTaskDispatcher Dispatcher(Q);
   bool TaskRan = false;
 
   Dispatcher.dispatch(makeGenericTask([&]() { TaskRan = true; }));
+  Dispatcher.shutdown();
 
-  auto Task = Dispatcher.pop_back();
+  auto Task = Q.takeFirstIn();
   EXPECT_NE(Task, nullptr);
-
   Task->run();
   EXPECT_TRUE(TaskRan);
 
-  // Should be empty now
-  EXPECT_EQ(Dispatcher.pop_back(), nullptr);
-  EXPECT_EQ(Dispatcher.pop_front(), nullptr);
-
-  Dispatcher.shutdown();
+  // Queue is shut down and drained — should return nullptr.
+  EXPECT_EQ(Q.takeFirstIn(), nullptr);
 }
 
 TEST(QueueingTaskDispatcherTest, MultipleTasks) {
-  // Test dispatching multiple tasks
-  QueueingTaskDispatcher Dispatcher;
+  // Test dispatching and running multiple tasks.
+  QueueingTaskDispatcher::TaskQueue Q;
+  QueueingTaskDispatcher Dispatcher(Q);
   int TaskCount = 0;
   constexpr int NumTasks = 5;
 
   for (int I = 0; I < NumTasks; ++I)
     Dispatcher.dispatch(makeGenericTask([&]() { ++TaskCount; }));
+  Dispatcher.shutdown();
 
-  // Pop all tasks and run them
+  // Take and run all tasks.
   for (int I = 0; I < NumTasks; ++I) {
-    auto Task = Dispatcher.pop_back();
+    auto Task = Q.takeFirstIn();
     EXPECT_NE(Task, nullptr);
     Task->run();
   }
 
   EXPECT_EQ(TaskCount, NumTasks);
-  EXPECT_EQ(Dispatcher.pop_back(), nullptr);
-
-  Dispatcher.shutdown();
+  EXPECT_EQ(Q.takeFirstIn(), nullptr);
 }
 
-TEST(QueueingTaskDispatcherTest, PopBackLIFOOrder) {
-  // Test that pop_back retrieves tasks in LIFO (Last-In-First-Out) order
-  QueueingTaskDispatcher Dispatcher;
+TEST(QueueingTaskDispatcherTest, TakeLastInLIFOOrder) {
+  // Test that takeLastIn retrieves tasks in LIFO order.
+  QueueingTaskDispatcher::TaskQueue Q;
+  QueueingTaskDispatcher Dispatcher(Q);
   std::vector<int> ExecutionOrder;
 
-  // Dispatch tasks with different values
   for (int I = 0; I < 3; ++I)
     Dispatcher.dispatch(makeGenericTask(
         [&ExecutionOrder, I]() { ExecutionOrder.push_back(I); }));
+  Dispatcher.shutdown();
 
-  // Pop from back should give us tasks in reverse order (LIFO)
-  while (auto Task = Dispatcher.pop_back())
+  while (auto Task = Q.takeLastIn())
     Task->run();
 
-  EXPECT_EQ(ExecutionOrder.size(), 3u);
-  EXPECT_EQ(ExecutionOrder[0], 2); // Last dispatched task
+  ASSERT_EQ(ExecutionOrder.size(), 3u);
+  EXPECT_EQ(ExecutionOrder[0], 2);
   EXPECT_EQ(ExecutionOrder[1], 1);
-  EXPECT_EQ(ExecutionOrder[2], 0); // First dispatched task
-
-  Dispatcher.shutdown();
+  EXPECT_EQ(ExecutionOrder[2], 0);
 }
 
-TEST(QueueingTaskDispatcherTest, PopFrontFIFOOrder) {
-  // Test that pop_front retrieves tasks in FIFO (First-In-First-Out) order
-  QueueingTaskDispatcher Dispatcher;
+TEST(QueueingTaskDispatcherTest, TakeFirstInFIFOOrder) {
+  // Test that takeFirstIn retrieves tasks in FIFO order.
+  QueueingTaskDispatcher::TaskQueue Q;
+  QueueingTaskDispatcher Dispatcher(Q);
   std::vector<int> ExecutionOrder;
 
-  // Dispatch tasks with different values
   for (int I = 0; I < 3; ++I)
     Dispatcher.dispatch(makeGenericTask(
         [&ExecutionOrder, I]() { ExecutionOrder.push_back(I); }));
+  Dispatcher.shutdown();
 
-  // Pop from front should give us tasks in original order (FIFO)
-  while (auto Task = Dispatcher.pop_front())
+  while (auto Task = Q.takeFirstIn())
     Task->run();
 
-  EXPECT_EQ(ExecutionOrder.size(), 3u);
-  EXPECT_EQ(ExecutionOrder[0], 0); // First dispatched task
+  ASSERT_EQ(ExecutionOrder.size(), 3u);
+  EXPECT_EQ(ExecutionOrder[0], 0);
   EXPECT_EQ(ExecutionOrder[1], 1);
-  EXPECT_EQ(ExecutionOrder[2], 2); // Last dispatched task
-
-  Dispatcher.shutdown();
+  EXPECT_EQ(ExecutionOrder[2], 2);
 }
 
 TEST(QueueingTaskDispatcherTest, RunLIFOUntilEmpty) {
-  // Test the runLIFOUntilEmpty method
-  QueueingTaskDispatcher Dispatcher;
+  // Test the runLIFOUntilEmpty convenience method.
+  QueueingTaskDispatcher::TaskQueue Q;
+  QueueingTaskDispatcher Dispatcher(Q);
   std::vector<int> ExecutionOrder;
 
-  // Dispatch tasks
   for (int I = 0; I < 3; ++I)
     Dispatcher.dispatch(makeGenericTask(
         [&ExecutionOrder, I]() { ExecutionOrder.push_back(I); }));
+  Dispatcher.shutdown();
 
-  // Run all tasks in LIFO order
-  Dispatcher.runLIFOUntilEmpty();
+  Q.runLIFOUntilEmpty();
 
-  EXPECT_EQ(ExecutionOrder.size(), 3u);
-  EXPECT_EQ(ExecutionOrder[0], 2); // Last dispatched task runs first
+  ASSERT_EQ(ExecutionOrder.size(), 3u);
+  EXPECT_EQ(ExecutionOrder[0], 2);
   EXPECT_EQ(ExecutionOrder[1], 1);
-  EXPECT_EQ(ExecutionOrder[2], 0); // First dispatched task runs last
-
-  // Should be empty now
-  EXPECT_EQ(Dispatcher.pop_back(), nullptr);
-  EXPECT_EQ(Dispatcher.pop_front(), nullptr);
-
-  Dispatcher.shutdown();
+  EXPECT_EQ(ExecutionOrder[2], 0);
 }
 
 TEST(QueueingTaskDispatcherTest, RunFIFOUntilEmpty) {
-  // Test the runFIFOUntilEmpty method
-  QueueingTaskDispatcher Dispatcher;
+  // Test the runFIFOUntilEmpty convenience method.
+  QueueingTaskDispatcher::TaskQueue Q;
+  QueueingTaskDispatcher Dispatcher(Q);
   std::vector<int> ExecutionOrder;
 
-  // Dispatch tasks
   for (int I = 0; I < 3; ++I)
     Dispatcher.dispatch(makeGenericTask(
         [&ExecutionOrder, I]() { ExecutionOrder.push_back(I); }));
+  Dispatcher.shutdown();
 
-  // Run all tasks in FIFO order
-  Dispatcher.runFIFOUntilEmpty();
+  Q.runFIFOUntilEmpty();
 
-  EXPECT_EQ(ExecutionOrder.size(), 3u);
-  EXPECT_EQ(ExecutionOrder[0], 0); // First dispatched task runs first
+  ASSERT_EQ(ExecutionOrder.size(), 3u);
+  EXPECT_EQ(ExecutionOrder[0], 0);
   EXPECT_EQ(ExecutionOrder[1], 1);
-  EXPECT_EQ(ExecutionOrder[2], 2); // Last dispatched task runs last
-
-  // Should be empty now
-  EXPECT_EQ(Dispatcher.pop_back(), nullptr);
-  EXPECT_EQ(Dispatcher.pop_front(), nullptr);
-
-  Dispatcher.shutdown();
+  EXPECT_EQ(ExecutionOrder[2], 2);
 }
 
-TEST(QueueingTaskDispatcherTest, MixedPopOperations) {
-  // Test mixing pop_front and pop_back operations
-  QueueingTaskDispatcher Dispatcher;
+TEST(QueueingTaskDispatcherTest, MixedTakeOperations) {
+  // Test mixing takeFirstIn and takeLastIn.
+  QueueingTaskDispatcher::TaskQueue Q;
+  QueueingTaskDispatcher Dispatcher(Q);
   std::vector<int> ExecutionOrder;
 
-  // Dispatch tasks 0, 1, 2
+  // Dispatch tasks 0, 1, 2.
   for (int I = 0; I < 3; ++I)
     Dispatcher.dispatch(makeGenericTask(
         [&ExecutionOrder, I]() { ExecutionOrder.push_back(I); }));
+  Dispatcher.shutdown();
 
-  // Pop from back (should get task 2)
-  auto Task1 = Dispatcher.pop_back();
-  EXPECT_NE(Task1, nullptr);
+  // takeLastIn should get task 2.
+  auto Task1 = Q.takeLastIn();
+  ASSERT_NE(Task1, nullptr);
   Task1->run();
 
-  // Pop from front (should get task 0)
-  auto Task2 = Dispatcher.pop_front();
-  EXPECT_NE(Task2, nullptr);
+  // takeFirstIn should get task 0.
+  auto Task2 = Q.takeFirstIn();
+  ASSERT_NE(Task2, nullptr);
   Task2->run();
 
-  // Pop from back again (should get task 1)
-  auto Task3 = Dispatcher.pop_back();
-  EXPECT_NE(Task3, nullptr);
+  // takeLastIn should get task 1 (only one left).
+  auto Task3 = Q.takeLastIn();
+  ASSERT_NE(Task3, nullptr);
   Task3->run();
 
-  // Should be empty now
-  EXPECT_EQ(Dispatcher.pop_back(), nullptr);
-  EXPECT_EQ(Dispatcher.pop_front(), nullptr);
-
-  EXPECT_EQ(ExecutionOrder.size(), 3u);
-  EXPECT_EQ(ExecutionOrder[0], 2); // Last task (pop_back)
-  EXPECT_EQ(ExecutionOrder[1], 0); // First task (pop_front)
-  EXPECT_EQ(ExecutionOrder[2], 1); // Middle task (pop_back)
+  EXPECT_EQ(Q.takeFirstIn(), nullptr);
 
-  Dispatcher.shutdown();
+  ASSERT_EQ(ExecutionOrder.size(), 3u);
+  EXPECT_EQ(ExecutionOrder[0], 2);
+  EXPECT_EQ(ExecutionOrder[1], 0);
+  EXPECT_EQ(ExecutionOrder[2], 1);
 }
 
-TEST(QueueingTaskDispatcherTest, ShutdownWithPendingTasks) {
-  // Test shutdown behavior when tasks remain in queue
-  QueueingTaskDispatcher Dispatcher;
+TEST(QueueingTaskDispatcherTest, ShutdownDrainsRemainingTasks) {
+  // Verify that tasks dispatched before shutdown can still be taken.
+  QueueingTaskDispatcher::TaskQueue Q;
+  QueueingTaskDispatcher Dispatcher(Q);
+  int TaskCount = 0;
 
-  // Dispatch some tasks but don't run them
   for (int I = 0; I < 3; ++I)
-    Dispatcher.dispatch(makeGenericTask([]() {
-      // These tasks won't be executed in this test
-    }));
+    Dispatcher.dispatch(makeGenericTask([&]() { ++TaskCount; }));
 
-  // Should be able to shutdown even with pending tasks
   Dispatcher.shutdown();
 
-  // After shutdown, no tasks should be available
-  EXPECT_EQ(Dispatcher.pop_back(), nullptr);
-  EXPECT_EQ(Dispatcher.pop_front(), nullptr);
+  // All pre-shutdown tasks should still be available.
+  while (auto Task = Q.takeFirstIn())
+    Task->run();
+
+  EXPECT_EQ(TaskCount, 3);
 }
 
 TEST(QueueingTaskDispatcherTest, DispatchAfterShutdown) {
-  // Test behavior of dispatch after shutdown
-  QueueingTaskDispatcher Dispatcher;
+  // Tasks dispatched after shutdown should be discarded.
+  QueueingTaskDispatcher::TaskQueue Q;
+  QueueingTaskDispatcher Dispatcher(Q);
   bool TaskRan = false;
 
   Dispatcher.shutdown();
 
-  // Dispatch should work even after shutdown (tasks are queued)
   Dispatcher.dispatch(makeGenericTask([&]() { TaskRan = true; }));
 
-  // Task should not be retrievable
-  EXPECT_EQ(Dispatcher.pop_back(), nullptr);
-  EXPECT_EQ(Dispatcher.pop_front(), nullptr);
-
+  EXPECT_EQ(Q.takeFirstIn(), nullptr);
   EXPECT_FALSE(TaskRan);
 }
 
-TEST(QueueingTaskDispatcherTest, RunMethodsOnEmptyDispatcher) {
-  // Test that run methods work correctly on empty dispatcher
-  QueueingTaskDispatcher Dispatcher;
+TEST(QueueingTaskDispatcherTest, TakeBlocksUntilTaskAvailable) {
+  // Verify that takeFirstIn blocks on an empty queue until a task arrives.
+  QueueingTaskDispatcher::TaskQueue Q;
+  QueueingTaskDispatcher Dispatcher(Q);
+  std::atomic<bool> TaskTaken = false;
 
-  // These should not crash or hang
-  Dispatcher.runLIFOUntilEmpty();
-  Dispatcher.runFIFOUntilEmpty();
+  std::thread Consumer([&]() {
+    auto Task = Q.takeFirstIn();
+    TaskTaken = true;
+    EXPECT_NE(Task, nullptr);
+    Task->run();
+  });
 
-  Dispatcher.shutdown();
-}
+  // Give the consumer a moment to block.
+  std::this_thread::sleep_for(std::chrono::milliseconds(50));
+  EXPECT_FALSE(TaskTaken);
 
-TEST(QueueingTaskDispatcherTest, InterleaveDispatchAndPop) {
-  // Test interleaving dispatch and pop operations
-  QueueingTaskDispatcher Dispatcher;
-  std::vector<int> ExecutionOrder;
+  // Dispatching a task should unblock the consumer.
+  std::atomic<bool> TaskRan = false;
+  Dispatcher.dispatch(makeGenericTask([&]() { TaskRan = true; }));
 
-  // Dispatch task 0
-  Dispatcher.dispatch(
-      makeGenericTask([&ExecutionOrder]() { ExecutionOrder.push_back(0); }));
+  Consumer.join();
 
-  // Pop and run task 0
-  auto Task1 = Dispatcher.pop_back();
-  EXPECT_NE(Task1, nullptr);
-  Task1->run();
+  EXPECT_TRUE(TaskTaken);
+  EXPECT_TRUE(TaskRan);
 
-  // Dispatch tasks 1 and 2
-  for (int I = 1; I < 3; ++I)
-    Dispatcher.dispatch(makeGenericTask(
-        [&ExecutionOrder, I]() { ExecutionOrder.push_back(I); }));
+  Dispatcher.shutdown();
+}
 
-  // Pop and run remaining tasks
-  while (auto Task = Dispatcher.pop_front())
-    Task->run();
+TEST(QueueingTaskDispatcherTest, TakeReturnsNullptrOnShutdown) {
+  // Verify that a blocked take returns nullptr when the queue is shut down.
+  QueueingTaskDispatcher::TaskQueue Q;
+  QueueingTaskDispatcher Dispatcher(Q);
+  std::atomic<bool> TakeReturned = false;
 
-  EXPECT_EQ(ExecutionOrder.size(), 3u);
-  EXPECT_EQ(ExecutionOrder[0], 0); // First task executed immediately
-  EXPECT_EQ(ExecutionOrder[1], 1); // Second task (FIFO order)
-  EXPECT_EQ(ExecutionOrder[2], 2); // Third task (FIFO order)
+  std::thread Consumer([&]() {
+    auto Task = Q.takeFirstIn();
+    EXPECT_EQ(Task, nullptr);
+    TakeReturned.store(true);
+  });
 
+  // Give the consumer a moment to block.
+  std::this_thread::sleep_for(std::chrono::milliseconds(50));
+  EXPECT_FALSE(TakeReturned);
+
+  // Shutting down should unblock the consumer with nullptr.
   Dispatcher.shutdown();
+
+  Consumer.join();
+  EXPECT_TRUE(TakeReturned);
 }
 
 TEST(QueueingTaskDispatcherTest, ThreadSafety) {
-  // Test thread safety of the dispatcher
-  QueueingTaskDispatcher Dispatcher;
-  constexpr int NumThreads = 4;
-  constexpr int TasksPerThread = 25;
+  // Test thread safety with concurrent dispatch and take.
+  QueueingTaskDispatcher::TaskQueue Q;
+  QueueingTaskDispatcher Dispatcher(Q);
+  constexpr int NumProducers = 4;
+  constexpr int TasksPerProducer = 25;
+  constexpr int TotalTasks = NumProducers * TasksPerProducer;
   std::atomic<int> TasksCompleted = 0;
 
-  std::vector<std::thread> DispatchThreads;
-  std::vector<std::thread> PopThreads;
-
-  // Create threads that dispatch tasks
-  for (int ThreadId = 0; ThreadId < NumThreads; ++ThreadId) {
-    DispatchThreads.emplace_back([&]() {
-      for (int I = 0; I < TasksPerThread; ++I) {
+  // Producer threads dispatch tasks.
+  std::vector<std::thread> Producers;
+  for (int I = 0; I < NumProducers; ++I) {
+    Producers.emplace_back([&]() {
+      for (int J = 0; J < TasksPerProducer; ++J)
         Dispatcher.dispatch(makeGenericTask([&]() { ++TasksCompleted; }));
-      }
     });
   }
 
-  // Create threads that pop and run tasks
-  for (int ThreadId = 0; ThreadId < NumThreads; ++ThreadId) {
-    PopThreads.emplace_back([&]() {
-      for (int I = 0; I < TasksPerThread; ++I) {
-        std::unique_ptr<Task> Task;
-
-        // Keep trying to pop a task
-        while (!Task) {
-          Task = Dispatcher.pop_back();
-          if (!Task) {
-            std::this_thread::yield();
-          }
-        }
-
-        Task->run();
-      }
-    });
-  }
+  // Consumer thread takes and runs tasks until shutdown.
+  std::thread Consumer([&]() {
+    while (auto Task = Q.takeFirstIn())
+      Task->run();
+  });
 
-  // Wait for all threads to complete
-  for (auto &T : DispatchThreads)
-    T.join();
-  for (auto &T : PopThreads)
+  // Wait for all producers to finish, then shut down.
+  for (auto &T : Producers)
     T.join();
-
-  EXPECT_EQ(TasksCompleted.load(), NumThreads * TasksPerThread);
-
   Dispatcher.shutdown();
-}
-
-TEST(QueueingTaskDispatcherTest, LargeNumberOfTasks) {
-  // Test with a large number of tasks to ensure no performance issues
-  QueueingTaskDispatcher Dispatcher;
-  constexpr int NumTasks = 1000;
-  int TasksRun = 0;
-
-  // Dispatch many tasks
-  for (int I = 0; I < NumTasks; ++I)
-    Dispatcher.dispatch(makeGenericTask([&TasksRun]() { ++TasksRun; }));
 
-  // Run all tasks using FIFO
-  Dispatcher.runFIFOUntilEmpty();
-
-  EXPECT_EQ(TasksRun, NumTasks);
-
-  Dispatcher.shutdown();
+  Consumer.join();
+  EXPECT_EQ(TasksCompleted, TotalTasks);
 }
 
 } // end anonymous namespace



More information about the llvm-commits mailing list