[ORC] Simplify JITLinkRedirectableSymbolManager, fix definition locations.

Redirectable stubs should be placed in the same JITDylib as their names, with
their lifetimes managed according to the ResourceTracker used when adding them.

The original implementation created a single pool of stubs in the JITDylib
that is passed to the JITLinkRedirectableSymbolManager constructor, but this
may cause the addresses of the redirectable symbols themselves (added to the
JITDylibs passed to createRedirectableSymbols) to appear outside the address
range of their defining JITDylib.

This patch fixes the issue by dropping the pool and emitting a new graph for
each set of requested redirectable symbols. We lose the ability to recycle
stubs, but gain the ability to free them entirely. Since we don't expect stub
reuse to be a common case this is likely the best trade-off.

If in the future we do need to return to a stub pool we should create a
separate one for each JITDylib to ensure that addresses behave as expected.

Author: lhames

llvm/include/llvm/ExecutionEngine/JITLink/JITLink.h

@@ -1077,6 +1077,16 @@ class LinkGraph {
     return MutableArrayRef<char>(AllocatedBuffer, SourceStr.size());
   }
 
+  /// Allocate a copy of the given string using the LinkGraph's allocator
+  /// and return it as a StringRef.
+  ///
+  /// This is a convenience wrapper around allocateContent(Twine) that is
+  /// handy when creating new symbol names within the graph.
+  StringRef allocateName(Twine Source) {
+    auto Buf = allocateContent(Source);
+    return {Buf.data(), Buf.size()};
+  }
+
   /// Allocate a copy of the given string using the LinkGraph's allocator.
   ///
   /// The allocated string will be terminated with a null character, and the

llvm/include/llvm/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.h

@@ -17,15 +17,16 @@
 #include "llvm/ExecutionEngine/Orc/RedirectionManager.h"
 #include "llvm/Support/StringSaver.h"
 
