[ORC] Add JITDylib argument to ResourceManager notify-removing/transferring ops.

In some cases it's helpful to group trackers by JITDylib. E.g. Platform classes
may want to track initializer symbols with a `JITDylib -> Tracker -> [ Symbol ]`
map. This makes it easy to collect all symbols for the JITDylib, while still
allowing efficient removal of a single tracker. Passing the JITDylib as an
argument to ResourceManager::notifyRemovingResources and
ResourceManager::notifyTransferringResources supports such use-cases.

16 files changed, 86 insertions, 70 deletions

diff --git a/llvm/examples/OrcV2Examples/LLJITWithObjectLinkingLayerPlugin/LLJITWithObjectLinkingLayerPlugin.cpp b/llvm/examples/OrcV2Examples/LLJITWithObjectLinkingLayerPlugin/LLJITWithObjectLinkingLayerPlugin.cpp
index 56da4e4..dd097a1 100644
--- a/llvm/examples/OrcV2Examples/LLJITWithObjectLinkingLayerPlugin/LLJITWithObjectLinkingLayerPlugin.cpp
+++ b/llvm/examples/OrcV2Examples/LLJITWithObjectLinkingLayerPlugin/LLJITWithObjectLinkingLayerPlugin.cpp
@@ -83,11 +83,11 @@ public:
     return Error::success();
   }
 
