2 files changed, 153 insertions, 51 deletions

diff --git a/llvm/lib/ExecutionEngine/JITLink/EHFrameSupport.cpp b/llvm/lib/ExecutionEngine/JITLink/EHFrameSupport.cpp
index c531df82446..b54fd378aa5 100644
--- a/llvm/lib/ExecutionEngine/JITLink/EHFrameSupport.cpp
+++ b/llvm/lib/ExecutionEngine/JITLink/EHFrameSupport.cpp
@@ -507,8 +507,9 @@ Error deregisterEHFrameSection(const void *EHFrameSectionAddr) {
 #endif
 }
 
-AtomGraphPassFunction createEHFrameRecorderPass(const Triple &TT,
-                                                JITTargetAddress &EHFrameAddr) {
+AtomGraphPassFunction
+createEHFrameRecorderPass(const Triple &TT,
+                          StoreFrameAddressFunction StoreFrameAddress) {
   const char *EHFrameSectionName = nullptr;
   if (TT.getObjectFormat() == Triple::MachO)
     EHFrameSectionName = "__eh_frame";
@@ -516,7 +517,7 @@ AtomGraphPassFunction createEHFrameRecorderPass(const Triple &TT,
     EHFrameSectionName = ".eh_frame";
 
   auto RecordEHFrame = [EHFrameSectionName,
-                        &EHFrameAddr](AtomGraph &G) -> Error {
+                        StoreFrameAddress](AtomGraph &G) -> Error {
     // Search for a non-empty eh-frame and record the address of the first atom
     // in it.
     JITTargetAddress Addr = 0;
@@ -529,7 +530,7 @@ AtomGraphPassFunction createEHFrameRecorderPass(const Triple &TT,
         break;
       }
 
-    EHFrameAddr = Addr;
+    StoreFrameAddress(Addr);
     return Error::success();
   };
 
diff --git a/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp b/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
index 154d7040b76..eb0ea47e2d4 100644
--- a/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
@@ -133,25 +133,19 @@ public:
 
     MR.resolve(InternedResult);
 
-    if (Layer.NotifyLoaded)
-      Layer.NotifyLoaded(MR.getVModuleKey());
+    Layer.notifyLoaded(MR);
   }
 
   void notifyFinalized(
       std::unique_ptr<JITLinkMemoryManager::Allocation> A) override {
 
-    if (EHFrameAddr) {
-      // If there is an eh-frame then try to register it.
-      if (auto Err = registerEHFrameSection((void *)EHFrameAddr)) {
-        Layer.getExecutionSession().reportError(std::move(Err));
-        MR.failMaterialization();
-        return;
-      }
-    }
+    if (auto Err = Layer.notifyEmitted(MR, std::move(A))) {
+      Layer.getExecutionSession().reportError(std::move(Err));
+      MR.failMaterialization();
 
+      return;
+    }
     MR.emit();
-    Layer.notifyFinalized(
-        ObjectLinkingLayer::ObjectResources(std::move(A), EHFrameAddr));
   }
 
   AtomGraphPassFunction getMarkLivePass(const Triple &TT) const override {
@@ -166,11 +160,7 @@ public:
     Config.PostPrunePasses.push_back(
         [this](AtomGraph &G) { return computeNamedSymbolDependencies(G); });
 
-    Config.PostFixupPasses.push_back(
-        createEHFrameRecorderPass(TT, EHFrameAddr));
-
-    if (Layer.ModifyPassConfig)
-      Layer.ModifyPassConfig(TT, Config);
+    Layer.modifyPassConfig(MR, TT, Config);
 
     return Error::success();
   }
@@ -328,16 +318,18 @@ private:
   MaterializationResponsibility MR;
   std::unique_ptr<MemoryBuffer> ObjBuffer;
   DenseMap<SymbolStringPtr, SymbolNameSet> NamedSymbolDeps;
-  JITTargetAddress EHFrameAddr = 0;
 };
 
-ObjectLinkingLayer::ObjectLinkingLayer(
-    ExecutionSession &ES, JITLinkMemoryManager &MemMgr,
-    NotifyLoadedFunction NotifyLoaded, NotifyEmittedFunction NotifyEmitted,
-    ModifyPassConfigFunction ModifyPassConfig)
-    : ObjectLayer(ES), MemMgr(MemMgr), NotifyLoaded(std::move(NotifyLoaded)),
-      NotifyEmitted(std::move(NotifyEmitted)),
-      ModifyPassConfig(std::move(ModifyPassConfig)) {}
+ObjectLinkingLayer::Plugin::~Plugin() {}
+
+ObjectLinkingLayer::ObjectLinkingLayer(ExecutionSession &ES,
+                                       JITLinkMemoryManager &MemMgr)
+    : ObjectLayer(ES), MemMgr(MemMgr) {}
+
+ObjectLinkingLayer::~ObjectLinkingLayer() {
+  if (auto Err = removeAllModules())
+    getExecutionSession().reportError(std::move(Err));
+}
 
 void ObjectLinkingLayer::emit(MaterializationResponsibility R,
                               std::unique_ptr<MemoryBuffer> O) {
@@ -346,36 +338,145 @@ void ObjectLinkingLayer::emit(MaterializationResponsibility R,
       *this, std::move(R), std::move(O)));
 }
 
