[Orc][RuntimeDyld] Prevent duplicate calls to finalizeMemory on shared memory

managers.

Prior to this patch, recursive finalization (where finalization of one
RuntimeDyld instance triggers finalization of another instance on which the
first depends) could trigger memory access failures: When the inner (dependent)
RuntimeDyld instance and its memory manager are finalized, memory allocated
(but not yet relocated) by the outer instance is locked, and relocation in the
outer instance fails with a memory access error.

This patch adds a latch to the RuntimeDyld::MemoryManager base class that is
checked by a new method: RuntimeDyld::finalizeWithMemoryManagerLocking, ensuring
that shared memory managers are only finalized by the outermost RuntimeDyld
instance.

This allows ORC clients to supply the same memory manager to multiple calls to
addModuleSet. In particular it enables the use of user-supplied memory managers
with the CompileOnDemandLayer which must reuse the supplied memory manager for
each function that is lazily compiled.

llvm-svn: 257263

llvm/include/llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h

@@ -108,9 +108,7 @@ private:
 
     void Finalize() override {
       State = Finalizing;
-      RTDyld->resolveRelocations();
-      RTDyld->registerEHFrames();
-      MemMgr->finalizeMemory();
+      RTDyld->finalizeWithMemoryManagerLocking();
       State = Finalized;
     }
 

llvm/include/llvm/ExecutionEngine/RuntimeDyld.h

@@ -95,7 +95,9 @@ public:
 
   /// \brief Memory Management.
   class MemoryManager {
+    friend class RuntimeDyld;
   public:
+    MemoryManager() : FinalizationLocked(false) {}
     virtual ~MemoryManager() {}
 
     /// Allocate a memory block of (at least) the given size suitable for
@@ -153,6 +155,7 @@ public:
 
   private:
     virtual void anchor();
+    bool FinalizationLocked;
   };
 
   /// \brief Symbol resolution.
@@ -241,6 +244,25 @@ public:
     this->ProcessAllSections = ProcessAllSections;
   }
 
+  /// Perform all actions needed to make the code owned by this RuntimeDyld
+  /// instance executable:
+  ///
+  /// 1) Apply relocations.
+  /// 2) Register EH frames.
+  /// 3) Update memory permissions*.
+  ///
+  /// * Finalization is potentially recursive**, and the 3rd step will only be
+  ///   applied by the outermost call to finalize. This allows different
+  ///   RuntimeDyld instances to share a memory manager without the innermost
+  ///   finalization locking the memory and causing relocation fixup errors in
+  ///   outer instances.
+  ///
+  /// ** Recursive finalization occurs when one RuntimeDyld instances needs the
+  ///   address of a symbol owned by some other instance in order to apply
+  ///   relocations.
+  ///
+  void finalizeWithMemoryManagerLocking();
+
 private:
   // RuntimeDyldImpl is the actual class. RuntimeDyld is just the public
   // interface.

llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp

@@ -967,6 +967,17 @@ bool RuntimeDyld::hasError() { return Dyld->hasError(); }
 
 StringRef RuntimeDyld::getErrorString() { return Dyld->getErrorString(); }
 
