[orc-rt] Add Session class. (#162590)

An orc-rt Session contains a JIT'd program, managing resources and
communicating with a remote JIT-controller instance (expected to be an
orc::ExecutionSession, though this is not required).

Author: lhames

orc-rt/include/CMakeLists.txt

@@ -15,6 +15,7 @@ set(ORC_RT_HEADERS
     orc-rt/ResourceManager.h
     orc-rt/RTTI.h
     orc-rt/ScopeExit.h
+    orc-rt/Session.h
     orc-rt/SimpleNativeMemoryMap.h
     orc-rt/SimplePackedSerialization.h
     orc-rt/SPSAllocAction.h

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

@@ -0,0 +1,74 @@
+//===-------- Session.h - Session class and related APIs  -------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Session class and related APIs.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef ORC_RT_SESSION_H
+#define ORC_RT_SESSION_H
+
+#include "orc-rt/Error.h"
+#include "orc-rt/ResourceManager.h"
+#include "orc-rt/move_only_function.h"
+
+#include <vector>
+
+namespace orc_rt {
+
+/// Represents an ORC executor Session.
+class Session {
+public:
+  using ErrorReporterFn = move_only_function<void(Error)>;
+  using OnShutdownCompleteFn = move_only_function<void()>;
+
+  /// 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
+  /// program are not generally visible to ORC-RT, but can optionally be
+  /// reported by calling orc_rc_Session_reportError function.
+  ///
+  /// Note that entry into the reporter is not synchronized: it may be
+  /// called from multiple threads concurrently.
+  Session(ErrorReporterFn ReportError) : ReportError(std::move(ReportError)) {}
+
+  // Sessions are not copyable or moveable.
+  Session(const Session &) = delete;
+  Session &operator=(const Session &) = delete;
+
+  ~Session();
+
+  /// Report an error via the ErrorReporter function.
+  void reportError(Error Err) { ReportError(std::move(Err)); }
+
+  /// Initiate session shutdown.
+  ///
+  /// Runs shutdown on registered resources in reverse order.
+  void shutdown(OnShutdownCompleteFn OnComplete);
+
+  /// Initiate session shutdown and block until complete.
+  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));
+  }
+
+private:
+  void shutdownNext(OnShutdownCompleteFn OnShutdownComplete, Error Err,
+                    std::vector<std::unique_ptr<ResourceManager>> RemainingRMs);
+
+  std::mutex M;
+  ErrorReporterFn ReportError;
+  std::vector<std::unique_ptr<ResourceManager>> ResourceMgrs;
+};
+
+} // namespace orc_rt
+
+#endif // ORC_RT_SESSION_H

orc-rt/lib/executor/CMakeLists.txt

@@ -3,6 +3,7 @@ set(files
   AllocAction.cpp
   ResourceManager.cpp
   RTTI.cpp
+  Session.cpp
   SimpleNativeMemoryMap.cpp
   )
 

orc-rt/lib/executor/Session.cpp

@@ -0,0 +1,60 @@
+//===- Session.cpp --------------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Contains the implementation of the Session class and related APIs.
+//
+//===----------------------------------------------------------------------===//
+
+#include "orc-rt/Session.h"
+
+#include <future>
+
+namespace orc_rt {
+
+Session::~Session() { waitForShutdown(); }
+
+void Session::shutdown(OnShutdownCompleteFn OnShutdownComplete) {
+  std::vector<std::unique_ptr<ResourceManager>> ToShutdown;
+
+  {
+    std::scoped_lock<std::mutex> Lock(M);
+    std::swap(ResourceMgrs, ToShutdown);
+  }
+
+  shutdownNext(std::move(OnShutdownComplete), Error::success(),
+               std::move(ToShutdown));
+}
+
+void Session::waitForShutdown() {
+  std::promise<void> P;
+  auto F = P.get_future();
+
+  shutdown([P = std::move(P)]() mutable { P.set_value(); });
+
+  F.wait();
+}
+
+void Session::shutdownNext(
+    OnShutdownCompleteFn OnComplete, Error Err,
+    std::vector<std::unique_ptr<ResourceManager>> RemainingRMs) {
+  if (Err)
+    reportError(std::move(Err));
+
+  if (RemainingRMs.empty())
+    return OnComplete();
+
+  auto NextRM = std::move(RemainingRMs.back());
+  RemainingRMs.pop_back();
+  NextRM->shutdown([this, RemainingRMs = std::move(RemainingRMs),
+                    OnComplete = std::move(OnComplete)](Error Err) mutable {
+    shutdownNext(std::move(OnComplete), std::move(Err),
+                 std::move(RemainingRMs));
+  });
+}
+
+} // namespace orc_rt

