Make RuntimeDyld::MemoryManager responsible for tracking EH frames.

[lhames]

Summary

RuntimeDyld is currently responsible for tracking allocated EH frames, but it makes more sense to have the RuntimeDyld::MemoryManager track them (since the frames are allocated through the memory manager, and written to memory owned by the memory manager). This patch moves the frame tracking into RTDyldMemoryManager, and changes the deregisterFrames method on RuntimeDyld::MemoryManager from:

void deregisterEHFrames(uint8_t *Addr, uint64_t LoadAddr, size_t Size);

to just

void deregisterEHFrames();

Separating this responsibility will allow ORC to continue to throw the RuntimeDyld instances away post-link (saving a few dozen bytes per lazy function) while properly deregistering frames when modules are unloaded.

This patch also updates ORC to call deregisterEHFrames when modules are unloaded. This fixes a bug where an exception that tears down the JIT can then unwind through dangling EH frames that have been deallocated but not deregistered, resulting in UB.

For people using SectionMemoryManager this should be pretty much a no-op. For people with custom allocators that override registerEHFrames/deregisterEHFrames, you will now be responsible for tracking allocated EH frames.

[lhames]

The bug that motivated this change is http://llvm.org/PR32901 .

[chfast]

include/llvm/ExecutionEngine/Orc/OrcRemoteTargetClient.h

147

How about emplace_back(LoadAddr, static_cast<uint32_t>(Size))?
BTW, why cast to uint32_t?

335

I'm not sure it will it more clear, but you write UnfinalizedEHFrames = {}.

394

Can you use size_t for Size? That might avoid the previous static_cast.

include/llvm/ExecutionEngine/RTDyldMemoryManager.h

141

And there we have uint64_t? Can it be size_t as well?

lib/ExecutionEngine/RuntimeDyld/RTDyldMemoryManager.cpp

140

emplace_back?

[lhames]

Will make the suggested changes and upload a new patch.

include/llvm/ExecutionEngine/Orc/OrcRemoteTargetClient.h

147

Yep - emplace_back will work.

The cast is there to suppress the implicit conversion warning. The situation with the Size argument is a bit weird: The method declares Size to be a size_t, but we want this to work even when we're cross-JITing from 32 to 64-bit, so we should probably fix it to uint64_t. I think I'll leave this as is for now and update the interface in a subsequent patch.

335

Oh - that's much neater.

include/llvm/ExecutionEngine/RTDyldMemoryManager.h

141

Oddly enough (after the discussion above), this specific size member could be a size_t, since RTDyldMemoryManager's default implementation for register/deregister is local-only. (That's kind of broken in its own way, but that's an issue for another patch).

lib/ExecutionEngine/RuntimeDyld/RTDyldMemoryManager.cpp

140

Yep.

[lhames]

Actually those push_backs should stay push_backs: To use emplace the structs would have to have non-default constructors, which seems like more effort than its worth for the conceptual neatness of emplace.

[lhames]

Switch EHFrames::Size to size_t, discard now-redundant static_cast, use UnfinalizedEHFrames = {}.

[chfast]

include/llvm/ExecutionEngine/Orc/OrcRemoteTargetClient.h

147

The Size argument of type size_t is used here, but EHFrame struct expects uint64_t. Can we clean it up?

392

It's hard to see from the diff why this is declared twice.

[lhames]

include/llvm/ExecutionEngine/Orc/OrcRemoteTargetClient.h

147

As mentioned in the previous comment the correct size for the EHFrame struct is uint64_t. The interface should also change from size_t to uint64_t, but I think that warrants a separate patch.

392

You mean in OrcRemoteTargetClient.h and RTDyldMemoryManager.h? They're doing slightly different things. For backwards compatibility RTDyldMemoryManager is registering frames in the current process only. The RemoteClient version is dedicated to registering frames on a remote. That's why one uses JITTargetAddress/uint64_t and the other can get away with a pointer/size_t.

[lhames]

