[llvm] c32a4f8 - [ORC] Allow removal of ObjectLinkingLayer Plugins.

[Lang Hames via llvm-commits]

Author: Lang Hames
Date: 2024-05-01T22:40:59-09:30
New Revision: c32a4f83b59c4293f0b91503ed217371ae1ab543

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

LOG: [ORC] Allow removal of ObjectLinkingLayer Plugins.

This adds a removePlugin operation to ObjectLinkingLayer. The removal of a
plugin will be visible in all links started after the removal (ongoing links
started before the removal will still use the removed plugin).

Coding my way home: 17.56037S, 149.61118W

Added: 
    

Modified: 
    llvm/include/llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h
    llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h b/llvm/include/llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h
index e452b90598a0ad..689daba8ae73b8 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h
@@ -121,13 +121,26 @@ class ObjectLinkingLayer : public RTTIExtends<ObjectLinkingLayer, ObjectLayer>,
     this->ReturnObjectBuffer = std::move(ReturnObjectBuffer);
   }
 
-  /// Add a pass-config modifier.
+  /// Add a plugin.
   ObjectLinkingLayer &addPlugin(std::shared_ptr<Plugin> P) {
     std::lock_guard<std::mutex> Lock(LayerMutex);
     Plugins.push_back(std::move(P));
     return *this;
   }
 
+  /// Remove a plugin. This remove applies only to subsequent links (links
+  /// already underway will continue to use the plugin), and does not of itself
+  /// destroy the plugin -- destruction will happen once all shared pointers
+  /// (including those held by in-progress links) are destroyed.
+  void removePlugin(Plugin &P) {
+    std::lock_guard<std::mutex> Lock(LayerMutex);
+    auto I = llvm::find_if(Plugins, [&](const std::shared_ptr<Plugin> &Elem) {
+      return Elem.get() == &P;
+    });
+    assert(I != Plugins.end() && "Plugin not present");
+    Plugins.erase(I);
+  }
+
   /// Add a LinkGraph to the JITDylib targeted by the given tracker.
   Error add(ResourceTrackerSP, std::unique_ptr<jitlink::LinkGraph> G);
 

diff  --git a/llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp b/llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp
index 7ab3e40df7459d..70570055fea9dd 100644
--- a/llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp
@@ -178,6 +178,82 @@ TEST_F(ObjectLinkingLayerTest, HandleErrorDuringPostAllocationPass) {
   EXPECT_THAT_EXPECTED(ES.lookup(&JD, "_anchor"), Failed());
 }
 
+TEST_F(ObjectLinkingLayerTest, AddAndRemovePlugins) {
+  class TestPlugin : public ObjectLinkingLayer::Plugin {
+  public:
+    TestPlugin(size_t &ActivationCount, bool &PluginDestroyed)
+        : ActivationCount(ActivationCount), PluginDestroyed(PluginDestroyed) {}
+
+    ~TestPlugin() { PluginDestroyed = true; }
+
+    void modifyPassConfig(MaterializationResponsibility &MR,
+                          jitlink::LinkGraph &G,
+                          jitlink::PassConfiguration &Config) override {
+      ++ActivationCount;
+    }
+
+    Error notifyFailed(MaterializationResponsibility &MR) override {
+      ADD_FAILURE() << "TestPlugin::notifyFailed called unexpectedly";
+      return Error::success();
+    }
+
+    Error notifyRemovingResources(JITDylib &JD, ResourceKey K) override {
+      return Error::success();
+    }
+
+    void notifyTransferringResources(JITDylib &JD, ResourceKey DstKey,
+                                     ResourceKey SrcKey) override {}
+
+  private:
+    size_t &ActivationCount;
+    bool &PluginDestroyed;
+  };
+
+  size_t ActivationCount = 0;
+  bool PluginDestroyed = false;
+
+  auto P = std::make_shared<TestPlugin>(ActivationCount, PluginDestroyed);
+
+  ObjLinkingLayer.addPlugin(P);
+
+  {
+    auto G1 = std::make_unique<LinkGraph>("G1", Triple("x86_64-apple-darwin"),
+                                          8, llvm::endianness::little,
+                                          x86_64::getEdgeKindName);
+
+    auto &DataSec = G1->createSection("__data", MemProt::Read | MemProt::Write);
+    auto &DataBlock = G1->createContentBlock(DataSec, BlockContent,
+                                             orc::ExecutorAddr(0x1000), 8, 0);
+    G1->addDefinedSymbol(DataBlock, 4, "_anchor1", 4, Linkage::Weak,
+                         Scope::Default, false, true);
+
+    EXPECT_THAT_ERROR(ObjLinkingLayer.add(JD, std::move(G1)), Succeeded());
+    EXPECT_THAT_EXPECTED(ES.lookup(&JD, "_anchor1"), Succeeded());
+    EXPECT_EQ(ActivationCount, 1U);
+  }
+
+  ObjLinkingLayer.removePlugin(*P);
+
+  {
+    auto G2 = std::make_unique<LinkGraph>("G2", Triple("x86_64-apple-darwin"),
+                                          8, llvm::endianness::little,
+                                          x86_64::getEdgeKindName);
+
+    auto &DataSec = G2->createSection("__data", MemProt::Read | MemProt::Write);
+    auto &DataBlock = G2->createContentBlock(DataSec, BlockContent,
+                                             orc::ExecutorAddr(0x1000), 8, 0);
+    G2->addDefinedSymbol(DataBlock, 4, "_anchor2", 4, Linkage::Weak,
+                         Scope::Default, false, true);
+
+    EXPECT_THAT_ERROR(ObjLinkingLayer.add(JD, std::move(G2)), Succeeded());
+    EXPECT_THAT_EXPECTED(ES.lookup(&JD, "_anchor2"), Succeeded());
+    EXPECT_EQ(ActivationCount, 1U);
+  }
+
+  P.reset();
+  EXPECT_TRUE(PluginDestroyed);
+}
+
 TEST(ObjectLinkingLayerSearchGeneratorTest, AbsoluteSymbolsObjectLayer) {
   class TestEPC : public UnsupportedExecutorProcessControl {
   public: