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 =
  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() { 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
 108 private:
 109   ExecutionSession &ES;
 110   HandleRemoveFunction HandleRemove;
 111   HandleTransferFunction HandleTransfer;
 112   RecordedResourcesMap Resources;
 113 };
 114
 115 TEST_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
 153 TEST_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
 191 TEST_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
 245 TEST_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
 289 TEST_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 implicitly
 335   // transferred to the default tracker triggering a second call to our
 336   // transfer function above (which expects only one call).
 337   cantFail(FooRT->remove());
 338 }
 339
 340 TEST_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
 391 TEST_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 } // namespace