Thanks! Committed in r302589.

include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h

       return nullptr;
     }
 
+    void removeModulesFromBaseLayer(BaseLayerT &BaseLayer) {
+      for (auto &BLH : BaseLayerHandles)
+        BaseLayer.removeModuleSet(BLH);
+    }
+
     std::unique_ptr<JITSymbolResolver> ExternalSymbolResolver;
     std::unique_ptr<ResourceOwner<RuntimeDyld::MemoryManager>> MemMgr;
     std::unique_ptr<IndirectStubsMgrT> StubsMgr;

         CreateIndirectStubsManager(std::move(CreateIndirectStubsManager)),
         CloneStubsIntoPartitions(CloneStubsIntoPartitions) {}
 
+  ~CompileOnDemandLayer() {
+    while (!LogicalDylibs.empty())
+      removeModuleSet(LogicalDylibs.begin());
+  }
+  
   /// @brief Add a module to the compile-on-demand layer.
   template <typename ModuleSetT, typename MemoryManagerPtrT,
             typename SymbolResolverPtrT>

   ///   This will remove all modules in the layers below that were derived from
   /// the module represented by H.
   void removeModuleSet(ModuleSetHandleT H) {
+    H->removeModulesFromBaseLayer(BaseLayer);
     LogicalDylibs.erase(H);
   }
 

         return 0;
     }
 
+    LD.BaseLayerHandles.push_back(PartH);
+
     return CalledAddr;
   }
 

include/llvm/ExecutionEngine/Orc/OrcRemoteTargetClient.h

 
     void registerEHFrames(uint8_t *Addr, uint64_t LoadAddr,
                           size_t Size) override {
-      UnfinalizedEHFrames.push_back(
-          std::make_pair(LoadAddr, static_cast<uint32_t>(Size)));
+      UnfinalizedEHFrames.push_back({LoadAddr, static_cast<uint32_t>(Size)});
     }
 
