[ORC] Re-apply 98f2bb4461, enable JITEventListeners in OrcV2, with fixes.

Updates the object buffer ownership scheme in jitLinkForOrc and related
functions: Ownership of both the object::ObjectFile and underlying
MemoryBuffer is passed into jitLinkForOrc and passed back to the onEmit
callback once linking is complete. This avoids the use-after-free errors
that were seen in 98f2bb4461.

llvm/examples/OrcV2Examples/CMakeLists.txt

@@ -1,6 +1,7 @@
 add_subdirectory(BasicOrcV2CBindings)
 add_subdirectory(LLJITDumpObjects)
-add_subdirectory(LLJITWithObjectCache)
 add_subdirectory(LLJITWithCustomObjectLinkingLayer)
+add_subdirectory(LLJITWithGDBRegistrationListener)
 add_subdirectory(LLJITWithLazyReexports)
+add_subdirectory(LLJITWithObjectCache)
 add_subdirectory(LLJITWithObjectLinkingLayerPlugin)

llvm/examples/OrcV2Examples/LLJITWithGDBRegistrationListener/CMakeLists.txt

@@ -0,0 +1,16 @@
+set(LLVM_LINK_COMPONENTS
+  Core
+  IRReader
+  JITLink
+  OrcJIT
+  Support
+  nativecodegen
+  )
+
+add_llvm_example(LLJITWithGDBRegistrationListener
+  LLJITWithGDBRegistrationListener.cpp
+  )
+
+# We want JIT'd code to be able to link against process symbols like printf
+# for this example, so make sure they're exported.
+export_executable_symbols(LLJITWithGDBRegistrationListener)

llvm/examples/OrcV2Examples/LLJITWithGDBRegistrationListener/LLJITWithGDBRegistrationListener.cpp

@@ -0,0 +1,109 @@
+//===--------------- LLJITWithCustomObjectLinkingLayer.cpp ----------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file shows how to switch LLJIT to use a custom object linking layer (we
+// use ObjectLinkingLayer, which is backed by JITLink, as an example).
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ADT/StringMap.h"
+#include "llvm/ExecutionEngine/JITEventListener.h"
+#include "llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h"
+#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
+#include "llvm/ExecutionEngine/Orc/LLJIT.h"
+#include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
+#include "llvm/ExecutionEngine/SectionMemoryManager.h"
+#include "llvm/Support/InitLLVM.h"
+#include "llvm/Support/TargetSelect.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include "../ExampleModules.h"
+
+using namespace llvm;
+using namespace llvm::orc;
+
+ExitOnError ExitOnErr;
+
+static cl::opt<std::string>
+    EntryPointName("entry", cl::desc("Symbol to call as main entry point"),
+                   cl::init("main"));
+
+static cl::list<std::string> InputFiles(cl::Positional, cl::OneOrMore,
+                                        cl::desc("input files"));
+
+static cl::list<std::string> InputArgv("args", cl::Positional,
+                                       cl::desc("<program arguments>..."),
+                                       cl::ZeroOrMore, cl::PositionalEatsArgs);
+
+int main(int argc, char *argv[]) {
+  // Initialize LLVM.
+  InitLLVM X(argc, argv);
+
+  InitializeNativeTarget();
+  InitializeNativeTargetAsmPrinter();
+
+  cl::ParseCommandLineOptions(argc, argv, "LLJITWithCustomObjectLinkingLayer");
+  ExitOnErr.setBanner(std::string(argv[0]) + ": ");
+
+  // Detect the host and set code model to small.
+  auto JTMB = ExitOnErr(JITTargetMachineBuilder::detectHost());
+  if (!JTMB.getTargetTriple().isOSLinux())
+    errs()
+        << "Warning: This demo may not work for platforms other than Linux.\n";
+
+  // Create an LLJIT instance and use a custom object linking layer creator to
+  // register the GDBRegistrationListener with our RTDyldObjectLinkingLayer.
+  auto J =
+      ExitOnErr(LLJITBuilder()
+                    .setJITTargetMachineBuilder(std::move(JTMB))
+                    .setObjectLinkingLayerCreator([&](ExecutionSession &ES,
+                                                      const Triple &TT) {
+                      auto GetMemMgr = []() {
+                        return std::make_unique<SectionMemoryManager>();
+                      };
+                      auto ObjLinkingLayer =
+                          std::make_unique<RTDyldObjectLinkingLayer>(
+                              ES, std::move(GetMemMgr));
+                      ObjLinkingLayer->registerJITEventListener(
+                          *JITEventListener::createGDBRegistrationListener());
+                      return ObjLinkingLayer;
+                    })
+                    .create());
+
+  // Make sure that our process symbols are visible to JIT'd code.
+  {
+    MangleAndInterner Mangle(J->getExecutionSession(), J->getDataLayout());
+    J->getMainJITDylib().addGenerator(
+        ExitOnErr(orc::DynamicLibrarySearchGenerator::GetForCurrentProcess(
+            J->getDataLayout().getGlobalPrefix(),
+            [MainName = Mangle("main")](const orc::SymbolStringPtr &Name) {
+              return Name != MainName;
+            })));
+  }
+
+  // Load the input modules.
+  for (auto &InputFile : InputFiles) {
+    auto Ctx = std::make_unique<LLVMContext>();
+    SMDiagnostic Err;
+    std::unique_ptr<Module> M = parseIRFile(InputFile, Err, *Ctx);
+    if (!M) {
+      Err.print(argv[0], errs());
+      return 1;
+    }
+
+    ExitOnErr(J->addIRModule(ThreadSafeModule(std::move(M), std::move(Ctx))));
+  }
+
+  // Look up the entry point, cast it to a C main function pointer, then use
+  // runAsMain to call it.
+  auto EntrySym = ExitOnErr(J->lookup(EntryPointName));
+  auto EntryFn =
+      jitTargetAddressToFunction<int (*)(int, char *[])>(EntrySym.getAddress());
+
+  return runAsMain(EntryFn, InputArgv, StringRef(InputFiles.front()));
+}

