[llvm] 8312876 - [ORC] Fix Task cleanup during DynamicThreadPoolTaskDispatcher::shutdown.

[Lang Hames via llvm-commits]

Author: Lang Hames
Date: 2025-01-09T00:46:05Z
New Revision: 831287620567559e7078cb7f4cd1962d35c49893

URL: https://github.com/llvm/llvm-project/commit/831287620567559e7078cb7f4cd1962d35c49893
DIFF: https://github.com/llvm/llvm-project/commit/831287620567559e7078cb7f4cd1962d35c49893.diff

LOG: [ORC] Fix Task cleanup during DynamicThreadPoolTaskDispatcher::shutdown.

Threads created by DynamicThreadPoolTaskDispatcher::dispatch had been holding a
unique_ptr to the most recent Task, meaning that the Task would be destroyed
when the thread object was destroyed, but this would happen *after* the thread
signaled the Dispatcher that it was finished. This could cause
DynamicThreadPoolTaskDispatcher::shutdown to return (and consequently
ExecutionSession to be destroyed) before all Tasks were destroyed, with Task
destructors accessing ExecutionSession and related objects after they were
freed.

The fix is to reset the Task pointer immediately after it is run to trigger
cleanup, *then* (if there are no other tasks to run) signal the Dispatcher that
the thread is finished.

This patch also updates DynamicThreadPoolTaskDispatcher::dispatch to reject any
new Tasks dispatched after DynamicThreadPoolTaskDispatcher::shutdown is called.

Added: 
    

Modified: 
    llvm/include/llvm/ExecutionEngine/Orc/TaskDispatch.h
    llvm/lib/ExecutionEngine/Orc/TaskDispatch.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ExecutionEngine/Orc/TaskDispatch.h b/llvm/include/llvm/ExecutionEngine/Orc/TaskDispatch.h
index 8c65677aae25a4..d7939864fd609b 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/TaskDispatch.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/TaskDispatch.h
@@ -122,7 +122,7 @@ class DynamicThreadPoolTaskDispatcher : public TaskDispatcher {
   void shutdown() override;
 private:
   std::mutex DispatchMutex;
-  bool Running = true;
+  bool Shutdown = false;
   size_t Outstanding = 0;
   std::condition_variable OutstandingCV;
 

diff  --git a/llvm/lib/ExecutionEngine/Orc/TaskDispatch.cpp b/llvm/lib/ExecutionEngine/Orc/TaskDispatch.cpp
index fbe4b093b0c642..1af17e85220db5 100644
--- a/llvm/lib/ExecutionEngine/Orc/TaskDispatch.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/TaskDispatch.cpp
@@ -31,6 +31,10 @@ void DynamicThreadPoolTaskDispatcher::dispatch(std::unique_ptr<Task> T) {
   {
     std::lock_guard<std::mutex> Lock(DispatchMutex);
 
+    // Reject new tasks if they're dispatched after a call to shutdown.
+    if (Shutdown)
+      return;
+
     if (IsMaterializationTask) {
 
       // If this is a materialization task and there are too many running
@@ -54,6 +58,14 @@ void DynamicThreadPoolTaskDispatcher::dispatch(std::unique_ptr<Task> T) {
       // Run the task.
       T->run();
 
+      // Reset the task to free any resources. We need this to happen *before*
+      // we notify anyone (via Outstanding) that this thread is done to ensure
+      // that we don't proceed with JIT shutdown while still holding resources.
+      // (E.g. this was causing "Dangling SymbolStringPtr" assertions).
+      T.reset();
+
+      // Check the work queue state and either proceed with the next task or
+      // end this thread.
       std::lock_guard<std::mutex> Lock(DispatchMutex);
       if (!MaterializationTaskQueue.empty()) {
         // If there are any materialization tasks running then steal that work.
@@ -64,7 +76,6 @@ void DynamicThreadPoolTaskDispatcher::dispatch(std::unique_ptr<Task> T) {
           IsMaterializationTask = true;
         }
       } else {
-        // Otherwise decrement work counters.
         if (IsMaterializationTask)
           --NumMaterializationThreads;
         --Outstanding;
@@ -78,7 +89,7 @@ void DynamicThreadPoolTaskDispatcher::dispatch(std::unique_ptr<Task> T) {
 
 void DynamicThreadPoolTaskDispatcher::shutdown() {
   std::unique_lock<std::mutex> Lock(DispatchMutex);
-  Running = false;
+  Shutdown = true;
   OutstandingCV.wait(Lock, [this]() { return Outstanding == 0; });
 }
 #endif