2 files changed, 44 insertions, 0 deletions
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/RPCUtils.h b/llvm/include/llvm/ExecutionEngine/Orc/RPCUtils.h
index 5267b50f744..b6945b50bd9 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/RPCUtils.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/RPCUtils.h
@@ -887,6 +887,24 @@ public:
     SequenceNumberMgr.reset();
   }
 
+  /// Remove the handler for the given function.
+  /// A handler must currently be registered for this function.
+  template <typename Func>
+  void removeHandler() {
+    auto IdItr = LocalFunctionIds.find(Func::getPrototype());
+    assert(IdItr != LocalFunctionIds.end() &&
+           "Function does not have a registered handler");
+    auto HandlerItr = Handlers.find(IdItr->second);
+    assert(HandlerItr != Handlers.end() &&
+           "Function does not have a registered handler");
+    Handlers.erase(HandlerItr);
+  }
+
+  /// Clear all handlers.
+  void clearHandlers() {
+    Handlers.clear();
+  }
+
 protected:
   // The LaunchPolicy type allows a launch policy to be specified when adding
   // a function handler. See addHandlerImpl.
diff --git a/llvm/unittests/ExecutionEngine/Orc/RPCUtilsTest.cpp b/llvm/unittests/ExecutionEngine/Orc/RPCUtilsTest.cpp
index 186c3d40848..9abf401af41 100644
--- a/llvm/unittests/ExecutionEngine/Orc/RPCUtilsTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/RPCUtilsTest.cpp
@@ -504,3 +504,29 @@ TEST(DummyRPC, TestAPICalls) {
 
   ServerThread.join();
 }
+
+TEST(DummyRPC, TestRemoveHandler) {
+  Queue Q1, Q2;
+  DummyRPCEndpoint Server(Q1, Q2);
+
+  Server.addHandler<DummyRPCAPI::VoidBool>(
+    [](bool B) {
+      EXPECT_EQ(B, true)
+        << "Server void(bool) received unexpected result";
+    });
+
+  Server.removeHandler<DummyRPCAPI::VoidBool>();
+}
+
+TEST(DummyRPC, TestClearHandlers) {
+  Queue Q1, Q2;
+  DummyRPCEndpoint Server(Q1, Q2);
+
+  Server.addHandler<DummyRPCAPI::VoidBool>(
+    [](bool B) {
+      EXPECT_EQ(B, true)
+        << "Server void(bool) received unexpected result";
+    });
+
+  Server.clearHandlers();
+}