[orc-rt] Introduce Task and TaskDispatcher APIs and implementations. (llvm#168514)

Introduces the Task and TaskDispatcher interfaces (TaskDispatcher.h),
ThreadPoolTaskDispatcher implementation (ThreadPoolTaskDispatch.h), and
updates Session to include a TaskDispatcher instance that can be used to
run tasks.

TaskDispatcher's introduction is motivated by the need to handle calls
to JIT'd code initiated from the controller process: Incoming calls will
be wrapped in Tasks and dispatched. Session shutdown will wait on
TaskDispatcher shutdown, ensuring that all Tasks are run or destroyed
prior to the Session being destroyed.

Author: lhames

orc-rt/include/CMakeLists.txt

@@ -22,6 +22,8 @@ set(ORC_RT_HEADERS
     orc-rt/SPSMemoryFlags.h
     orc-rt/SPSWrapperFunction.h
     orc-rt/SPSWrapperFunctionBuffer.h
+    orc-rt/TaskDispatcher.h
+    orc-rt/ThreadPoolTaskDispatcher.h
     orc-rt/WrapperFunction.h
     orc-rt/bind.h
     orc-rt/bit.h

orc-rt/include/orc-rt/Session.h

@@ -15,10 +15,12 @@
 
 #include "orc-rt/Error.h"
 #include "orc-rt/ResourceManager.h"
+#include "orc-rt/TaskDispatcher.h"
 #include "orc-rt/move_only_function.h"
 
 #include "orc-rt-c/CoreTypes.h"
 
+#include <condition_variable>
 #include <memory>
 #include <mutex>
 #include <vector>
@@ -39,7 +41,10 @@ class Session {
   ///
   /// Note that entry into the reporter is not synchronized: it may be
   /// called from multiple threads concurrently.
-  Session(ErrorReporterFn ReportError) : ReportError(std::move(ReportError)) {}
+  Session(std::unique_ptr<TaskDispatcher> Dispatcher,
+          ErrorReporterFn ReportError)
+      : Dispatcher(std::move(Dispatcher)), ReportError(std::move(ReportError)) {
+  }
 
   // Sessions are not copyable or moveable.
   Session(const Session &) = delete;
@@ -49,6 +54,9 @@ class Session {
 
   ~Session();
 
+  /// Dispatch a task using the Session's TaskDispatcher.
+  void dispatch(std::unique_ptr<Task> T) { Dispatcher->dispatch(std::move(T)); }
+
   /// Report an error via the ErrorReporter function.
   void reportError(Error Err) { ReportError(std::move(Err)); }
 
@@ -67,12 +75,21 @@ class Session {
   }
 
 private:
-  void shutdownNext(OnShutdownCompleteFn OnShutdownComplete, Error Err,
+  void shutdownNext(Error Err,
                     std::vector<std::unique_ptr<ResourceManager>> RemainingRMs);
 
-  std::mutex M;
+  void shutdownComplete();
+
+  std::unique_ptr<TaskDispatcher> Dispatcher;
   ErrorReporterFn ReportError;
+
+  enum class SessionState { Running, ShuttingDown, Shutdown };
+
+  std::mutex M;
+  SessionState State = SessionState::Running;
+  std::condition_variable StateCV;
   std::vector<std::unique_ptr<ResourceManager>> ResourceMgrs;
+  std::vector<OnShutdownCompleteFn> ShutdownCallbacks;
 };
 
 inline orc_rt_SessionRef wrap(Session *S) noexcept {

orc-rt/include/orc-rt/TaskDispatcher.h

@@ -0,0 +1,64 @@
+//===----------- TaskDispatcher.h - Task dispatch utils ---------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Task and TaskDispatcher classes.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef ORC_RT_TASKDISPATCHER_H
+#define ORC_RT_TASKDISPATCHER_H
+
+#include "orc-rt/RTTI.h"
+
+#include <memory>
+#include <utility>
+
+namespace orc_rt {
+
+/// Represents an abstract task to be run.
+class Task : public RTTIExtends<Task, RTTIRoot> {
+public:
+  virtual ~Task();
+  virtual void run() = 0;
+};
+
+/// Base class for generic tasks.
+class GenericTask : public RTTIExtends<GenericTask, Task> {};
+
+/// Generic task implementation.
+template <typename FnT> class GenericTaskImpl : public GenericTask {
+public:
+  GenericTaskImpl(FnT &&Fn) : Fn(std::forward<FnT>(Fn)) {}
+  void run() override { Fn(); }
+
+private:
+  FnT Fn;
+};
+
+/// Create a generic task from a function object.
+template <typename FnT> std::unique_ptr<GenericTask> makeGenericTask(FnT &&Fn) {
+  return std::make_unique<GenericTaskImpl<std::decay_t<FnT>>>(
+      std::forward<FnT>(Fn));
+}
+
+/// Abstract base for classes that dispatch Tasks.
+class TaskDispatcher {
+public:
+  virtual ~TaskDispatcher();
+
+  /// Run the given task.
+  virtual void dispatch(std::unique_ptr<Task> T) = 0;
+
+  /// Called by Session. Should cause further dispatches to be rejected, and
+  /// wait until all previously dispatched tasks have completed.
+  virtual void shutdown() = 0;
+};
+
+} // End namespace orc_rt
+
+#endif // ORC_RT_TASKDISPATCHER_H

orc-rt/include/orc-rt/ThreadPoolTaskDispatcher.h

@@ -0,0 +1,48 @@
+//===--- ThreadPoolTaskDispatcher.h - Run tasks in thread pool --*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// ThreadPoolTaskDispatcher implementation.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef ORC_RT_THREADPOOLTASKDISPATCHER_H
+#define ORC_RT_THREADPOOLTASKDISPATCHER_H
+
+#include "orc-rt/TaskDispatcher.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+#include <vector>
+
+namespace orc_rt {
+
+/// Thread-pool based TaskDispatcher.
+///
+/// Will spawn NumThreads threads to run dispatched Tasks.
+class ThreadPoolTaskDispatcher : public TaskDispatcher {
+public:
+  ThreadPoolTaskDispatcher(size_t NumThreads);
+  ~ThreadPoolTaskDispatcher() override;
+  void dispatch(std::unique_ptr<Task> T) override;
+  void shutdown() override;
+
+private:
+  void taskLoop();
+
+  std::vector<std::thread> Threads;
+
+  std::mutex M;
+  bool AcceptingTasks = true;
+  std::condition_variable CV;
+  std::vector<std::unique_ptr<Task>> PendingTasks;
+};
+
+} // End namespace orc_rt
+
+#endif // ORC_RT_THREADPOOLTASKDISPATCHER_H

orc-rt/lib/executor/CMakeLists.txt

@@ -4,6 +4,8 @@ set(files
   RTTI.cpp
   Session.cpp
   SimpleNativeMemoryMap.cpp
+  TaskDispatcher.cpp
+  ThreadPoolTaskDispatcher.cpp
   )
 
 add_library(orc-rt-executor STATIC ${files})

orc-rt/lib/executor/Session.cpp

@@ -12,8 +12,6 @@
 
 #include "orc-rt/Session.h"
 
-#include <future>
-
 namespace orc_rt {
 
 Session::~Session() { waitForShutdown(); }
@@ -23,38 +21,62 @@ void Session::shutdown(OnShutdownCompleteFn OnShutdownComplete) {
 
   {
     std::scoped_lock<std::mutex> Lock(M);
+    ShutdownCallbacks.push_back(std::move(OnShutdownComplete));
+
+    // If somebody else has already called shutdown then there's nothing further
+    // for us to do here.
+    if (State >= SessionState::ShuttingDown)
+      return;
+
+    State = SessionState::ShuttingDown;
     std::swap(ResourceMgrs, ToShutdown);
   }
 
-  shutdownNext(std::move(OnShutdownComplete), Error::success(),
-               std::move(ToShutdown));
+  shutdownNext(Error::success(), std::move(ToShutdown));
 }
 
 void Session::waitForShutdown() {
-  std::promise<void> P;
-  auto F = P.get_future();
-
-  shutdown([P = std::move(P)]() mutable { P.set_value(); });
-
-  F.wait();
+  shutdown([]() {});
+  std::unique_lock<std::mutex> Lock(M);
+  StateCV.wait(Lock, [&]() { return State == SessionState::Shutdown; });
 }
 
 void Session::shutdownNext(
-    OnShutdownCompleteFn OnComplete, Error Err,
-    std::vector<std::unique_ptr<ResourceManager>> RemainingRMs) {
+    Error Err, std::vector<std::unique_ptr<ResourceManager>> RemainingRMs) {
   if (Err)
     reportError(std::move(Err));
 
   if (RemainingRMs.empty())
-    return OnComplete();
+    return shutdownComplete();
 
   auto NextRM = std::move(RemainingRMs.back());
   RemainingRMs.pop_back();
-  NextRM->shutdown([this, RemainingRMs = std::move(RemainingRMs),
-                    OnComplete = std::move(OnComplete)](Error Err) mutable {
-    shutdownNext(std::move(OnComplete), std::move(Err),
-                 std::move(RemainingRMs));
-  });
+  NextRM->shutdown(
+      [this, RemainingRMs = std::move(RemainingRMs)](Error Err) mutable {
+        shutdownNext(std::move(Err), std::move(RemainingRMs));
+      });
+}
+
+void Session::shutdownComplete() {
+
+  std::unique_ptr<TaskDispatcher> TmpDispatcher;
+  std::vector<OnShutdownCompleteFn> TmpShutdownCallbacks;
+  {
+    std::lock_guard<std::mutex> Lock(M);
+    TmpDispatcher = std::move(Dispatcher);
+    TmpShutdownCallbacks = std::move(ShutdownCallbacks);
+  }
+
+  TmpDispatcher->shutdown();
+
+  for (auto &OnShutdownComplete : TmpShutdownCallbacks)
+    OnShutdownComplete();
+
+  {
+    std::lock_guard<std::mutex> Lock(M);
+    State = SessionState::Shutdown;
+  }
+  StateCV.notify_all();
 }
 
 } // namespace orc_rt

orc-rt/lib/executor/TaskDispatcher.cpp

@@ -0,0 +1,20 @@
+//===- TaskDispatch.cpp ---------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Contains the implementation of APIs in the orc-rt/TaskDispatch.h header.
+//
+//===----------------------------------------------------------------------===//
+
+#include "orc-rt/TaskDispatcher.h"
+
+namespace orc_rt {
+
+Task::~Task() = default;
+TaskDispatcher::~TaskDispatcher() = default;
+
+} // namespace orc_rt

orc-rt/lib/executor/ThreadPoolTaskDispatcher.cpp

@@ -0,0 +1,70 @@
+//===- ThreadPoolTaskDispatch.cpp -----------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Contains the implementation of APIs in the orc-rt/ThreadPoolTaskDispatch.h
+// header.
+//
+//===----------------------------------------------------------------------===//
+
+#include "orc-rt/ThreadPoolTaskDispatcher.h"
+
+#include <cassert>
+
+namespace orc_rt {
+
+ThreadPoolTaskDispatcher::~ThreadPoolTaskDispatcher() {
+  assert(!AcceptingTasks && "shutdown was not run");
+}
+
+ThreadPoolTaskDispatcher::ThreadPoolTaskDispatcher(size_t NumThreads) {
+  Threads.reserve(NumThreads);
+  for (size_t I = 0; I < NumThreads; ++I)
+    Threads.emplace_back([this]() { taskLoop(); });
+}
+
+void ThreadPoolTaskDispatcher::dispatch(std::unique_ptr<Task> T) {
+  {
+    std::scoped_lock<std::mutex> Lock(M);
+    if (!AcceptingTasks)
+      return;
+    PendingTasks.push_back(std::move(T));
+  }
+  CV.notify_one();
+}
+
+void ThreadPoolTaskDispatcher::shutdown() {
+  {
+    std::scoped_lock<std::mutex> Lock(M);
+    assert(AcceptingTasks && "ThreadPoolTaskDispatcher already shut down?");
+    AcceptingTasks = false;
+  }
+  CV.notify_all();
+  for (auto &Thread : Threads)
+    Thread.join();
+}
+
+void ThreadPoolTaskDispatcher::taskLoop() {
+  while (true) {
+    std::unique_ptr<Task> T;
+    {
+      std::unique_lock<std::mutex> Lock(M);
+      CV.wait(Lock,
+              [this]() { return !PendingTasks.empty() || !AcceptingTasks; });
+
+      if (!AcceptingTasks && PendingTasks.empty())
+        return;
+
+      T = std::move(PendingTasks.back());
+      PendingTasks.pop_back();
+    }
+
+    T->run();
+  }
+}
+
+} // namespace orc_rt

orc-rt/unittests/CMakeLists.txt

@@ -31,6 +31,7 @@ add_orc_rt_unittest(CoreTests
   SPSMemoryFlagsTest.cpp
   SPSWrapperFunctionTest.cpp
   SPSWrapperFunctionBufferTest.cpp
+  ThreadPoolTaskDispatcherTest.cpp
   WrapperFunctionBufferTest.cpp
   bind-test.cpp
   bit-test.cpp