[orc-rt] Add Controller Interface (CI) symbol table to Session. (#186747)

The Controller Interface is the extended set of symbols (mostly wrapper
functions) that the controller can call prior to loading any JIT'd code.
It is expected that it will be used to inspect the process and create /
configure services to enable JITing.

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

@@ -14,6 +14,7 @@
 #define ORC_RT_SESSION_H
 
 #include "orc-rt/Error.h"
+#include "orc-rt/LockedAccess.h"
 #include "orc-rt/Service.h"
 #include "orc-rt/TaskDispatcher.h"
 #include "orc-rt/WrapperFunction.h"
@@ -50,6 +51,8 @@ public:
   using OnCallHandlerCompleteFn =
       move_only_function<void(WrapperFunctionBuffer)>;
 
+  using SymbolMap = std::unordered_map<std::string, void *>;
+
   /// Provides access to the controller.
   class ControllerAccess {
     friend class Session;
@@ -116,9 +119,7 @@ public:
   /// Note that entry into the reporter is not synchronized: it may be
   /// called from multiple threads concurrently.
   Session(std::unique_ptr<TaskDispatcher> Dispatcher,
-          ErrorReporterFn ReportError)
-      : Dispatcher(std::move(Dispatcher)), ReportError(std::move(ReportError)) {
-  }
+          ErrorReporterFn ReportError);
 
   // Sessions are not copyable or moveable.
   Session(const Session &) = delete;
@@ -134,6 +135,12 @@ public:
   /// Report an error via the ErrorReporter function.
   void reportError(Error Err) { ReportError(std::move(Err)); }
 
+  /// Controller interface symbols map.
+  auto controllerInterface() { return LockedAccess(ControllerInterface, M); }
+  auto controllerInterface() const {
+    return LockedAccess(ControllerInterface, M);
+  }
+
   /// Initiate session shutdown.
   ///
   /// Runs shutdown on registered resources in reverse order.
@@ -204,8 +211,9 @@ private:
   std::shared_ptr<ControllerAccess> CA;
   ErrorReporterFn ReportError;
 
-  std::mutex M;
+  mutable std::mutex M;
   std::vector<std::unique_ptr<Service>> Services;
+  SymbolMap ControllerInterface;
   std::unique_ptr<ShutdownInfo> SI;
 };
 

orc-rt/lib/executor/Session.cpp

@@ -16,6 +16,12 @@ namespace orc_rt {
 
 Session::ControllerAccess::~ControllerAccess() = default;
 
+Session::Session(std::unique_ptr<TaskDispatcher> Dispatcher,
+                 ErrorReporterFn ReportError)
+    : Dispatcher(std::move(Dispatcher)), ReportError(std::move(ReportError)) {
+  ControllerInterface["orc_rt_SessionInstance"] = static_cast<void *>(this);
+}
+
 Session::~Session() { waitForShutdown(); }
 
 void Session::shutdown(OnShutdownCompleteFn OnShutdownComplete) {

orc-rt/unittests/SessionTest.cpp

@@ -387,6 +387,35 @@ TEST(SessionTest, CreateServiceAndUseRef) {
   CS.doMoreConfig(1);
 }
 
+TEST(SessionTest, ControllerInterfaceContainsSessionByDefault) {
+  Session S(std::make_unique<NoDispatcher>(), noErrors);
+  ASSERT_TRUE(S.controllerInterface()->count("orc_rt_SessionInstance"));
+  EXPECT_EQ(S.controllerInterface()->at("orc_rt_SessionInstance"),
+            static_cast<void *>(&S));
+}
+
+TEST(SessionTest, ControllerInterfaceWithRef) {
+  Session S(std::make_unique<NoDispatcher>(), noErrors);
+  int X = 0, Y = 0;
+  S.controllerInterface().with_ref([&](Session::SymbolMap &Syms) {
+    Syms["orc_rt_A"] = &X;
+    Syms["orc_rt_B"] = &Y;
+  });
+
+  EXPECT_EQ(S.controllerInterface()->at("orc_rt_A"), &X);
+  EXPECT_EQ(S.controllerInterface()->at("orc_rt_B"), &Y);
+}
+
+TEST(SessionTest, ControllerInterfaceConstAccess) {
+  Session S(std::make_unique<NoDispatcher>(), noErrors);
+  int X = 0;
+  S.controllerInterface()->emplace("orc_rt_X", &X);
+
+  const Session &CS = S;
+  ASSERT_TRUE(CS.controllerInterface()->count("orc_rt_X"));
+  EXPECT_EQ(CS.controllerInterface()->at("orc_rt_X"), &X);
+}
+
 TEST(ControllerAccessTest, Basics) {
   // Test that we can set the ControllerAccess implementation and still shut
   // down as expected.