[JITLink] Don't try to abandon non-existent allocations.

If JITLinkGeneric::linkPhase2 receives an Error rather than an InFlightAlloc
then we need to call JITLinkContext::notifyFailed, rather than calling
abandonAllocAndBailOut -- the latter asserts that there is an allocation to
abandon, and this was turning allocation errors into assertion failures in
debug mode.

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 abandonAllocAndBailOut(std::move(Self), AR.takeError());
+    return Ctx->notifyFailed(AR.takeError());
 
   LLVM_DEBUG({
     dbgs() << "Link graph \"" << G->getName()

llvm/unittests/ExecutionEngine/JITLink/CMakeLists.txt

@@ -10,7 +10,9 @@ set(LLVM_LINK_COMPONENTS
 add_llvm_unittest(JITLinkTests
     AArch32Tests.cpp
     EHFrameSupportTests.cpp
+    JITLinkMocks.cpp
     LinkGraphTests.cpp
+    MemoryManagerErrorTests.cpp
   )
 
 target_link_libraries(JITLinkTests PRIVATE LLVMTestingSupport)

llvm/unittests/ExecutionEngine/JITLink/JITLinkMocks.cpp

@@ -0,0 +1,72 @@
+//===--------- JITLinkMocks.cpp - Mock APIs for JITLink unit tests --------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "JITLinkMocks.h"
+#include "llvm/ExecutionEngine/JITLink/MachO_x86_64.h"
+
+#include "llvm/Testing/Support/Error.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+using namespace llvm::orc;
+using namespace llvm::jitlink;
+
+void lookupResolveEverythingToNull(
+    const llvm::jitlink::JITLinkContext::LookupMap &Symbols,
+    std::unique_ptr<llvm::jitlink::JITLinkAsyncLookupContinuation> LC) {
+  llvm::orc::ExecutorAddr Null;
+  llvm::jitlink::AsyncLookupResult Result;
+  for (auto &KV : Symbols)
+    Result[KV.first] = {Null, llvm::JITSymbolFlags::Exported};
+  LC->run(std::move(Result));
+}
+
+void lookupErrorOut(
+    const llvm::jitlink::JITLinkContext::LookupMap &Symbols,
+    std::unique_ptr<llvm::jitlink::JITLinkAsyncLookupContinuation> LC) {
+  LC->run(llvm::make_error<llvm::StringError>("Lookup failed",
+                                              llvm::inconvertibleErrorCode()));
+}
+
+std::unique_ptr<MockJITLinkContext> makeMockContext(
+    llvm::unique_function<void(llvm::Error)> HandleFailed,
+    llvm::unique_function<void(MockJITLinkMemoryManager &)> SetupMemMgr,
+    llvm::unique_function<void(MockJITLinkContext &)> SetupContext) {
+  auto MemMgr = std::make_unique<MockJITLinkMemoryManager>();
+  SetupMemMgr(*MemMgr);
+  auto Ctx = std::make_unique<MockJITLinkContext>(std::move(MemMgr),
+                                                  std::move(HandleFailed));
+  SetupContext(*Ctx);
+  return Ctx;
+}
+
+void defaultMemMgrSetup(MockJITLinkMemoryManager &) {}
+void defaultCtxSetup(MockJITLinkContext &) {}
+
+TEST(JITLinkMocks, SmokeTest) {
+  // Check that the testing infrastructure defaults can "link" a graph
+  // successfully.
+  auto G = std::make_unique<LinkGraph>("foo", Triple("x86_64-apple-darwin"), 8,
+                                       support::little, getGenericEdgeKindName);
+
+  ArrayRef<char> Content = "hello, world!";
+  auto &Sec =
+      G->createSection("__data", orc::MemProt::Read | orc::MemProt::Write);
+  orc::ExecutorAddr B1Addr(0x1000);
+  auto &B = G->createContentBlock(Sec, Content, B1Addr, 8, 0);
+  G->addDefinedSymbol(B, 4, "S", 4, Linkage::Strong, Scope::Default, false,
+                      false);
+
+  Error Err = Error::success();
+  auto Ctx =
+      makeMockContext(JoinErrorsInto(Err), defaultMemMgrSetup, defaultCtxSetup);
+
+  link_MachO_x86_64(std::move(G), std::move(Ctx));
+
+  EXPECT_THAT_ERROR(std::move(Err), Succeeded());
+}

llvm/unittests/ExecutionEngine/JITLink/JITLinkMocks.h

@@ -0,0 +1,228 @@
+//===----- JITLinkMocks.h - Mock APIs for JITLink unit tests ----*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Mock APIs for JITLink unit tests.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_UNITTESTS_EXECUTIONENGINE_JITLINK_JITLINKMOCKS_H
+#define LLVM_UNITTESTS_EXECUTIONENGINE_JITLINK_JITLINKMOCKS_H
+
+#include "llvm/ExecutionEngine/JITLink/JITLink.h"
+
+class MockJITLinkMemoryManager : public llvm::jitlink::JITLinkMemoryManager {
+public:
+  class Alloc {
+  public:
+    virtual ~Alloc() {}
+  };
+
+  class SimpleAlloc : public Alloc {
+  public:
+    SimpleAlloc(const llvm::jitlink::JITLinkDylib *JD,
+                llvm::jitlink::LinkGraph &G) {
+      for (auto *B : G.blocks())
+        (void)B->getMutableContent(G);
+    }
+  };
+
+  class MockInFlightAlloc : public InFlightAlloc {
+  public:
+    MockInFlightAlloc(MockJITLinkMemoryManager &MJMM, std::unique_ptr<Alloc> A)
+        : MJMM(MJMM), A(std::move(A)) {}
+
+    void abandon(OnAbandonedFunction OnAbandoned) override {
+      OnAbandoned(MJMM.Abandon(std::move(A)));
+    }
+
+    void finalize(OnFinalizedFunction OnFinalized) override {
+      OnFinalized(MJMM.Finalize(std::move(A)));
+    }
+
+  private:
+    MockJITLinkMemoryManager &MJMM;
+    std::unique_ptr<Alloc> A;
+  };
+
+  MockJITLinkMemoryManager() {
+    Allocate = [this](const llvm::jitlink::JITLinkDylib *JD,
+                      llvm::jitlink::LinkGraph &G) {
+      return defaultAllocate(JD, G);
+    };
+
+    Deallocate = [this](std::vector<FinalizedAlloc> Allocs) {
+      return defaultDeallocate(std::move(Allocs));
+    };
+
+    Abandon = [this](std::unique_ptr<Alloc> A) {
+      return defaultAbandon(std::move(A));
+    };
+
+    Finalize = [this](std::unique_ptr<Alloc> A) {
+      return defaultFinalize(std::move(A));
+    };
+  }
+
+  void allocate(const llvm::jitlink::JITLinkDylib *JD,
+                llvm::jitlink::LinkGraph &G,
+                OnAllocatedFunction OnAllocated) override {
+    auto A = Allocate(JD, G);
+    if (!A)
+      OnAllocated(A.takeError());
+    else
+      OnAllocated(std::make_unique<MockInFlightAlloc>(*this, std::move(*A)));
+  }
+
+  void deallocate(std::vector<FinalizedAlloc> Allocs,
+                  OnDeallocatedFunction OnDeallocated) override {
+    OnDeallocated(Deallocate(std::move(Allocs)));
+  }
+
+  using JITLinkMemoryManager::allocate;
+  using JITLinkMemoryManager::deallocate;
+
+  llvm::Expected<std::unique_ptr<Alloc>>
+  defaultAllocate(const llvm::jitlink::JITLinkDylib *JD,
+                  llvm::jitlink::LinkGraph &G) {
+    return std::make_unique<SimpleAlloc>(JD, G);
+  }
+
+  llvm::Error defaultDeallocate(std::vector<FinalizedAlloc> Allocs) {
+    for (auto &A : Allocs)
+      delete A.release().toPtr<Alloc *>();
+    return llvm::Error::success();
+  }
+
+  llvm::Error defaultAbandon(std::unique_ptr<Alloc> A) {
+    return llvm::Error::success();
+  }
+
+  llvm::Expected<FinalizedAlloc> defaultFinalize(std::unique_ptr<Alloc> A) {
+    return FinalizedAlloc(llvm::orc::ExecutorAddr::fromPtr(A.release()));
+  }
+
+  llvm::unique_function<llvm::Expected<std::unique_ptr<Alloc>>(
+      const llvm::jitlink::JITLinkDylib *, llvm::jitlink::LinkGraph &)>
+      Allocate;
+  llvm::unique_function<llvm::Error(std::vector<FinalizedAlloc>)> Deallocate;
+  llvm::unique_function<llvm::Error(std::unique_ptr<Alloc>)> Abandon;
+  llvm::unique_function<llvm::Expected<FinalizedAlloc>(std::unique_ptr<Alloc>)>
+      Finalize;
+};
+
+void lookupResolveEverythingToNull(
+    const llvm::jitlink::JITLinkContext::LookupMap &Symbols,
+    std::unique_ptr<llvm::jitlink::JITLinkAsyncLookupContinuation> LC);
+
+void lookupErrorOut(
+    const llvm::jitlink::JITLinkContext::LookupMap &Symbols,
+    std::unique_ptr<llvm::jitlink::JITLinkAsyncLookupContinuation> LC);
+
+class MockJITLinkContext : public llvm::jitlink::JITLinkContext {
+public:
+  using HandleFailedFn = llvm::unique_function<void(llvm::Error)>;
+
+  MockJITLinkContext(std::unique_ptr<MockJITLinkMemoryManager> MJMM,
+                     HandleFailedFn HandleFailed)
+      : JITLinkContext(&JD), MJMM(std::move(MJMM)),
+        HandleFailed(std::move(HandleFailed)) {}
+
+  ~MockJITLinkContext() {
+    if (auto Err = MJMM->deallocate(std::move(FinalizedAllocs)))
+      notifyFailed(std::move(Err));
+  }
+
+  llvm::jitlink::JITLinkMemoryManager &getMemoryManager() override {
+    return *MJMM;
+  }
+
+  void notifyFailed(llvm::Error Err) override { HandleFailed(std::move(Err)); }
+
+  void lookup(const LookupMap &Symbols,
+              std::unique_ptr<llvm::jitlink::JITLinkAsyncLookupContinuation> LC)
+      override {
+    Lookup(Symbols, std::move(LC));
+  }
+
+  llvm::Error notifyResolved(llvm::jitlink::LinkGraph &G) override {
+    return NotifyResolved ? NotifyResolved(G) : llvm::Error::success();
+  }
+
+  void notifyFinalized(
+      llvm::jitlink::JITLinkMemoryManager::FinalizedAlloc Alloc) override {
+    if (NotifyFinalized)
+      NotifyFinalized(std::move(Alloc));
+    else
+      FinalizedAllocs.push_back(std::move(Alloc));
+  }
+
+  bool shouldAddDefaultTargetPasses(const llvm::Triple &TT) const override {
+    return true;
+  }
+
+  llvm::jitlink::LinkGraphPassFunction
+  getMarkLivePass(const llvm::Triple &TT) const override {
+    return MarkLivePass ? llvm::jitlink::LinkGraphPassFunction(
+                              [this](llvm::jitlink::LinkGraph &G) {
+                                return MarkLivePass(G);
+                              })
+                        : llvm::jitlink::LinkGraphPassFunction(
+                              [](llvm::jitlink::LinkGraph &G) {
+                                return markAllSymbolsLive(G);
+                              });
+  }
+
+  llvm::Error
+  modifyPassConfig(llvm::jitlink::LinkGraph &G,
+                   llvm::jitlink::PassConfiguration &Config) override {
+    if (ModifyPassConfig)
+      return ModifyPassConfig(G, Config);
+    return llvm::Error::success();
+  }
+
+  llvm::jitlink::JITLinkDylib JD{"JD"};
+  std::unique_ptr<MockJITLinkMemoryManager> MJMM;
+  HandleFailedFn HandleFailed;
+  llvm::unique_function<void(
+      const LookupMap &,
+      std::unique_ptr<llvm::jitlink::JITLinkAsyncLookupContinuation>)>
+      Lookup;
+  llvm::unique_function<llvm::Error(llvm::jitlink::LinkGraph &)> NotifyResolved;
+  llvm::unique_function<void(
+      llvm::jitlink::JITLinkMemoryManager::FinalizedAlloc)>
+      NotifyFinalized;
+  mutable llvm::unique_function<llvm::Error(llvm::jitlink::LinkGraph &)>
+      MarkLivePass;
+  llvm::unique_function<llvm::Error(llvm::jitlink::LinkGraph &,
+                                    llvm::jitlink::PassConfiguration &)>
+      ModifyPassConfig;
+
+  std::vector<llvm::jitlink::JITLinkMemoryManager::FinalizedAlloc>
+      FinalizedAllocs;
+};
+
+std::unique_ptr<MockJITLinkContext> makeMockContext(
+    llvm::unique_function<void(llvm::Error)> HandleFailed,
+    llvm::unique_function<void(MockJITLinkMemoryManager &)> SetupMemMgr,
+    llvm::unique_function<void(MockJITLinkContext &)> SetupContext);
+
+void defaultMemMgrSetup(MockJITLinkMemoryManager &);
+void defaultCtxSetup(MockJITLinkContext &);
+
+class JoinErrorsInto {
+public:
+  JoinErrorsInto(llvm::Error &Err) : Err(Err) {}
+  void operator()(llvm::Error E2) {
+    Err = llvm::joinErrors(std::move(Err), std::move(E2));
+  }
+
+private:
+  llvm::Error &Err;
+};
+
+#endif // LLVM_UNITTESTS_EXECUTIONENGINE_JITLINK_JITLINKMOCKS_H

llvm/unittests/ExecutionEngine/JITLink/MemoryManagerErrorTests.cpp

@@ -0,0 +1,46 @@
+//===---- MemoryManagerErrorTests.cpp - Test memory manager error paths ---===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "JITLinkMocks.h"
+#include "llvm/ExecutionEngine/JITLink/MachO_x86_64.h"
+
+#include "llvm/Testing/Support/Error.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+using namespace llvm::orc;
+using namespace llvm::jitlink;
+
+TEST(MemoryManagerErrorTest, ErrorOnFirstAllocate) {
+  // Check that we can get addresses for blocks, symbols, and edges.
+  auto G = std::make_unique<LinkGraph>("foo", Triple("x86_64-apple-darwin"), 8,
+                                       support::little, getGenericEdgeKindName);
+
+  ArrayRef<char> Content = "hello, world!";
+  auto &Sec =
+      G->createSection("__data", orc::MemProt::Read | orc::MemProt::Write);
+  orc::ExecutorAddr B1Addr(0x1000);
+  auto &B = G->createContentBlock(Sec, Content, B1Addr, 8, 0);
+  G->addDefinedSymbol(B, 4, "S", 4, Linkage::Strong, Scope::Default, false,
+                      false);
+
+  Error Err = Error::success();
+  auto Ctx = makeMockContext(
+      JoinErrorsInto(Err),
+      [](MockJITLinkMemoryManager &MemMgr) {
+        MemMgr.Allocate = [](const JITLinkDylib *JD, LinkGraph &G) {
+          return make_error<StringError>("Failed to allocate",
+                                         inconvertibleErrorCode());
+        };
+      },
+      defaultCtxSetup);
+
+  link_MachO_x86_64(std::move(G), std::move(Ctx));
+
+  EXPECT_THAT_ERROR(std::move(Err), Failed());
+}