[ORC] Introduce deferred allocation-actions scheme for MachOPlatform bootstrap.

This patch modifies the MachOPlatform bootstrap process to record allocation
actions for ORC runtime platform support code in a "deferred actions" vector
rather than attaching it to the corresponding LinkGraphs up-front. The deferred
allocation-actions are run after all the platform support code has been loaded
by attaching them to a separate "bootstrap-complete" graph.

This change should allow the mach-o platform support code in the ORC runtime to
use advanced mach-o platform features (e.g. static inits, TLVs), provided that
the support code does not use these features at runtime before the bootstrap
process completes, or after the shutdown process starts. This is a nice
improvement in and of itself but is motivated by specific future plans: we
want to start recording unwind info in the mach-o platform state object*, and
the recording functions will have their own frame info that needs registering.
The deferred allocation-actions scheme allows for this.

* The plan is to add a new unwind-info-lookup path to libunwind to allow it to
  call back to the ORC runtime to find unwind sections. This will simplify the
  implementation of support for JIT'd compact-unwind info.

Author: lhames

compiler-rt/lib/orc/CMakeLists.txt

@@ -15,7 +15,6 @@ set(ALL_ORC_SOURCES
   coff_platform.cpp
   coff_platform.per_jd.cpp
   elfnix_platform.cpp
-  macho_ehframe_registration.cpp
   macho_platform.cpp
   )
 
@@ -96,7 +95,6 @@ if (APPLE)
 
   set(ORC_SOURCES
     ${ORC_COMMON_SOURCES}
-    macho_ehframe_registration.cpp
     macho_platform.cpp
     )
 

compiler-rt/lib/orc/macho_ehframe_registration.cpp

