Lang Hames d47aa6b62b
[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.
2026-03-16 20:45:51 +11:00

124 lines
3.9 KiB
C++

//===- 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"
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) {
assert(OnShutdownComplete && "OnShutdownComplete must be set");
// Safe to call concurrently / redundantly.
detachFromController();
{
std::scoped_lock<std::mutex> Lock(M);
if (SI) {
// SI exists: someone called shutdown already. If the shutdown is not yet
// complete then just add OnShutdownComplete to the list of pending
// callbacks for the in-progress shutdown, then return.
// (If the shutdown is already complete then we'll run the handler
// directly below).
if (!SI->Complete)
return SI->OnCompletes.push_back(std::move(OnShutdownComplete));
} else {
// SI does not exist: We're the first to call shutdown. Create a
// ShutdownInfo struct and add OnShutdownComplete to the list of pending
// callbacks, then call shutdownNext below (outside the lock).
SI = std::make_unique<ShutdownInfo>();
SI->OnCompletes.push_back(std::move(OnShutdownComplete));
std::swap(SI->Services, Services);
}
}
// OnShutdownComplete is set (i.e. not moved into the list of pending
// callbacks). This can only happen if shutdown is already complete. Call
// OnComplete directly and return.
if (OnShutdownComplete)
return OnShutdownComplete();
// OnShutdownComplete is _not_ set (i.e. was moved into the list of pending
// handlers), and we didn't return under the lock above, so we must be
// responsible for the shutdown. Call shutdownNext.
shutdownNext();
}
void Session::waitForShutdown() {
std::promise<void> P;
auto F = P.get_future();
shutdown([P = std::move(P)]() mutable { P.set_value(); });
F.get();
}
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() {
if (SI->Services.empty())
return shutdownComplete();
// Get the next Service to shut down.
auto NextSrv = std::move(SI->Services.back());
SI->Services.pop_back();
NextSrv->onShutdown([this]() { shutdownNext(); });
}
void Session::shutdownComplete() {
std::unique_ptr<TaskDispatcher> TmpDispatcher;
{
std::scoped_lock<std::mutex> Lock(M);
TmpDispatcher = std::move(Dispatcher);
}
TmpDispatcher->shutdown();
std::vector<OnShutdownCompleteFn> OnCompletes;
{
std::scoped_lock<std::mutex> Lock(M);
SI->Complete = true;
OnCompletes = std::move(SI->OnCompletes);
}
for (auto &OnShutdownComplete : OnCompletes)
OnShutdownComplete();
}
void Session::wrapperReturn(orc_rt_SessionRef S, uint64_t CallId,
orc_rt_WrapperFunctionBuffer ResultBytes) {
unwrap(S)->sendWrapperResult(CallId, WrapperFunctionBuffer(ResultBytes));
}
} // namespace orc_rt