[OrcJIT] distinguish normal and abrupt TaskDispatcher shutdown

Add a `run_to_complete` complement to `shutdown`, which does nearly the
same operations, but does not cancel future / pending work. The
distinction is often mainly academic, but split for clarity of intent of
each caller.

Annotate each and every tests that assumes that any previously scheduled
tasks have finished running with a call of `run_to_completion` before
reading the variables written by those tasks. This more clearly
emphasizes the expectations of each of these tests and makes them robust
against data-races or changes to the default scheduler.

These should have no impact currently, but have been tested and shown to
pass with a scheduler which defers all work (as if being run on a
thread) until an explicit synchronization point "acquires" that state,
demonstrating that this is each and every test that would be affected by
changes to the default scheduler.

Author: vtjnash

llvm/include/llvm/ExecutionEngine/Orc/TaskDispatch.h

@@ -134,8 +134,14 @@ class LLVM_ABI TaskDispatcher {
     dispatch(std::move(T));
   }
 
-  /// Called by ExecutionSession. Waits until all tasks have completed.
-  virtual void shutdown() = 0;
+  /// The TaskDispatcher may be reused immediatly afterwards.
+  void run_to_complete() { run(false); }
+
+  /// Called by ExecutionSession. Halts all in-progress work as soon as
+  /// possible. May cause deadlocks since promises will not be set, so this
+  /// should only be used immediately before exiting.
+  /// The TaskDispatcher should not be reused afterwards.
+  void shutdown() { run(true); }
 
 protected:
   friend class future_base;
@@ -147,6 +153,9 @@ class LLVM_ABI TaskDispatcher {
   /// Notify all task dispatchers that a future with have_waiter became ready
   LLVM_ABI static void notifyWaiters();
 
+  /// Called by ExecutionSession. Waits until all tasks have completed.
+  virtual void run(bool cancel) = 0;
+
 #if LLVM_ENABLE_THREADS
   /// Shared synchronization primitives for all dispatchers
   static std::mutex DispatchMutex;
@@ -163,9 +172,9 @@ class LLVM_ABI InPlaceTaskDispatcher : public TaskDispatcher {
 public:
   void dispatch(std::unique_ptr<Task> T) override;
   void dispatch_elsewhere(std::unique_ptr<Task> T) override;
-  void shutdown() override;
 
 private:
+  void run(bool cancel) override;
   void work_until(future_base &F) override;
 
   /// C++ does not support non-static thread_local variables, so this needs to
@@ -188,9 +197,9 @@ class LLVM_ABI DynamicThreadPoolTaskDispatcher : public TaskDispatcher {
       : MaxMaterializationThreads(MaxMaterializationThreads) {}
 
   void dispatch(std::unique_ptr<Task> T) override;
-  void shutdown() override;
 
 private:
+  void run(bool cancel) override;
   void work_until(future_base &F) override;
   bool canRunMaterializationTaskNow();
   bool canRunIdleTaskNow();

llvm/lib/ExecutionEngine/Orc/TaskDispatch.cpp

@@ -58,7 +58,7 @@ void InPlaceTaskDispatcher::dispatch_elsewhere(std::unique_ptr<Task> T) {
   FutureReadyCV.notify_one();
 }
 
-void InPlaceTaskDispatcher::shutdown() {
+void InPlaceTaskDispatcher::run(bool cancel) {
   // Keep processing until no tasks belonging to this dispatcher remain
   while (true) {
     // Check if any task belongs to this dispatcher
@@ -212,9 +212,10 @@ void DynamicThreadPoolTaskDispatcher::dispatch(std::unique_ptr<Task> T) {
   }).detach();
 }
 
-void DynamicThreadPoolTaskDispatcher::shutdown() {
+void DynamicThreadPoolTaskDispatcher::run(bool cancel) {
   std::unique_lock<std::mutex> Lock(DispatchMutex);
-  Shutdown = true;
+  if (cancel)
+    Shutdown = true;
   OutstandingCV.wait(Lock, [this]() { return Outstanding == 0; });
 }