[ORC-RT][ORC][MachO] Fix some issues with executor-side symbol tables.

1. Prevent deadlock by unlocking JDStatesMutex when calling back to the
   controller to request a push of new symbols. (If JDStatesMutex is locked
   then the push operation can't register the new symbols, and so can't
   complete).

2. Record MachOPlatform runtime symbols during bootstrap and attach their
   registration to the bootstrap-completion graph, similar to the way that
   deferred allocation actions are handled. We can't register the symbols
   the normal way during bootstrap since the symbol registration function is
   itself in the process of being materialized.

3. Add dlsym testcases to exercise these fixes.

Author: lhames

compiler-rt/lib/orc/macho_platform.cpp

@@ -830,7 +830,7 @@ int MachOPlatformRuntimeState::dlclose(void *DSOHandle) {
 }
 
 void *MachOPlatformRuntimeState::dlsym(void *DSOHandle, const char *Symbol) {
-  std::lock_guard<std::mutex> Lock(JDStatesMutex);
+  std::unique_lock<std::mutex> Lock(JDStatesMutex);
   auto *JDS = getJITDylibStateByHeader(DSOHandle);
   if (!JDS) {
     std::ostringstream ErrStream;
@@ -860,10 +860,12 @@ void *MachOPlatformRuntimeState::dlsym(void *DSOHandle, const char *Symbol) {
 
   // Otherwise call back to the controller to try to request that the symbol
   // be materialized.
+  Lock.unlock();
   if (auto Err = requestPushSymbols(*JDS, {Symbols.data(), Symbols.size()})) {
     DLFcnError = toString(std::move(Err));
     return nullptr;
   }
+  Lock.lock();
 
   // Try another local resolution.
   visitSymbolAddrs(*JDS, Symbols, [&](size_t Idx, ElemResult E) {

compiler-rt/test/orc/TestCases/Darwin/arm64/Inputs/ret_self.S

@@ -0,0 +1,11 @@
+	.section	__TEXT,__text,regular,pure_instructions
+	.build_version macos, 14, 0	sdk_version 14, 4
+	.globl	_ret_self
+	.p2align	2
+_ret_self:
+	adrp	x0, _ret_self@PAGE
+	add	x0, x0, _ret_self@PAGEOFF
+	ret
+
+.subsections_via_symbols
+

compiler-rt/test/orc/TestCases/Darwin/arm64/trivial-dlsym.c

@@ -0,0 +1,44 @@
+// Test that __orc_rt_macho_jit_dlsym works as expected.
+//
+// RUN: %clang -c -o %t.sym.o %p/Inputs/ret_self.S
+// RUN: %clang -c -o %t.test.o %s
+// RUN: %llvm_jitlink \
+// RUN:   -alias Platform:_dlopen=___orc_rt_macho_jit_dlopen \
+// RUN:   -alias Platform:_dlsym=___orc_rt_macho_jit_dlsym \
+// RUN:   -alias Platform:_dlclose=___orc_rt_macho_jit_dlclose \
+// RUN:   %t.test.o -lextra_sym -jd extra_sym %t.sym.o | FileCheck %s
+
+// CHECK: entering main
+// CHECK-NEXT: found "ret_self" at
+// CHECK-NEXT: address of "ret_self" is consistent
+// CHECK-NEXT: leaving main
+
+int printf(const char *restrict format, ...);
+void *dlopen(const char *path, int mode);
+void *dlsym(void *handle, const char *symbol);
+int dlclose(void *handle);
+
+int main(int argc, char *argv[]) {
+  printf("entering main\n");
+  void *H = dlopen("extra_sym", 0);
+  if (!H) {
+    printf("failed\n");
+    return -1;
+  }
+
+  void *(*ret_self)(void) = (void *(*)(void))dlsym(H, "ret_self");
+  if (ret_self)
+    printf("found \"ret_self\" at %p\n", ret_self);
+  else
+    printf("failed to find \"ret_self\" via dlsym\n");
+
+  printf("address of \"ret_self\" is %s\n",
+         ret_self() == ret_self ? "consistent" : "inconsistent");
+
+  if (dlclose(H) == -1) {
+    printf("failed\n");
+    return -1;
+  }
+  printf("leaving main\n");
+  return 0;
+}

compiler-rt/test/orc/TestCases/Darwin/x86-64/Inputs/ret_self.S

@@ -0,0 +1,12 @@
+// A function that returns its own address. Handy for testing whether JIT'd code
+// and JIT symbol tables agree on addresses.
+
+        .section	__TEXT,__text,regular,pure_instructions
+	.build_version macos, 14, 0
+	.globl	_ret_self
+	.p2align	4, 0x90
+_ret_self:
+	leaq	_ret_self(%rip), %rax
+	retq
+
+.subsections_via_symbols

compiler-rt/test/orc/TestCases/Darwin/x86-64/trivial-dlsym.c

@@ -0,0 +1,44 @@
+// Test that __orc_rt_macho_jit_dlsym works as expected.
+//
+// RUN: %clang -c -o %t.sym.o %p/Inputs/ret_self.S
+// RUN: %clang -c -o %t.test.o %s
+// RUN: %llvm_jitlink \
+// RUN:   -alias Platform:_dlopen=___orc_rt_macho_jit_dlopen \
+// RUN:   -alias Platform:_dlsym=___orc_rt_macho_jit_dlsym \
+// RUN:   -alias Platform:_dlclose=___orc_rt_macho_jit_dlclose \
+// RUN:   %t.test.o -lextra_sym -jd extra_sym %t.sym.o | FileCheck %s
+
+// CHECK: entering main
+// CHECK-NEXT: found "ret_self" at
+// CHECK-NEXT: address of "ret_self" is consistent
+// CHECK-NEXT: leaving main
+
+int printf(const char *restrict format, ...);
+void *dlopen(const char *path, int mode);
+void *dlsym(void *handle, const char *symbol);
+int dlclose(void *handle);
+
+int main(int argc, char *argv[]) {
+  printf("entering main\n");
+  void *H = dlopen("extra_sym", 0);
+  if (!H) {
+    printf("failed\n");
+    return -1;
+  }
+
+  void *(*ret_self)(void) = (void *(*)(void))dlsym(H, "ret_self");
+  if (ret_self)
+    printf("found \"ret_self\" at %p\n", ret_self);
+  else
+    printf("failed to find \"ret_self\" via dlsym\n");
+
+  printf("address of \"ret_self\" is %s\n",
+         ret_self() == ret_self ? "consistent" : "inconsistent");
+
+  if (dlclose(H) == -1) {
+    printf("failed\n");
+    return -1;
+  }
+  printf("leaving main\n");
+  return 0;
+}

llvm/include/llvm/ExecutionEngine/Orc/MachOPlatform.h

@@ -118,13 +118,17 @@ class MachOPlatform : public Platform {
   standardRuntimeUtilityAliases();
 
 private:
+  using SymbolTableVector = SmallVector<
+      std::tuple<ExecutorAddr, ExecutorAddr, MachOExecutorSymbolFlags>>;
+
   // Data needed for bootstrap only.
   struct BootstrapInfo {
     std::mutex Mutex;
     std::condition_variable CV;
     size_t ActiveGraphs = 0;
     shared::AllocActions DeferredAAs;
     ExecutorAddr MachOHeaderAddr;
+    SymbolTableVector SymTab;
   };
 
   // The MachOPlatformPlugin scans/modifies LinkGraphs to support MachO

llvm/lib/ExecutionEngine/Orc/MachOPlatform.cpp

@@ -89,6 +89,11 @@ class SPSSerializationTraits<SPSMachOExecutorSymbolFlags,
 
 namespace {
 
+using SPSRegisterSymbolsArgs =
+    SPSArgList<SPSExecutorAddr,
+               SPSSequence<SPSTuple<SPSExecutorAddr, SPSExecutorAddr,
+                                    SPSMachOExecutorSymbolFlags>>>;
+
 std::unique_ptr<jitlink::LinkGraph> createPlatformGraph(MachOPlatform &MOP,
                                                         std::string Name) {
   unsigned PointerSize;
@@ -208,24 +213,28 @@ constexpr MachOHeaderMaterializationUnit::HeaderSymbol
 class MachOPlatformCompleteBootstrapMaterializationUnit
     : public MaterializationUnit {
 public:
+  using SymbolTableVector =
+      SmallVector<std::tuple<ExecutorAddr, ExecutorAddr,
+                             MachOPlatform::MachOExecutorSymbolFlags>>;
+
   MachOPlatformCompleteBootstrapMaterializationUnit(
       MachOPlatform &MOP, StringRef PlatformJDName,
-      SymbolStringPtr CompleteBootstrapSymbol, shared::AllocActions DeferredAAs,
+      SymbolStringPtr CompleteBootstrapSymbol, SymbolTableVector SymTab,
+      shared::AllocActions DeferredAAs, ExecutorAddr MachOHeaderAddr,
       ExecutorAddr PlatformBootstrap, ExecutorAddr PlatformShutdown,
       ExecutorAddr RegisterJITDylib, ExecutorAddr DeregisterJITDylib,
       ExecutorAddr RegisterObjectSymbolTable,
-      ExecutorAddr DeregisterObjectSymbolTable, ExecutorAddr MachOHeaderAddr)
+      ExecutorAddr DeregisterObjectSymbolTable)
       : MaterializationUnit(
             {{{CompleteBootstrapSymbol, JITSymbolFlags::None}}, nullptr}),
         MOP(MOP), PlatformJDName(PlatformJDName),
         CompleteBootstrapSymbol(std::move(CompleteBootstrapSymbol)),
-        DeferredAAs(std::move(DeferredAAs)),
-        PlatformBootstrap(PlatformBootstrap),
+        SymTab(std::move(SymTab)), DeferredAAs(std::move(DeferredAAs)),
+        MachOHeaderAddr(MachOHeaderAddr), PlatformBootstrap(PlatformBootstrap),
         PlatformShutdown(PlatformShutdown), RegisterJITDylib(RegisterJITDylib),
         DeregisterJITDylib(DeregisterJITDylib),
         RegisterObjectSymbolTable(RegisterObjectSymbolTable),
-        DeregisterObjectSymbolTable(DeregisterObjectSymbolTable),
-        MachOHeaderAddr(MachOHeaderAddr) {}
+        DeregisterObjectSymbolTable(DeregisterObjectSymbolTable) {}
 
   StringRef getName() const override {
     return "MachOPlatformCompleteBootstrap";
@@ -242,7 +251,7 @@ class MachOPlatformCompleteBootstrapMaterializationUnit
                         Linkage::Strong, Scope::Hidden, false, true);
 
     // Reserve space for the stolen actions, plus two extras.
-    G->allocActions().reserve(DeferredAAs.size() + 2);
+    G->allocActions().reserve(DeferredAAs.size() + 3);
 
     // 1. Bootstrap the platform support code.
     G->allocActions().push_back(
@@ -258,7 +267,14 @@ class MachOPlatformCompleteBootstrapMaterializationUnit
          cantFail(WrapperFunctionCall::Create<SPSArgList<SPSExecutorAddr>>(
              DeregisterJITDylib, MachOHeaderAddr))});
 
-    // 3. Add the deferred actions to the graph.
+    // 3. Register deferred symbols.
+    G->allocActions().push_back(
+        {cantFail(WrapperFunctionCall::Create<SPSRegisterSymbolsArgs>(
+             RegisterObjectSymbolTable, MachOHeaderAddr, SymTab)),
+         cantFail(WrapperFunctionCall::Create<SPSRegisterSymbolsArgs>(
+             DeregisterObjectSymbolTable, MachOHeaderAddr, SymTab))});
+
+    // 4. Add the deferred actions to the graph.
     std::move(DeferredAAs.begin(), DeferredAAs.end(),
               std::back_inserter(G->allocActions()));
 
@@ -271,14 +287,15 @@ class MachOPlatformCompleteBootstrapMaterializationUnit
   MachOPlatform &MOP;
   StringRef PlatformJDName;
   SymbolStringPtr CompleteBootstrapSymbol;
+  SymbolTableVector SymTab;
   shared::AllocActions DeferredAAs;
+  ExecutorAddr MachOHeaderAddr;
   ExecutorAddr PlatformBootstrap;
   ExecutorAddr PlatformShutdown;
   ExecutorAddr RegisterJITDylib;
   ExecutorAddr DeregisterJITDylib;
   ExecutorAddr RegisterObjectSymbolTable;
   ExecutorAddr DeregisterObjectSymbolTable;
-  ExecutorAddr MachOHeaderAddr;
 };
 
 static StringRef ObjCRuntimeObjectSectionsData[] = {
@@ -601,10 +618,11 @@ MachOPlatform::MachOPlatform(
   if ((Err = PlatformJD.define(
            std::make_unique<MachOPlatformCompleteBootstrapMaterializationUnit>(
                *this, PlatformJD.getName(), BootstrapCompleteSymbol,
-               std::move(BI.DeferredAAs), PlatformBootstrap.Addr,
+               std::move(BI.SymTab), std::move(BI.DeferredAAs),
+               BI.MachOHeaderAddr, PlatformBootstrap.Addr,
                PlatformShutdown.Addr, RegisterJITDylib.Addr,
                DeregisterJITDylib.Addr, RegisterObjectSymbolTable.Addr,
-               DeregisterObjectSymbolTable.Addr, BI.MachOHeaderAddr))))
+               DeregisterObjectSymbolTable.Addr))))
     return;
   if ((Err = ES.lookup(makeJITDylibSearchOrder(
                            &PlatformJD, JITDylibLookupFlags::MatchAllSymbols),
@@ -1714,16 +1732,17 @@ Error MachOPlatform::MachOPlatformPlugin::addSymbolTableRegistration(
     HeaderAddr = I->second;
   }
 
-  SmallVector<std::tuple<ExecutorAddr, ExecutorAddr, MachOExecutorSymbolFlags>>
-      SymTab;
+  SymbolTableVector LocalSymTab;
+  auto &SymTab = LLVM_LIKELY(!InBootstrapPhase) ? LocalSymTab
+                                                : MP.Bootstrap.load()->SymTab;
   for (auto &[OriginalSymbol, NameSym] : JITSymTabInfo)
     SymTab.push_back({NameSym->getAddress(), OriginalSymbol->getAddress(),
                       flagsForSymbol(*OriginalSymbol)});
 
-  using SPSRegisterSymbolsArgs =
-      SPSArgList<SPSExecutorAddr,
-                 SPSSequence<SPSTuple<SPSExecutorAddr, SPSExecutorAddr,
-                                      SPSMachOExecutorSymbolFlags>>>;
+  // Bail out if we're in the bootstrap phase -- registration of thees symbols
+  // will be attached to the bootstrap graph.
+  if (LLVM_UNLIKELY(InBootstrapPhase))
+    return Error::success();
 
   shared::AllocActions &allocActions = LLVM_LIKELY(!InBootstrapPhase)
                                            ? G.allocActions()