[ORC] Update ELF debugger support plugin to use AllocActions (#167866)

The `DebugObjectManagerPlugin` implements debugger support for ELF
platforms with the GDB JIT Interface. It emits a separate debug object
allocation in addition to the LinkGraph's own allocation. This used to
happen in the plugin's `notifyEmitted()` callback, i.e. after the
LinkGraph's allocation was finalized. In the meantime, it had to block
finalization of the corresponding materialization unit to make sure that
the debugger can register the object before the code runs.

This patch switches the plugin to use an allocation action instead. We
can remove the `notifyEmitted()` hook and implement all steps as JITLink
passes.

Author: weliveindetail

llvm/include/llvm/ExecutionEngine/Orc/DebugObjectManagerPlugin.h

@@ -15,8 +15,8 @@
 
 #include "llvm/ExecutionEngine/JITLink/JITLink.h"
 #include "llvm/ExecutionEngine/Orc/Core.h"
-#include "llvm/ExecutionEngine/Orc/EPCDebugObjectRegistrar.h"
 #include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
+#include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/Memory.h"
@@ -48,10 +48,6 @@ class DebugObject;
 ///
 class LLVM_ABI DebugObjectManagerPlugin : public ObjectLinkingLayer::Plugin {
 public:
-  // DEPRECATED - Please specify options explicitly
-  DebugObjectManagerPlugin(ExecutionSession &ES,
-                           std::unique_ptr<DebugObjectRegistrar> Target);
-
   /// Create the plugin to submit DebugObjects for JITLink artifacts. For all
   /// options the recommended setting is true.
   ///
@@ -67,16 +63,14 @@ class LLVM_ABI DebugObjectManagerPlugin : public ObjectLinkingLayer::Plugin {
   ///   sequence. When turning this off, the user has to issue the call to
   ///   __jit_debug_register_code() on the executor side manually.
   ///
-  DebugObjectManagerPlugin(ExecutionSession &ES,
-                           std::unique_ptr<DebugObjectRegistrar> Target,
-                           bool RequireDebugSections, bool AutoRegisterCode);
+  DebugObjectManagerPlugin(ExecutionSession &ES, bool RequireDebugSections,
+                           bool AutoRegisterCode, Error &Err);
   ~DebugObjectManagerPlugin() override;
 
   void notifyMaterializing(MaterializationResponsibility &MR,
                            jitlink::LinkGraph &G, jitlink::JITLinkContext &Ctx,
                            MemoryBufferRef InputObject) override;
 
-  Error notifyEmitted(MaterializationResponsibility &MR) override;
   Error notifyFailed(MaterializationResponsibility &MR) override;
   Error notifyRemovingResources(JITDylib &JD, ResourceKey K) override;
 
@@ -97,7 +91,7 @@ class LLVM_ABI DebugObjectManagerPlugin : public ObjectLinkingLayer::Plugin {
   std::mutex PendingObjsLock;
   std::mutex RegisteredObjsLock;
 
-  std::unique_ptr<DebugObjectRegistrar> Target;
+  ExecutorAddr RegistrationAction;
   bool RequireDebugSections;
   bool AutoRegisterCode;
 };

llvm/include/llvm/ExecutionEngine/Orc/Shared/OrcRTBridge.h

@@ -59,6 +59,8 @@ LLVM_ABI extern const char *MemoryReadStringsWrapperName;
 LLVM_ABI extern const char *RegisterEHFrameSectionAllocActionName;
 LLVM_ABI extern const char *DeregisterEHFrameSectionAllocActionName;
 
+LLVM_ABI extern const char *RegisterJITLoaderGDBAllocActionName;
+
 LLVM_ABI extern const char *RunAsMainWrapperName;
 LLVM_ABI extern const char *RunAsVoidFunctionWrapperName;
 LLVM_ABI extern const char *RunAsIntFunctionWrapperName;

llvm/lib/ExecutionEngine/Orc/DebugObjectManagerPlugin.cpp

@@ -19,6 +19,7 @@
 #include "llvm/BinaryFormat/ELF.h"
 #include "llvm/ExecutionEngine/JITLink/JITLinkDylib.h"
 #include "llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h"
+#include "llvm/ExecutionEngine/Orc/Shared/OrcRTBridge.h"
 #include "llvm/Object/ELFObjectFile.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/MSVCErrorWorkarounds.h"
@@ -115,7 +116,9 @@ class DebugObject {
 public:
   DebugObject(JITLinkMemoryManager &MemMgr, const JITLinkDylib *JD,
               ExecutionSession &ES)
-      : MemMgr(MemMgr), JD(JD), ES(ES), Flags(DebugObjectFlags{}) {}
+      : MemMgr(MemMgr), JD(JD), ES(ES), Flags(DebugObjectFlags{}) {
+    FinalizeFuture = FinalizePromise.get_future();
+  }
 
   bool hasFlags(DebugObjectFlags F) const { return Flags & F; }
   void setFlags(DebugObjectFlags F) {
@@ -126,8 +129,17 @@ class DebugObject {
   }
 
   using FinalizeContinuation = std::function<void(Expected<ExecutorAddrRange>)>;
+  void finalizeAsync(FinalizeContinuation OnAsync);
+
+  void failMaterialization(Error Err) {
+    FinalizePromise.set_value(std::move(Err));
+  }
 
-  void finalizeAsync(FinalizeContinuation OnFinalize);
+  void reportTargetMem(ExecutorAddrRange TargetMem) {
+    FinalizePromise.set_value(TargetMem);
+  }
+
+  Expected<ExecutorAddrRange> awaitTargetMem() { return FinalizeFuture.get(); }
 
   virtual ~DebugObject() {
     if (Alloc) {
@@ -151,6 +163,9 @@ class DebugObject {
   const JITLinkDylib *JD = nullptr;
   ExecutionSession &ES;
 
+  std::promise<MSVCPExpected<ExecutorAddrRange>> FinalizePromise;
+  std::future<MSVCPExpected<ExecutorAddrRange>> FinalizeFuture;
+
 private:
   DebugObjectFlags Flags;
   FinalizedAlloc Alloc;
@@ -160,22 +175,27 @@ class DebugObject {
 // copying memory over to the target and pass on the result once we're done.
 // Ownership of the allocation remains with us for the rest of our lifetime.
 void DebugObject::finalizeAsync(FinalizeContinuation OnFinalize) {
-  assert(!Alloc && "Cannot finalize more than once");
-
+  assert(!this->Alloc && "Cannot finalize more than once");
   if (auto SimpleSegAlloc = finalizeWorkingMemory()) {
     auto ROSeg = SimpleSegAlloc->getSegInfo(MemProt::Read);
     ExecutorAddrRange DebugObjRange(ROSeg.Addr, ROSeg.WorkingMem.size());
     SimpleSegAlloc->finalize(
         [this, DebugObjRange,
          OnFinalize = std::move(OnFinalize)](Expected<FinalizedAlloc> FA) {
           if (FA) {
-            Alloc = std::move(*FA);
+            // Note: FA->getAddress() is supposed to be the address of the
+            // memory range on the target, but InProcessMemoryManager returns
+            // the address of a FinalizedAllocInfo helper instead.
+            this->Alloc = std::move(*FA);
             OnFinalize(DebugObjRange);
           } else
             OnFinalize(FA.takeError());
         });
-  } else
+  } else {
+    // We could report this error synchronously, but it's easier this way,
+    // because the FinalizePromise will be triggered unconditionally.
     OnFinalize(SimpleSegAlloc.takeError());
+  }
 }
 
 /// The current implementation of ELFDebugObject replicates the approach used in
@@ -386,16 +406,17 @@ createDebugObjectFromBuffer(ExecutionSession &ES, LinkGraph &G,
   }
 }
 
-DebugObjectManagerPlugin::DebugObjectManagerPlugin(
-    ExecutionSession &ES, std::unique_ptr<DebugObjectRegistrar> Target,
-    bool RequireDebugSections, bool AutoRegisterCode)
-    : ES(ES), Target(std::move(Target)),
-      RequireDebugSections(RequireDebugSections),
-      AutoRegisterCode(AutoRegisterCode) {}
-
-DebugObjectManagerPlugin::DebugObjectManagerPlugin(
-    ExecutionSession &ES, std::unique_ptr<DebugObjectRegistrar> Target)
-    : DebugObjectManagerPlugin(ES, std::move(Target), true, true) {}
+DebugObjectManagerPlugin::DebugObjectManagerPlugin(ExecutionSession &ES,
+                                                   bool RequireDebugSections,
+                                                   bool AutoRegisterCode,
+                                                   Error &Err)
+    : ES(ES), RequireDebugSections(RequireDebugSections),
+      AutoRegisterCode(AutoRegisterCode) {
+  // Pass bootstrap symbol for registration function to enable debugging
+  ErrorAsOutParameter _(&Err);
+  Err = ES.getExecutorProcessControl().getBootstrapSymbols(
+      {{RegistrationAction, rt::RegisterJITLoaderGDBAllocActionName}});
+}
 
 DebugObjectManagerPlugin::~DebugObjectManagerPlugin() = default;
 
@@ -440,48 +461,50 @@ void DebugObjectManagerPlugin::modifyPassConfig(
                                                     SectionRange(GraphSection));
           return Error::success();
         });
-  }
-}
 
-Error DebugObjectManagerPlugin::notifyEmitted(
-    MaterializationResponsibility &MR) {
-  std::lock_guard<std::mutex> Lock(PendingObjsLock);
-  auto It = PendingObjs.find(&MR);
-  if (It == PendingObjs.end())
-    return Error::success();
-
-  // During finalization the debug object is registered with the target.
-  // Materialization must wait for this process to finish. Otherwise we might
-  // start running code before the debugger processed the corresponding debug
-  // info.
-  std::promise<MSVCPError> FinalizePromise;
-  std::future<MSVCPError> FinalizeErr = FinalizePromise.get_future();
-
-  It->second->finalizeAsync(
-      [this, &FinalizePromise, &MR](Expected<ExecutorAddrRange> TargetMem) {
-        // Any failure here will fail materialization.
-        if (!TargetMem) {
-          FinalizePromise.set_value(TargetMem.takeError());
-          return;
-        }
-        if (Error Err =
-                Target->registerDebugObject(*TargetMem, AutoRegisterCode)) {
-          FinalizePromise.set_value(std::move(Err));
-          return;
-        }
-
-        // Once our tracking info is updated, notifyEmitted() can return and
-        // finish materialization.
-        FinalizePromise.set_value(MR.withResourceKeyDo([&](ResourceKey K) {
-          assert(PendingObjs.count(&MR) && "We still hold PendingObjsLock");
-          std::lock_guard<std::mutex> Lock(RegisteredObjsLock);
-          auto It = PendingObjs.find(&MR);
-          RegisteredObjs[K].push_back(std::move(It->second));
-          PendingObjs.erase(It);
-        }));
-      });
-
-  return FinalizeErr.get();
+    PassConfig.PreFixupPasses.push_back(
+        [this, &DebugObj, &MR](LinkGraph &G) -> Error {
+          DebugObj.finalizeAsync([this, &DebugObj,
+                                  &MR](Expected<ExecutorAddrRange> TargetMem) {
+            if (!TargetMem) {
+              DebugObj.failMaterialization(TargetMem.takeError());
+              return;
+            }
+            // Update tracking info
+            Error Err = MR.withResourceKeyDo([&](ResourceKey K) {
+              std::lock_guard<std::mutex> LockPending(PendingObjsLock);
+              std::lock_guard<std::mutex> LockRegistered(RegisteredObjsLock);
+              auto It = PendingObjs.find(&MR);
+              RegisteredObjs[K].push_back(std::move(It->second));
+              PendingObjs.erase(It);
+            });
+
+            if (Err)
+              DebugObj.failMaterialization(std::move(Err));
+
+            // Unblock post-fixup pass
+            DebugObj.reportTargetMem(*TargetMem);
+          });
+          return Error::success();
+        });
+
+    PassConfig.PostFixupPasses.push_back(
+        [this, &DebugObj](LinkGraph &G) -> Error {
+          Expected<ExecutorAddrRange> R = DebugObj.awaitTargetMem();
+          if (!R)
+            return R.takeError();
+          if (R->empty())
+            return Error::success();
+
+          using namespace shared;
+          G.allocActions().push_back(
+              {cantFail(WrapperFunctionCall::Create<
+                        SPSArgList<SPSExecutorAddrRange, bool>>(
+                   RegistrationAction, *R, AutoRegisterCode)),
+               {/* no deregistration */}});
+          return Error::success();
+        });
+  }
 }
 
 Error DebugObjectManagerPlugin::notifyFailed(

llvm/lib/ExecutionEngine/Orc/Debugging/DebuggerSupport.cpp

@@ -35,12 +35,10 @@ Error enableDebuggerSupport(LLJIT &J) {
 
   switch (TT.getObjectFormat()) {
   case Triple::ELF: {
-    auto Registrar = createJITLoaderGDBRegistrar(ES);
-    if (!Registrar)
-      return Registrar.takeError();
+    Error TargetSymErr = Error::success();
     ObjLinkingLayer->addPlugin(std::make_unique<DebugObjectManagerPlugin>(
-        ES, std::move(*Registrar), false, true));
-    return Error::success();
+        ES, false, true, TargetSymErr));
+    return TargetSymErr;
   }
   case Triple::MachO: {
     auto DS = GDBJITDebugInfoRegistrationPlugin::Create(ES, *ProcessSymsJD, TT);

llvm/lib/ExecutionEngine/Orc/Shared/OrcRTBridge.cpp

@@ -74,6 +74,9 @@ const char *RegisterEHFrameSectionAllocActionName =
 const char *DeregisterEHFrameSectionAllocActionName =
     "llvm_orc_deregisterEHFrameAllocAction";
 
+const char *RegisterJITLoaderGDBAllocActionName =
+    "llvm_orc_registerJITLoaderGDBAllocAction";
+
 const char *RunAsMainWrapperName = "__llvm_orc_bootstrap_run_as_main_wrapper";
 const char *RunAsVoidFunctionWrapperName =
     "__llvm_orc_bootstrap_run_as_void_function_wrapper";

llvm/lib/ExecutionEngine/Orc/TargetProcess/DefaultHostBootstrapValues.cpp

@@ -9,6 +9,7 @@
 #include "llvm/ExecutionEngine/Orc/TargetProcess/DefaultHostBootstrapValues.h"
 
 #include "llvm/ExecutionEngine/Orc/Shared/OrcRTBridge.h"
+#include "llvm/ExecutionEngine/Orc/TargetProcess/JITLoaderGDB.h"
 #include "llvm/ExecutionEngine/Orc/TargetProcess/RegisterEHFrames.h"
 
 #ifdef __APPLE__
@@ -27,6 +28,9 @@ void addDefaultBootstrapValuesForHostProcess(
   BootstrapSymbols[rt::DeregisterEHFrameSectionAllocActionName] =
       ExecutorAddr::fromPtr(&llvm_orc_deregisterEHFrameSectionAllocAction);
 
+  BootstrapSymbols[rt::RegisterJITLoaderGDBAllocActionName] =
+      ExecutorAddr::fromPtr(&llvm_orc_registerJITLoaderGDBAllocAction);
+
 #ifdef __APPLE__
   if (!dlsym(RTLD_DEFAULT, "__unw_add_find_dynamic_unwind_sections"))
     BootstrapMap["darwin-use-ehframes-only"].push_back(1);

llvm/lib/ExecutionEngine/Orc/TargetProcess/SimpleRemoteEPCServer.cpp

@@ -9,6 +9,7 @@
 #include "llvm/ExecutionEngine/Orc/TargetProcess/SimpleRemoteEPCServer.h"
 
 #include "llvm/ExecutionEngine/Orc/Shared/OrcRTBridge.h"
+#include "llvm/ExecutionEngine/Orc/TargetProcess/DefaultHostBootstrapValues.h"
 #include "llvm/ExecutionEngine/Orc/TargetProcess/RegisterEHFrames.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/Process.h"
@@ -206,10 +207,7 @@ Error SimpleRemoteEPCServer::sendSetupMessage(
          "Dispatch function name should not be set");
   EI.BootstrapSymbols[ExecutorSessionObjectName] = ExecutorAddr::fromPtr(this);
   EI.BootstrapSymbols[DispatchFnName] = ExecutorAddr::fromPtr(jitDispatchEntry);
-  EI.BootstrapSymbols[rt::RegisterEHFrameSectionAllocActionName] =
-      ExecutorAddr::fromPtr(&llvm_orc_registerEHFrameSectionAllocAction);
-  EI.BootstrapSymbols[rt::DeregisterEHFrameSectionAllocActionName] =
-      ExecutorAddr::fromPtr(&llvm_orc_deregisterEHFrameSectionAllocAction);
+  addDefaultBootstrapValuesForHostProcess(EI.BootstrapMap, EI.BootstrapSymbols);
 
   using SPSSerialize =
       shared::SPSArgList<shared::SPSSimpleRemoteEPCExecutorInfo>;

llvm/tools/llvm-jitlink/llvm-jitlink.cpp

@@ -1297,9 +1297,16 @@ Session::Session(std::unique_ptr<ExecutorProcessControl> EPC, Error &Err)
   } else if (TT.isOSBinFormatELF()) {
     if (!NoExec)
       ObjLayer.addPlugin(ExitOnErr(EHFrameRegistrationPlugin::Create(ES)));
-    if (DebuggerSupport)
-      ObjLayer.addPlugin(std::make_unique<DebugObjectManagerPlugin>(
-          ES, ExitOnErr(createJITLoaderGDBRegistrar(this->ES)), true, true));
+    if (DebuggerSupport) {
+      Error TargetSymErr = Error::success();
+      auto Plugin = std::make_unique<DebugObjectManagerPlugin>(ES, true, true,
+                                                               TargetSymErr);
+      if (!TargetSymErr)
+        ObjLayer.addPlugin(std::move(Plugin));
+      else
+        logAllUnhandledErrors(std::move(TargetSymErr), errs(),
+                              "Debugger support not available: ");
+    }
   }
 
   if (auto MainJDOrErr = ES.createJITDylib("main"))