467 lines · cpp
1//===- SessionTest.cpp ----------------------------------------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8//9// Tests for orc-rt's Session.h APIs.10//11//===----------------------------------------------------------------------===//12 13#include "orc-rt/Session.h"14#include "orc-rt/SPSWrapperFunction.h"15#include "orc-rt/ThreadPoolTaskDispatcher.h"16 17#include "gmock/gmock.h"18#include "gtest/gtest.h"19 20#include <chrono>21#include <deque>22#include <future>23#include <optional>24 25using namespace orc_rt;26using ::testing::Eq;27using ::testing::Optional;28 29class MockResourceManager : public ResourceManager {30public:31 enum class Op { Detach, Shutdown };32 33 static Error alwaysSucceed(Op) { return Error::success(); }34 35 MockResourceManager(std::optional<size_t> &DetachOpIdx,36 std::optional<size_t> &ShutdownOpIdx, size_t &OpIdx,37 move_only_function<Error(Op)> GenResult = alwaysSucceed)38 : DetachOpIdx(DetachOpIdx), ShutdownOpIdx(ShutdownOpIdx), OpIdx(OpIdx),39 GenResult(std::move(GenResult)) {}40 41 void detach(OnCompleteFn OnComplete) override {42 DetachOpIdx = OpIdx++;43 OnComplete(GenResult(Op::Detach));44 }45 46 void shutdown(OnCompleteFn OnComplete) override {47 ShutdownOpIdx = OpIdx++;48 OnComplete(GenResult(Op::Shutdown));49 }50 51private:52 std::optional<size_t> &DetachOpIdx;53 std::optional<size_t> &ShutdownOpIdx;54 size_t &OpIdx;55 move_only_function<Error(Op)> GenResult;56};57 58class NoDispatcher : public TaskDispatcher {59public:60 void dispatch(std::unique_ptr<Task> T) override {61 assert(false && "strictly no dispatching!");62 }63 void shutdown() override {}64};65 66class EnqueueingDispatcher : public TaskDispatcher {67public:68 using OnShutdownRunFn = move_only_function<void()>;69 EnqueueingDispatcher(std::deque<std::unique_ptr<Task>> &Tasks,70 OnShutdownRunFn OnShutdownRun = {})71 : Tasks(Tasks), OnShutdownRun(std::move(OnShutdownRun)) {}72 void dispatch(std::unique_ptr<Task> T) override {73 Tasks.push_back(std::move(T));74 }75 void shutdown() override {76 if (OnShutdownRun)77 OnShutdownRun();78 }79 80 /// Run up to NumTasks (arbitrarily many if NumTasks == std::nullopt) tasks81 /// from the front of the queue, returning the number actually run.82 static size_t83 runTasksFromFront(std::deque<std::unique_ptr<Task>> &Tasks,84 std::optional<size_t> NumTasks = std::nullopt) {85 size_t NumRun = 0;86 87 while (!Tasks.empty() && (!NumTasks || NumRun != *NumTasks)) {88 auto T = std::move(Tasks.front());89 Tasks.pop_front();90 T->run();91 ++NumRun;92 }93 94 return NumRun;95 }96 97private:98 std::deque<std::unique_ptr<Task>> &Tasks;99 OnShutdownRunFn OnShutdownRun;100};101 102class MockControllerAccess : public Session::ControllerAccess {103public:104 MockControllerAccess(Session &SS) : Session::ControllerAccess(SS), SS(SS) {}105 106 void disconnect() override {107 std::unique_lock<std::mutex> Lock(M);108 Shutdown = true;109 ShutdownCV.wait(Lock, [this]() { return Shutdown && Outstanding == 0; });110 }111 112 void callController(OnCallHandlerCompleteFn OnComplete, HandlerTag T,113 WrapperFunctionBuffer ArgBytes) override {114 // Simulate a call to the controller by dispatching a task to run the115 // requested function.116 size_t CId;117 {118 std::scoped_lock<std::mutex> Lock(M);119 if (Shutdown)120 return;121 CId = CallId++;122 Pending[CId] = std::move(OnComplete);123 ++Outstanding;124 }125 126 SS.dispatch(makeGenericTask([this, CId, OnComplete = std::move(OnComplete),127 T, ArgBytes = std::move(ArgBytes)]() mutable {128 auto Fn = reinterpret_cast<orc_rt_WrapperFunction>(T);129 Fn(reinterpret_cast<orc_rt_SessionRef>(this), CId, wfReturn,130 ArgBytes.release());131 }));132 133 bool Notify = false;134 {135 std::scoped_lock<std::mutex> Lock(M);136 if (--Outstanding == 0 && Shutdown)137 Notify = true;138 }139 if (Notify)140 ShutdownCV.notify_all();141 }142 143 void sendWrapperResult(uint64_t CallId,144 WrapperFunctionBuffer ResultBytes) override {145 // Respond to a simulated call by the controller.146 OnCallHandlerCompleteFn OnComplete;147 {148 std::scoped_lock<std::mutex> Lock(M);149 if (Shutdown) {150 assert(Pending.empty() && "Shut down but results still pending?");151 return;152 }153 auto I = Pending.find(CallId);154 assert(I != Pending.end());155 OnComplete = std::move(I->second);156 Pending.erase(I);157 ++Outstanding;158 }159 160 SS.dispatch(161 makeGenericTask([OnComplete = std::move(OnComplete),162 ResultBytes = std::move(ResultBytes)]() mutable {163 OnComplete(std::move(ResultBytes));164 }));165 166 bool Notify = false;167 {168 std::scoped_lock<std::mutex> Lock(M);169 if (--Outstanding == 0 && Shutdown)170 Notify = true;171 }172 if (Notify)173 ShutdownCV.notify_all();174 }175 176 void callFromController(OnCallHandlerCompleteFn OnComplete,177 orc_rt_WrapperFunction Fn,178 WrapperFunctionBuffer ArgBytes) {179 size_t CId = 0;180 bool BailOut = false;181 {182 std::scoped_lock<std::mutex> Lock(M);183 if (!Shutdown) {184 CId = CallId++;185 Pending[CId] = std::move(OnComplete);186 ++Outstanding;187 } else188 BailOut = true;189 }190 if (BailOut)191 return OnComplete(WrapperFunctionBuffer::createOutOfBandError(192 "Controller disconnected"));193 194 handleWrapperCall(CId, Fn, std::move(ArgBytes));195 196 bool Notify = false;197 {198 std::scoped_lock<std::mutex> Lock(M);199 if (--Outstanding == 0 && Shutdown)200 Notify = true;201 }202 203 if (Notify)204 ShutdownCV.notify_all();205 }206 207 /// Simulate start of outstanding operation.208 void incOutstanding() {209 std::scoped_lock<std::mutex> Lock(M);210 ++Outstanding;211 }212 213 /// Simulate end of outstanding operation.214 void decOutstanding() {215 bool Notify = false;216 {217 std::scoped_lock<std::mutex> Lock(M);218 if (--Outstanding == 0 && Shutdown)219 Notify = true;220 }221 if (Notify)222 ShutdownCV.notify_all();223 }224 225private:226 static void wfReturn(orc_rt_SessionRef S, uint64_t CallId,227 orc_rt_WrapperFunctionBuffer ResultBytes) {228 // Abuse "session" to refer to the ControllerAccess object.229 // We can just re-use sendFunctionResult for this.230 reinterpret_cast<MockControllerAccess *>(S)->sendWrapperResult(CallId,231 ResultBytes);232 }233 234 Session &SS;235 236 std::mutex M;237 bool Shutdown = false;238 size_t Outstanding = 0;239 size_t CallId = 0;240 std::unordered_map<size_t, OnCallHandlerCompleteFn> Pending;241 std::condition_variable ShutdownCV;242};243 244class CallViaMockControllerAccess {245public:246 CallViaMockControllerAccess(MockControllerAccess &CA,247 orc_rt_WrapperFunction Fn)248 : CA(CA), Fn(Fn) {}249 void operator()(Session::OnCallHandlerCompleteFn OnComplete,250 WrapperFunctionBuffer ArgBytes) {251 CA.callFromController(std::move(OnComplete), Fn, std::move(ArgBytes));252 }253 254private:255 MockControllerAccess &CA;256 orc_rt_WrapperFunction Fn;257};258 259// Non-overloaded version of cantFail: allows easy construction of260// move_only_functions<void(Error)>s.261static void noErrors(Error Err) { cantFail(std::move(Err)); }262 263TEST(SessionTest, TrivialConstructionAndDestruction) {264 Session S(std::make_unique<NoDispatcher>(), noErrors);265}266 267TEST(SessionTest, ReportError) {268 Error E = Error::success();269 cantFail(std::move(E)); // Force error into checked state.270 271 Session S(std::make_unique<NoDispatcher>(),272 [&](Error Err) { E = std::move(Err); });273 S.reportError(make_error<StringError>("foo"));274 275 if (E)276 EXPECT_EQ(toString(std::move(E)), "foo");277 else278 ADD_FAILURE() << "Missing error value";279}280 281TEST(SessionTest, DispatchTask) {282 int X = 0;283 std::deque<std::unique_ptr<Task>> Tasks;284 Session S(std::make_unique<EnqueueingDispatcher>(Tasks), noErrors);285 286 EXPECT_EQ(Tasks.size(), 0U);287 S.dispatch(makeGenericTask([&]() { ++X; }));288 EXPECT_EQ(Tasks.size(), 1U);289 auto T = std::move(Tasks.front());290 Tasks.pop_front();291 T->run();292 EXPECT_EQ(X, 1);293}294 295TEST(SessionTest, SingleResourceManager) {296 size_t OpIdx = 0;297 std::optional<size_t> DetachOpIdx;298 std::optional<size_t> ShutdownOpIdx;299 300 {301 Session S(std::make_unique<NoDispatcher>(), noErrors);302 S.addResourceManager(std::make_unique<MockResourceManager>(303 DetachOpIdx, ShutdownOpIdx, OpIdx));304 }305 306 EXPECT_EQ(OpIdx, 1U);307 EXPECT_EQ(DetachOpIdx, std::nullopt);308 EXPECT_THAT(ShutdownOpIdx, Optional(Eq(0)));309}310 311TEST(SessionTest, MultipleResourceManagers) {312 size_t OpIdx = 0;313 std::optional<size_t> DetachOpIdx[3];314 std::optional<size_t> ShutdownOpIdx[3];315 316 {317 Session S(std::make_unique<NoDispatcher>(), noErrors);318 for (size_t I = 0; I != 3; ++I)319 S.addResourceManager(std::make_unique<MockResourceManager>(320 DetachOpIdx[I], ShutdownOpIdx[I], OpIdx));321 }322 323 EXPECT_EQ(OpIdx, 3U);324 // Expect shutdown in reverse order.325 for (size_t I = 0; I != 3; ++I) {326 EXPECT_EQ(DetachOpIdx[I], std::nullopt);327 EXPECT_THAT(ShutdownOpIdx[I], Optional(Eq(2 - I)));328 }329}330 331TEST(SessionTest, ExpectedShutdownSequence) {332 // Check that Session shutdown results in...333 // 1. ResourceManagers being shut down.334 // 2. The TaskDispatcher being shut down.335 // 3. A call to OnShutdownComplete.336 337 size_t OpIdx = 0;338 std::optional<size_t> DetachOpIdx;339 std::optional<size_t> ShutdownOpIdx;340 341 bool DispatcherShutDown = false;342 bool SessionShutdownComplete = false;343 std::deque<std::unique_ptr<Task>> Tasks;344 Session S(std::make_unique<EnqueueingDispatcher>(345 Tasks,346 [&]() {347 EXPECT_TRUE(ShutdownOpIdx);348 EXPECT_EQ(*ShutdownOpIdx, 0);349 EXPECT_FALSE(SessionShutdownComplete);350 DispatcherShutDown = true;351 }),352 noErrors);353 S.addResourceManager(354 std::make_unique<MockResourceManager>(DetachOpIdx, ShutdownOpIdx, OpIdx));355 356 S.shutdown([&]() {357 EXPECT_TRUE(DispatcherShutDown);358 SessionShutdownComplete = true;359 });360 S.waitForShutdown();361 362 EXPECT_TRUE(SessionShutdownComplete);363}364 365TEST(ControllerAccessTest, Basics) {366 // Test that we can set the ControllerAccess implementation and still shut367 // down as expected.368 std::deque<std::unique_ptr<Task>> Tasks;369 Session S(std::make_unique<EnqueueingDispatcher>(Tasks), noErrors);370 auto CA = std::make_shared<MockControllerAccess>(S);371 S.setController(CA);372 373 EnqueueingDispatcher::runTasksFromFront(Tasks);374 375 S.waitForShutdown();376}377 378static void add_sps_wrapper(orc_rt_SessionRef S, uint64_t CallId,379 orc_rt_WrapperFunctionReturn Return,380 orc_rt_WrapperFunctionBuffer ArgBytes) {381 SPSWrapperFunction<int32_t(int32_t, int32_t)>::handle(382 S, CallId, Return, ArgBytes,383 [](move_only_function<void(int32_t)> Return, int32_t X, int32_t Y) {384 Return(X + Y);385 });386}387 388TEST(ControllerAccessTest, ValidCallToController) {389 // Simulate a call to a controller handler.390 std::deque<std::unique_ptr<Task>> Tasks;391 Session S(std::make_unique<EnqueueingDispatcher>(Tasks), noErrors);392 auto CA = std::make_shared<MockControllerAccess>(S);393 S.setController(CA);394 395 int32_t Result = 0;396 SPSWrapperFunction<int32_t(int32_t, int32_t)>::call(397 CallViaSession(S, reinterpret_cast<Session::HandlerTag>(add_sps_wrapper)),398 [&](Expected<int32_t> R) { Result = cantFail(std::move(R)); }, 41, 1);399 400 EnqueueingDispatcher::runTasksFromFront(Tasks);401 402 EXPECT_EQ(Result, 42);403 404 S.waitForShutdown();405}406 407TEST(ControllerAccessTest, CallToControllerBeforeAttach) {408 // Expect calls to the controller prior to attaching to fail.409 std::deque<std::unique_ptr<Task>> Tasks;410 Session S(std::make_unique<EnqueueingDispatcher>(Tasks), noErrors);411 412 Error Err = Error::success();413 SPSWrapperFunction<int32_t(int32_t, int32_t)>::call(414 CallViaSession(S, reinterpret_cast<Session::HandlerTag>(add_sps_wrapper)),415 [&](Expected<int32_t> R) {416 ErrorAsOutParameter _(Err);417 Err = R.takeError();418 },419 41, 1);420 421 EXPECT_EQ(toString(std::move(Err)), "no controller attached");422 423 S.waitForShutdown();424}425 426TEST(ControllerAccessTest, CallToControllerAfterDetach) {427 // Expect calls to the controller prior to attaching to fail.428 std::deque<std::unique_ptr<Task>> Tasks;429 Session S(std::make_unique<EnqueueingDispatcher>(Tasks), noErrors);430 auto CA = std::make_shared<MockControllerAccess>(S);431 S.setController(CA);432 433 S.detachFromController();434 435 Error Err = Error::success();436 SPSWrapperFunction<int32_t(int32_t, int32_t)>::call(437 CallViaSession(S, reinterpret_cast<Session::HandlerTag>(add_sps_wrapper)),438 [&](Expected<int32_t> R) {439 ErrorAsOutParameter _(Err);440 Err = R.takeError();441 },442 41, 1);443 444 EXPECT_EQ(toString(std::move(Err)), "no controller attached");445 446 S.waitForShutdown();447}448 449TEST(ControllerAccessTest, CallFromController) {450 // Simulate a call from the controller.451 std::deque<std::unique_ptr<Task>> Tasks;452 Session S(std::make_unique<EnqueueingDispatcher>(Tasks), noErrors);453 auto CA = std::make_shared<MockControllerAccess>(S);454 S.setController(CA);455 456 int32_t Result = 0;457 SPSWrapperFunction<int32_t(int32_t, int32_t)>::call(458 CallViaMockControllerAccess(*CA, add_sps_wrapper),459 [&](Expected<int32_t> R) { Result = cantFail(std::move(R)); }, 41, 1);460 461 EnqueueingDispatcher::runTasksFromFront(Tasks);462 463 EXPECT_EQ(Result, 42);464 465 S.waitForShutdown();466}467