-  Error notifyRemovingResources(ResourceKey K) override {
+  Error notifyRemovingResources(JITDylib &JD, ResourceKey K) override {
     return Error::success();
   }
 
-  void notifyTransferringResources(ResourceKey DstKey,
+  void notifyTransferringResources(JITDylib &JD, ResourceKey DstKey,
                                    ResourceKey SrcKey) override {}
 
 private:
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/COFFPlatform.h b/llvm/include/llvm/ExecutionEngine/Orc/COFFPlatform.h
index 79358ab..0a9e08f 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/COFFPlatform.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/COFFPlatform.h
@@ -104,11 +104,11 @@ private:
       return Error::success();
     }
 
-    Error notifyRemovingResources(ResourceKey K) override {
+    Error notifyRemovingResources(JITDylib &JD, ResourceKey K) override {
       return Error::success();
     }
 
-    void notifyTransferringResources(ResourceKey DstKey,
+    void notifyTransferringResources(JITDylib &JD, ResourceKey DstKey,
                                      ResourceKey SrcKey) override {}
 
   private:
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/Core.h b/llvm/include/llvm/ExecutionEngine/Orc/Core.h
index b6d1bf0..daa0fa2 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/Core.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/Core.h
@@ -101,8 +101,9 @@ private:
 class ResourceManager {
 public:
   virtual ~ResourceManager();
-  virtual Error handleRemoveResources(ResourceKey K) = 0;
-  virtual void handleTransferResources(ResourceKey DstK, ResourceKey SrcK) = 0;
+  virtual Error handleRemoveResources(JITDylib &JD, ResourceKey K) = 0;
+  virtual void handleTransferResources(JITDylib &JD, ResourceKey DstK,
+                                       ResourceKey SrcK) = 0;
 };
 
 /// A set of symbol names (represented by SymbolStringPtrs for
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/DebugObjectManagerPlugin.h b/llvm/include/llvm/ExecutionEngine/Orc/DebugObjectManagerPlugin.h
index 455f5a8..9f10a77 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/DebugObjectManagerPlugin.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/DebugObjectManagerPlugin.h
@@ -57,9 +57,9 @@ public:
 
   Error notifyEmitted(MaterializationResponsibility &MR) override;
   Error notifyFailed(MaterializationResponsibility &MR) override;
-  Error notifyRemovingResources(ResourceKey K) override;
+  Error notifyRemovingResources(JITDylib &JD, ResourceKey K) override;
 
-  void notifyTransferringResources(ResourceKey DstKey,
+  void notifyTransferringResources(JITDylib &JD, ResourceKey DstKey,
                                    ResourceKey SrcKey) override;
 
   void modifyPassConfig(MaterializationResponsibility &MR,
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/DebuggerSupportPlugin.h b/llvm/include/llvm/ExecutionEngine/Orc/DebuggerSupportPlugin.h
index 253b1c87..e21d582 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/DebuggerSupportPlugin.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/DebuggerSupportPlugin.h
@@ -41,9 +41,9 @@ public:
       : RegisterActionAddr(RegisterActionAddr) {}
 
   Error notifyFailed(MaterializationResponsibility &MR) override;
-  Error notifyRemovingResources(ResourceKey K) override;
+  Error notifyRemovingResources(JITDylib &JD, ResourceKey K) override;
 
-  void notifyTransferringResources(ResourceKey DstKey,
+  void notifyTransferringResources(JITDylib &JD, ResourceKey DstKey,
                                    ResourceKey SrcKey) override;
 
   void modifyPassConfig(MaterializationResponsibility &MR,
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h b/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h
index 3fcacde..758c001 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h
@@ -143,11 +143,11 @@ private:
       return Error::success();
     }
 
-    Error notifyRemovingResources(ResourceKey K) override {
+    Error notifyRemovingResources(JITDylib &JD, ResourceKey K) override {
       return Error::success();
     }
 
-    void notifyTransferringResources(ResourceKey DstKey,
+    void notifyTransferringResources(JITDylib &JD, ResourceKey DstKey,
                                      ResourceKey SrcKey) override {}
 
   private:
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/MachOPlatform.h b/llvm/include/llvm/ExecutionEngine/Orc/MachOPlatform.h
index 6ca47d1..197905e 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/MachOPlatform.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/MachOPlatform.h
@@ -127,11 +127,11 @@ private:
       return Error::success();
     }
 
-    Error notifyRemovingResources(ResourceKey K) override {
+    Error notifyRemovingResources(JITDylib &JD, ResourceKey K) override {
       return Error::success();
     }
 
-    void notifyTransferringResources(ResourceKey DstKey,
+    void notifyTransferringResources(JITDylib &JD, ResourceKey DstKey,
                                      ResourceKey SrcKey) override {}
 
   private:
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h b/llvm/include/llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h
index 12505fc..c1f6240 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h
@@ -80,8 +80,8 @@ public:
       return Error::success();
     }
     virtual Error notifyFailed(MaterializationResponsibility &MR) = 0;
-    virtual Error notifyRemovingResources(ResourceKey K) = 0;
-    virtual void notifyTransferringResources(ResourceKey DstKey,
+    virtual Error notifyRemovingResources(JITDylib &JD, ResourceKey K) = 0;
+    virtual void notifyTransferringResources(JITDylib &JD, ResourceKey DstKey,
                                              ResourceKey SrcKey) = 0;
 
     /// Return any dependencies that synthetic symbols (e.g. init symbols)
@@ -188,8 +188,9 @@ private:
   void notifyLoaded(MaterializationResponsibility &MR);
   Error notifyEmitted(MaterializationResponsibility &MR, FinalizedAlloc FA);
 
-  Error handleRemoveResources(ResourceKey K) override;
-  void handleTransferResources(ResourceKey DstKey, ResourceKey SrcKey) override;
+  Error handleRemoveResources(JITDylib &JD, ResourceKey K) override;
+  void handleTransferResources(JITDylib &JD, ResourceKey DstKey,
+                               ResourceKey SrcKey) override;
 
   mutable std::mutex LayerMutex;
   jitlink::JITLinkMemoryManager &MemMgr;
@@ -211,8 +212,8 @@ public:
                         jitlink::PassConfiguration &PassConfig) override;
   Error notifyEmitted(MaterializationResponsibility &MR) override;
   Error notifyFailed(MaterializationResponsibility &MR) override;
-  Error notifyRemovingResources(ResourceKey K) override;
-  void notifyTransferringResources(ResourceKey DstKey,
+  Error notifyRemovingResources(JITDylib &JD, ResourceKey K) override;
+  void notifyTransferringResources(JITDylib &JD, ResourceKey DstKey,
                                    ResourceKey SrcKey) override;
 
 private:
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h b/llvm/include/llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h
index d5a5518..5d8c4e6 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h
@@ -140,8 +140,9 @@ private:
                  std::unique_ptr<RuntimeDyld::LoadedObjectInfo> LoadedObjInfo,
                  Error Err);
 
-  Error handleRemoveResources(ResourceKey K) override;
-  void handleTransferResources(ResourceKey DstKey, ResourceKey SrcKey) override;
+  Error handleRemoveResources(JITDylib &JD, ResourceKey K) override;
+  void handleTransferResources(JITDylib &JD, ResourceKey DstKey,
+                               ResourceKey SrcKey) override;
 
   mutable std::mutex RTDyldLayerMutex;
   GetMemoryManagerFunction GetMemoryManager;
diff --git a/llvm/lib/ExecutionEngine/Orc/Core.cpp b/llvm/lib/ExecutionEngine/Orc/Core.cpp
index d058da2..4a9d0d4 100644
--- a/llvm/lib/ExecutionEngine/Orc/Core.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/Core.cpp
@@ -2285,9 +2285,10 @@ Error ExecutionSession::removeResourceTracker(ResourceTracker &RT) {
 
   Error Err = Error::success();
 
+  auto &JD = RT.getJITDylib();
   for (auto *L : reverse(CurrentResourceManagers))
-    Err =
-        joinErrors(std::move(Err), L->handleRemoveResources(RT.getKeyUnsafe()));
+    Err = joinErrors(std::move(Err),
+                     L->handleRemoveResources(JD, RT.getKeyUnsafe()));
 
   for (auto &Q : QueriesToFail)
     Q->handleFailed(
@@ -2316,7 +2317,8 @@ void ExecutionSession::transferResourceTracker(ResourceTracker &DstRT,
     auto &JD = DstRT.getJITDylib();
     JD.transferTracker(DstRT, SrcRT);
     for (auto *L : reverse(ResourceManagers))
-      L->handleTransferResources(DstRT.getKeyUnsafe(), SrcRT.getKeyUnsafe());
+      L->handleTransferResources(JD, DstRT.getKeyUnsafe(),
+                                 SrcRT.getKeyUnsafe());
   });
 }
 
diff --git a/llvm/lib/ExecutionEngine/Orc/DebugObjectManagerPlugin.cpp b/llvm/lib/ExecutionEngine/Orc/DebugObjectManagerPlugin.cpp
index 418dc39..02c3e61 100644
--- a/llvm/lib/ExecutionEngine/Orc/DebugObjectManagerPlugin.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/DebugObjectManagerPlugin.cpp
@@ -492,7 +492,8 @@ Error DebugObjectManagerPlugin::notifyFailed(
   return Error::success();
 }
 
-void DebugObjectManagerPlugin::notifyTransferringResources(ResourceKey DstKey,
+void DebugObjectManagerPlugin::notifyTransferringResources(JITDylib &JD,
+                                                           ResourceKey DstKey,
                                                            ResourceKey SrcKey) {
   // Debug objects are stored by ResourceKey only after registration.
   // Thus, pending objects don't need to be updated here.
@@ -507,7 +508,8 @@ void DebugObjectManagerPlugin::notifyTransferringResources(ResourceKey DstKey,
   }
 }
 
-Error DebugObjectManagerPlugin::notifyRemovingResources(ResourceKey Key) {
+Error DebugObjectManagerPlugin::notifyRemovingResources(JITDylib &JD,
+                                                        ResourceKey Key) {
   // Removing the resource for a pending object fails materialization, so they
   // get cleaned up in the notifyFailed() handler.
   std::lock_guard<std::mutex> Lock(RegisteredObjsLock);
diff --git a/llvm/lib/ExecutionEngine/Orc/DebuggerSupportPlugin.cpp b/llvm/lib/ExecutionEngine/Orc/DebuggerSupportPlugin.cpp
index 009e2d7..15e7ffb 100644
--- a/llvm/lib/ExecutionEngine/Orc/DebuggerSupportPlugin.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/DebuggerSupportPlugin.cpp
@@ -390,12 +390,12 @@ Error GDBJITDebugInfoRegistrationPlugin::notifyFailed(
 }
 
 Error GDBJITDebugInfoRegistrationPlugin::notifyRemovingResources(
-    ResourceKey K) {
+    JITDylib &JD, ResourceKey K) {
   return Error::success();
 }
 
 void GDBJITDebugInfoRegistrationPlugin::notifyTransferringResources(
-    ResourceKey DstKey, ResourceKey SrcKey) {}
+    JITDylib &JD, ResourceKey DstKey, ResourceKey SrcKey) {}
 
 void GDBJITDebugInfoRegistrationPlugin::modifyPassConfig(
     MaterializationResponsibility &MR, LinkGraph &LG,
diff --git a/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp b/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
index 1a4a451..2b11c47 100644
--- a/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
@@ -680,12 +680,12 @@ Error ObjectLinkingLayer::notifyEmitted(MaterializationResponsibility &MR,
       [&](ResourceKey K) { Allocs[K].push_back(std::move(FA)); });
 }
 
-Error ObjectLinkingLayer::handleRemoveResources(ResourceKey K) {
+Error ObjectLinkingLayer::handleRemoveResources(JITDylib &JD, ResourceKey K) {
 
   {
     Error Err = Error::success();
     for (auto &P : Plugins)
-      Err = joinErrors(std::move(Err), P->notifyRemovingResources(K));
+      Err = joinErrors(std::move(Err), P->notifyRemovingResources(JD, K));
     if (Err)
       return Err;
   }
@@ -705,7 +705,8 @@ Error ObjectLinkingLayer::handleRemoveResources(ResourceKey K) {
   return MemMgr.deallocate(std::move(AllocsToRemove));
 }
 
-void ObjectLinkingLayer::handleTransferResources(ResourceKey DstKey,
+void ObjectLinkingLayer::handleTransferResources(JITDylib &JD,
+                                                 ResourceKey DstKey,
                                                  ResourceKey SrcKey) {
   auto I = Allocs.find(SrcKey);
   if (I != Allocs.end()) {
@@ -721,7 +722,7 @@ void ObjectLinkingLayer::handleTransferResources(ResourceKey DstKey,
   }
 
   for (auto &P : Plugins)
-    P->notifyTransferringResources(DstKey, SrcKey);
+    P->notifyTransferringResources(JD, DstKey, SrcKey);
 }
 
 EHFrameRegistrationPlugin::EHFrameRegistrationPlugin(
@@ -773,7 +774,8 @@ Error EHFrameRegistrationPlugin::notifyFailed(
   return Error::success();
 }
 
-Error EHFrameRegistrationPlugin::notifyRemovingResources(ResourceKey K) {
+Error EHFrameRegistrationPlugin::notifyRemovingResources(JITDylib &JD,
+                                                         ResourceKey K) {
   std::vector<ExecutorAddrRange> RangesToRemove;
 
   ES.runSessionLocked([&] {
@@ -797,7 +799,7 @@ Error EHFrameRegistrationPlugin::notifyRemovingResources(ResourceKey K) {
 }
 
 void EHFrameRegistrationPlugin::notifyTransferringResources(
-    ResourceKey DstKey, ResourceKey SrcKey) {
+    JITDylib &JD, ResourceKey DstKey, ResourceKey SrcKey) {
   auto SI = EHFrameRanges.find(SrcKey);
   if (SI == EHFrameRanges.end())
     return;
diff --git a/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp b/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp
index 8e9d4cd..9e781e9 100644
--- a/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp
@@ -350,7 +350,8 @@ void RTDyldObjectLinkingLayer::onObjEmit(
   }
 }
 
-Error RTDyldObjectLinkingLayer::handleRemoveResources(ResourceKey K) {
+Error RTDyldObjectLinkingLayer::handleRemoveResources(JITDylib &JD,
+                                                      ResourceKey K) {
 
   std::vector<MemoryManagerUP> MemMgrsToRemove;
 
@@ -374,7 +375,8 @@ Error RTDyldObjectLinkingLayer::handleRemoveResources(ResourceKey K) {
   return Error::success();
 }
 
-void RTDyldObjectLinkingLayer::handleTransferResources(ResourceKey DstKey,
+void RTDyldObjectLinkingLayer::handleTransferResources(JITDylib &JD,
+                                                       ResourceKey DstKey,
                                                        ResourceKey SrcKey) {
   auto I = MemMgrs.find(SrcKey);
   if (I != MemMgrs.end()) {
diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
index 9184945..4180be4 100644
--- a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
+++ b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
@@ -945,10 +945,10 @@ Session::Session(std::unique_ptr<ExecutorProcessControl> EPC, Error &Err)
     Error notifyFailed(MaterializationResponsibility &MR) override {
       return Error::success();
     }
-    Error notifyRemovingResources(ResourceKey K) override {
+    Error notifyRemovingResources(JITDylib &JD, ResourceKey K) override {
       return Error::success();
     }
-    void notifyTransferringResources(ResourceKey DstKey,
+    void notifyTransferringResources(JITDylib &JD, ResourceKey DstKey,
                                      ResourceKey SrcKey) override {}
 
   private:
diff --git a/llvm/unittests/ExecutionEngine/Orc/ResourceTrackerTest.cpp b/llvm/unittests/ExecutionEngine/Orc/ResourceTrackerTest.cpp
index 56362cd..fa0c374 100644
--- a/llvm/unittests/ExecutionEngine/Orc/ResourceTrackerTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/ResourceTrackerTest.cpp
@@ -23,21 +23,23 @@ namespace {
 template <typename ResourceT = unsigned>
 class SimpleResourceManager : public ResourceManager {
 public:
-  using HandleRemoveFunction = unique_function<Error(ResourceKey)>;
+  using HandleRemoveFunction =
+      unique_function<Error(JITDylib &JD, ResourceKey)>;
 
   using HandleTransferFunction =
-      unique_function<void(ResourceKey, ResourceKey)>;
+      unique_function<void(JITDylib &JD, ResourceKey, ResourceKey)>;
 
   using RecordedResourcesMap = DenseMap<ResourceKey, ResourceT>;
 
   SimpleResourceManager(ExecutionSession &ES) : ES(ES) {
-    HandleRemove = [&](ResourceKey K) -> Error {
-      ES.runSessionLocked([&] { removeResource(K); });
+    HandleRemove = [&](JITDylib &JD, ResourceKey K) -> Error {
+      ES.runSessionLocked([&] { removeResource(JD, K); });
       return Error::success();
     };
 
-    HandleTransfer = [this](ResourceKey DstKey, ResourceKey SrcKey) {
-      transferResources(DstKey, SrcKey);
+    HandleTransfer = [this](JITDylib &JD, ResourceKey DstKey,
+                            ResourceKey SrcKey) {
+      transferResources(JD, DstKey, SrcKey);
     };
 
     ES.registerResourceManager(*this);
@@ -69,11 +71,11 @@ public:
   }
 
   /// Remove the resource associated with K from the map if present.
-  void removeResource(ResourceKey K) { Resources.erase(K); }
+  void removeResource(JITDylib &JD, ResourceKey K) { Resources.erase(K); }
 
   /// Transfer resources from DstKey to SrcKey.
   template <typename MergeOp = std::plus<ResourceT>>
-  void transferResources(ResourceKey DstKey, ResourceKey SrcKey,
+  void transferResources(JITDylib &JD, ResourceKey DstKey, ResourceKey SrcKey,
                          MergeOp Merge = MergeOp()) {
     auto &DstResourceRef = Resources[DstKey];
     ResourceT DstResources;
@@ -90,13 +92,13 @@ public:
   RecordedResourcesMap &getRecordedResources() { return Resources; }
   const RecordedResourcesMap &getRecordedResources() const { return Resources; }
 
-  Error handleRemoveResources(ResourceKey K) override {
-    return HandleRemove(K);
+  Error handleRemoveResources(JITDylib &JD, ResourceKey K) override {
+    return HandleRemove(JD, K);
   }
 
-  void handleTransferResources(ResourceKey DstKey,
+  void handleTransferResources(JITDylib &JD, ResourceKey DstKey,
                                ResourceKey SrcKey) override {
-    HandleTransfer(DstKey, SrcKey);
+    HandleTransfer(JD, DstKey, SrcKey);
   }
 
   static void transferNotAllowed(ResourceKey DstKey, ResourceKey SrcKey) {
@@ -115,11 +117,11 @@ TEST_F(ResourceTrackerStandardTest,
 
   bool ResourceManagerGotRemove = false;
   SimpleResourceManager<> SRM(ES);
-  SRM.setHandleRemove([&](ResourceKey K) -> Error {
+  SRM.setHandleRemove([&](JITDylib &JD, ResourceKey K) -> Error {
     ResourceManagerGotRemove = true;
     EXPECT_EQ(SRM.getRecordedResources().size(), 0U)
         << "Unexpected resources recorded";
-    SRM.removeResource(K);
+    SRM.removeResource(JD, K);
     return Error::success();
   });
 
@@ -152,13 +154,13 @@ TEST_F(ResourceTrackerStandardTest, BasicDefineAndRemoveAllAfterMaterializing) {
 
   bool ResourceManagerGotRemove = false;
   SimpleResourceManager<> SRM(ES);
-  SRM.setHandleRemove([&](ResourceKey K) -> Error {
+  SRM.setHandleRemove([&](JITDylib &JD, ResourceKey K) -> Error {
     ResourceManagerGotRemove = true;
     EXPECT_EQ(SRM.getRecordedResources().size(), 1U)
         << "Unexpected number of resources recorded";
     EXPECT_EQ(SRM.getRecordedResources().count(K), 1U)
         << "Unexpected recorded resource";
-    SRM.removeResource(K);
+    SRM.removeResource(JD, K);
     return Error::success();
   });
 
@@ -190,11 +192,11 @@ TEST_F(ResourceTrackerStandardTest, BasicDefineAndRemoveAllWhileMaterializing) {
 
   bool ResourceManagerGotRemove = false;
   SimpleResourceManager<> SRM(ES);
-  SRM.setHandleRemove([&](ResourceKey K) -> Error {
+  SRM.setHandleRemove([&](JITDylib &JD, ResourceKey K) -> Error {
     ResourceManagerGotRemove = true;
     EXPECT_EQ(SRM.getRecordedResources().size(), 0U)
         << "Unexpected resources recorded";
-    SRM.removeResource(K);
+    SRM.removeResource(JD, K);
     return Error::success();
   });
 
@@ -289,11 +291,12 @@ TEST_F(ResourceTrackerStandardTest,
 
   bool ResourceManagerGotTransfer = false;
   SimpleResourceManager<> SRM(ES);
-  SRM.setHandleTransfer([&](ResourceKey DstKey, ResourceKey SrcKey) {
-    ResourceManagerGotTransfer = true;
-    auto &RR = SRM.getRecordedResources();
-    EXPECT_EQ(RR.size(), 0U) << "Expected no resources recorded yet";
-  });
+  SRM.setHandleTransfer(
+      [&](JITDylib &JD, ResourceKey DstKey, ResourceKey SrcKey) {
+        ResourceManagerGotTransfer = true;
+        auto &RR = SRM.getRecordedResources();
+        EXPECT_EQ(RR.size(), 0U) << "Expected no resources recorded yet";
+      });
 
   auto MakeMU = [&](SymbolStringPtr Name, JITEvaluatedSymbol Sym) {
     return std::make_unique<SimpleMaterializationUnit>(
@@ -339,10 +342,11 @@ TEST_F(ResourceTrackerStandardTest,
 
   bool ResourceManagerGotTransfer = false;
   SimpleResourceManager<> SRM(ES);
-  SRM.setHandleTransfer([&](ResourceKey DstKey, ResourceKey SrcKey) {
-    ResourceManagerGotTransfer = true;
-    SRM.transferResources(DstKey, SrcKey);
-  });
+  SRM.setHandleTransfer(
+      [&](JITDylib &JD, ResourceKey DstKey, ResourceKey SrcKey) {
+        ResourceManagerGotTransfer = true;
+        SRM.transferResources(JD, DstKey, SrcKey);
+      });
 
   auto MakeMU = [&](SymbolStringPtr Name, JITEvaluatedSymbol Sym) {
     return std::make_unique<SimpleMaterializationUnit>(
@@ -389,10 +393,11 @@ TEST_F(ResourceTrackerStandardTest,
 
   bool ResourceManagerGotTransfer = false;
   SimpleResourceManager<> SRM(ES);
-  SRM.setHandleTransfer([&](ResourceKey DstKey, ResourceKey SrcKey) {
-    ResourceManagerGotTransfer = true;
-    SRM.transferResources(DstKey, SrcKey);
-  });
+  SRM.setHandleTransfer(
+      [&](JITDylib &JD, ResourceKey DstKey, ResourceKey SrcKey) {
+        ResourceManagerGotTransfer = true;
+        SRM.transferResources(JD, DstKey, SrcKey);
+      });
 
   auto FooRT = JD.createResourceTracker();
   std::unique_ptr<MaterializationResponsibility> FooMR;