[llvm] 4288fb8 - [ORC][MachO] Fix JITDylib header-addr tracking in MachOPlatform.

[Lang Hames via llvm-commits]

Author: Lang Hames
Date: 2023-12-04T18:30:33-08:00
New Revision: 4288fb8c26447280b9e8ddb22160ebbfc082a815

URL: https://github.com/llvm/llvm-project/commit/4288fb8c26447280b9e8ddb22160ebbfc082a815
DIFF: https://github.com/llvm/llvm-project/commit/4288fb8c26447280b9e8ddb22160ebbfc082a815.diff

LOG: [ORC][MachO] Fix JITDylib header-addr tracking in MachOPlatform.

HeaderAddr shouldn't be a member variable of MachOPlatformPlugin: there's only
one plugin instance shared between all JITDylibs, so the shared HeaderAddr will
be overwritten in an unpredictable and unsafe way. We haven't seen any issues
due to this yet, but it triggered failures during testing of an upcoming
llvm-jitlink patch (e.g. ORC-RT test Darwin/x86-64/jit-re-dlopen-trivial.S).
This patch pre-fixes the issue in advance of the llvm-jitlink patch landing.

This patch also removes some stale debugging output in MachOPlatform.

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ExecutionEngine/Orc/MachOPlatform.h b/llvm/include/llvm/ExecutionEngine/Orc/MachOPlatform.h
index 37c044a7415d6..7203b80052b5f 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/MachOPlatform.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/MachOPlatform.h
@@ -207,12 +207,12 @@ class MachOPlatform : public Platform {
     Error prepareSymbolTableRegistration(jitlink::LinkGraph &G,
                                          JITSymTabVector &JITSymTabInfo);
     Error addSymbolTableRegistration(jitlink::LinkGraph &G,
+                                     MaterializationResponsibility &MR,
                                      JITSymTabVector &JITSymTabInfo,
                                      bool InBootstrapPhase);
 
     std::mutex PluginMutex;
     MachOPlatform ∓
-    ExecutorAddr HeaderAddr;
 
     // FIXME: ObjCImageInfos and HeaderAddrs need to be cleared when
     // JITDylibs are removed.

diff  --git a/llvm/lib/ExecutionEngine/Orc/MachOPlatform.cpp b/llvm/lib/ExecutionEngine/Orc/MachOPlatform.cpp
index 850e29cf9952d..a0bd9b6266ff1 100644
--- a/llvm/lib/ExecutionEngine/Orc/MachOPlatform.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/MachOPlatform.cpp
@@ -754,12 +754,6 @@ void MachOPlatform::rt_pushInitializers(PushInitializersSendResultFn SendResult,
 void MachOPlatform::rt_pushSymbols(
     PushSymbolsInSendResultFn SendResult, ExecutorAddr Handle,
     const std::vector<std::pair<StringRef, bool>> &SymbolNames) {
-  LLVM_DEBUG({
-    dbgs() << "MachOPlatform::rt_pushSymbols(" << Handle << ", [ ";
-    for (auto &Name : SymbolNames)
-      dbgs() << "\"" << Name.first << "\" ";
-    dbgs() << "])\n";
-  });
 
   JITDylib *JD = nullptr;
 
@@ -769,6 +763,16 @@ void MachOPlatform::rt_pushSymbols(
     if (I != HeaderAddrToJITDylib.end())
       JD = I->second;
   }
+  LLVM_DEBUG({
+    dbgs() << "MachOPlatform::rt_pushSymbols(";
+    if (JD)
+      dbgs() << "\"" << JD->getName() << "\", [ ";
+    else
+      dbgs() << "<invalid handle " << Handle << ">, [ ";
+    for (auto &Name : SymbolNames)
+      dbgs() << "\"" << Name.first << "\" ";
+    dbgs() << "])\n";
+  });
 
   if (!JD) {
     SendResult(make_error<StringError>("No JITDylib associated with handle " +
@@ -787,7 +791,6 @@ void MachOPlatform::rt_pushSymbols(
       LookupKind::DLSym, {{JD, JITDylibLookupFlags::MatchExportedSymbolsOnly}},
       std::move(LS), SymbolState::Ready,
       [SendResult = std::move(SendResult)](Expected<SymbolMap> Result) mutable {
-        dbgs() << "Sending result pushSymbols result...\n";
         SendResult(Result.takeError());
       },
       NoDependenciesToRegister);
@@ -813,14 +816,6 @@ void MachOPlatform::MachOPlatformPlugin::modifyPassConfig(
 
   using namespace jitlink;
 
-  // Check for a header address.
-  {
-    std::lock_guard<std::mutex> Lock(MP.PlatformMutex);
-    auto I = MP.JITDylibToHeaderAddr.find(&MR.getTargetJITDylib());
-    if (I != MP.JITDylibToHeaderAddr.end())
-      HeaderAddr = I->second;
-  }
-
   bool InBootstrapPhase =
       &MR.getTargetJITDylib() == &MP.PlatformJD && MP.Bootstrap;
 
@@ -875,10 +870,10 @@ void MachOPlatform::MachOPlatformPlugin::modifyPassConfig(
   Config.PostPrunePasses.push_back([this, JITSymTabInfo](LinkGraph &G) {
     return prepareSymbolTableRegistration(G, *JITSymTabInfo);
   });
-  Config.PostFixupPasses.push_back(
-      [this, JITSymTabInfo, InBootstrapPhase](LinkGraph &G) {
-        return addSymbolTableRegistration(G, *JITSymTabInfo, InBootstrapPhase);
-      });
+  Config.PostFixupPasses.push_back([this, &MR, JITSymTabInfo,
+                                    InBootstrapPhase](LinkGraph &G) {
+    return addSymbolTableRegistration(G, MR, *JITSymTabInfo, InBootstrapPhase);
+  });
 
   // Add a pass to register the final addresses of any special sections in the
   // object with the runtime.
@@ -1427,7 +1422,15 @@ Error MachOPlatform::MachOPlatformPlugin::registerObjectPlatformSections(
                                              ? G.allocActions()
                                              : MP.Bootstrap.load()->DeferredAAs;
 
-    assert(HeaderAddr && "No HeaderAddr for JITDylib");
+    ExecutorAddr HeaderAddr;
+    {
+      std::lock_guard<std::mutex> Lock(MP.PlatformMutex);
+      auto I = MP.JITDylibToHeaderAddr.find(&JD);
+      assert(I != MP.JITDylibToHeaderAddr.end() &&
+             "No header registered for JD");
+      assert(I->second && "Null header registered for JD");
+      HeaderAddr = I->second;
+    }
     allocActions.push_back(
         {cantFail(
              WrapperFunctionCall::Create<SPSRegisterObjectPlatformSectionsArgs>(
@@ -1699,16 +1702,23 @@ Error MachOPlatform::MachOPlatformPlugin::prepareSymbolTableRegistration(
 }
 
 Error MachOPlatform::MachOPlatformPlugin::addSymbolTableRegistration(
-    jitlink::LinkGraph &G, JITSymTabVector &JITSymTabInfo,
-    bool InBootstrapPhase) {
+    jitlink::LinkGraph &G, MaterializationResponsibility &MR,
+    JITSymTabVector &JITSymTabInfo, bool InBootstrapPhase) {
+
+  ExecutorAddr HeaderAddr;
+  {
+    std::lock_guard<std::mutex> Lock(MP.PlatformMutex);
+    auto I = MP.JITDylibToHeaderAddr.find(&MR.getTargetJITDylib());
+    assert(I != MP.JITDylibToHeaderAddr.end() && "No header registered for JD");
+    assert(I->second && "Null header registered for JD");
+    HeaderAddr = I->second;
+  }
 
   SmallVector<std::tuple<ExecutorAddr, ExecutorAddr, MachOExecutorSymbolFlags>>
       SymTab;
-  for (auto &[OriginalSymbol, NameSym] : JITSymTabInfo) {
-    // dbgs() << "Original symbol: \"" << OriginalSymbol->getName() << "\"\n";
+  for (auto &[OriginalSymbol, NameSym] : JITSymTabInfo)
     SymTab.push_back({NameSym->getAddress(), OriginalSymbol->getAddress(),
                       flagsForSymbol(*OriginalSymbol)});
-  }
 
   using SPSRegisterSymbolsArgs =
       SPSArgList<SPSExecutorAddr,