[orc-rt] Add ControllerAccess interface. (#169598)

ControllerAccess provides an abstract interface for bidirectional RPC
between the executor (running JIT'd code) and the controller (containing
the llvm::orc::ExecutionSession). ControllerAccess implementations are
expected to implement IPC / RPC using a concrete communication method
(shared memory, pipes, sockets, native system IPC, etc).

Calls from executor to controller are made via callController, with
"handler tags" (addresses in the executor) specifying the target handler
in the controller. A handler must be associated in the controller with
the given tag for the call to succeed. This ensures that only registered
entry points in the controller can be used, and avoids leaking
controller addresses into the executor.

Calls in both directions are to "wrapper functions" that take a buffer
of bytes as input and return a buffer of bytes as output. In the ORC
runtime these must be `orc_rt_WrapperFunction`s (see
Session::handleWrapperCall). The interpretation of the byte buffers is
up to the wrapper functions: the ORC runtime imposes no restrictions on
how the bytes are to be interpreted.

ControllerAccess objects may be detached from the Session prior to
Session shutdown, in which case no further calls may be made in either
direction, and any pending results (from calls made that haven't
returned yet) should return errors. If the ControllerAccess class is
still attached at Session shutdown time it will be detached as part of
the shutdown process. The ControllerAccess::disconnect method must
support concurrent entry on multiple threads, and all callers must block
until they can guarantee that no further calls will be received or
accepted.

Author: lhames

orc-rt/include/orc-rt/Session.h

@@ -16,23 +16,97 @@
 #include "orc-rt/Error.h"
 #include "orc-rt/ResourceManager.h"
 #include "orc-rt/TaskDispatcher.h"
+#include "orc-rt/WrapperFunction.h"
 #include "orc-rt/move_only_function.h"
 
 #include "orc-rt-c/CoreTypes.h"
+#include "orc-rt-c/WrapperFunction.h"
 
+#include <cassert>
 #include <condition_variable>
 #include <memory>
 #include <mutex>
 #include <vector>
 
 namespace orc_rt {
 
+class Session;
+
+inline orc_rt_SessionRef wrap(Session *S) noexcept {
+  return reinterpret_cast<orc_rt_SessionRef>(S);
+}
+
+inline Session *unwrap(orc_rt_SessionRef S) noexcept {
+  return reinterpret_cast<Session *>(S);
+}
+
 /// Represents an ORC executor Session.
 class Session {
 public:
   using ErrorReporterFn = move_only_function<void(Error)>;
   using OnShutdownCompleteFn = move_only_function<void()>;
 
+  using HandlerTag = void *;
+  using OnCallHandlerCompleteFn =
+      move_only_function<void(WrapperFunctionBuffer)>;
+
+  /// Provides access to the controller.
+  class ControllerAccess {
+    friend class Session;
+
+  public:
+    virtual ~ControllerAccess();
+
+  protected:
+    using HandlerTag = Session::HandlerTag;
+    using OnCallHandlerCompleteFn = Session::OnCallHandlerCompleteFn;
+
+    ControllerAccess(Session &S) : S(&S) {}
+
+    /// Called by the Session to disconnect the session with the Controller.
+    ///
+    /// disconnect implementations must support concurrent entry on multiple
+    /// threads, and all calls must block until the disconnect operation is
+    /// complete.
+    ///
+    /// Once disconnect completes, implementations should make no further
+    /// calls to the Session, and should ignore any calls from the session
+    /// (implementations are free to ignore any calls from the Session after
+    /// disconnect is called).
+    virtual void disconnect() = 0;
+
+    /// Report an error to the session.
+    void reportError(Error Err) {
+      assert(S && "Already disconnected");
+      S->reportError(std::move(Err));
+    }
+
+    /// Call the handler in the controller associated with the given tag.
+    virtual void callController(OnCallHandlerCompleteFn OnComplete,
+                                HandlerTag T,
+                                WrapperFunctionBuffer ArgBytes) = 0;
+
+    /// Send the result of the given wrapper function call to the controller.
+    virtual void sendWrapperResult(uint64_t CallId,
+                                   WrapperFunctionBuffer ResultBytes) = 0;
+
+    /// Ask the Session to run the given wrapper function.
+    ///
+    /// Subclasses must not call this method after disconnect returns.
+    void handleWrapperCall(uint64_t CallId, orc_rt_WrapperFunction Fn,
+                           WrapperFunctionBuffer ArgBytes) {
+      assert(S && "Already disconnected");
+      S->handleWrapperCall(CallId, Fn, std::move(ArgBytes));
+    }
+
+  private:
+    void doDisconnect() {
+      disconnect();
+      S = nullptr;
+    }
+    Session *S;
+  };
+
   /// Create a session object. The ReportError function will be called to
   /// report errors generated while serving JIT'd code, e.g. if a memory
   /// management request cannot be fulfilled. (Error's within the JIT'd
@@ -69,9 +143,22 @@ class Session {
   void waitForShutdown();
 
   /// Add a ResourceManager to the session.
-  void addResourceManager(std::unique_ptr<ResourceManager> RM) {
-    std::scoped_lock<std::mutex> Lock(M);
-    ResourceMgrs.push_back(std::move(RM));
+  void addResourceManager(std::unique_ptr<ResourceManager> RM);
+
+  /// Set the ControllerAccess object.
+  void setController(std::shared_ptr<ControllerAccess> CA);
+
+  /// Disconnect the ControllerAccess object.
+  void detachFromController();
+
+  void callController(OnCallHandlerCompleteFn OnComplete, HandlerTag T,
+                      WrapperFunctionBuffer ArgBytes) {
+    if (auto TmpCA = CA)
+      CA->callController(std::move(OnComplete), T, std::move(ArgBytes));
+    else
+      OnComplete(
+          WrapperFunctionBuffer::createOutOfBandError("no controller attached")
+              .release());
   }
 
 private:
@@ -85,21 +172,43 @@ class Session {
   void shutdownNext(Error Err);
   void shutdownComplete();
 
+  void handleWrapperCall(uint64_t CallId, orc_rt_WrapperFunction Fn,
+                         WrapperFunctionBuffer ArgBytes) {
+    dispatch(makeGenericTask([=, ArgBytes = std::move(ArgBytes)]() mutable {
+      Fn(wrap(this), CallId, wrapperReturn, ArgBytes.release());
+    }));
+  }
+
+  void sendWrapperResult(uint64_t CallId, WrapperFunctionBuffer ResultBytes) {
+    if (auto TmpCA = CA)
+      TmpCA->sendWrapperResult(CallId, std::move(ResultBytes));
+  }
+
+  static void wrapperReturn(orc_rt_SessionRef S, uint64_t CallId,
+                            orc_rt_WrapperFunctionBuffer ResultBytes);
+
   std::unique_ptr<TaskDispatcher> Dispatcher;
+  std::shared_ptr<ControllerAccess> CA;
   ErrorReporterFn ReportError;
 
   std::mutex M;
   std::vector<std::unique_ptr<ResourceManager>> ResourceMgrs;
   std::unique_ptr<ShutdownInfo> SI;
 };
 
-inline orc_rt_SessionRef wrap(Session *S) noexcept {
-  return reinterpret_cast<orc_rt_SessionRef>(S);
-}
+class CallViaSession {
+public:
+  CallViaSession(Session &S, Session::HandlerTag T) : S(S), T(T) {}
 
-inline Session *unwrap(orc_rt_SessionRef S) noexcept {
-  return reinterpret_cast<Session *>(S);
-}
+  void operator()(Session::OnCallHandlerCompleteFn &&HandleResult,
+                  WrapperFunctionBuffer ArgBytes) {
+    S.callController(std::move(HandleResult), T, std::move(ArgBytes));
+  }
+
+private:
+  Session &S;
+  Session::HandlerTag T;
+};
 
 } // namespace orc_rt
 

orc-rt/lib/executor/Session.cpp

@@ -14,9 +14,14 @@
 
 namespace orc_rt {
 
+Session::ControllerAccess::~ControllerAccess() = default;
+
 Session::~Session() { waitForShutdown(); }
 
 void Session::shutdown(OnShutdownCompleteFn OnShutdownComplete) {
+  // Safe to call concurrently / redundantly.
+  detachFromController();
+
   {
     std::scoped_lock<std::mutex> Lock(M);
     if (SI) {
@@ -38,6 +43,27 @@ void Session::waitForShutdown() {
   SI->CompleteCV.wait(Lock, [&]() { return SI->Complete; });
 }
 
+void Session::addResourceManager(std::unique_ptr<ResourceManager> RM) {
+  std::scoped_lock<std::mutex> Lock(M);
+  assert(!SI && "addResourceManager called after shutdown");
+  ResourceMgrs.push_back(std::move(RM));
+}
+
+void Session::setController(std::shared_ptr<ControllerAccess> CA) {
+  assert(CA && "Cannot attach null controller");
+  std::scoped_lock<std::mutex> Lock(M);
+  assert(!this->CA && "Cannot re-attach controller");
+  assert(!SI && "Cannot attach controller after shutdown");
+  this->CA = std::move(CA);
+}
+
+void Session::detachFromController() {
+  if (auto TmpCA = CA) {
+    TmpCA->doDisconnect();
+    CA = nullptr;
+  }
+}
+
 void Session::shutdownNext(Error Err) {
   if (Err)
     reportError(std::move(Err));
@@ -72,4 +98,9 @@ void Session::shutdownComplete() {
   SI->CompleteCV.notify_all();
 }
 
+void Session::wrapperReturn(orc_rt_SessionRef S, uint64_t CallId,
+                            orc_rt_WrapperFunctionBuffer ResultBytes) {
+  unwrap(S)->sendWrapperResult(CallId, ResultBytes);
+}
+
 } // namespace orc_rt