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