orc-rt/unittests/CMakeLists.txt

@@ -23,6 +23,7 @@ add_orc_rt_unittest(CoreTests
   MemoryFlagsTest.cpp
   RTTITest.cpp
   ScopeExitTest.cpp
+  SessionTest.cpp
   SimpleNativeMemoryMapTest.cpp
   SimplePackedSerializationTest.cpp
   SPSAllocActionTest.cpp

orc-rt/unittests/SessionTest.cpp

@@ -0,0 +1,105 @@
+//===- SessionTest.cpp ----------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Tests for orc-rt's Session.h APIs.
+//
+//===----------------------------------------------------------------------===//
+
+#include "orc-rt/Session.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+#include <optional>
+
+using namespace orc_rt;
+using ::testing::Eq;
+using ::testing::Optional;
+
+class MockResourceManager : public ResourceManager {
+public:
+  enum class Op { Detach, Shutdown };
+
+  static Error alwaysSucceed(Op) { return Error::success(); }
+
+  MockResourceManager(std::optional<size_t> &DetachOpIdx,
+                      std::optional<size_t> &ShutdownOpIdx, size_t &OpIdx,
+                      move_only_function<Error(Op)> GenResult = alwaysSucceed)
+      : DetachOpIdx(DetachOpIdx), ShutdownOpIdx(ShutdownOpIdx), OpIdx(OpIdx),
+        GenResult(std::move(GenResult)) {}
+
+  void detach(OnCompleteFn OnComplete) override {
+    DetachOpIdx = OpIdx++;
+    OnComplete(GenResult(Op::Detach));
+  }
+
+  void shutdown(OnCompleteFn OnComplete) override {
+    ShutdownOpIdx = OpIdx++;
+    OnComplete(GenResult(Op::Shutdown));
+  }
+
+private:
+  std::optional<size_t> &DetachOpIdx;
+  std::optional<size_t> &ShutdownOpIdx;
+  size_t &OpIdx;
+  move_only_function<Error(Op)> GenResult;
+};
+
+// Non-overloaded version of cantFail: allows easy construction of
+// move_only_functions<void(Error)>s.
+static void noErrors(Error Err) { cantFail(std::move(Err)); }
+
+TEST(SessionTest, TrivialConstructionAndDestruction) { Session S(noErrors); }
+
+TEST(SessionTest, ReportError) {
+  Error E = Error::success();
+  cantFail(std::move(E)); // Force error into checked state.
+
+  Session S([&](Error Err) { E = std::move(Err); });
+  S.reportError(make_error<StringError>("foo"));
+
+  if (E)
+    EXPECT_EQ(toString(std::move(E)), "foo");
+  else
+    ADD_FAILURE() << "Missing error value";
+}
+
+TEST(SessionTest, SingleResourceManager) {
+  size_t OpIdx = 0;
+  std::optional<size_t> DetachOpIdx;
+  std::optional<size_t> ShutdownOpIdx;
+
+  {
+    Session S(noErrors);
+    S.addResourceManager(std::make_unique<MockResourceManager>(
+        DetachOpIdx, ShutdownOpIdx, OpIdx));
+  }
+
+  EXPECT_EQ(OpIdx, 1U);
+  EXPECT_EQ(DetachOpIdx, std::nullopt);
+  EXPECT_THAT(ShutdownOpIdx, Optional(Eq(0)));
+}
+
+TEST(SessionTest, MultipleResourceManagers) {
+  size_t OpIdx = 0;
+  std::optional<size_t> DetachOpIdx[3];
+  std::optional<size_t> ShutdownOpIdx[3];
+
+  {
+    Session S(noErrors);
+    for (size_t I = 0; I != 3; ++I)
+      S.addResourceManager(std::make_unique<MockResourceManager>(
+          DetachOpIdx[I], ShutdownOpIdx[I], OpIdx));
+  }
+
+  EXPECT_EQ(OpIdx, 3U);
+  // Expect shutdown in reverse order.
+  for (size_t I = 0; I != 3; ++I) {
+    EXPECT_EQ(DetachOpIdx[I], std::nullopt);
+    EXPECT_THAT(ShutdownOpIdx[I], Optional(Eq(2 - I)));
+  }
+}