[llvm] llvm-jitlink fixes (PR #122067)

[Lang Hames via llvm-commits]

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

None

>From bec1e787368f962d822540c5a96b5a49e67da84a Mon Sep 17 00:00:00 2001
From: Lang Hames <lhames at gmail.com>
Date: Wed, 8 Jan 2025 08:09:12 +0000
Subject: [PATCH 1/2] [llvm-jitlink] Shut down the session on an error return
 path.

Ensures cleanup of task dispatcher threads. This may address some of the
nondeterministic failures seen in llvm-jitlink regression tests recently.
---
 llvm/tools/llvm-jitlink/llvm-jitlink.cpp | 1 +
 1 file changed, 1 insertion(+)

diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
index 431b86a27e16b1..963c36322c8ab8 100644
--- a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
+++ b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
@@ -2562,6 +2562,7 @@ int main(int argc, char *argv[]) {
     if (Timers)
       Timers->JITLinkTG.printAll(errs());
     reportLLVMJITLinkError(EntryPoint.takeError());
+    ExitOnErr(S->ES.endSession());
     exit(1);
   }
 

>From 5ea1b1ff8733f76f385ee2dfe07bbd49fbbe353e Mon Sep 17 00:00:00 2001
From: Lang Hames <lhames at gmail.com>
Date: Wed, 8 Jan 2025 08:12:11 +0000
Subject: [PATCH 2/2] [ORC] Fix task cleanup during
 DynamicThreadPoolTaskDispatcher::shutdown.

Dispatcher shutdown must ensure that any dispatched Tasks have completed *and
been destroyed* prior to returning, so that the ExecutionSession (which Tasks
may access) can be safely destroyed.

DynamicThreadPoolTaskDispatcher::dispatch was holding a unique pointer to the
most recent Task past the point where it notified the shutdown method (via
OutstandingCV) that the thread was done. This could lead to the
ExecutionSession being destroyed before the Task, resulting in use-after-free
style errors (e.g. assertions about dangling SymbolStringPtrs in the
SymbolStringPool destructor).

The solution is just to destroy the Task object immediately after running
it.

This patch also updates the dispatch method to reject any Task dispatched after
the shutdown flag is raised.
---
 .../llvm/ExecutionEngine/Orc/TaskDispatch.h       |  2 +-
 llvm/lib/ExecutionEngine/Orc/TaskDispatch.cpp     | 15 +++++++++++++--
 2 files changed, 14 insertions(+), 3 deletions(-)

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