+#include <atomic>
+
 namespace llvm {
 namespace orc {
 
-class JITLinkRedirectableSymbolManager : public RedirectableSymbolManager,
-                                         public ResourceManager {
+class JITLinkRedirectableSymbolManager : public RedirectableSymbolManager {
 public:
   /// Create redirection manager that uses JITLink based implementaion.
   static Expected<std::unique_ptr<RedirectableSymbolManager>>
-  Create(ObjectLinkingLayer &ObjLinkingLayer, JITDylib &JD) {
+  Create(ObjectLinkingLayer &ObjLinkingLayer) {
     auto AnonymousPtrCreator(jitlink::getAnonymousPointerCreator(
         ObjLinkingLayer.getExecutionSession().getTargetTriple()));
     auto PtrJumpStubCreator(jitlink::getPointerJumpStubCreator(
@@ -35,69 +36,27 @@ class JITLinkRedirectableSymbolManager : public RedirectableSymbolManager,
                                      inconvertibleErrorCode());
     return std::unique_ptr<RedirectableSymbolManager>(
         new JITLinkRedirectableSymbolManager(
-            ObjLinkingLayer, JD, AnonymousPtrCreator, PtrJumpStubCreator));
+            ObjLinkingLayer, AnonymousPtrCreator, PtrJumpStubCreator));
   }
 
   void emitRedirectableSymbols(std::unique_ptr<MaterializationResponsibility> R,
                                const SymbolAddrMap &InitialDests) override;
 
   Error redirect(JITDylib &TargetJD, const SymbolAddrMap &NewDests) override;
 
-  Error handleRemoveResources(JITDylib &TargetJD, ResourceKey K) override;
-
-  void handleTransferResources(JITDylib &TargetJD, ResourceKey DstK,
-                               ResourceKey SrcK) override;
-
 private:
-  using StubHandle = unsigned;
-  constexpr static unsigned StubBlockSize = 256;
-  constexpr static StringRef JumpStubPrefix = "$__IND_JUMP_STUBS";
-  constexpr static StringRef StubPtrPrefix = "$IND_JUMP_PTR_";
-  constexpr static StringRef JumpStubTableName = "$IND_JUMP_";
-  constexpr static StringRef StubPtrTableName = "$__IND_JUMP_PTRS";
-
   JITLinkRedirectableSymbolManager(
-      ObjectLinkingLayer &ObjLinkingLayer, JITDylib &JD,
+      ObjectLinkingLayer &ObjLinkingLayer,
       jitlink::AnonymousPointerCreator &AnonymousPtrCreator,
       jitlink::PointerJumpStubCreator &PtrJumpStubCreator)
-      : ObjLinkingLayer(ObjLinkingLayer), JD(JD),
+      : ObjLinkingLayer(ObjLinkingLayer),
         AnonymousPtrCreator(std::move(AnonymousPtrCreator)),
-        PtrJumpStubCreator(std::move(PtrJumpStubCreator)) {
-    ObjLinkingLayer.getExecutionSession().registerResourceManager(*this);
-  }
-
-  ~JITLinkRedirectableSymbolManager() {
-    ObjLinkingLayer.getExecutionSession().deregisterResourceManager(*this);
-  }
-
-  StringRef JumpStubSymbolName(unsigned I) {
-    return *ObjLinkingLayer.getExecutionSession().intern(
-        (JumpStubPrefix + Twine(I)).str());
-  }
-
-  StringRef StubPtrSymbolName(unsigned I) {
-    return *ObjLinkingLayer.getExecutionSession().intern(
-        (StubPtrPrefix + Twine(I)).str());
-  }
-
-  unsigned GetNumAvailableStubs() const { return AvailableStubs.size(); }
-
-  Error redirectInner(JITDylib &TargetJD, const SymbolAddrMap &NewDests);
-  Error grow(unsigned Need);
+        PtrJumpStubCreator(std::move(PtrJumpStubCreator)) {}
 
   ObjectLinkingLayer &ObjLinkingLayer;
-  JITDylib &JD;
   jitlink::AnonymousPointerCreator AnonymousPtrCreator;
   jitlink::PointerJumpStubCreator PtrJumpStubCreator;
-
-  std::vector<StubHandle> AvailableStubs;
-  using SymbolToStubMap = DenseMap<SymbolStringPtr, StubHandle>;
-  DenseMap<JITDylib *, SymbolToStubMap> SymbolToStubs;
-  std::vector<ExecutorSymbolDef> JumpStubs;
-  std::vector<ExecutorSymbolDef> StubPointers;
-  DenseMap<ResourceKey, std::vector<SymbolStringPtr>> TrackedResources;
-
-  std::mutex Mutex;
+  std::atomic_size_t StubGraphIdx{0};
 };
 
 } // namespace orc

llvm/lib/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.cpp

@@ -9,171 +9,88 @@
 #include "llvm/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.h"
 #include "llvm/ExecutionEngine/Orc/Core.h"
 
+#include "llvm/ExecutionEngine/Orc/DebugUtils.h"
+
 #define DEBUG_TYPE "orc"
 
 using namespace llvm;
 using namespace llvm::orc;
 
+namespace {
+constexpr StringRef JumpStubSectionName = "__orc_stubs";
+constexpr StringRef StubPtrSectionName = "__orc_stub_ptrs";
+constexpr StringRef StubSuffix = "$__stub_ptr";
+} // namespace
+
 void JITLinkRedirectableSymbolManager::emitRedirectableSymbols(
     std::unique_ptr<MaterializationResponsibility> R,
     const SymbolAddrMap &InitialDests) {
-  auto &ES = ObjLinkingLayer.getExecutionSession();
-  std::unique_lock<std::mutex> Lock(Mutex);
-  if (GetNumAvailableStubs() < InitialDests.size())
-    if (auto Err = grow(InitialDests.size() - GetNumAvailableStubs())) {
-      ES.reportError(std::move(Err));
-      R->failMaterialization();
-      return;
-    }
 
-  JITDylib &TargetJD = R->getTargetJITDylib();
-  SymbolMap NewSymbolDefs;
-  std::vector<SymbolStringPtr> Symbols;
-  for (auto &[K, V] : InitialDests) {
-    StubHandle StubID = AvailableStubs.back();
-    if (SymbolToStubs[&TargetJD].count(K)) {
-      ES.reportError(make_error<StringError>(
-          "Tried to create duplicate redirectable symbols",
-          inconvertibleErrorCode()));
-      R->failMaterialization();
-      return;
-    }
-    SymbolToStubs[&TargetJD][K] = StubID;
-    NewSymbolDefs[K] = JumpStubs[StubID];
-    NewSymbolDefs[K].setFlags(V.getFlags());
-    Symbols.push_back(K);
-    AvailableStubs.pop_back();
-  }
+  auto &ES = ObjLinkingLayer.getExecutionSession();
+  Triple TT = ES.getTargetTriple();
 
-  // FIXME: when this fails we can return stubs to the pool
-  if (auto Err = redirectInner(TargetJD, InitialDests)) {
-    ES.reportError(std::move(Err));
-    R->failMaterialization();
-    return;
+  auto G = std::make_unique<jitlink::LinkGraph>(
+      ("<INDIRECT STUBS #" + Twine(++StubGraphIdx) + ">").str(), TT,
+      TT.isArch64Bit() ? 8 : 4,
+      TT.isLittleEndian() ? endianness::little : endianness::big,
+      jitlink::getGenericEdgeKindName);
+  auto &PointerSection =
+      G->createSection(StubPtrSectionName, MemProt::Write | MemProt::Read);
+  auto &StubsSection =
+      G->createSection(JumpStubSectionName, MemProt::Exec | MemProt::Read);
+
+  SymbolFlagsMap NewSymbols;
+  for (auto &[Name, Def] : InitialDests) {
+    jitlink::Symbol *TargetSym = nullptr;
+    if (Def.getAddress())
+      TargetSym = &G->addAbsoluteSymbol(
+          G->allocateName(*Name + "$__init_tgt"), Def.getAddress(), 0,
+          jitlink::Linkage::Strong, jitlink::Scope::Local, false);
+
+    auto PtrName = ES.intern((*Name + StubSuffix).str());
+    auto &Ptr = AnonymousPtrCreator(*G, PointerSection, TargetSym, 0);
+    Ptr.setName(*PtrName);
+    Ptr.setScope(jitlink::Scope::Hidden);
+    auto &Stub = PtrJumpStubCreator(*G, StubsSection, Ptr);
+    Stub.setName(*Name);
+    Stub.setScope(jitlink::Scope::Default);
+    NewSymbols[std::move(PtrName)] = JITSymbolFlags();
   }
 
-  // FIXME: return stubs to the pool here too.
-  if (auto Err = R->replace(absoluteSymbols(NewSymbolDefs))) {
+  // Try to claim responsibility for the new stub symbols.
+  if (auto Err = R->defineMaterializing(std::move(NewSymbols))) {
     ES.reportError(std::move(Err));
-    R->failMaterialization();
-    return;
+    return R->failMaterialization();
   }
 
-  // FIXME: return stubs to the pool here too.
-  if (auto Err = R->withResourceKeyDo([&](ResourceKey Key) {
-        TrackedResources[Key].insert(TrackedResources[Key].end(),
-                                     Symbols.begin(), Symbols.end());
-      })) {
-    ES.reportError(std::move(Err));
-    R->failMaterialization();
-    return;
-  }
+  ObjLinkingLayer.emit(std::move(R), std::move(G));
 }
 
 Error JITLinkRedirectableSymbolManager::redirect(
     JITDylib &TargetJD, const SymbolAddrMap &NewDests) {
-  std::unique_lock<std::mutex> Lock(Mutex);
-  return redirectInner(TargetJD, NewDests);
-}
+  auto &ES = ObjLinkingLayer.getExecutionSession();
+  SymbolLookupSet LS;
+  DenseMap<NonOwningSymbolStringPtr, SymbolStringPtr> PtrToStub;
+  for (auto &[StubName, Sym] : NewDests) {
+    auto PtrName = ES.intern((*StubName + StubSuffix).str());
+    PtrToStub[NonOwningSymbolStringPtr(PtrName)] = StubName;
+    LS.add(std::move(PtrName));
+  }
+  auto PtrSyms = ES.lookup({{&TargetJD, JITDylibLookupFlags::MatchAllSymbols}},
+                           std::move(LS));
+  if (!PtrSyms)
+    return PtrSyms.takeError();
 
-Error JITLinkRedirectableSymbolManager::redirectInner(
-    JITDylib &TargetJD, const SymbolAddrMap &NewDests) {
   std::vector<tpctypes::PointerWrite> PtrWrites;
-  for (auto &[K, V] : NewDests) {
-    if (!SymbolToStubs[&TargetJD].count(K))
-      return make_error<StringError>(
-          "Tried to redirect non-existent redirectalbe symbol",
-          inconvertibleErrorCode());
-    StubHandle StubID = SymbolToStubs[&TargetJD].at(K);
-    PtrWrites.push_back({StubPointers[StubID].getAddress(), V.getAddress()});
+  for (auto &[PtrName, PtrSym] : *PtrSyms) {
+    auto DestSymI = NewDests.find(PtrToStub[NonOwningSymbolStringPtr(PtrName)]);
+    assert(DestSymI != NewDests.end() && "Bad ptr -> stub mapping");
+    auto &DestSym = DestSymI->second;
+    PtrWrites.push_back({PtrSym.getAddress(), DestSym.getAddress()});
   }
+
   return ObjLinkingLayer.getExecutionSession()
       .getExecutorProcessControl()
       .getMemoryAccess()
       .writePointers(PtrWrites);
 }
-
-Error JITLinkRedirectableSymbolManager::grow(unsigned Need) {
-  unsigned OldSize = JumpStubs.size();
-  unsigned NumNewStubs = alignTo(Need, StubBlockSize);
-  unsigned NewSize = OldSize + NumNewStubs;
-
-  JumpStubs.resize(NewSize);
-  StubPointers.resize(NewSize);
-  AvailableStubs.reserve(NewSize);
-
-  SymbolLookupSet LookupSymbols;
-  DenseMap<SymbolStringPtr, ExecutorSymbolDef *> NewDefsMap;
-
-  auto &ES = ObjLinkingLayer.getExecutionSession();
-  Triple TT = ES.getTargetTriple();
-  auto G = std::make_unique<jitlink::LinkGraph>(
-      "<INDIRECT STUBS>", TT, TT.isArch64Bit() ? 8 : 4,
-      TT.isLittleEndian() ? endianness::little : endianness::big,
-      jitlink::getGenericEdgeKindName);
-  auto &PointerSection =
-      G->createSection(StubPtrTableName, MemProt::Write | MemProt::Read);
-  auto &StubsSection =
-      G->createSection(JumpStubTableName, MemProt::Exec | MemProt::Read);
-
-  // FIXME: We can batch the stubs into one block and use address to access them
-  for (size_t I = OldSize; I < NewSize; I++) {
-    auto &Pointer = AnonymousPtrCreator(*G, PointerSection, nullptr, 0);
-
-    StringRef PtrSymName = StubPtrSymbolName(I);
-    Pointer.setName(PtrSymName);
-    Pointer.setScope(jitlink::Scope::Default);
-    LookupSymbols.add(ES.intern(PtrSymName));
-    NewDefsMap[ES.intern(PtrSymName)] = &StubPointers[I];
-
-    auto &Stub = PtrJumpStubCreator(*G, StubsSection, Pointer);
-
-    StringRef JumpStubSymName = JumpStubSymbolName(I);
-    Stub.setName(JumpStubSymName);
-    Stub.setScope(jitlink::Scope::Default);
-    LookupSymbols.add(ES.intern(JumpStubSymName));
-    NewDefsMap[ES.intern(JumpStubSymName)] = &JumpStubs[I];
-  }
-
-  if (auto Err = ObjLinkingLayer.add(JD, std::move(G)))
-    return Err;
-
-  auto LookupResult = ES.lookup(makeJITDylibSearchOrder(&JD), LookupSymbols);
-  if (auto Err = LookupResult.takeError())
-    return Err;
-
-  for (auto &[K, V] : *LookupResult)
-    *NewDefsMap.at(K) = V;
-
-  for (size_t I = OldSize; I < NewSize; I++)
-    AvailableStubs.push_back(I);
-
-  return Error::success();
-}
-
-Error JITLinkRedirectableSymbolManager::handleRemoveResources(
-    JITDylib &TargetJD, ResourceKey K) {
-  std::unique_lock<std::mutex> Lock(Mutex);
-  for (auto &Symbol : TrackedResources[K]) {
-    if (!SymbolToStubs[&TargetJD].count(Symbol))
-      return make_error<StringError>(
-          "Tried to remove non-existent redirectable symbol",
-          inconvertibleErrorCode());
-    AvailableStubs.push_back(SymbolToStubs[&TargetJD].at(Symbol));
-    SymbolToStubs[&TargetJD].erase(Symbol);
-    if (SymbolToStubs[&TargetJD].empty())
-      SymbolToStubs.erase(&TargetJD);
-  }
-  TrackedResources.erase(K);
-
-  return Error::success();
-}
-
-void JITLinkRedirectableSymbolManager::handleTransferResources(
-    JITDylib &TargetJD, ResourceKey DstK, ResourceKey SrcK) {
-  std::unique_lock<std::mutex> Lock(Mutex);
-  TrackedResources[DstK].insert(TrackedResources[DstK].end(),
-                                TrackedResources[SrcK].begin(),
-                                TrackedResources[SrcK].end());
-  TrackedResources.erase(SrcK);
-}