From 98ccac607a9ff044ec7d024e119bc975eb03ed4c Mon Sep 17 00:00:00 2001
From: Lang Hames <lhames@gmail.com>
Date: Sun, 15 Mar 2026 15:07:16 +1100
Subject: [PATCH] [orc-rt] Return ref from Session::addService, add
 createService. (#186640)

Session::addService now returns a reference to the added Service. This
allows clients to hold a reference for further direct interaction with
the Service object.

This commit also introduces a new Session::createService convenience
method that creates the service and returns a reference to it.
---
 orc-rt/include/orc-rt/Session.h  | 17 ++++++++++++++++-
 orc-rt/lib/executor/Session.cpp  |  6 ------
 orc-rt/unittests/SessionTest.cpp | 27 +++++++++++++++++++++++++++
 3 files changed, 43 insertions(+), 7 deletions(-)

diff --git a/orc-rt/include/orc-rt/Session.h b/orc-rt/include/orc-rt/Session.h
index 258a63945c98..f6848dae09c1 100644
--- a/orc-rt/include/orc-rt/Session.h
+++ b/orc-rt/include/orc-rt/Session.h
@@ -143,7 +143,22 @@ public:
   void waitForShutdown();
 
   /// Add a Service to the session.
-  void addService(std::unique_ptr<Service> Srv);
+  template <typename ServiceT>
+  ServiceT &addService(std::unique_ptr<ServiceT> Srv) {
+    assert(Srv && "addService called with null value");
+    ServiceT &Ref = *Srv;
+    std::scoped_lock<std::mutex> Lock(M);
+    assert(!SI && "addService called after shutdown");
+    Services.push_back(std::move(Srv));
+    return Ref;
+  }
+
+  /// Construct an instance of ServiceT from the given arguments and add it to
+  /// the Session.
+  template <typename ServiceT, typename... ArgTs>
+  ServiceT &createService(ArgTs &&...Args) {
+    return addService(std::make_unique<ServiceT>(std::forward<ArgTs>(Args)...));
+  }
 
   /// Set the ControllerAccess object.
   void setController(std::shared_ptr<ControllerAccess> CA);
diff --git a/orc-rt/lib/executor/Session.cpp b/orc-rt/lib/executor/Session.cpp
index fe023064073a..f8b66c122424 100644
--- a/orc-rt/lib/executor/Session.cpp
+++ b/orc-rt/lib/executor/Session.cpp
@@ -63,12 +63,6 @@ void Session::waitForShutdown() {
   F.get();
 }
 
-void Session::addService(std::unique_ptr<Service> Srv) {
-  std::scoped_lock<std::mutex> Lock(M);
-  assert(!SI && "addService called after shutdown");
-  Services.push_back(std::move(Srv));
-}
-
 void Session::setController(std::shared_ptr<ControllerAccess> CA) {
   assert(CA && "Cannot attach null controller");
   std::scoped_lock<std::mutex> Lock(M);
diff --git a/orc-rt/unittests/SessionTest.cpp b/orc-rt/unittests/SessionTest.cpp
index ca1a25b0928d..ecf632d4eb4e 100644
--- a/orc-rt/unittests/SessionTest.cpp
+++ b/orc-rt/unittests/SessionTest.cpp
@@ -55,6 +55,21 @@ private:
   move_only_function<Error(Op)> GenResult;
 };
 
+class ConfigurableService : public Service {
+public:
+  ConfigurableService(int ConstructorOption) {}
+
+  void onDetach(OnCompleteFn OnComplete) override {
+    OnComplete(Error::success());
+  }
+
+  void onShutdown(OnCompleteFn OnComplete) override {
+    OnComplete(Error::success());
+  }
+
+  void doMoreConfig(int) noexcept {}
+};
+
 class NoDispatcher : public TaskDispatcher {
 public:
   void dispatch(std::unique_ptr<Task> T) override {
@@ -362,6 +377,18 @@ TEST(SessionTest, ExpectedShutdownSequence) {
   EXPECT_TRUE(SessionShutdownComplete);
 }
 
+TEST(SessionTest, AddServiceAndUseRef) {
+  Session S(std::make_unique<NoDispatcher>(), noErrors);
+  auto &CS = S.addService(std::make_unique<ConfigurableService>(42));
+  CS.doMoreConfig(1);
+}
+
+TEST(SessionTest, CreateServiceAndUseRef) {
+  Session S(std::make_unique<NoDispatcher>(), noErrors);
+  auto &CS = S.createService<ConfigurableService>(42);
+  CS.doMoreConfig(1);
+}
+
 TEST(ControllerAccessTest, Basics) {
   // Test that we can set the ControllerAccess implementation and still shut
   // down as expected.