+void RuntimeDyld::finalizeWithMemoryManagerLocking() {
+  bool MemoryFinalizationLocked = MemMgr.FinalizationLocked;
+  MemMgr.FinalizationLocked = true;
+  resolveRelocations();
+  registerEHFrames();
+  if (!MemoryFinalizationLocked) {
+    MemMgr.finalizeMemory();
+    MemMgr.FinalizationLocked = false;
+  }
+}
+
 void RuntimeDyld::registerEHFrames() {
   if (Dyld)
     Dyld->registerEHFrames();

llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp

@@ -7,6 +7,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "OrcTestCommon.h"
 #include "llvm/ExecutionEngine/ExecutionEngine.h"
 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
@@ -21,6 +22,10 @@ using namespace llvm::orc;
 
 namespace {
 
+class ObjectLinkingLayerExecutionTest : public testing::Test,
+                                        public OrcExecutionTest {
+};
+
 TEST(ObjectLinkingLayerTest, TestSetProcessAllSections) {
 
   class SectionMemoryManagerWrapper : public SectionMemoryManager {
@@ -91,4 +96,81 @@ TEST(ObjectLinkingLayerTest, TestSetProcessAllSections) {
   }
 }
 
+
+TEST_F(ObjectLinkingLayerExecutionTest, NoDuplicateFinalization) {
+
+  if (!TM)
+    return;
+
+  class SectionMemoryManagerWrapper : public SectionMemoryManager {
+  public:
+    int FinalizationCount = 0;
+    bool finalizeMemory(std::string *ErrMsg = 0) override {
+      ++FinalizationCount;
+      return SectionMemoryManager::finalizeMemory(ErrMsg);
+    }
+  };
+
+  ObjectLinkingLayer<> ObjLayer;
+  SimpleCompiler Compile(*TM);
+
+  // Create a pair of modules that will trigger recursive finalization:
+  // Module 1:
+  //   int bar() { return 42; }
+  // Module 2:
+  //   int bar();
+  //   int foo() { return bar(); }
+
+  ModuleBuilder MB1(getGlobalContext(), "", "dummy");
+  {
+    MB1.getModule()->setDataLayout(TM->createDataLayout());
+    Function *BarImpl = MB1.createFunctionDecl<int32_t(void)>("bar");
+    BasicBlock *BarEntry = BasicBlock::Create(getGlobalContext(), "entry",
+                                              BarImpl);
+    IRBuilder<> Builder(BarEntry);
+    IntegerType *Int32Ty = IntegerType::get(getGlobalContext(), 32);
+    Value *FourtyTwo = ConstantInt::getSigned(Int32Ty, 42);
+    Builder.CreateRet(FourtyTwo);
+  }
+
+  auto Obj1 = Compile(*MB1.getModule());
+  std::vector<object::ObjectFile*> Obj1Set;
+  Obj1Set.push_back(Obj1.getBinary());
+
+  ModuleBuilder MB2(getGlobalContext(), "", "dummy");
+  {
+    MB2.getModule()->setDataLayout(TM->createDataLayout());
+    Function *BarDecl = MB2.createFunctionDecl<int32_t(void)>("bar");
+    Function *FooImpl = MB2.createFunctionDecl<int32_t(void)>("foo");
+    BasicBlock *FooEntry = BasicBlock::Create(getGlobalContext(), "entry",
+                                              FooImpl);
+    IRBuilder<> Builder(FooEntry);
+    Builder.CreateRet(Builder.CreateCall(BarDecl));
+  }
+  auto Obj2 = Compile(*MB2.getModule());
+  std::vector<object::ObjectFile*> Obj2Set;
+  Obj2Set.push_back(Obj2.getBinary());
+
+  auto Resolver =
+    createLambdaResolver(
+      [&](const std::string &Name) {
+        if (auto Sym = ObjLayer.findSymbol(Name, true))
+          return RuntimeDyld::SymbolInfo(Sym.getAddress(), Sym.getFlags());
+        return RuntimeDyld::SymbolInfo(nullptr);
+      },
+      [](const std::string &Name) {
+        return RuntimeDyld::SymbolInfo(nullptr);
+      });
+
+  SectionMemoryManagerWrapper SMMW;
+  ObjLayer.addObjectSet(std::move(Obj1Set), &SMMW, &*Resolver);
+  auto H = ObjLayer.addObjectSet(std::move(Obj2Set), &SMMW, &*Resolver);
+  ObjLayer.emitAndFinalize(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)
+      << "Extra call to finalize";
+}
+
 }

llvm/unittests/ExecutionEngine/Orc/OrcTestCommon.cpp

@@ -19,7 +19,7 @@ bool OrcExecutionTest::NativeTargetInitialized = false;
 
 ModuleBuilder::ModuleBuilder(LLVMContext &Context, StringRef Triple,
                              StringRef Name)
-  : M(new Module(Name, Context)),
-    Builder(Context) {
-  M->setTargetTriple(Triple);
+  : M(new Module(Name, Context)) {
+  if (Triple != "")
+    M->setTargetTriple(Triple);
 }

llvm/unittests/ExecutionEngine/Orc/OrcTestCommon.h

@@ -20,6 +20,7 @@
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/TypeBuilder.h"
+#include "llvm/Object/ObjectFile.h"
 #include "llvm/ExecutionEngine/ExecutionEngine.h"
 #include "llvm/ExecutionEngine/Orc/JITSymbol.h"
 #include "llvm/Support/TargetSelect.h"
@@ -74,7 +75,6 @@ public:
 
 private:
   std::unique_ptr<Module> M;
-  IRBuilder<> Builder;
 };
 
 // Dummy struct type.