453 lines · cpp
1//===------ ResourceTrackerTest.cpp - Unit tests ResourceTracker API ------===//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#include "OrcTestCommon.h"10#include "llvm/ADT/FunctionExtras.h"11#include "llvm/Config/llvm-config.h"12#include "llvm/ExecutionEngine/Orc/Core.h"13#include "llvm/ExecutionEngine/Orc/Shared/OrcError.h"14#include "llvm/Testing/Support/Error.h"15 16using namespace llvm;17using namespace llvm::orc;18 19class ResourceTrackerStandardTest : public CoreAPIsBasedStandardTest {};20 21namespace {22 23template <typename ResourceT = unsigned>24class SimpleResourceManager : public ResourceManager {25public:26 using HandleRemoveFunction =27 unique_function<Error(JITDylib &JD, ResourceKey)>;28 29 using HandleTransferFunction =30 unique_function<void(JITDylib &JD, ResourceKey, ResourceKey)>;31 32 using RecordedResourcesMap = DenseMap<ResourceKey, ResourceT>;33 34 SimpleResourceManager(ExecutionSession &ES) : ES(ES) {35 HandleRemove = [&](JITDylib &JD, ResourceKey K) -> Error {36 ES.runSessionLocked([&] { removeResource(JD, K); });37 return Error::success();38 };39 40 HandleTransfer = [this](JITDylib &JD, ResourceKey DstKey,41 ResourceKey SrcKey) {42 transferResources(JD, DstKey, SrcKey);43 };44 45 ES.registerResourceManager(*this);46 }47 48 SimpleResourceManager(const SimpleResourceManager &) = delete;49 SimpleResourceManager &operator=(const SimpleResourceManager &) = delete;50 SimpleResourceManager(SimpleResourceManager &&) = delete;51 SimpleResourceManager &operator=(SimpleResourceManager &&) = delete;52 53 ~SimpleResourceManager() override { ES.deregisterResourceManager(*this); }54 55 /// Set the HandleRemove function object.56 void setHandleRemove(HandleRemoveFunction HandleRemove) {57 this->HandleRemove = std::move(HandleRemove);58 }59 60 /// Set the HandleTransfer function object.61 void setHandleTransfer(HandleTransferFunction HandleTransfer) {62 this->HandleTransfer = std::move(HandleTransfer);63 }64 65 /// Create an association between the given key and resource.66 template <typename MergeOp = std::plus<ResourceT>>67 void recordResource(ResourceKey K, ResourceT Val = ResourceT(),68 MergeOp Merge = MergeOp()) {69 auto Tmp = std::move(Resources[K]);70 Resources[K] = Merge(std::move(Tmp), std::move(Val));71 }72 73 /// Remove the resource associated with K from the map if present.74 void removeResource(JITDylib &JD, ResourceKey K) { Resources.erase(K); }75 76 /// Transfer resources from DstKey to SrcKey.77 template <typename MergeOp = std::plus<ResourceT>>78 void transferResources(JITDylib &JD, ResourceKey DstKey, ResourceKey SrcKey,79 MergeOp Merge = MergeOp()) {80 auto &DstResourceRef = Resources[DstKey];81 ResourceT DstResources;82 std::swap(DstResourceRef, DstResources);83 84 auto SI = Resources.find(SrcKey);85 assert(SI != Resources.end() && "No resource associated with SrcKey");86 87 DstResourceRef = Merge(std::move(DstResources), std::move(SI->second));88 Resources.erase(SI);89 }90 91 /// Return a reference to the Resources map.92 RecordedResourcesMap &getRecordedResources() { return Resources; }93 const RecordedResourcesMap &getRecordedResources() const { return Resources; }94 95 Error handleRemoveResources(JITDylib &JD, ResourceKey K) override {96 return HandleRemove(JD, K);97 }98 99 void handleTransferResources(JITDylib &JD, ResourceKey DstKey,100 ResourceKey SrcKey) override {101 HandleTransfer(JD, DstKey, SrcKey);102 }103 104 static void transferNotAllowed(ResourceKey DstKey, ResourceKey SrcKey) {105 llvm_unreachable("Resource transfer not allowed");106 }107 108private:109 ExecutionSession &ES;110 HandleRemoveFunction HandleRemove;111 HandleTransferFunction HandleTransfer;112 RecordedResourcesMap Resources;113};114 115TEST_F(ResourceTrackerStandardTest,116 BasicDefineAndRemoveAllBeforeMaterializing) {117 118 bool ResourceManagerGotRemove = false;119 SimpleResourceManager<> SRM(ES);120 SRM.setHandleRemove([&](JITDylib &JD, ResourceKey K) -> Error {121 ResourceManagerGotRemove = true;122 EXPECT_EQ(SRM.getRecordedResources().size(), 0U)123 << "Unexpected resources recorded";124 SRM.removeResource(JD, K);125 return Error::success();126 });127 128 bool MaterializationUnitDestroyed = false;129 auto MU = std::make_unique<SimpleMaterializationUnit>(130 SymbolFlagsMap({{Foo, FooSym.getFlags()}}),131 [&](std::unique_ptr<MaterializationResponsibility> R) {132 llvm_unreachable("Never called");133 },134 nullptr, SimpleMaterializationUnit::DiscardFunction(),135 [&]() { MaterializationUnitDestroyed = true; });136 137 auto RT = JD.createResourceTracker();138 cantFail(JD.define(std::move(MU), RT));139 cantFail(RT->remove());140 auto SymFlags = cantFail(ES.lookupFlags(141 LookupKind::Static,142 {{&JD, JITDylibLookupFlags::MatchExportedSymbolsOnly}},143 SymbolLookupSet(Foo, SymbolLookupFlags::WeaklyReferencedSymbol)));144 145 EXPECT_EQ(SymFlags.size(), 0U)146 << "Symbols should have been removed from the symbol table";147 EXPECT_TRUE(ResourceManagerGotRemove)148 << "ResourceManager did not receive handleRemoveResources";149 EXPECT_TRUE(MaterializationUnitDestroyed)150 << "MaterializationUnit not destroyed in response to removal";151}152 153TEST_F(ResourceTrackerStandardTest, BasicDefineAndRemoveAllAfterMaterializing) {154 155 bool ResourceManagerGotRemove = false;156 SimpleResourceManager<> SRM(ES);157 SRM.setHandleRemove([&](JITDylib &JD, ResourceKey K) -> Error {158 ResourceManagerGotRemove = true;159 EXPECT_EQ(SRM.getRecordedResources().size(), 1U)160 << "Unexpected number of resources recorded";161 EXPECT_EQ(SRM.getRecordedResources().count(K), 1U)162 << "Unexpected recorded resource";163 SRM.removeResource(JD, K);164 return Error::success();165 });166 167 auto MU = std::make_unique<SimpleMaterializationUnit>(168 SymbolFlagsMap({{Foo, FooSym.getFlags()}}),169 [&](std::unique_ptr<MaterializationResponsibility> R) {170 cantFail(R->withResourceKeyDo(171 [&](ResourceKey K) { SRM.recordResource(K); }));172 cantFail(R->notifyResolved({{Foo, FooSym}}));173 cantFail(R->notifyEmitted({}));174 });175 176 auto RT = JD.createResourceTracker();177 cantFail(JD.define(std::move(MU), RT));178 cantFail(ES.lookup({&JD}, Foo));179 cantFail(RT->remove());180 auto SymFlags = cantFail(ES.lookupFlags(181 LookupKind::Static,182 {{&JD, JITDylibLookupFlags::MatchExportedSymbolsOnly}},183 SymbolLookupSet(Foo, SymbolLookupFlags::WeaklyReferencedSymbol)));184 185 EXPECT_EQ(SymFlags.size(), 0U)186 << "Symbols should have been removed from the symbol table";187 EXPECT_TRUE(ResourceManagerGotRemove)188 << "ResourceManager did not receive handleRemoveResources";189}190 191TEST_F(ResourceTrackerStandardTest, BasicDefineAndRemoveAllWhileMaterializing) {192 193 bool ResourceManagerGotRemove = false;194 SimpleResourceManager<> SRM(ES);195 SRM.setHandleRemove([&](JITDylib &JD, ResourceKey K) -> Error {196 ResourceManagerGotRemove = true;197 EXPECT_EQ(SRM.getRecordedResources().size(), 0U)198 << "Unexpected resources recorded";199 SRM.removeResource(JD, K);200 return Error::success();201 });202 203 std::unique_ptr<MaterializationResponsibility> MR;204 auto MU = std::make_unique<SimpleMaterializationUnit>(205 SymbolFlagsMap({{Foo, FooSym.getFlags()}}),206 [&](std::unique_ptr<MaterializationResponsibility> R) {207 MR = std::move(R);208 });209 210 auto RT = JD.createResourceTracker();211 cantFail(JD.define(std::move(MU), RT));212 213 ES.lookup(214 LookupKind::Static, makeJITDylibSearchOrder(&JD), SymbolLookupSet(Foo),215 SymbolState::Ready,216 [](Expected<SymbolMap> Result) {217 EXPECT_THAT_EXPECTED(Result, Failed<FailedToMaterialize>())218 << "Lookup failed unexpectedly";219 },220 NoDependenciesToRegister);221 222 cantFail(RT->remove());223 auto SymFlags = cantFail(ES.lookupFlags(224 LookupKind::Static,225 {{&JD, JITDylibLookupFlags::MatchExportedSymbolsOnly}},226 SymbolLookupSet(Foo, SymbolLookupFlags::WeaklyReferencedSymbol)));227 228 EXPECT_EQ(SymFlags.size(), 0U)229 << "Symbols should have been removed from the symbol table";230 EXPECT_TRUE(ResourceManagerGotRemove)231 << "ResourceManager did not receive handleRemoveResources";232 233 EXPECT_THAT_ERROR(MR->withResourceKeyDo([](ResourceKey K) {234 ADD_FAILURE() << "Should not reach withResourceKeyDo body for removed key";235 }),236 Failed<ResourceTrackerDefunct>())237 << "withResourceKeyDo on MR with removed tracker should have failed";238 EXPECT_THAT_ERROR(MR->notifyResolved({{Foo, FooSym}}),239 Failed<ResourceTrackerDefunct>())240 << "notifyResolved on MR with removed tracker should have failed";241 242 MR->failMaterialization();243}244 245TEST_F(ResourceTrackerStandardTest, JITDylibClear) {246 SimpleResourceManager<> SRM(ES);247 248 // Add materializer for Foo.249 cantFail(JD.define(std::make_unique<SimpleMaterializationUnit>(250 SymbolFlagsMap({{Foo, FooSym.getFlags()}}),251 [&](std::unique_ptr<MaterializationResponsibility> R) {252 cantFail(R->withResourceKeyDo(253 [&](ResourceKey K) { ++SRM.getRecordedResources()[K]; }));254 cantFail(R->notifyResolved({{Foo, FooSym}}));255 cantFail(R->notifyEmitted({}));256 })));257 258 // Add materializer for Bar.259 cantFail(JD.define(std::make_unique<SimpleMaterializationUnit>(260 SymbolFlagsMap({{Bar, BarSym.getFlags()}}),261 [&](std::unique_ptr<MaterializationResponsibility> R) {262 cantFail(R->withResourceKeyDo(263 [&](ResourceKey K) { ++SRM.getRecordedResources()[K]; }));264 cantFail(R->notifyResolved({{Bar, BarSym}}));265 cantFail(R->notifyEmitted({}));266 })));267 268 EXPECT_TRUE(SRM.getRecordedResources().empty())269 << "Expected no resources recorded yet.";270 271 cantFail(272 ES.lookup(makeJITDylibSearchOrder(&JD), SymbolLookupSet({Foo, Bar})));273 274 auto JDResourceKey = JD.getDefaultResourceTracker()->getKeyUnsafe();275 EXPECT_EQ(SRM.getRecordedResources().size(), 1U)276 << "Expected exactly one entry (for JD's ResourceKey)";277 EXPECT_EQ(SRM.getRecordedResources().count(JDResourceKey), 1U)278 << "Expected an entry for JD's ResourceKey";279 EXPECT_EQ(SRM.getRecordedResources()[JDResourceKey], 2U)280 << "Expected value of 2 for JD's ResourceKey "281 "(+1 for each of Foo and Bar)";282 283 cantFail(JD.clear());284 285 EXPECT_TRUE(SRM.getRecordedResources().empty())286 << "Expected no resources recorded after clear";287}288 289TEST_F(ResourceTrackerStandardTest,290 BasicDefineAndExplicitTransferBeforeMaterializing) {291 292 bool ResourceManagerGotTransfer = false;293 SimpleResourceManager<> SRM(ES);294 SRM.setHandleTransfer(295 [&](JITDylib &JD, ResourceKey DstKey, ResourceKey SrcKey) {296 ResourceManagerGotTransfer = true;297 auto &RR = SRM.getRecordedResources();298 EXPECT_EQ(RR.size(), 0U) << "Expected no resources recorded yet";299 });300 301 auto MakeMU = [&](SymbolStringPtr Name, ExecutorSymbolDef Sym) {302 return std::make_unique<SimpleMaterializationUnit>(303 SymbolFlagsMap({{Name, Sym.getFlags()}}),304 [=, &SRM](std::unique_ptr<MaterializationResponsibility> R) {305 cantFail(R->withResourceKeyDo(306 [&](ResourceKey K) { SRM.recordResource(K); }));307 cantFail(R->notifyResolved({{Name, Sym}}));308 cantFail(R->notifyEmitted({}));309 });310 };311 312 auto FooRT = JD.createResourceTracker();313 cantFail(JD.define(MakeMU(Foo, FooSym), FooRT));314 315 auto BarRT = JD.createResourceTracker();316 cantFail(JD.define(MakeMU(Bar, BarSym), BarRT));317 318 BarRT->transferTo(*FooRT);319 320 EXPECT_TRUE(ResourceManagerGotTransfer)321 << "ResourceManager did not receive transfer";322 EXPECT_TRUE(BarRT->isDefunct()) << "BarRT should now be defunct";323 324 cantFail(325 ES.lookup(makeJITDylibSearchOrder({&JD}), SymbolLookupSet({Foo, Bar})));326 327 EXPECT_EQ(SRM.getRecordedResources().size(), 1U)328 << "Expected exactly one entry (for FooRT's Key)";329 EXPECT_EQ(SRM.getRecordedResources().count(FooRT->getKeyUnsafe()), 1U)330 << "Expected an entry for FooRT's ResourceKey";331 EXPECT_EQ(SRM.getRecordedResources().count(BarRT->getKeyUnsafe()), 0U)332 << "Expected no entry for BarRT's ResourceKey";333 334 // We need to explicitly destroy FooRT or its resources will be implicitly335 // transferred to the default tracker triggering a second call to our336 // transfer function above (which expects only one call).337 cantFail(FooRT->remove());338}339 340TEST_F(ResourceTrackerStandardTest,341 BasicDefineAndExplicitTransferAfterMaterializing) {342 343 bool ResourceManagerGotTransfer = false;344 SimpleResourceManager<> SRM(ES);345 SRM.setHandleTransfer(346 [&](JITDylib &JD, ResourceKey DstKey, ResourceKey SrcKey) {347 ResourceManagerGotTransfer = true;348 SRM.transferResources(JD, DstKey, SrcKey);349 });350 351 auto MakeMU = [&](SymbolStringPtr Name, ExecutorSymbolDef Sym) {352 return std::make_unique<SimpleMaterializationUnit>(353 SymbolFlagsMap({{Name, Sym.getFlags()}}),354 [=, &SRM](std::unique_ptr<MaterializationResponsibility> R) {355 cantFail(R->withResourceKeyDo(356 [&](ResourceKey K) { SRM.recordResource(K, 1); }));357 cantFail(R->notifyResolved({{Name, Sym}}));358 cantFail(R->notifyEmitted({}));359 });360 };361 362 auto FooRT = JD.createResourceTracker();363 cantFail(JD.define(MakeMU(Foo, FooSym), FooRT));364 365 auto BarRT = JD.createResourceTracker();366 cantFail(JD.define(MakeMU(Bar, BarSym), BarRT));367 368 EXPECT_EQ(SRM.getRecordedResources().size(), 0U)369 << "Expected no recorded resources yet";370 371 cantFail(372 ES.lookup(makeJITDylibSearchOrder({&JD}), SymbolLookupSet({Foo, Bar})));373 374 EXPECT_EQ(SRM.getRecordedResources().size(), 2U)375 << "Expected recorded resources for both Foo and Bar";376 377 BarRT->transferTo(*FooRT);378 379 EXPECT_TRUE(ResourceManagerGotTransfer)380 << "ResourceManager did not receive transfer";381 EXPECT_TRUE(BarRT->isDefunct()) << "BarRT should now be defunct";382 383 EXPECT_EQ(SRM.getRecordedResources().size(), 1U)384 << "Expected recorded resources for Foo only";385 EXPECT_EQ(SRM.getRecordedResources().count(FooRT->getKeyUnsafe()), 1U)386 << "Expected recorded resources for Foo";387 EXPECT_EQ(SRM.getRecordedResources()[FooRT->getKeyUnsafe()], 2U)388 << "Expected resources value for for Foo to be '2'";389}390 391TEST_F(ResourceTrackerStandardTest,392 BasicDefineAndExplicitTransferWhileMaterializing) {393 394 bool ResourceManagerGotTransfer = false;395 SimpleResourceManager<> SRM(ES);396 SRM.setHandleTransfer(397 [&](JITDylib &JD, ResourceKey DstKey, ResourceKey SrcKey) {398 ResourceManagerGotTransfer = true;399 SRM.transferResources(JD, DstKey, SrcKey);400 });401 402 auto FooRT = JD.createResourceTracker();403 std::unique_ptr<MaterializationResponsibility> FooMR;404 cantFail(JD.define(std::make_unique<SimpleMaterializationUnit>(405 SymbolFlagsMap({{Foo, FooSym.getFlags()}}),406 [&](std::unique_ptr<MaterializationResponsibility> R) {407 FooMR = std::move(R);408 }),409 FooRT));410 411 auto BarRT = JD.createResourceTracker();412 413 ES.lookup(414 LookupKind::Static, makeJITDylibSearchOrder(&JD), SymbolLookupSet(Foo),415 SymbolState::Ready,416 [](Expected<SymbolMap> Result) { cantFail(Result.takeError()); },417 NoDependenciesToRegister);418 419 cantFail(FooMR->withResourceKeyDo([&](ResourceKey K) {420 EXPECT_EQ(FooRT->getKeyUnsafe(), K)421 << "Expected FooRT's ResourceKey for Foo here";422 SRM.recordResource(K, 1);423 }));424 425 EXPECT_EQ(SRM.getRecordedResources().size(), 1U)426 << "Expected one recorded resource here";427 EXPECT_EQ(SRM.getRecordedResources()[FooRT->getKeyUnsafe()], 1U)428 << "Expected Resource value for FooRT to be '1' here";429 430 FooRT->transferTo(*BarRT);431 432 EXPECT_TRUE(ResourceManagerGotTransfer)433 << "Expected resource manager to receive handleTransferResources call";434 435 cantFail(FooMR->withResourceKeyDo([&](ResourceKey K) {436 EXPECT_EQ(BarRT->getKeyUnsafe(), K)437 << "Expected BarRT's ResourceKey for Foo here";438 SRM.recordResource(K, 1);439 }));440 441 EXPECT_EQ(SRM.getRecordedResources().size(), 1U)442 << "Expected one recorded resource here";443 EXPECT_EQ(SRM.getRecordedResources().count(BarRT->getKeyUnsafe()), 1U)444 << "Expected RecordedResources to contain an entry for BarRT";445 EXPECT_EQ(SRM.getRecordedResources()[BarRT->getKeyUnsafe()], 2U)446 << "Expected Resource value for BarRT to be '2' here";447 448 cantFail(FooMR->notifyResolved({{Foo, FooSym}}));449 cantFail(FooMR->notifyEmitted({}));450}451 452} // namespace453