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