llvm/include/llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h

@@ -16,6 +16,7 @@
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/ExecutionEngine/JITEventListener.h"
 #include "llvm/ExecutionEngine/JITSymbol.h"
 #include "llvm/ExecutionEngine/Orc/Core.h"
 #include "llvm/ExecutionEngine/Orc/Layer.h"
@@ -115,15 +116,23 @@ public:
     return *this;
   }
 
+  /// Register a JITEventListener.
+  void registerJITEventListener(JITEventListener &L);
+
+  /// Unregister a JITEventListener.
+  void unregisterJITEventListener(JITEventListener &L);
+
 private:
   Error onObjLoad(VModuleKey K, MaterializationResponsibility &R,
-                  object::ObjectFile &Obj,
+                  const object::ObjectFile &Obj,
+                  RuntimeDyld::MemoryManager *MemMgr,
                   std::unique_ptr<RuntimeDyld::LoadedObjectInfo> LoadedObjInfo,
                   std::map<StringRef, JITEvaluatedSymbol> Resolved,
                   std::set<StringRef> &InternalSymbols);
 
-  void onObjEmit(VModuleKey K, std::unique_ptr<MemoryBuffer> ObjBuffer,
+  void onObjEmit(VModuleKey K, MaterializationResponsibility &R,
-                 MaterializationResponsibility &R, Error Err);
+                 object::OwningBinary<object::ObjectFile> O,
+                 RuntimeDyld::MemoryManager *MemMgr, Error Err);
 
   mutable std::mutex RTDyldLayerMutex;
   GetMemoryManagerFunction GetMemoryManager;
@@ -133,6 +142,10 @@ private:
   bool OverrideObjectFlags = false;
   bool AutoClaimObjectSymbols = false;
   std::vector<std::unique_ptr<RuntimeDyld::MemoryManager>> MemMgrs;
+  std::vector<JITEventListener *> EventListeners;
+  DenseMap<RuntimeDyld::MemoryManager *,
+           std::unique_ptr<RuntimeDyld::LoadedObjectInfo>>
+      LoadedObjInfos;
 };
 
 class LegacyRTDyldObjectLinkingLayerBase {

llvm/include/llvm/ExecutionEngine/RuntimeDyld.h

@@ -267,15 +267,16 @@ public:
   void finalizeWithMemoryManagerLocking();
 
 private:
-  friend void
+  friend void jitLinkForORC(
-  jitLinkForORC(object::ObjectFile &Obj,
+      object::OwningBinary<object::ObjectFile> O,
-                std::unique_ptr<MemoryBuffer> UnderlyingBuffer,
+      RuntimeDyld::MemoryManager &MemMgr, JITSymbolResolver &Resolver,
-                RuntimeDyld::MemoryManager &MemMgr, JITSymbolResolver &Resolver,
+      bool ProcessAllSections,
-                bool ProcessAllSections,
+      unique_function<Error(const object::ObjectFile &Obj,
-                unique_function<Error(std::unique_ptr<LoadedObjectInfo>,
+                            std::unique_ptr<LoadedObjectInfo>,
-                                      std::map<StringRef, JITEvaluatedSymbol>)>
+                            std::map<StringRef, JITEvaluatedSymbol>)>
-                    OnLoaded,
+          OnLoaded,
-                unique_function<void(Error)> OnEmitted);
+      unique_function<void(object::OwningBinary<object::ObjectFile> O, Error)>
+          OnEmitted);
 
   // RuntimeDyldImpl is the actual class. RuntimeDyld is just the public
   // interface.
@@ -293,13 +294,15 @@ private:
 // instance and uses continuation passing to perform the fix-up and finalize
 // steps asynchronously.
 void jitLinkForORC(
-    object::ObjectFile &Obj, std::unique_ptr<MemoryBuffer> UnderlyingBuffer,
+    object::OwningBinary<object::ObjectFile> O,
     RuntimeDyld::MemoryManager &MemMgr, JITSymbolResolver &Resolver,
     bool ProcessAllSections,
-    unique_function<Error(std::unique_ptr<RuntimeDyld::LoadedObjectInfo>,
+    unique_function<Error(const object::ObjectFile &Obj,
+                          std::unique_ptr<RuntimeDyld::LoadedObjectInfo>,
                           std::map<StringRef, JITEvaluatedSymbol>)>
         OnLoaded,
-    unique_function<void(Error)> OnEmitted);
+    unique_function<void(object::OwningBinary<object::ObjectFile>, Error)>
+        OnEmitted);
 
 } // end namespace llvm
 

llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp

@@ -81,8 +81,12 @@ RTDyldObjectLinkingLayer::RTDyldObjectLinkingLayer(
 
 RTDyldObjectLinkingLayer::~RTDyldObjectLinkingLayer() {
   std::lock_guard<std::mutex> Lock(RTDyldLayerMutex);
-  for (auto &MemMgr : MemMgrs)
+  for (auto &MemMgr : MemMgrs) {
+    for (auto *L : EventListeners)
+      L->notifyFreeingObject(
+          static_cast<uint64_t>(reinterpret_cast<uintptr_t>(MemMgr.get())));
     MemMgr->deregisterEHFrames();
+  }
 }
 
 void RTDyldObjectLinkingLayer::emit(MaterializationResponsibility R,
@@ -97,13 +101,7 @@ void RTDyldObjectLinkingLayer::emit(MaterializationResponsibility R,
 
   auto &ES = getExecutionSession();
 
-  // Create a MemoryBufferRef backed MemoryBuffer (i.e. shallow) copy of the
+  auto Obj = object::ObjectFile::createObjectFile(*O);
-  // the underlying buffer to pass into RuntimeDyld. This allows us to hold
-  // ownership of the real underlying buffer and return it to the user once
-  // the object has been emitted.
-  auto ObjBuffer = MemoryBuffer::getMemBuffer(O->getMemBufferRef(), false);
-
-  auto Obj = object::ObjectFile::createObjectFile(*ObjBuffer);
 
   if (!Obj) {
     getExecutionSession().reportError(Obj.takeError());
@@ -154,20 +152,39 @@ void RTDyldObjectLinkingLayer::emit(MaterializationResponsibility R,
   JITDylibSearchOrderResolver Resolver(*SharedR);
 
   jitLinkForORC(
-      **Obj, std::move(O), *MemMgr, Resolver, ProcessAllSections,
+      object::OwningBinary<object::ObjectFile>(std::move(*Obj), std::move(O)),
-      [this, K, SharedR, &Obj, InternalSymbols](
+      *MemMgr, Resolver, ProcessAllSections,
+      [this, K, SharedR, MemMgr, InternalSymbols](
+          const object::ObjectFile &Obj,
           std::unique_ptr<RuntimeDyld::LoadedObjectInfo> LoadedObjInfo,
           std::map<StringRef, JITEvaluatedSymbol> ResolvedSymbols) {
-        return onObjLoad(K, *SharedR, **Obj, std::move(LoadedObjInfo),
+        return onObjLoad(K, *SharedR, Obj, MemMgr, std::move(LoadedObjInfo),
                          ResolvedSymbols, *InternalSymbols);
       },
-      [this, K, SharedR, O = std::move(O)](Error Err) mutable {
+      [this, K, SharedR, MemMgr](object::OwningBinary<object::ObjectFile> Obj,
-        onObjEmit(K, std::move(O), *SharedR, std::move(Err));
+                                 Error Err) mutable {
+        onObjEmit(K, *SharedR, std::move(Obj), MemMgr, std::move(Err));
       });
 }
 
+void RTDyldObjectLinkingLayer::registerJITEventListener(JITEventListener &L) {
+  std::lock_guard<std::mutex> Lock(RTDyldLayerMutex);
+  assert(llvm::none_of(EventListeners,
+                       [&](JITEventListener *O) { return O == &L; }) &&
+         "Listener has already been registered");
+  EventListeners.push_back(&L);
+}
+
+void RTDyldObjectLinkingLayer::unregisterJITEventListener(JITEventListener &L) {
+  std::lock_guard<std::mutex> Lock(RTDyldLayerMutex);
+  auto I = llvm::find(EventListeners, &L);
+  assert(I != EventListeners.end() && "Listener not registered");
+  EventListeners.erase(I);
+}
+
 Error RTDyldObjectLinkingLayer::onObjLoad(
-    VModuleKey K, MaterializationResponsibility &R, object::ObjectFile &Obj,
+    VModuleKey K, MaterializationResponsibility &R,
+    const object::ObjectFile &Obj, RuntimeDyld::MemoryManager *MemMgr,
     std::unique_ptr<RuntimeDyld::LoadedObjectInfo> LoadedObjInfo,
     std::map<StringRef, JITEvaluatedSymbol> Resolved,
     std::set<StringRef> &InternalSymbols) {
@@ -252,12 +269,17 @@ Error RTDyldObjectLinkingLayer::onObjLoad(
   if (NotifyLoaded)
     NotifyLoaded(K, Obj, *LoadedObjInfo);
 
+  std::lock_guard<std::mutex> Lock(RTDyldLayerMutex);
+  assert(!LoadedObjInfos.count(MemMgr) && "Duplicate loaded info for MemMgr");
+  LoadedObjInfos[MemMgr] = std::move(LoadedObjInfo);
+
   return Error::success();
 }
 
 void RTDyldObjectLinkingLayer::onObjEmit(
-    VModuleKey K, std::unique_ptr<MemoryBuffer> ObjBuffer,
+    VModuleKey K, MaterializationResponsibility &R,
-    MaterializationResponsibility &R, Error Err) {
+    object::OwningBinary<object::ObjectFile> O,
+    RuntimeDyld::MemoryManager *MemMgr, Error Err) {
   if (Err) {
     getExecutionSession().reportError(std::move(Err));
     R.failMaterialization();
@@ -270,6 +292,22 @@ void RTDyldObjectLinkingLayer::onObjEmit(
     return;
   }
 
+  std::unique_ptr<object::ObjectFile> Obj;
+  std::unique_ptr<MemoryBuffer> ObjBuffer;
+  std::tie(Obj, ObjBuffer) = O.takeBinary();
+
+  // Run EventListener notifyLoaded callbacks.
+  {
+    std::lock_guard<std::mutex> Lock(RTDyldLayerMutex);
+    auto LOIItr = LoadedObjInfos.find(MemMgr);
+    assert(LOIItr != LoadedObjInfos.end() && "LoadedObjInfo missing");
+    for (auto *L : EventListeners)
+      L->notifyObjectLoaded(
+          static_cast<uint64_t>(reinterpret_cast<uintptr_t>(MemMgr)), *Obj,
+          *LOIItr->second);
+    LoadedObjInfos.erase(MemMgr);
+  }
+
   if (NotifyEmitted)
     NotifyEmitted(K, std::move(ObjBuffer));
 }

llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp

@@ -1190,16 +1190,16 @@ Error RuntimeDyldImpl::resolveExternalSymbols() {
 
 void RuntimeDyldImpl::finalizeAsync(
     std::unique_ptr<RuntimeDyldImpl> This,
-    unique_function<void(Error)> OnEmitted,
+    unique_function<void(object::OwningBinary<object::ObjectFile>, Error)>
-    std::unique_ptr<MemoryBuffer> UnderlyingBuffer) {
+        OnEmitted,
+    object::OwningBinary<object::ObjectFile> O) {
 
   auto SharedThis = std::shared_ptr<RuntimeDyldImpl>(std::move(This));
   auto PostResolveContinuation =
-      [SharedThis, OnEmitted = std::move(OnEmitted),
+      [SharedThis, OnEmitted = std::move(OnEmitted), O = std::move(O)](
-       UnderlyingBuffer = std::move(UnderlyingBuffer)](
           Expected<JITSymbolResolver::LookupResult> Result) mutable {
         if (!Result) {
-          OnEmitted(Result.takeError());
+          OnEmitted(std::move(O), Result.takeError());
           return;
         }
 
@@ -1213,10 +1213,11 @@ void RuntimeDyldImpl::finalizeAsync(
         SharedThis->registerEHFrames();
         std::string ErrMsg;
         if (SharedThis->MemMgr.finalizeMemory(&ErrMsg))
-          OnEmitted(make_error<StringError>(std::move(ErrMsg),
+          OnEmitted(std::move(O),
+                    make_error<StringError>(std::move(ErrMsg),
                                             inconvertibleErrorCode()));
         else
-          OnEmitted(Error::success());
+          OnEmitted(std::move(O), Error::success());
       };
 
   JITSymbolResolver::LookupSet Symbols;
@@ -1403,32 +1404,35 @@ void RuntimeDyld::deregisterEHFrames() {
 // FIXME: Kill this with fire once we have a new JIT linker: this is only here
 // so that we can re-use RuntimeDyld's implementation without twisting the
 // interface any further for ORC's purposes.
-void jitLinkForORC(object::ObjectFile &Obj,
+void jitLinkForORC(
-                   std::unique_ptr<MemoryBuffer> UnderlyingBuffer,
+    object::OwningBinary<object::ObjectFile> O,
-                   RuntimeDyld::MemoryManager &MemMgr,
+    RuntimeDyld::MemoryManager &MemMgr, JITSymbolResolver &Resolver,
-                   JITSymbolResolver &Resolver, bool ProcessAllSections,
+    bool ProcessAllSections,
-                   unique_function<Error(
+    unique_function<
-                       std::unique_ptr<RuntimeDyld::LoadedObjectInfo> LoadedObj,
+        Error(const object::ObjectFile &Obj,
-                       std::map<StringRef, JITEvaluatedSymbol>)>
+              std::unique_ptr<RuntimeDyld::LoadedObjectInfo> LoadedObj,
-                       OnLoaded,
+              std::map<StringRef, JITEvaluatedSymbol>)>
-                   unique_function<void(Error)> OnEmitted) {
+        OnLoaded,
+    unique_function<void(object::OwningBinary<object::ObjectFile>, Error)>
+        OnEmitted) {
 
   RuntimeDyld RTDyld(MemMgr, Resolver);
   RTDyld.setProcessAllSections(ProcessAllSections);
 
-  auto Info = RTDyld.loadObject(Obj);
+  auto Info = RTDyld.loadObject(*O.getBinary());
 
   if (RTDyld.hasError()) {
-    OnEmitted(make_error<StringError>(RTDyld.getErrorString(),
+    OnEmitted(std::move(O), make_error<StringError>(RTDyld.getErrorString(),
-                                      inconvertibleErrorCode()));
+                                                    inconvertibleErrorCode()));
     return;
   }
 
-  if (auto Err = OnLoaded(std::move(Info), RTDyld.getSymbolTable()))
+  if (auto Err =
-    OnEmitted(std::move(Err));
+          OnLoaded(*O.getBinary(), std::move(Info), RTDyld.getSymbolTable()))
+    OnEmitted(std::move(O), std::move(Err));
 
   RuntimeDyldImpl::finalizeAsync(std::move(RTDyld.Dyld), std::move(OnEmitted),
-                                 std::move(UnderlyingBuffer));
+                                 std::move(O));
 }
 
 } // end namespace llvm

llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldImpl.h

@@ -549,9 +549,11 @@ public:
 
   void resolveLocalRelocations();
 
-  static void finalizeAsync(std::unique_ptr<RuntimeDyldImpl> This,
+  static void finalizeAsync(
-                            unique_function<void(Error)> OnEmitted,
+      std::unique_ptr<RuntimeDyldImpl> This,
-                            std::unique_ptr<MemoryBuffer> UnderlyingBuffer);
+      unique_function<void(object::OwningBinary<object::ObjectFile>, Error)>
+          OnEmitted,
+      object::OwningBinary<object::ObjectFile> O);
 
   void reassignSectionAddress(unsigned SectionID, uint64_t Addr);
 

llvm/tools/lli/lli.cpp

@@ -30,6 +30,7 @@
 #include "llvm/ExecutionEngine/Orc/LLJIT.h"
 #include "llvm/ExecutionEngine/Orc/MachOPlatform.h"
 #include "llvm/ExecutionEngine/Orc/OrcRemoteTargetClient.h"
+#include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
 #include "llvm/ExecutionEngine/OrcMCJITReplacement.h"
 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
 #include "llvm/IR/IRBuilder.h"
@@ -891,6 +892,11 @@ int runOrcLazyJIT(const char *ProgName) {
 
   auto J = ExitOnErr(Builder.create());
 
+  if (TT->isOSBinFormatELF())
+    static_cast<llvm::orc::RTDyldObjectLinkingLayer &>(J->getObjLinkingLayer())
+        .registerJITEventListener(
+            *JITEventListener::createGDBRegistrationListener());
+
   if (PerModuleLazy)
     J->setPartitionFunction(orc::CompileOnDemandLayer::compileWholeModule);