@@ -284,6 +284,9 @@ class MachOPlatformRuntimeState {
   Expected<ExecutorAddr> lookupSymbolInJITDylib(void *DSOHandle,
                                                 std::string_view Symbol);
 
+  static Error registerEHFrames(span<const char> EHFrameSection);
+  static Error deregisterEHFrames(span<const char> EHFrameSection);
+
   static Error registerObjCSelectors(JITDylibState &JDS);
   static Error registerObjCClasses(JITDylibState &JDS);
   static Error registerSwift5Protocols(JITDylibState &JDS);
@@ -438,7 +441,10 @@ Error MachOPlatformRuntimeState::registerObjectPlatformSections(
 
   for (auto &KV : Secs) {
     // FIXME: Validate section ranges?
-    if (KV.first == "__DATA,__data") {
+    if (KV.first == "__TEXT,__eh_frame") {
+      if (auto Err = registerEHFrames(KV.second.toSpan<const char>()))
+        return Err;
+    } else if (KV.first == "__DATA,__data") {
       assert(!JDS->DataSectionContent.count(KV.second.Start.toPtr<char *>()) &&
              "Address already registered.");
       auto S = KV.second.toSpan<char>();
@@ -505,7 +511,10 @@ Error MachOPlatformRuntimeState::deregisterObjectPlatformSections(
 
   for (auto &KV : Secs) {
     // FIXME: Validate section ranges?
-    if (KV.first == "__DATA,__data") {
+    if (KV.first == "__TEXT,__eh_frame") {
+      if (auto Err = deregisterEHFrames(KV.second.toSpan<const char>()))
+        return Err;
+    } else if (KV.first == "__DATA,__data") {
       JDS->DataSectionContent.erase(KV.second.Start.toPtr<char *>());
     } else if (KV.first == "__DATA,__common") {
       JDS->ZeroInitRanges.erase(KV.second.Start.toPtr<char *>());
@@ -673,6 +682,45 @@ MachOPlatformRuntimeState::lookupSymbolInJITDylib(void *DSOHandle,
   return Result;
 }
 
+// eh-frame registration functions.
+// We expect these to be available for all processes.
+extern "C" void __register_frame(const void *);
+extern "C" void __deregister_frame(const void *);
+
+template <typename HandleFDEFn>
+void walkEHFrameSection(span<const char> EHFrameSection,
+                        HandleFDEFn HandleFDE) {
+  const char *CurCFIRecord = EHFrameSection.data();
+  uint64_t Size = *reinterpret_cast<const uint32_t *>(CurCFIRecord);
+
+  while (CurCFIRecord != EHFrameSection.end() && Size != 0) {
+    const char *OffsetField = CurCFIRecord + (Size == 0xffffffff ? 12 : 4);
+    if (Size == 0xffffffff)
+      Size = *reinterpret_cast<const uint64_t *>(CurCFIRecord + 4) + 12;
+    else
+      Size += 4;
+    uint32_t Offset = *reinterpret_cast<const uint32_t *>(OffsetField);
+
+    if (Offset != 0)
+      HandleFDE(CurCFIRecord);
+
+    CurCFIRecord += Size;
+    Size = *reinterpret_cast<const uint32_t *>(CurCFIRecord);
+  }
+}
+
+Error MachOPlatformRuntimeState::registerEHFrames(
+    span<const char> EHFrameSection) {
+  walkEHFrameSection(EHFrameSection, __register_frame);
+  return Error::success();
+}
+
+Error MachOPlatformRuntimeState::deregisterEHFrames(
+    span<const char> EHFrameSection) {
+  walkEHFrameSection(EHFrameSection, __deregister_frame);
+  return Error::success();
+}
+
 Error MachOPlatformRuntimeState::registerObjCSelectors(JITDylibState &JDS) {
   if (!JDS.ObjCSelRefsSections.hasNewSections())
     return Error::success();
@@ -1050,14 +1098,24 @@ Error runWrapperFunctionCalls(std::vector<WrapperFunctionCall> WFCs) {
 
 ORC_RT_INTERFACE __orc_rt_CWrapperFunctionResult
 __orc_rt_macho_platform_bootstrap(char *ArgData, size_t ArgSize) {
-  MachOPlatformRuntimeState::initialize();
-  return WrapperFunctionResult().release();
+  return WrapperFunction<SPSError()>::handle(
+             ArgData, ArgSize,
+             []() -> Error {
+               MachOPlatformRuntimeState::initialize();
+               return Error::success();
+             })
+      .release();
 }
 
 ORC_RT_INTERFACE __orc_rt_CWrapperFunctionResult
 __orc_rt_macho_platform_shutdown(char *ArgData, size_t ArgSize) {
-  MachOPlatformRuntimeState::destroy();
-  return WrapperFunctionResult().release();
+  return WrapperFunction<SPSError()>::handle(
+             ArgData, ArgSize,
+             []() -> Error {
+               MachOPlatformRuntimeState::destroy();
+               return Error::success();
+             })
+      .release();
 }
 
 ORC_RT_INTERFACE __orc_rt_CWrapperFunctionResult

compiler-rt/lib/orc/macho_platform.cpp

@@ -107,6 +107,15 @@ class MachOPlatform : public Platform {
   static bool isInitializerSection(StringRef SegName, StringRef SectName);
 
 private:
+  // Data needed for bootstrap only.
+  struct BootstrapInfo {
+    std::mutex Mutex;
+    std::condition_variable CV;
+    size_t ActiveGraphs = 0;
+    shared::AllocActions DeferredAAs;
+    ExecutorAddr MachOHeaderAddr;
+  };
+
   // The MachOPlatformPlugin scans/modifies LinkGraphs to support MachO
   // platform features including initializers, exceptions, TLV, and language
   // runtime registration.
@@ -138,6 +147,12 @@ class MachOPlatform : public Platform {
     using InitSymbolDepMap =
         DenseMap<MaterializationResponsibility *, JITLinkSymbolSet>;
 
+    Error bootstrapPipelineStart(jitlink::LinkGraph &G);
+    Error bootstrapPipelineRecordRuntimeFunctions(jitlink::LinkGraph &G);
+    Error bootstrapPipelineEnd(jitlink::LinkGraph &G);
+
+    Error recordRuntimeRegistrationFunctions(jitlink::LinkGraph &G);
+
     Error associateJITDylibHeaderSymbol(jitlink::LinkGraph &G,
                                         MaterializationResponsibility &MR);
 
@@ -149,9 +164,8 @@ class MachOPlatform : public Platform {
 
     Error fixTLVSectionsAndEdges(jitlink::LinkGraph &G, JITDylib &JD);
 
-    Error registerObjectPlatformSections(jitlink::LinkGraph &G, JITDylib &JD);
-
-    Error registerEHSectionsPhase1(jitlink::LinkGraph &G);
+    Error registerObjectPlatformSections(jitlink::LinkGraph &G, JITDylib &JD,
+                                         bool InBootstrapPhase);
 
     std::mutex PluginMutex;
     MachOPlatform &MP;
@@ -179,7 +193,7 @@ class MachOPlatform : public Platform {
                 Error &Err);
 
   // Associate MachOPlatform JIT-side runtime support functions with handlers.
-  Error associateRuntimeSupportFunctions(JITDylib &PlatformJD);
+  Error associateRuntimeSupportFunctions();
 
   // Implements rt_pushInitializers by making repeat async lookups for
   // initializer symbols (each lookup may spawn more initializer symbols if
@@ -195,19 +209,14 @@ class MachOPlatform : public Platform {
   void rt_lookupSymbol(SendSymbolAddressFn SendResult, ExecutorAddr Handle,
                        StringRef SymbolName);
 
-  // Records the addresses of runtime symbols used by the platform.
-  Error bootstrapMachORuntime(JITDylib &PlatformJD);
-
   // Call the ORC runtime to create a pthread key.
   Expected<uint64_t> createPThreadKey();
 
-  enum PlatformState { BootstrapPhase1, BootstrapPhase2, Initialized };
-
   ExecutionSession &ES;
+  JITDylib &PlatformJD;
   ObjectLinkingLayer &ObjLinkingLayer;
 
-  SymbolStringPtr MachOHeaderStartSymbol;
-  std::atomic<PlatformState> State{BootstrapPhase1};
+  SymbolStringPtr MachOHeaderStartSymbol = ES.intern("___dso_handle");
 
   struct RuntimeFunction {
     RuntimeFunction(SymbolStringPtr Name) : Name(std::move(Name)) {}
@@ -240,6 +249,8 @@ class MachOPlatform : public Platform {
   DenseMap<JITDylib *, ExecutorAddr> JITDylibToHeaderAddr;
   DenseMap<ExecutorAddr, JITDylib *> HeaderAddrToJITDylib;
   DenseMap<JITDylib *, uint64_t> JITDylibToPThreadKey;
+
+  std::atomic<BootstrapInfo *> Bootstrap;
 };
 
 namespace shared {