-    void deregisterEHFrames(uint8_t *Addr, uint64_t LoadAddr,
-                            size_t Size) override {
-      auto Err = Client.deregisterEHFrames(LoadAddr, Size);
+    void deregisterEHFrames() override {
+      for (auto &Frame : RegisteredEHFrames) {
+        auto Err = Client.deregisterEHFrames(Frame.Addr, Frame.Size);
         // FIXME: Add error poll.
         assert(!Err && "Failed to register remote EH frames.");
         (void)Err;
+      }
     }
 
     void notifyObjectLoaded(RuntimeDyld &Dyld,

       Unfinalized.clear();
 
       for (auto &EHFrame : UnfinalizedEHFrames) {
-        if (auto Err = Client.registerEHFrames(EHFrame.first, EHFrame.second)) {
+        if (auto Err = Client.registerEHFrames(EHFrame.Addr, EHFrame.Size)) {
           // FIXME: Replace this once finalizeMemory can return an Error.
           handleAllErrors(std::move(Err), [&](ErrorInfoBase &EIB) {
             if (ErrMsg) {

           return false;
         }
       }
-      UnfinalizedEHFrames.clear();
+      RegisteredEHFrames = std::move(UnfinalizedEHFrames);
+      UnfinalizedEHFrames = std::vector<EHFrame>();
 
       return false;
     }

     ResourceIdMgr::ResourceId Id;
     std::vector<ObjectAllocs> Unmapped;
     std::vector<ObjectAllocs> Unfinalized;
-    std::vector<std::pair<uint64_t, uint32_t>> UnfinalizedEHFrames;
+
+    struct EHFrame {
+      JITTargetAddress Addr;
+      uint32_t Size;
+    };
+    std::vector<EHFrame> UnfinalizedEHFrames;
+    std::vector<EHFrame> RegisteredEHFrames;
   };
 
   /// Remote indirect stubs manager.

include/llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h

       buildInitialSymbolTable(PFC->Objects);
     }
 
+    ~ConcreteLinkedObjectSet() override {
+      MemMgr->deregisterEHFrames();
+    }
+    
     void setHandle(ObjSetHandleT H) {
       PFC->Handle = H;
     }

include/llvm/ExecutionEngine/RTDyldMemoryManager.h

   /// Deregister EH frames in the current proces.
   static void deregisterEHFramesInProcess(uint8_t *Addr, size_t Size);
 
-  void registerEHFrames(uint8_t *Addr, uint64_t LoadAddr, size_t Size) override {
-    registerEHFramesInProcess(Addr, Size);
-  }
+  void registerEHFrames(uint8_t *Addr, uint64_t LoadAddr, size_t Size) override;
+  void deregisterEHFrames() override;
 
-  void deregisterEHFrames(uint8_t *Addr, uint64_t LoadAddr, size_t Size) override {
-    deregisterEHFramesInProcess(Addr, Size);
-  }
-
   /// This method returns the address of the specified function or variable in
   /// the current process.
   static uint64_t getSymbolAddressInProcess(const std::string &Name);

   /// MCJIT or RuntimeDyld.  Use getSymbolAddress instead.
   virtual void *getPointerToNamedFunction(const std::string &Name,
                                           bool AbortOnFailure = true);
+
+private:
+  struct EHFrame {
+    uint8_t *Addr;
+    uint64_t Size;
+  };
+  std::vector<EHFrame> EHFrames;
 };
 
 // Create wrappers for C Binding types (see CBindingWrapping.h).

include/llvm/ExecutionEngine/RuntimeDyld.h

     /// be the case for local execution) these two values will be the same.
     virtual void registerEHFrames(uint8_t *Addr, uint64_t LoadAddr,
                                   size_t Size) = 0;
-    virtual void deregisterEHFrames(uint8_t *addr, uint64_t LoadAddr,
-                                    size_t Size) = 0;
+    virtual void deregisterEHFrames() = 0;
 
     /// This method is called when object loading is complete and section page
     /// permissions can be applied.  It is up to the memory manager implementation

lib/ExecutionEngine/Orc/OrcMCJITReplacement.h

       return ClientMM->registerEHFrames(Addr, LoadAddr, Size);
     }
 
-    void deregisterEHFrames(uint8_t *Addr, uint64_t LoadAddr,
-                            size_t Size) override {
-      return ClientMM->deregisterEHFrames(Addr, LoadAddr, Size);
+    void deregisterEHFrames() override {
+      return ClientMM->deregisterEHFrames();
     }
 
     void notifyObjectLoaded(RuntimeDyld &RTDyld,

lib/ExecutionEngine/RuntimeDyld/RTDyldMemoryManager.cpp

 
 #endif
 
+void RTDyldMemoryManager::registerEHFrames(uint8_t *Addr, uint64_t LoadAddr,
+                                          size_t Size) {
+  registerEHFramesInProcess(Addr, Size);
+  EHFrames.push_back({Addr, Size});
+}
+
+void RTDyldMemoryManager::deregisterEHFrames() {
+  for (auto &Frame : EHFrames)
+    deregisterEHFramesInProcess(Frame.Addr, Frame.Size);
+  EHFrames.clear();
+}
+
 static int jit_noop() {
   return 0;
 }

lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp

 
 void RuntimeDyldImpl::registerEHFrames() {}
 
-void RuntimeDyldImpl::deregisterEHFrames() {}
+void RuntimeDyldImpl::deregisterEHFrames() {
+  MemMgr.deregisterEHFrames();
+}
 
 #ifndef NDEBUG
 static void dumpSectionMemory(const SectionEntry &S, StringRef State) {

lib/ExecutionEngine/RuntimeDyld/RuntimeDyldELF.h

   // in a table until we receive a request to register all unregistered
   // EH frame sections with the memory manager.
   SmallVector<SID, 2> UnregisteredEHFrameSections;
-  SmallVector<SID, 2> RegisteredEHFrameSections;
 
   // Map between GOT relocation value and corresponding GOT offset
   std::map<RelocationValueRef, uint64_t> GOTOffsetMap;

                        StubMap &Stubs) override;
   bool isCompatibleFile(const object::ObjectFile &Obj) const override;
   void registerEHFrames() override;
-  void deregisterEHFrames() override;
   Error finalizeLoad(const ObjectFile &Obj,
                      ObjSectionToIDMap &SectionMap) override;
 };

lib/ExecutionEngine/RuntimeDyld/RuntimeDyldELF.cpp

     uint64_t EHFrameLoadAddr = Sections[EHFrameSID].getLoadAddress();
     size_t EHFrameSize = Sections[EHFrameSID].getSize();
     MemMgr.registerEHFrames(EHFrameAddr, EHFrameLoadAddr, EHFrameSize);
-    RegisteredEHFrameSections.push_back(EHFrameSID);
   }
   UnregisteredEHFrameSections.clear();
 }
 