-ObjectLinkingLayer::ObjectResources::ObjectResources(
-    AllocPtr Alloc, JITTargetAddress EHFrameAddr)
-    : Alloc(std::move(Alloc)), EHFrameAddr(EHFrameAddr) {}
+void ObjectLinkingLayer::modifyPassConfig(MaterializationResponsibility &MR,
+                                          const Triple &TT,
+                                          PassConfiguration &PassConfig) {
+  for (auto &P : Plugins)
+    P->modifyPassConfig(MR, TT, PassConfig);
+}
+
+void ObjectLinkingLayer::notifyLoaded(MaterializationResponsibility &MR) {
+  for (auto &P : Plugins)
+    P->notifyLoaded(MR);
+}
+
+Error ObjectLinkingLayer::notifyEmitted(MaterializationResponsibility &MR,
+                                        AllocPtr Alloc) {
+  Error Err = Error::success();
+  for (auto &P : Plugins)
+    Err = joinErrors(std::move(Err), P->notifyEmitted(MR));
+
+  if (Err)
+    return Err;
+
+  {
+    std::lock_guard<std::mutex> Lock(LayerMutex);
+    UntrackedAllocs.push_back(std::move(Alloc));
+  }
 
-ObjectLinkingLayer::ObjectResources::ObjectResources(ObjectResources &&Other)
-    : Alloc(std::move(Other.Alloc)), EHFrameAddr(Other.EHFrameAddr) {
-  Other.EHFrameAddr = 0;
+  return Error::success();
 }
 
-ObjectLinkingLayer::ObjectResources &
-ObjectLinkingLayer::ObjectResources::operator=(ObjectResources &&Other) {
-  std::swap(Alloc, Other.Alloc);
-  std::swap(EHFrameAddr, Other.EHFrameAddr);
-  return *this;
+Error ObjectLinkingLayer::removeModule(VModuleKey K) {
+  Error Err = Error::success();
+
+  for (auto &P : Plugins)
+    Err = joinErrors(std::move(Err), P->notifyRemovingModule(K));
+
+  AllocPtr Alloc;
+
+  {
+    std::lock_guard<std::mutex> Lock(LayerMutex);
+    auto AllocItr = TrackedAllocs.find(K);
+    Alloc = std::move(AllocItr->second);
+    TrackedAllocs.erase(AllocItr);
+  }
+
+  assert(Alloc && "No allocation for key K");
+
+  return joinErrors(std::move(Err), Alloc->deallocate());
 }
 
-ObjectLinkingLayer::ObjectResources::~ObjectResources() {
-  const char *ErrBanner =
-      "ObjectLinkingLayer received error deallocating object resources:";
+Error ObjectLinkingLayer::removeAllModules() {
 
-  assert((EHFrameAddr == 0 || Alloc) &&
-         "Non-null EHFrameAddr must have an associated allocation");
+  Error Err = Error::success();
 
-  if (EHFrameAddr)
-    if (auto Err = deregisterEHFrameSection((void *)EHFrameAddr))
-      logAllUnhandledErrors(std::move(Err), llvm::errs(), ErrBanner);
+  for (auto &P : Plugins)
+    Err = joinErrors(std::move(Err), P->notifyRemovingAllModules());
+
+  std::vector<AllocPtr> Allocs;
+  {
+    std::lock_guard<std::mutex> Lock(LayerMutex);
+    Allocs = std::move(UntrackedAllocs);
+
+    for (auto &KV : TrackedAllocs)
+      Allocs.push_back(std::move(KV.second));
+
+    TrackedAllocs.clear();
+  }
+
+  while (!Allocs.empty()) {
+    Err = joinErrors(std::move(Err), Allocs.back()->deallocate());
+    Allocs.pop_back();
+  }
+
+  return Err;
+}
+
+void LocalEHFrameRegistrationPlugin::modifyPassConfig(
+    MaterializationResponsibility &MR, const Triple &TT,
+    PassConfiguration &PassConfig) {
+  assert(!InProcessLinks.count(&MR) && "Link for MR already being tracked?");
+
+  PassConfig.PostFixupPasses.push_back(
+      createEHFrameRecorderPass(TT, [this, &MR](JITTargetAddress Addr) {
+        if (Addr)
+          InProcessLinks[&MR] = jitTargetAddressToPointer<void *>(Addr);
+      }));
+}
+
+Error LocalEHFrameRegistrationPlugin::notifyEmitted(
+    MaterializationResponsibility &MR) {
+
+  auto EHFrameAddrItr = InProcessLinks.find(&MR);
+  if (EHFrameAddrItr == InProcessLinks.end())
+    return Error::success();
+
+  const void *EHFrameAddr = EHFrameAddrItr->second;
+  assert(EHFrameAddr && "eh-frame addr to register can not be null");
+
+  InProcessLinks.erase(EHFrameAddrItr);
+  if (auto Key = MR.getVModuleKey())
+    TrackedEHFrameAddrs[Key] = EHFrameAddr;
+  else
+    UntrackedEHFrameAddrs.push_back(EHFrameAddr);
+
+  return registerEHFrameSection(EHFrameAddr);
+}
+
+Error LocalEHFrameRegistrationPlugin::notifyRemovingModule(VModuleKey K) {
+  auto EHFrameAddrItr = TrackedEHFrameAddrs.find(K);
+  if (EHFrameAddrItr == TrackedEHFrameAddrs.end())
+    return Error::success();
+
+  const void *EHFrameAddr = EHFrameAddrItr->second;
+  assert(EHFrameAddr && "Tracked eh-frame addr must not be null");
+
+  TrackedEHFrameAddrs.erase(EHFrameAddrItr);
+
+  return deregisterEHFrameSection(EHFrameAddr);
+}
+
+Error LocalEHFrameRegistrationPlugin::notifyRemovingAllModules() {
+
+  std::vector<const void *> EHFrameAddrs = std::move(UntrackedEHFrameAddrs);
+  EHFrameAddrs.reserve(EHFrameAddrs.size() + TrackedEHFrameAddrs.size());
+
+  for (auto &KV : TrackedEHFrameAddrs)
+    EHFrameAddrs.push_back(KV.second);
+
+  TrackedEHFrameAddrs.clear();
+
+  Error Err = Error::success();
+
+  while (!EHFrameAddrs.empty()) {
+    const void *EHFrameAddr = EHFrameAddrs.back();
+    assert(EHFrameAddr && "Untracked eh-frame addr must not be null");
+    EHFrameAddrs.pop_back();
+    Err = joinErrors(std::move(Err), deregisterEHFrameSection(EHFrameAddr));
+  }
 
-  if (Alloc)
-    if (auto Err = Alloc->deallocate())
-      logAllUnhandledErrors(std::move(Err), llvm::errs(), ErrBanner);
+  return Err;
 }
 
 } // End namespace orc.