[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

6 files changed, 120 insertions, 7 deletions

diff --git a/llvm/include/llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h b/llvm/include/llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h
index 2acfecfb94d..4dc48f11488 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h
+++ b/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;
     }
 
diff --git a/llvm/include/llvm/ExecutionEngine/RuntimeDyld.h b/llvm/include/llvm/ExecutionEngine/RuntimeDyld.h
index 385b8d0a30b..afd8de64913 100644
--- a/llvm/include/llvm/ExecutionEngine/RuntimeDyld.h
+++ b/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.
diff --git a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp
index a95f3bbe417..24b911c190a 100644
--- a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp
+++ b/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();
diff --git a/llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp b/llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp
index a37177c5d1d..d7c0b74670c 100644
--- a/llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp
+++ b/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";
+}
+
 }
diff --git a/llvm/unittests/ExecutionEngine/Orc/OrcTestCommon.cpp b/llvm/unittests/ExecutionEngine/Orc/OrcTestCommon.cpp
index 1b5485d3b33..17d1e9c9276 100644
--- a/llvm/unittests/ExecutionEngine/Orc/OrcTestCommon.cpp
+++ b/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);
 }
diff --git a/llvm/unittests/ExecutionEngine/Orc/OrcTestCommon.h b/llvm/unittests/ExecutionEngine/Orc/OrcTestCommon.h
index 15d9f54a37f..f480e0789ae 100644
--- a/llvm/unittests/ExecutionEngine/Orc/OrcTestCommon.h
+++ b/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.