-void RuntimeDyldELF::deregisterEHFrames() {
-  for (int i = 0, e = RegisteredEHFrameSections.size(); i != e; ++i) {
-    SID EHFrameSID = RegisteredEHFrameSections[i];
-    uint8_t *EHFrameAddr = Sections[EHFrameSID].getAddress();
-    uint64_t EHFrameLoadAddr = Sections[EHFrameSID].getLoadAddress();
-    size_t EHFrameSize = Sections[EHFrameSID].getSize();
-    MemMgr.deregisterEHFrames(EHFrameAddr, EHFrameLoadAddr, EHFrameSize);
-  }
-  RegisteredEHFrameSections.clear();
-}
-
 std::unique_ptr<RuntimeDyldELF>
 llvm::RuntimeDyldELF::create(Triple::ArchType Arch,
                              RuntimeDyld::MemoryManager &MemMgr,

lib/ExecutionEngine/RuntimeDyld/RuntimeDyldImpl.h

 
   virtual void registerEHFrames();
 
-  virtual void deregisterEHFrames();
+  void deregisterEHFrames();
 
   virtual Error finalizeLoad(const ObjectFile &ObjImg,
                              ObjSectionToIDMap &SectionMap) {

lib/ExecutionEngine/RuntimeDyld/Targets/RuntimeDyldCOFFI386.h

   }
 
   void registerEHFrames() override {}
-  void deregisterEHFrames() override {}
 };
 
 }

lib/ExecutionEngine/RuntimeDyld/Targets/RuntimeDyldCOFFThumb.h

   }
 
   void registerEHFrames() override {}
-  void deregisterEHFrames() override {}
 };
 
 }

lib/ExecutionEngine/RuntimeDyld/Targets/RuntimeDyldCOFFX86_64.h

     }
     UnregisteredEHFrameSections.clear();
   }
-  void deregisterEHFrames() override {
-    // Stub
-  }
   Error finalizeLoad(const ObjectFile &Obj,
                      ObjSectionToIDMap &SectionMap) override {
     // Look for and record the EH frame section IDs.

tools/lli/RemoteJITUtils.h

     MemMgr->registerEHFrames(Addr, LoadAddr, Size);
   }
 
