[llvm] 7bd481d - [ORC] Add ExecutionSession::removeJITDylibs (plural), use it in endSession.

[Lang Hames via llvm-commits]

Author: Lang Hames
Date: 2023-07-30T08:51:42-07:00
New Revision: 7bd481d9afc83a99fbb5d0a92484ab9993942784

URL: https://github.com/llvm/llvm-project/commit/7bd481d9afc83a99fbb5d0a92484ab9993942784
DIFF: https://github.com/llvm/llvm-project/commit/7bd481d9afc83a99fbb5d0a92484ab9993942784.diff

LOG: [ORC] Add ExecutionSession::removeJITDylibs (plural), use it in endSession.

The ExecutionSession::removeJITDylibs operation will remove all JITDylibs in
the given list (i.e. first clear them, then remove them from the session).

ExecutionSession::endSession is updated to remove JITDylibs rather than just
clearing them. This prevents new code from being added to any JITDylib once
endSession has been called.

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ExecutionEngine/Orc/Core.h b/llvm/include/llvm/ExecutionEngine/Orc/Core.h
index 9554a219594803..2c755968c45b5f 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/Core.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/Core.h
@@ -1465,17 +1465,29 @@ class ExecutionSession {
   /// If no Platform is attached this call is equivalent to createBareJITDylib.
   Expected<JITDylib &> createJITDylib(std::string Name);
 
-  /// Closes the given JITDylib.
+  /// Removes the given JITDylibs from the ExecutionSession.
   ///
-  /// This method clears all resources held for the JITDylib, puts it in the
-  /// closed state, and clears all references held by the ExecutionSession and
-  /// other JITDylibs. No further code can be added to the JITDylib, and the
-  /// object will be freed once any remaining JITDylibSPs to it are destroyed.
+  /// This method clears all resources held for the JITDylibs, puts them in the
+  /// closed state, and clears all references to them that are held by the
+  /// ExecutionSession or other JITDylibs. No further code can be added to the
+  /// removed JITDylibs, and the JITDylib objects will be freed once any
+  /// remaining JITDylibSPs pointing to them are destroyed.
   ///
-  /// This method does *not* run static destructors.
+  /// This method does *not* run static destructors for code contained in the
+  /// JITDylibs, and each JITDylib can only be removed once.
   ///
-  /// This method can only be called once for each JITDylib.
-  Error removeJITDylib(JITDylib &JD);
+  /// JITDylibs will be removed in the order given. Teardown is usually
+  /// independent for each JITDylib, but not always. In particular, where the
+  /// ORC runtime is used it is expected that teardown off all JITDylibs will
+  /// depend on it, so the JITDylib containing the ORC runtime must be removed
+  /// last. If the client has introduced any other dependencies they should be
+  /// accounted for in the removal order too.
+  Error removeJITDylibs(std::vector<JITDylibSP> JDsToRemove);
+
+  /// Calls removeJTIDylibs on the gives JITDylib.
+  Error removeJITDylib(JITDylib &JD) {
+    return removeJITDylibs(std::vector<JITDylibSP>({&JD}));
+  }
 
   /// Set the error reporter function.
   ExecutionSession &setErrorReporter(ErrorReporter ReportError) {

diff  --git a/llvm/lib/ExecutionEngine/Orc/Core.cpp b/llvm/lib/ExecutionEngine/Orc/Core.cpp
index 0c23f2b2521932..f069c58d2b7004 100644
--- a/llvm/lib/ExecutionEngine/Orc/Core.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/Core.cpp
@@ -1918,16 +1918,14 @@ ExecutionSession::~ExecutionSession() {
 Error ExecutionSession::endSession() {
   LLVM_DEBUG(dbgs() << "Ending ExecutionSession " << this << "\n");
 
-  std::vector<JITDylibSP> JITDylibsToClose = runSessionLocked([&] {
+  auto JDsToRemove = runSessionLocked([&] {
     SessionOpen = false;
-    return std::move(JDs);
+    return JDs;
   });
 
-  // TODO: notifiy platform? run static deinits?
+  std::reverse(JDsToRemove.begin(), JDsToRemove.end());
 
-  Error Err = Error::success();
-  for (auto &JD : reverse(JITDylibsToClose))
-    Err = joinErrors(std::move(Err), JD->clear());
+  auto Err = removeJITDylibs(std::move(JDsToRemove));
 
   Err = joinErrors(std::move(Err), EPC->disconnect());
 
@@ -1977,42 +1975,45 @@ Expected<JITDylib &> ExecutionSession::createJITDylib(std::string Name) {
   return JD;
 }
 
-Error ExecutionSession::removeJITDylib(JITDylib &JD) {
-  // Keep JD alive throughout this routine, even if all other references
-  // have been dropped.
-  JITDylibSP JDKeepAlive = &JD;
+Error ExecutionSession::removeJITDylibs(std::vector<JITDylibSP> JDsToRemove) {
 
   // Set JD to 'Closing' state and remove JD from the ExecutionSession.
   runSessionLocked([&] {
-    assert(JD.State == JITDylib::Open && "JD already closed");
-    JD.State = JITDylib::Closing;
-    auto I = llvm::find(JDs, &JD);
-    assert(I != JDs.end() && "JD does not appear in session JDs");
-    JDs.erase(I);
+    for (auto &JD : JDsToRemove) {
+      assert(JD->State == JITDylib::Open && "JD already closed");
+      JD->State = JITDylib::Closing;
+      auto I = llvm::find(JDs, JD);
+      assert(I != JDs.end() && "JD does not appear in session JDs");
+      JDs.erase(I);
+    }
   });
 
-  // Clear the JITDylib. Hold on to any error while we clean up the
-  // JITDylib members below.
-  auto Err = JD.clear();
-
-  // Notify the platform of the teardown.
-  if (P)
-    Err = joinErrors(std::move(Err), P->teardownJITDylib(JD));
+  // Clear JITDylibs and notify the platform.
+  Error Err = Error::success();
+  for (auto JD : JDsToRemove) {
+    dbgs() << "---REMOVING--- " << JD->getName() << "\n";
+    Err = joinErrors(std::move(Err), JD->clear());
+    if (P)
+      Err = joinErrors(std::move(Err), P->teardownJITDylib(*JD));
+  }
 
   // Set JD to closed state. Clear remaining data structures.
   runSessionLocked([&] {
-    assert(JD.State == JITDylib::Closing && "JD should be closing");
-    JD.State = JITDylib::Closed;
-    assert(JD.Symbols.empty() && "JD.Symbols is not empty after clear");
-    assert(JD.UnmaterializedInfos.empty() &&
-           "JD.UnmaterializedInfos is not empty after clear");
-    assert(JD.MaterializingInfos.empty() &&
-           "JD.MaterializingInfos is not empty after clear");
-    assert(JD.TrackerSymbols.empty() &&
-           "TrackerSymbols is not empty after clear");
-    JD.DefGenerators.clear();
-    JD.LinkOrder.clear();
+    for (auto &JD : JDsToRemove) {
+      assert(JD->State == JITDylib::Closing && "JD should be closing");
+      JD->State = JITDylib::Closed;
+      assert(JD->Symbols.empty() && "JD.Symbols is not empty after clear");
+      assert(JD->UnmaterializedInfos.empty() &&
+             "JD.UnmaterializedInfos is not empty after clear");
+      assert(JD->MaterializingInfos.empty() &&
+             "JD.MaterializingInfos is not empty after clear");
+      assert(JD->TrackerSymbols.empty() &&
+             "TrackerSymbols is not empty after clear");
+      JD->DefGenerators.clear();
+      JD->LinkOrder.clear();
+    }
   });
+
   return Err;
 }