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 Srv); + template + ServiceT &addService(std::unique_ptr Srv) { + assert(Srv && "addService called with null value"); + ServiceT &Ref = *Srv; + std::scoped_lock 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 + ServiceT &createService(ArgTs &&...Args) { + return addService(std::make_unique(std::forward(Args)...)); + } /// Set the ControllerAccess object. void setController(std::shared_ptr 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 Srv) { - std::scoped_lock Lock(M); - assert(!SI && "addService called after shutdown"); - Services.push_back(std::move(Srv)); -} - void Session::setController(std::shared_ptr CA) { assert(CA && "Cannot attach null controller"); std::scoped_lock 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 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 T) override { @@ -362,6 +377,18 @@ TEST(SessionTest, ExpectedShutdownSequence) { EXPECT_TRUE(SessionShutdownComplete); } +TEST(SessionTest, AddServiceAndUseRef) { + Session S(std::make_unique(), noErrors); + auto &CS = S.addService(std::make_unique(42)); + CS.doMoreConfig(1); +} + +TEST(SessionTest, CreateServiceAndUseRef) { + Session S(std::make_unique(), noErrors); + auto &CS = S.createService(42); + CS.doMoreConfig(1); +} + TEST(ControllerAccessTest, Basics) { // Test that we can set the ControllerAccess implementation and still shut // down as expected.