xref: /llvm-project/llvm/unittests/ExecutionEngine/Orc/ResourceTrackerTest.cpp (revision b66b73be295c00a11472566f52b80044384f2fb2) 
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
448