-  void deregisterEHFrames(uint8_t *Addr, uint64_t LoadAddr,
-                          size_t Size) override {
-    MemMgr->deregisterEHFrames(Addr, LoadAddr, Size);
+  void deregisterEHFrames() override {
+    MemMgr->deregisterEHFrames();
   }
 
   bool finalizeMemory(std::string *ErrMsg = nullptr) override {

tools/llvm-rtdyld/llvm-rtdyld.cpp

 
   void registerEHFrames(uint8_t *Addr, uint64_t LoadAddr,
                         size_t Size) override {}
-  void deregisterEHFrames(uint8_t *Addr, uint64_t LoadAddr,
-                          size_t Size) override {}
+  void deregisterEHFrames() override {}
 
   void preallocateSlab(uint64_t Size) {
     std::string Err;

unittests/ExecutionEngine/Orc/ObjectTransformLayerTest.cpp

       return nullptr;
     }
     void registerEHFrames(uint8_t *, uint64_t, size_t) override {}
-    void deregisterEHFrames(uint8_t *, uint64_t, size_t) override {}
+    void deregisterEHFrames() override {}
     bool finalizeMemory(std::string *) override { return false; }
   };
 

unittests/ExecutionEngine/Orc/RTDyldObjectLinkingLayerTest.cpp

   Objs.push_back(OwningObj.getBinary());
 
   bool DebugSectionSeen = false;
-  SectionMemoryManagerWrapper SMMW(DebugSectionSeen);
+  auto SMMW =
+    std::make_shared<SectionMemoryManagerWrapper>(DebugSectionSeen);
   auto Resolver =
     createLambdaResolver(
       [](const std::string &Name) {

 
   {
     // Test with ProcessAllSections = false (the default).
-    auto H = ObjLayer.addObjectSet(Objs, &SMMW, &*Resolver);
+    auto H = ObjLayer.addObjectSet(Objs, SMMW, &*Resolver);
     ObjLayer.emitAndFinalize(H);
     EXPECT_EQ(DebugSectionSeen, false)
       << "Unexpected debug info section";

   {
     // Test with ProcessAllSections = true.
     ObjLayer.setProcessAllSections(true);
-    auto H = ObjLayer.addObjectSet(Objs, &SMMW, &*Resolver);
+    auto H = ObjLayer.addObjectSet(Objs, SMMW, &*Resolver);
     ObjLayer.emitAndFinalize(H);
     EXPECT_EQ(DebugSectionSeen, true)
       << "Expected debug info section not seen";

         return JITSymbol(nullptr);
       });
 
-  SectionMemoryManagerWrapper SMMW;
-  ObjLayer.addObjectSet(std::move(Obj1Set), &SMMW, &*Resolver);
-  auto H = ObjLayer.addObjectSet(std::move(Obj2Set), &SMMW, &*Resolver);
+  auto SMMW = std::make_shared<SectionMemoryManagerWrapper>();
+  ObjLayer.addObjectSet(std::move(Obj1Set), SMMW, &*Resolver);
+  auto H = ObjLayer.addObjectSet(std::move(Obj2Set), SMMW, &*Resolver);
   ObjLayer.emitAndFinalize(H);
-
+  ObjLayer.removeObjectSet(H);
+  
   // Finalization of module 2 should trigger finalization of module 1.
   // Verify that finalize on SMMW is only called once.
-  EXPECT_EQ(SMMW.FinalizationCount, 1)
+  EXPECT_EQ(SMMW->FinalizationCount, 1)
       << "Extra call to finalize";
 }
 

   std::vector<object::ObjectFile*> Obj2Set;
   Obj2Set.push_back(Obj2.getBinary());
 
-  SectionMemoryManagerWrapper SMMW;
+  auto SMMW = std::make_shared<SectionMemoryManagerWrapper>();
   NullResolver NR;
-  auto H = ObjLayer.addObjectSet(std::move(Obj1Set), &SMMW, &NR);
-  ObjLayer.addObjectSet(std::move(Obj2Set), &SMMW, &NR);
+  auto H = ObjLayer.addObjectSet(std::move(Obj1Set), SMMW, &NR);
+  ObjLayer.addObjectSet(std::move(Obj2Set), SMMW, &NR);
   ObjLayer.emitAndFinalize(H);
-
+  ObjLayer.removeObjectSet(H);
+  
   // Only one call to needsToReserveAllocationSpace should have been made.
-  EXPECT_EQ(SMMW.NeedsToReserveAllocationSpaceCount, 1)
+  EXPECT_EQ(SMMW->NeedsToReserveAllocationSpaceCount, 1)
       << "More than one call to needsToReserveAllocationSpace "
          "(multiple unrelated objects loaded prior to finalization)";
 }