[llvm] [orc-rt] Replace RunWrapperCall with generalized Dispatch (PR #210591)

Lang Hames via llvm-commits llvm-commits at lists.llvm.org
Sun Jul 19 05:59:19 PDT 2026


https://github.com/lhames created https://github.com/llvm/llvm-project/pull/210591

RunWrapperCall dispatched wrapper-function calls, but the Session also needs to dispatch continuations for results returned by the controller. Replace it with a generalized Dispatch mechanism (DispatchFn) that is handed an opaque Task for each unit of work, covering both.

QueueingRunner is correspondingly simplified: it no longer knows anything about wrapper functions and just enqueues the Task it is given. Tests and the noDispatch helper are updated to the new interface.

>From fc3817f9a2b09067a45eb6df88c357d77eeb8d11 Mon Sep 17 00:00:00 2001
From: Lang Hames <lhames at gmail.com>
Date: Sat, 18 Jul 2026 23:57:26 +1000
Subject: [PATCH] [orc-rt] Replace RunWrapperCall with generalized Dispatch

RunWrapperCall dispatched wrapper-function calls, but the Session also needs to
dispatch continuations for results returned by the controller. Replace it with a
generalized Dispatch mechanism (DispatchFn) that is handed an opaque Task for
each unit of work, covering both.

QueueingRunner is correspondingly simplified: it no longer knows anything about
wrapper functions and just enqueues the Task it is given. Tests and the
noDispatch helper are updated to the new interface.
---
 orc-rt/include/orc-rt/QueueingRunner.h        |  22 ++--
 orc-rt/include/orc-rt/Session.h               |  35 +++---
 orc-rt/lib/executor/Session.cpp               |   4 +-
 orc-rt/test/unit/CommonTestUtils.h            |  21 ++--
 .../unit/InProcessControllerAccessTest.cpp    |   2 +-
 orc-rt/test/unit/QueueingRunnerTest.cpp       | 118 +++++++-----------
 6 files changed, 80 insertions(+), 122 deletions(-)

diff --git a/orc-rt/include/orc-rt/QueueingRunner.h b/orc-rt/include/orc-rt/QueueingRunner.h
index 86798a92eaf60..0e9941a783f87 100644
--- a/orc-rt/include/orc-rt/QueueingRunner.h
+++ b/orc-rt/include/orc-rt/QueueingRunner.h
@@ -13,12 +13,8 @@
 #ifndef ORC_RT_QUEUEINGRUNNER_H
 #define ORC_RT_QUEUEINGRUNNER_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 <cstdint>
 #include <deque>
 #include <mutex>
 #include <optional>
@@ -59,9 +55,9 @@ template <typename T> class SynchronizedDeque {
 
 } // namespace detail
 
-/// A wrapper-call runner that pushes incoming calls onto a caller-owned work
-/// queue, leaving the caller free to drain the queue however and whenever
-/// they choose.
+/// A task runner that pushes dispatched tasks onto a caller-owned work queue,
+/// leaving the caller free to drain the queue however and whenever they
+/// choose.
 ///
 /// QueueingRunner is intended for use on systems where threads are not
 /// available, and for unit tests. For most uses of the ORC runtime,
@@ -82,13 +78,9 @@ class QueueingRunner {
 
   QueueingRunner(WorkQueueT &Pending) : Pending(Pending) {}
 
-  /// Enqueue a wrapper-function call to be run later.
-  void operator()(orc_rt_SessionRef S, uint64_t CallId,
-                  orc_rt_WrapperFunctionReturn Return,
-                  orc_rt_WrapperFunction Fn, WrapperFunctionBuffer ArgBytes) {
-    Pending.push_back([=, ArgBytes = std::move(ArgBytes)]() mutable {
-      Fn(S, CallId, Return, ArgBytes.release());
-    });
+  /// Enqueue a task to be run later.
+  void operator()(move_only_function<void()> Task) {
+    Pending.push_back(std::move(Task));
   }
 
   /// Run all currently-queued calls in last-in-first-out order, returning when
diff --git a/orc-rt/include/orc-rt/Session.h b/orc-rt/include/orc-rt/Session.h
index d2971cae04250..0b1670e4e6514 100644
--- a/orc-rt/include/orc-rt/Session.h
+++ b/orc-rt/include/orc-rt/Session.h
@@ -79,17 +79,17 @@ class Session {
   using OnControllerCallReturnFn =
       move_only_function<void(WrapperFunctionBuffer)>;
 
-  /// Callback used by the Session to run incoming wrapper-function calls.
+  /// A unit of work handed to the Session's DispatchFn for execution.
+  using Task = move_only_function<void()>;
+
+  /// Callback used by the Session to dispatch tasks for execution.
   ///
-  /// A ManagedCodeTaskGroup token is created for each call to this callback,
-  /// and implementations must eventually call either Fn (typically as
-  /// Fn(S, CallId, Return, ArgBytes.release())), or call Return directly to
-  /// bail out of the call (typically with
-  /// WrapperFunctionBuffer::createOutOfBandError(...)). Failing to do either
-  /// will block Session shutdown indefinitely.
-  using RunWrapperCall = move_only_function<void(
-      orc_rt_SessionRef S, uint64_t CallId, orc_rt_WrapperFunctionReturn Return,
-      orc_rt_WrapperFunction Fn, WrapperFunctionBuffer ArgBytes)>;
+  /// The Session builds a Task for each unit of work it needs run -- an
+  /// incoming wrapper-function call, or a continuation for a result returned
+  /// by the controller -- and hands it to this callback, which is responsible
+  /// for arranging the task to be run inline, queued, or posted to a thread
+  /// pool.
+  using DispatchFn = move_only_function<void(Task)>;
 
   /// Tag used to identify executor-callable functions in the controller.
   /// See callController.
@@ -192,13 +192,14 @@ class Session {
   /// program are not generally visible to ORC-RT, but can optionally be
   /// reported by calling the orc_rt_Session_reportError function.)
   ///
-  /// The RunCall callback will be invoked for every incoming wrapper-function
-  /// call, and is responsible for arranging the call to be run (inline,
-  /// queued, or posted to a thread pool, at the caller's discretion).
+  /// The Dispatch callback is invoked to run tasks generated by the Session
+  /// (incoming wrapper-function calls, and continuations for results returned
+  /// by the controller), and is responsible for arranging each task to be run
+  /// inline, queued, or posted to a thread pool.
   ///
   /// Note that entry into the reporter is not synchronized: it may be
   /// called from multiple threads concurrently.
-  Session(ExecutorProcessInfo EPI, RunWrapperCall RunCall,
+  Session(ExecutorProcessInfo EPI, DispatchFn Dispatch,
           ErrorReporterFn ReportError);
 
   // Sessions are not copyable or moveable.
@@ -470,7 +471,9 @@ class Session {
       return;
     }
 
-    RunCall(wrap(this), CallId, &wrapperReturn, Fn, std::move(ArgBytes));
+    Dispatch([this, CallId, Fn, ArgBytes = std::move(ArgBytes)]() mutable {
+      Fn(wrap(this), CallId, &wrapperReturn, ArgBytes.release());
+    });
   }
 
   void sendWrapperResult(uint64_t CallId, WrapperFunctionBuffer ResultBytes);
@@ -478,7 +481,7 @@ class Session {
                             orc_rt_WrapperFunctionBuffer ResultBytes);
 
   ExecutorProcessInfo EPI;
-  RunWrapperCall RunCall;
+  DispatchFn Dispatch;
   std::shared_ptr<TaskGroup> ManagedCodeTaskGroup = TaskGroup::Create();
   std::shared_ptr<ControllerAccess> CA;
   ErrorReporterFn ReportError;
diff --git a/orc-rt/lib/executor/Session.cpp b/orc-rt/lib/executor/Session.cpp
index e9e7207819f24..6bc536882966f 100644
--- a/orc-rt/lib/executor/Session.cpp
+++ b/orc-rt/lib/executor/Session.cpp
@@ -49,9 +49,9 @@ class Session::NotificationService : public Service {
 
 Session::ControllerAccess::~ControllerAccess() = default;
 
-Session::Session(ExecutorProcessInfo EPI, RunWrapperCall RunCall,
+Session::Session(ExecutorProcessInfo EPI, DispatchFn Dispatch,
                  ErrorReporterFn ReportError)
-    : EPI(std::move(EPI)), RunCall(std::move(RunCall)),
+    : EPI(std::move(EPI)), Dispatch(std::move(Dispatch)),
       ReportError(std::move(ReportError)),
       Notifiers(createService<NotificationService>()) {}
 
diff --git a/orc-rt/test/unit/CommonTestUtils.h b/orc-rt/test/unit/CommonTestUtils.h
index fa62549b6ee0f..f4cc5dc6d33b4 100644
--- a/orc-rt/test/unit/CommonTestUtils.h
+++ b/orc-rt/test/unit/CommonTestUtils.h
@@ -46,19 +46,14 @@ inline orc_rt::ExecutorProcessInfo mockExecutorProcessInfo() noexcept {
   return orc_rt::ExecutorProcessInfo("arm64-apple-darwin", 16384);
 }
 
-/// RunWrapperCall callback for tests that should never dispatch a wrapper
-/// call. Records a test failure on invocation, then completes the call with an
-/// out-of-band error so that any caller awaiting the result unblocks and fails
-/// too (rather than hanging), even in -Asserts builds or when the dispatch
-/// arrives on a non-test thread.
-inline void noDispatch(orc_rt_SessionRef S, uint64_t CallId,
-                       orc_rt_WrapperFunctionReturn Return,
-                       orc_rt_WrapperFunction, orc_rt::WrapperFunctionBuffer) {
-  ADD_FAILURE() << "unexpected wrapper-call dispatch in a no-dispatch session";
-  Return(S, CallId,
-         orc_rt::WrapperFunctionBuffer::createOutOfBandError(
-             "unexpected wrapper-call dispatch in a no-dispatch session")
-             .release());
+/// DispatchFn for tests that should never dispatch a task. Records a test
+/// failure on invocation, then runs the task inline so that any caller
+/// awaiting a result unblocks (rather than hanging) and the managed-code token
+/// is released, even in -Asserts builds or when the dispatch arrives on a
+/// non-test thread.
+inline void noDispatch(orc_rt::move_only_function<void()> Task) {
+  ADD_FAILURE() << "unexpected dispatch in a no-dispatch session";
+  Task();
 }
 
 template <size_t Idx = 0> class OpCounter {
diff --git a/orc-rt/test/unit/InProcessControllerAccessTest.cpp b/orc-rt/test/unit/InProcessControllerAccessTest.cpp
index 9ea5dcb2b3110..a9bf1278080dd 100644
--- a/orc-rt/test/unit/InProcessControllerAccessTest.cpp
+++ b/orc-rt/test/unit/InProcessControllerAccessTest.cpp
@@ -287,7 +287,7 @@ static void echoWrapper(orc_rt_SessionRef S, uint64_t CallId,
 
 TEST(InProcessControllerAccessTest, CallFromControllerSuccess) {
   // The mock IPEPC initiates a wrapper call into IPCA. The Session's
-  // RunWrapperCall hook (a QueueingRunner over `Tasks`) enqueues the
+  // dispatch hook (a QueueingRunner over `Tasks`) enqueues the
   // invocation; draining the queue runs the wrapper, which echoes its
   // arguments back. Verify the mock receives the echoed bytes via
   // ReturnWrapperResult.
diff --git a/orc-rt/test/unit/QueueingRunnerTest.cpp b/orc-rt/test/unit/QueueingRunnerTest.cpp
index 9df5328b69102..415ce1e229ed6 100644
--- a/orc-rt/test/unit/QueueingRunnerTest.cpp
+++ b/orc-rt/test/unit/QueueingRunnerTest.cpp
@@ -11,7 +11,6 @@
 #include "gtest/gtest.h"
 
 #include <cstdint>
-#include <deque>
 #include <thread>
 #include <vector>
 
@@ -19,78 +18,60 @@ using namespace orc_rt;
 
 namespace {
 
-// A dummy SessionRef value used purely to thread an opaque pointer through
-// the runner's enqueue path.
-inline orc_rt_SessionRef dummySession() noexcept {
-  return reinterpret_cast<orc_rt_SessionRef>(uintptr_t{0xABCD});
-}
-
-inline orc_rt_WrapperFunctionReturn dummyReturn() noexcept {
-  return [](orc_rt_SessionRef, uint64_t, orc_rt_WrapperFunctionBuffer) {};
-}
-
-// Test wrapper-function that records each invocation in a globally-accessible
-// log via its CallId.
-struct CallRecord {
-  orc_rt_SessionRef Session;
-  uint64_t CallId;
-};
-
-static std::vector<CallRecord> *RecordingLog = nullptr;
-
-static void recordingFn(orc_rt_SessionRef S, uint64_t CallId,
-                        orc_rt_WrapperFunctionReturn,
-                        orc_rt_WrapperFunctionBuffer ArgBytes) {
-  WrapperFunctionBuffer Owned(ArgBytes);
-  RecordingLog->push_back({S, CallId});
-}
+// Log of task ids, in the order the tasks ran. Tasks record their own id here
+// when run, letting tests observe both whether and in what order tasks ran.
+static std::vector<uint64_t> *RecordingLog = nullptr;
 
 class QueueingRunnerTest : public ::testing::Test {
 protected:
   void SetUp() override { RecordingLog = &Log; }
   void TearDown() override { RecordingLog = nullptr; }
 
-  std::vector<CallRecord> Log;
+  // Build a task that records the given id in the log when run.
+  static move_only_function<void()> recordingTask(uint64_t Id) {
+    return [Id]() { RecordingLog->push_back(Id); };
+  }
+
+  std::vector<uint64_t> Log;
   QueueingRunner<>::WorkQueue Q;
 };
 
 TEST_F(QueueingRunnerTest, EnqueueDoesNotRunImmediately) {
   QueueingRunner<> R(Q);
-  R(dummySession(), /*CallId=*/0, dummyReturn(), recordingFn,
-    WrapperFunctionBuffer());
-  EXPECT_EQ(Log.size(), 0u) << "Enqueue should not run the call";
-  // Pop initial call.
+  R(recordingTask(0));
+  EXPECT_EQ(Log.size(), 0u) << "Enqueue should not run the task";
+  // Pop initial task.
   EXPECT_TRUE(Q.pop_back())
-      << "At least one call should be sitting in the queue";
+      << "At least one task should be sitting in the queue";
   EXPECT_FALSE(Q.pop_back())
-      << "Exactly one call should have been sitting in the queue";
+      << "Exactly one task should have been sitting in the queue";
 }
 
 TEST_F(QueueingRunnerTest, RunFIFOUntilEmpty) {
   QueueingRunner R(Q);
   for (uint64_t I = 0; I < 3; ++I)
-    R(dummySession(), I, dummyReturn(), recordingFn, WrapperFunctionBuffer());
+    R(recordingTask(I));
 
   QueueingRunner<>::runFIFOUntilEmpty(Q);
 
   ASSERT_EQ(Log.size(), 3u);
-  EXPECT_EQ(Log[0].CallId, 0u);
-  EXPECT_EQ(Log[1].CallId, 1u);
-  EXPECT_EQ(Log[2].CallId, 2u);
+  EXPECT_EQ(Log[0], 0u);
+  EXPECT_EQ(Log[1], 1u);
+  EXPECT_EQ(Log[2], 2u);
   EXPECT_FALSE(Q.pop_back()); // Expect queue to be empty.
 }
 
 TEST_F(QueueingRunnerTest, RunLIFOUntilEmpty) {
   QueueingRunner R(Q);
   for (uint64_t I = 0; I < 3; ++I)
-    R(dummySession(), I, dummyReturn(), recordingFn, WrapperFunctionBuffer());
+    R(recordingTask(I));
 
   QueueingRunner<>::runLIFOUntilEmpty(Q);
 
   ASSERT_EQ(Log.size(), 3u);
-  EXPECT_EQ(Log[0].CallId, 2u);
-  EXPECT_EQ(Log[1].CallId, 1u);
-  EXPECT_EQ(Log[2].CallId, 0u);
+  EXPECT_EQ(Log[0], 2u);
+  EXPECT_EQ(Log[1], 1u);
+  EXPECT_EQ(Log[2], 0u);
   EXPECT_FALSE(Q.pop_back()); // Expect queue to be empty.
 }
 
@@ -102,58 +83,45 @@ TEST_F(QueueingRunnerTest, DrainOnEmptyQueueIsNoOp) {
   EXPECT_EQ(Log.size(), 0u);
 }
 
-TEST_F(QueueingRunnerTest, DrainPicksUpCallsEnqueuedDuringDrain) {
-  // A call enqueued by a running call should also be drained in the same
+TEST_F(QueueingRunnerTest, DrainPicksUpTasksEnqueuedDuringDrain) {
+  // A task enqueued by a running task should also be drained in the same
   // runFIFOUntilEmpty call.
   QueueingRunner R(Q);
 
-  // First call enqueues a second call from inside its body. We use a custom
-  // wrapper-function (not recordingFn) to do that, since recordingFn doesn't
-  // know about the queue.
-  static QueueingRunner<> *PendingR = nullptr;
-  PendingR = &R;
-  static auto reentrantFn = [](orc_rt_SessionRef S, uint64_t CallId,
-                               orc_rt_WrapperFunctionReturn,
-                               orc_rt_WrapperFunctionBuffer ArgBytes) {
-    WrapperFunctionBuffer Owned(ArgBytes);
-    RecordingLog->push_back({S, CallId});
-    if (CallId == 0)
-      (*PendingR)(S, /*CallId=*/1, dummyReturn(), recordingFn,
-                  WrapperFunctionBuffer());
-  };
-
-  R(dummySession(), /*CallId=*/0, dummyReturn(), reentrantFn,
-    WrapperFunctionBuffer());
+  // The first task enqueues a second task from inside its body.
+  R([&]() {
+    RecordingLog->push_back(0);
+    R(recordingTask(1));
+  });
 
   QueueingRunner<>::runFIFOUntilEmpty(Q);
 
   ASSERT_EQ(Log.size(), 2u);
-  EXPECT_EQ(Log[0].CallId, 0u);
-  EXPECT_EQ(Log[1].CallId, 1u);
-  PendingR = nullptr;
+  EXPECT_EQ(Log[0], 0u);
+  EXPECT_EQ(Log[1], 1u);
 }
 
 TEST_F(QueueingRunnerTest, ConcurrentProducerAndDrainer) {
   // Verify that QueueingRunner's default WorkQueue (SynchronizedDeque)
   // tolerates concurrent push from one thread and drain from another.
   //
-  // A producer thread enqueues NumCalls wrapper-function invocations while
-  // the main thread spins draining the queue. Once the producer has finished
-  // enqueueing, the main thread joins it and then performs a final drain to
-  // pick up any tail of calls enqueued after its last loop iteration.
-  constexpr uint64_t NumCalls = 1024;
+  // A producer thread enqueues NumTasks tasks while the main thread spins
+  // draining the queue. Once the producer has finished enqueueing, the main
+  // thread joins it and then performs a final drain to pick up any tail of
+  // tasks enqueued after its last loop iteration.
+  constexpr uint64_t NumTasks = 1024;
 
   QueueingRunner<> R(Q);
 
   std::thread Producer([&]() {
-    for (uint64_t I = 0; I < NumCalls; ++I)
-      R(dummySession(), I, dummyReturn(), recordingFn, WrapperFunctionBuffer());
+    for (uint64_t I = 0; I < NumTasks; ++I)
+      R(recordingTask(I));
   });
 
   // Drain concurrently with the producer. The drainer doesn't know when the
   // producer is done, so we just spin until the producer thread has joined
   // (after which a final drain will be definitive).
-  while (Log.size() < NumCalls) {
+  while (Log.size() < NumTasks) {
     QueueingRunner<>::runFIFOUntilEmpty(Q);
     std::this_thread::yield();
   }
@@ -161,12 +129,12 @@ TEST_F(QueueingRunnerTest, ConcurrentProducerAndDrainer) {
   Producer.join();
   QueueingRunner<>::runFIFOUntilEmpty(Q); // pick up any tail.
 
-  ASSERT_EQ(Log.size(), NumCalls);
-  // Producer enqueues in order 0..NumCalls; FIFO drain must observe the same
+  ASSERT_EQ(Log.size(), NumTasks);
+  // Producer enqueues in order 0..NumTasks; FIFO drain must observe the same
   // order. (Concurrent draining doesn't reorder per-producer enqueues for a
   // single producer.)
-  for (uint64_t I = 0; I < NumCalls; ++I)
-    EXPECT_EQ(Log[I].CallId, I);
+  for (uint64_t I = 0; I < NumTasks; ++I)
+    EXPECT_EQ(Log[I], I);
 
   EXPECT_FALSE(Q.pop_back()) << "Queue should be empty after final drain";
 }



More information about the llvm-commits mailing list