[llvm] [orc-rt] Drop callManagedCodeAsync; tokens guard execution (PR #210508)
Lang Hames via llvm-commits
llvm-commits at lists.llvm.org
Sat Jul 18 06:01:53 PDT 2026
https://github.com/lhames created https://github.com/llvm/llvm-project/pull/210508
callManagedCodeSync is renamed to callManagedCode, and callManagedCodeAsync is removed.
callManagedCodeAsync held a ManagedCodeTaskGroup token from the initial call until the async continuation ran -- across the gap in which the call is suspended and no managed code is executing. A token is meant to guard managed code that is executing on a stack, so holding one across that gap guards nothing while blocking shutdown.
callManagedCode holds a token only for the synchronous call to the wrapped function, including anything the function runs inline on the same thread before returning. An async function is called the same way, with its continuation passed as an ordinary argument. Work deferred past the function's return -- stashed, or handed to another thread -- runs on a stack the token no longer guards; whoever later runs it is responsible for ensuring a token covers it, acquiring one and aborting if the acquire is denied.
See the new "Managed code execution and shutdown" section in docs/Design.md for the full model (what a token guards, nested entries, context-free acquire, and the abortability obligation).
>From 004bef3dfdeca90e6a9392c324f9e18c3e87df21 Mon Sep 17 00:00:00 2001
From: Lang Hames <lhames at gmail.com>
Date: Fri, 17 Jul 2026 14:56:31 +1000
Subject: [PATCH] [orc-rt] Drop callManagedCodeAsync; tokens guard execution
callManagedCodeSync is renamed to callManagedCode, and
callManagedCodeAsync is removed.
callManagedCodeAsync held a ManagedCodeTaskGroup token from the initial
call until the async continuation ran -- across the gap in which the
call is suspended and no managed code is executing. A token is meant to
guard managed code that is executing on a stack, so holding one across
that gap guards nothing while blocking shutdown.
callManagedCode holds a token only for the synchronous call to the
wrapped function, including anything the function runs inline on the
same thread before returning. An async function is called the same way,
with its continuation passed as an ordinary argument. Work deferred past
the function's return -- stashed, or handed to another thread -- runs on
a stack the token no longer guards; whoever later runs it is responsible
for ensuring a token covers it, acquiring one and aborting if the
acquire is denied.
See the new "Managed code execution and shutdown" section in
docs/Design.md for the full model (what a token guards, nested entries,
context-free acquire, and the abortability obligation).
---
orc-rt/docs/Design.md | 25 +++++++
orc-rt/include/orc-rt/Session.h | 112 +++++++++-------------------
orc-rt/test/unit/SessionTest.cpp | 124 +++++--------------------------
3 files changed, 79 insertions(+), 182 deletions(-)
diff --git a/orc-rt/docs/Design.md b/orc-rt/docs/Design.md
index e444f189bfa03..a2dcdf4b7f920 100644
--- a/orc-rt/docs/Design.md
+++ b/orc-rt/docs/Design.md
@@ -67,6 +67,31 @@ The `onShutdown` operation will be called at `Session` destruction time, after
all in-flight managed code calls have completed. Services should release all
held resources during `onShutdown`.
+### Managed code execution and shutdown
+
+Session teardown must not proceed while JIT'd code is still live on a stack:
+freeing the JIT'd code (and the resources it runs against) out from under those
+frames would crash the moment control returns to them.
+
+To help enforce this the Session carries a managed-code `TaskGroup`. Code that
+is about to run JIT'd code obtains a `TaskGroup::Token` from the group to
+bracket that execution, and the group delays Session shutdown until every Token
+has been released. `Session::callManagedCode` acquires and holds a Token for you
+around a synchronous call; Tokens can also be acquired manually from the
+`TokenSource` returned by `Session::managedCodeTokenSource`.
+
+A Token brackets a single span of execution on a stack -- *not* a whole chain of
+asynchronous operations. When an asynchronous operation captures a continuation,
+whoever later invokes that continuation must obtain a fresh Token to bracket the
+invocation (the continuation can't do it itself: its entry point may already be
+JIT'd code).
+
+Token acquisition fails once Session shutdown has been requested -- and it can
+fail even for a nested or resumed call whose caller already holds a Token. So
+every caller of JIT'd code needs a way to abort and unwind when acquisition is
+denied; `callManagedCode` reports denial through its return value, which callers
+must check.
+
### TaskDispatcher
Runs Tasks within the ORC runtime. In particular, calls originating from the
diff --git a/orc-rt/include/orc-rt/Session.h b/orc-rt/include/orc-rt/Session.h
index 380a3a6d33fd5..76f6b2b117954 100644
--- a/orc-rt/include/orc-rt/Session.h
+++ b/orc-rt/include/orc-rt/Session.h
@@ -48,8 +48,8 @@ inline Session *unwrap(orc_rt_SessionRef S) noexcept {
/// Represents an ORC executor Session.
class Session {
private:
- // Implementation helper for callManagedCodeSync (non-void version).
- template <typename RetT> struct ManagedCodeSyncCaller {
+ // Implementation helper for callManagedCode (non-void version).
+ template <typename RetT> struct ManagedCodeCaller {
template <typename FnT, typename... ArgTs>
static std::optional<RetT> call(TaskGroup::Token Tok, FnT &&Fn,
ArgTs &&...Args) {
@@ -59,8 +59,8 @@ class Session {
}
};
- // Implementation helper for callManagedCodeSync (void version).
- template <> struct ManagedCodeSyncCaller<void> {
+ // Implementation helper for callManagedCode (void version).
+ template <> struct ManagedCodeCaller<void> {
template <typename FnT, typename... ArgTs>
static bool call(TaskGroup::Token Tok, FnT &&Fn, ArgTs &&...Args) {
if (!Tok)
@@ -70,37 +70,6 @@ class Session {
}
};
- template <typename ReturnArgTupleT> struct ManagedCodeAsyncCaller;
-
- // Implementation helper for callManagedCodeAsync (non-void version).
- template <typename T>
- struct ManagedCodeAsyncCaller<std::tuple<std::optional<T>>> {
- template <typename ReturnT, typename FnT, typename... ArgTs>
- static void call(TaskGroup::Token Tok, ReturnT &&Return, FnT &&Fn,
- ArgTs &&...Args) {
- if (!Tok)
- return std::forward<ReturnT>(Return)(std::nullopt);
-
- std::forward<FnT>(Fn)([Tok = std::move(Tok), R = std::move(Return)](
- T Value) { R(std::move(Value)); },
- std::forward<ArgTs>(Args)...);
- }
- };
-
- // Implementation helper for callManagedCodeAsync (void version).
- template <> struct ManagedCodeAsyncCaller<std::tuple<bool>> {
- template <typename ReturnT, typename FnT, typename... ArgTs>
- static void call(TaskGroup::Token Tok, ReturnT &&Return, FnT &&Fn,
- ArgTs &&...Args) {
- if (!Tok)
- return std::forward<ReturnT>(Return)(false);
-
- std::forward<FnT>(Fn)(
- [Tok = std::move(Tok), R = std::move(Return)]() { R(true); },
- std::forward<ArgTs>(Args)...);
- }
- };
-
public:
using ErrorReporterFn = move_only_function<void(Error)>;
using OnDetachFn = move_only_function<void()>;
@@ -361,60 +330,51 @@ class Session {
/// When calling code managed by a Session (e.g. JIT'd code, or library code
/// loaded on behalf of JIT'd code), clients should hold a token for this
/// group, which can be constructed from the returned TokenSource. That token
- /// will prevent the Session from shutting down any Services (and the Session
- /// itself) until tasks accessing managed code have completed.
+ /// will delay Session teardown until all tasks accessing managed code have
+ /// completed.
///
- /// Clients should prefer using the callManagedCodeSync and
- /// callManagedCodeAsync helpers to automatically acquire and hold a token
- /// for the duration of a call.
+ /// Clients should prefer using callManagedCode to automatically acquire
+ /// and hold a token for the duration of a call.
TaskGroup::TokenSource managedCodeTokenSource() const {
return ManagedCodeTaskGroup;
}
- /// Synchronously call managed code.
+ /// Call managed code.
+ ///
+ /// This helper tries to acquire a ManagedCodeTaskGroup token and, if
+ /// successful, calls the given function object with the given arguments
+ /// while holding the token.
+ ///
+ /// The token is held only for the duration of the (synchronous) call to Fn,
+ /// and released as soon as Fn returns; anything Fn runs inline on this thread
+ /// before returning is covered by it. Work Fn defers past its return --
+ /// stashed to run later, or handed to another thread -- is NOT covered: it
+ /// runs on a stack the token no longer guards. Whoever runs such deferred
+ /// work is responsible for ensuring a token covers it, typically by acquiring
+ /// one (e.g. from a TokenSource, see managedCodeTokenSource) and aborting the
+ /// work if the acquire is denied, exactly as for any entry into managed code.
+ /// (A resumed continuation cannot bracket its own entry, since its landing
+ /// point may itself be managed code.)
///
- /// This helper tries to acquire a ManagedCodeTaskGroup token and then call
- /// the given function object with the given arguments while holding the
- /// token.
+ /// If the token is successfully acquired then this function returns the
+ /// result of the call to Fn as a std::optional<T> (for a non-void return
+ /// type T), or boolean true (for void returns). Note that for asynchronous
+ /// functions (which typically return void) this reflects only that Fn was
+ /// invoked, not the result of the asynchronous operation.
///
- /// If the token is successfully acquired then this function will return the
- /// call result as a std::optional<T> (for a non-void return type T), or
- /// boolean true (for void returns).
+ /// If the token is not successfully acquired then Fn is not called and this
+ /// function returns std::nullopt (for a non-void return type) or boolean
+ /// false (for void returns).
///
- /// If the token is not successfully acquired then this function will return
- /// std::nullopt (for non-void return type) or boolean false (for void
- /// returns).
+ /// See the "Managed code execution and shutdown" section of docs/Design.md
+ /// for the model behind managed-code tokens and shutdown.
template <typename FnT, typename... ArgTs>
- decltype(auto) callManagedCodeSync(FnT &&Fn, ArgTs &&...Args) {
- return ManagedCodeSyncCaller<std::invoke_result_t<FnT, ArgTs...>>::call(
+ decltype(auto) callManagedCode(FnT &&Fn, ArgTs &&...Args) {
+ return ManagedCodeCaller<std::invoke_result_t<FnT, ArgTs...>>::call(
TaskGroup::Token(ManagedCodeTaskGroup), std::forward<FnT>(Fn),
std::forward<ArgTs>(Args)...);
}
- /// Asynchronously call managed code.
- ///
- /// ReturnT must be a function object that takes either a boolean or a
- /// std::optional<T>.
- ///
- /// callManagedCodeAsync tries to acquire a ManagedCodeTaskGroup token and
- /// then call the given async function object while holding that token.
- ///
- /// If the token is successfully acquired then this function will call Fn,
- /// passing in a wrapped version of Return that takes a T (if Return takes a
- /// std::optional<T>), or a wrapped version of Return that takes no arguments
- /// (if Return takes a bool).
- ///
- /// If the token is not successfully acquired then this function will not
- /// call Fn, but instead immediately call Return with std::nullopt (if Return
- /// takes a std::optional<T>), or false (if Return takes a boolean).
- template <typename ReturnT, typename FnT, typename... ArgTs>
- void callManagedCodeAsync(ReturnT &&Return, FnT &&Fn, ArgTs &&...Args) {
- ManagedCodeAsyncCaller<typename CallableArgInfo<ReturnT>::args_tuple_type>::
- call(TaskGroup::Token(ManagedCodeTaskGroup),
- std::forward<ReturnT>(Return), std::forward<FnT>(Fn),
- std::forward<ArgTs>(Args)...);
- }
-
/// Call a tagged handler in the Controller.
///
/// This method can be called directly, but is expected to be more commonly
diff --git a/orc-rt/test/unit/SessionTest.cpp b/orc-rt/test/unit/SessionTest.cpp
index 4ff48ca06fd81..c123e88d4056b 100644
--- a/orc-rt/test/unit/SessionTest.cpp
+++ b/orc-rt/test/unit/SessionTest.cpp
@@ -467,18 +467,17 @@ TEST(SessionTest, ActiveManagedCallsDelayShutdown) {
EXPECT_TRUE(ShutdownComplete);
}
-static void managedSyncVoidFunction(int *P) { *P = 42; }
+static void managedVoidFunction(int *P) { *P = 42; }
-TEST(SessionTest, SyncCallManagedCodeVoidFn) {
- // Test synchronous calls to a void function while holding a
- // ManagedCodeTaskGroup token.
+TEST(SessionTest, CallManagedCodeVoidFn) {
+ // Test calls to a void function while holding a ManagedCodeTaskGroup token.
Session S(mockExecutorProcessInfo(), noDispatch, noErrors);
{
// Pre-shutdown we expect token acquisition to succeed and the function to
// run.
int X = 0;
- bool CallSucceeded = S.callManagedCodeSync(managedSyncVoidFunction, &X);
+ bool CallSucceeded = S.callManagedCode(managedVoidFunction, &X);
EXPECT_TRUE(CallSucceeded);
EXPECT_EQ(X, 42U);
@@ -488,25 +487,25 @@ TEST(SessionTest, SyncCallManagedCodeVoidFn) {
{
// Post-shutdown we expect token acquisition to fail, and
- // callManagedCodeSync to return false.
+ // callManagedCode to return false.
int X = 0;
- bool CallSucceeded = S.callManagedCodeSync(managedSyncVoidFunction, &X);
+ bool CallSucceeded = S.callManagedCode(managedVoidFunction, &X);
EXPECT_FALSE(CallSucceeded);
}
}
-static int managedSyncNonVoidFunction(int N) { return N + 1; }
+static int managedNonVoidFunction(int N) { return N + 1; }
-TEST(SessionTest, SyncCallManagedCodeNonVoidFn) {
- // Test synchronous calls to a non-void function while holding a
- // ManagedCodeTaskGroup token.
+TEST(SessionTest, CallManagedCodeNonVoidFn) {
+ // Test calls to a non-void function while holding a ManagedCodeTaskGroup
+ // token.
Session S(mockExecutorProcessInfo(), noDispatch, noErrors);
{
// Pre-shutdown we expect token acquisition to succeed, the function to be
// run, and the result to be returned.
- auto Result = S.callManagedCodeSync(managedSyncNonVoidFunction, 41);
+ auto Result = S.callManagedCode(managedNonVoidFunction, 41);
EXPECT_TRUE(Result);
EXPECT_EQ(*Result, 42U);
@@ -516,57 +515,19 @@ TEST(SessionTest, SyncCallManagedCodeNonVoidFn) {
{
// Post-shutdown we expect token acquisition to fail, and
- // callManagedCodeSync to return std::nullopt.
- auto Result = S.callManagedCodeSync(managedSyncNonVoidFunction, 41);
+ // callManagedCode to return std::nullopt.
+ auto Result = S.callManagedCode(managedNonVoidFunction, 41);
EXPECT_EQ(Result, std::nullopt);
}
}
-static void managedAsyncVoidFunction(move_only_function<void()> Return,
- int *P) {
- *P = 42;
- Return();
-}
-
-TEST(SessionTest, AsyncCallManagedCodeVoidFn) {
- // Test asynchronous calls to a void function while holding a
- // ManagedCodeTaskGroup token.
- Session S(mockExecutorProcessInfo(), noDispatch, noErrors);
-
- {
- // Pre-shutdown we expect token acquisition to succeed, and the function
- // and Return callback to be run.
- int X = 0;
- bool ReturnSucceeded = false;
- S.callManagedCodeAsync([&](bool B) { ReturnSucceeded = B; },
- managedAsyncVoidFunction, &X);
- EXPECT_TRUE(ReturnSucceeded);
- EXPECT_EQ(X, 42U);
- }
-
- waitForShutdown(S);
-
- {
- // Post-shutdown we expect token acquisition to fail. Return should be
- // with `false` and the function should not be called.
- int X = 0;
- bool ReturnSucceeded = false;
- S.callManagedCodeAsync([&](bool B) { ReturnSucceeded = B; },
- managedAsyncVoidFunction, &X);
- EXPECT_FALSE(ReturnSucceeded);
- EXPECT_EQ(X, 0U);
- }
-}
-
-static void managedAsyncNonVoidFunction(move_only_function<void(int)> Return,
- int *P) {
+static void managedAsyncFunction(move_only_function<void(int)> Return, int *P) {
Return(++*P);
}
-TEST(SessionTest, AsyncCallManagedCodeNonVoidFn) {
- // Test asynchronous calls to a non-void function while holding a
- // ManagedCodeTaskGroup token.
+TEST(SessionTest, CallManagedCodeAsyncFn) {
+ // Test that calls to managed async functions via callManagedCode work.
Session S(mockExecutorProcessInfo(), noDispatch, noErrors);
{
@@ -574,8 +535,7 @@ TEST(SessionTest, AsyncCallManagedCodeNonVoidFn) {
// and Return callback to be run.
int N = 41;
std::optional<int> Result;
- S.callManagedCodeAsync([&](std::optional<int> N) { Result = N; },
- managedAsyncNonVoidFunction, &N);
+ S.callManagedCode(managedAsyncFunction, [&](int N) { Result = N; }, &N);
EXPECT_TRUE(Result);
EXPECT_EQ(*Result, 42U);
EXPECT_EQ(N, 42U);
@@ -588,60 +548,12 @@ TEST(SessionTest, AsyncCallManagedCodeNonVoidFn) {
// with `std::nullopt` and the function should not be called.
int N = 41;
std::optional<int> Result;
- S.callManagedCodeAsync([&](std::optional<int> N) { Result = N; },
- managedAsyncNonVoidFunction, &N);
+ S.callManagedCode(managedAsyncFunction, [&](int N) { Result = N; }, &N);
EXPECT_EQ(Result, std::nullopt);
EXPECT_EQ(N, 41U);
}
}
-TEST(SessionTest, AsyncCallManagedCodeHoldsTokenAcrossAsyncGap) {
- // Verify that the ManagedCodeTaskGroup token is held until the async
- // continuation runs, not just until callManagedCodeAsync returns. This
- // ensures shutdown blocks for the duration of the actual async work.
- Session S(mockExecutorProcessInfo(), noDispatch, noErrors);
-
- size_t OpIdx = 0;
- std::optional<size_t> DetachOpIdx;
- std::optional<size_t> ShutdownOpIdx;
- S.createService<MockService>(DetachOpIdx, ShutdownOpIdx, OpIdx);
-
- // The managed code function stashes its continuation instead of calling it.
- std::optional<int> Result;
- move_only_function<void(int)> StashedContinuation;
- S.callManagedCodeAsync([&](std::optional<int> N) { Result = std::move(N); },
- [&](move_only_function<void(int)> Return, int N) {
- // Stash the continuation and return without calling
- // it.
- StashedContinuation = std::move(Return);
- },
- 41);
-
- // callManagedCodeAsync has returned, but the continuation hasn't been
- // called yet. The token should still be held inside StashedContinuation.
- ASSERT_TRUE(StashedContinuation);
-
- // Request shutdown. It should detach but block on the outstanding token.
- bool ShutdownComplete = false;
- S.shutdown([&]() { ShutdownComplete = true; });
-
- EXPECT_EQ(DetachOpIdx, 0U);
- EXPECT_FALSE(ShutdownOpIdx);
- EXPECT_FALSE(ShutdownComplete);
-
- // Now invoke the stashed continuation and then destroy it, releasing the
- // token.
- StashedContinuation(42);
- StashedContinuation = {};
-
- // Check result.
- EXPECT_EQ(Result, 42);
-
- // Shutdown should now have completed.
- EXPECT_EQ(ShutdownOpIdx, 1U);
- EXPECT_TRUE(ShutdownComplete);
-}
-
TEST(SessionTest, AddServiceAndUseRef) {
Session S(mockExecutorProcessInfo(), noDispatch, noErrors);
auto &CS = S.addService(std::make_unique<ConfigurableService>(42));
More information about the llvm-commits
mailing list