[llvm] d42c235 - [JITLink] Ensure that in-flight alloc is abandoned on error in post-alloc phase.

Lang Hames via llvm-commits llvm-commits at lists.llvm.org
Wed Feb 1 18:05:20 PST 2023


Author: Lang Hames
Date: 2023-02-01T18:04:57-08:00
New Revision: d42c2352aa3a77fb8bf217191868c3eb6c53a5a2

URL: https://github.com/llvm/llvm-project/commit/d42c2352aa3a77fb8bf217191868c3eb6c53a5a2
DIFF: https://github.com/llvm/llvm-project/commit/d42c2352aa3a77fb8bf217191868c3eb6c53a5a2.diff

LOG: [JITLink] Ensure that in-flight alloc is abandoned on error in post-alloc phase.

If an error occurs during the post-allocation phase (JITLinkerBase::linkPhase2)
then we need to call JITLinkMemoryManager::InFlightAlloc::abandon so that the
allocation has a chance to clean up. This was already handled for later phases,
but we were skipping the abandon step when we bailed out of phase 2.

Added: 
    

Modified: 
    llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp
    llvm/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp
    llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp b/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp
index 17de84fa6e115..4ef3c3ada12de 100644
--- a/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp
+++ b/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp
@@ -65,7 +65,7 @@ void JITLinkerBase::linkPhase2(std::unique_ptr<JITLinkerBase> Self,
   if (AR)
     Alloc = std::move(*AR);
   else
-    return Ctx->notifyFailed(AR.takeError());
+    return abandonAllocAndBailOut(std::move(Self), AR.takeError());
 
   LLVM_DEBUG({
     dbgs() << "Link graph \"" << G->getName()
@@ -75,13 +75,13 @@ void JITLinkerBase::linkPhase2(std::unique_ptr<JITLinkerBase> Self,
 
   // Run post-allocation passes.
   if (auto Err = runPasses(Passes.PostAllocationPasses))
-    return Ctx->notifyFailed(std::move(Err));
+    return abandonAllocAndBailOut(std::move(Self), std::move(Err));
 
   // Notify client that the defined symbols have been assigned addresses.
   LLVM_DEBUG(dbgs() << "Resolving symbols defined in " << G->getName() << "\n");
 
   if (auto Err = Ctx->notifyResolved(*G))
-    return Ctx->notifyFailed(std::move(Err));
+    return abandonAllocAndBailOut(std::move(Self), std::move(Err));
 
   auto ExternalSymbols = getExternalSymbolNames();
 

diff  --git a/llvm/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp b/llvm/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp
index bd44b86f30819..e74aa059f4054 100644
--- a/llvm/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp
+++ b/llvm/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp
@@ -236,10 +236,14 @@ class InProcessMemoryManager::IPInFlightAlloc
   IPInFlightAlloc(InProcessMemoryManager &MemMgr, LinkGraph &G, BasicLayout BL,
                   sys::MemoryBlock StandardSegments,
                   sys::MemoryBlock FinalizationSegments)
-      : MemMgr(MemMgr), G(G), BL(std::move(BL)),
+      : MemMgr(MemMgr), G(&G), BL(std::move(BL)),
         StandardSegments(std::move(StandardSegments)),
         FinalizationSegments(std::move(FinalizationSegments)) {}
 
+  ~IPInFlightAlloc() {
+    assert(!G && "InFlight alloc neither abandoned nor finalized");
+  }
+
   void finalize(OnFinalizedFunction OnFinalized) override {
 
     // Apply memory protections to all segments.
@@ -249,7 +253,7 @@ class InProcessMemoryManager::IPInFlightAlloc
     }
 
     // Run finalization actions.
-    auto DeallocActions = runFinalizeActions(G.allocActions());
+    auto DeallocActions = runFinalizeActions(G->allocActions());
     if (!DeallocActions) {
       OnFinalized(DeallocActions.takeError());
       return;
@@ -261,6 +265,13 @@ class InProcessMemoryManager::IPInFlightAlloc
       return;
     }
 
+#ifndef NDEBUG
+    // Set 'G' to null to flag that we've been successfully finalized.
+    // This allows us to assert at destruction time that a call has been made
+    // to either finalize or abandon.
+    G = nullptr;
+#endif
+
     // Continue with finalized allocation.
     OnFinalized(MemMgr.createFinalizedAlloc(std::move(StandardSegments),
                                             std::move(*DeallocActions)));
@@ -272,6 +283,14 @@ class InProcessMemoryManager::IPInFlightAlloc
       Err = joinErrors(std::move(Err), errorCodeToError(EC));
     if (auto EC = sys::Memory::releaseMappedMemory(StandardSegments))
       Err = joinErrors(std::move(Err), errorCodeToError(EC));
+
+#ifndef NDEBUG
+    // Set 'G' to null to flag that we've been successfully finalized.
+    // This allows us to assert at destruction time that a call has been made
+    // to either finalize or abandon.
+    G = nullptr;
+#endif
+
     OnAbandoned(std::move(Err));
   }
 
@@ -295,7 +314,7 @@ class InProcessMemoryManager::IPInFlightAlloc
   }
 
   InProcessMemoryManager &MemMgr;
-  LinkGraph &G;
+  LinkGraph *G;
   BasicLayout BL;
   sys::MemoryBlock StandardSegments;
   sys::MemoryBlock FinalizationSegments;

diff  --git a/llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp b/llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp
index 605fb2524bf3e..c943eaf3bd7f1 100644
--- a/llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp
@@ -119,4 +119,58 @@ TEST_F(ObjectLinkingLayerTest, ClaimLateDefinedWeakSymbols) {
   EXPECT_THAT_EXPECTED(ES.lookup(&JD, "_anchor"), Succeeded());
 }
 
+TEST_F(ObjectLinkingLayerTest, HandleErrorDuringPostAllocationPass) {
+  // We want to confirm that Errors in post allocation passes correctly
+  // abandon the in-flight allocation and report an error.
+  class TestPlugin : public ObjectLinkingLayer::Plugin {
+  public:
+    ~TestPlugin() { EXPECT_TRUE(ErrorReported); }
+
+    void modifyPassConfig(MaterializationResponsibility &MR,
+                          jitlink::LinkGraph &G,
+                          jitlink::PassConfiguration &Config) override {
+      Config.PostAllocationPasses.insert(
+          Config.PostAllocationPasses.begin(), [](LinkGraph &G) {
+            return make_error<StringError>("Kaboom", inconvertibleErrorCode());
+          });
+    }
+
+    Error notifyFailed(MaterializationResponsibility &MR) override {
+      ErrorReported = true;
+      return Error::success();
+    }
+
+    Error notifyRemovingResources(JITDylib &JD, ResourceKey K) override {
+      return Error::success();
+    }
+    void notifyTransferringResources(JITDylib &JD, ResourceKey DstKey,
+                                     ResourceKey SrcKey) override {
+      llvm_unreachable("unexpected resource transfer");
+    }
+
+  private:
+    bool ErrorReported = false;
+  };
+
+  // We expect this test to generate errors. Consume them so that we don't
+  // add noise to the test logs.
+  ES.setErrorReporter(consumeError);
+
+  ObjLinkingLayer.addPlugin(std::make_unique<TestPlugin>());
+
+  auto G =
+      std::make_unique<LinkGraph>("foo", Triple("x86_64-apple-darwin"), 8,
+                                  support::little, x86_64::getEdgeKindName);
+
+  auto &DataSec = G->createSection("__data", MemProt::Read | MemProt::Write);
+  auto &DataBlock = G->createContentBlock(DataSec, BlockContent,
+                                          orc::ExecutorAddr(0x1000), 8, 0);
+  G->addDefinedSymbol(DataBlock, 4, "_anchor", 4, Linkage::Weak, Scope::Default,
+                      false, true);
+
+  EXPECT_THAT_ERROR(ObjLinkingLayer.add(JD, std::move(G)), Succeeded());
+
+  EXPECT_THAT_EXPECTED(ES.lookup(&JD, "_anchor"), Failed());
+}
+
 } // end anonymous namespace


        


More information about the llvm-commits mailing list