xref: /llvm-project/llvm/unittests/ExecutionEngine/Orc/ResourceTrackerTest.cpp (revision 0aec49c8531bc5282b095730d34681455826bc2c) 
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/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(
35       ExecutionSession &ES,
36       HandleRemoveFunction HandleRemove = HandleRemoveFunction(),
37       HandleTransferFunction HandleTransfer = HandleTransferFunction())
38       : ES(ES), HandleRemove(std::move(HandleRemove)),
39         HandleTransfer(std::move(HandleTransfer)) {
40 
41     // If HandleRemvoe is not supplied then use the default.
42     if (!this->HandleRemove)
43       this->HandleRemove = [&](ResourceKey K) -> Error {
44         ES.runSessionLocked([&] { removeResource(K); });
45         return Error::success();
46       };
47 
48     // If HandleTransfer is not supplied then use the default.
49     if (!this->HandleTransfer)
50       this->HandleTransfer = [this](ResourceKey DstKey, ResourceKey SrcKey) {
51         transferResources(DstKey, SrcKey);
52       };
53 
54     ES.registerResourceManager(*this);
55   }
56 
57   SimpleResourceManager(const SimpleResourceManager &) = delete;
58   SimpleResourceManager &operator=(const SimpleResourceManager &) = delete;
59   SimpleResourceManager(SimpleResourceManager &&) = delete;
60   SimpleResourceManager &operator=(SimpleResourceManager &&) = delete;
61 
62   ~SimpleResourceManager() { ES.deregisterResourceManager(*this); }
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, [&](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(JD.lookupFlags(
139       LookupKind::Static, JITDylibLookupFlags::MatchExportedSymbolsOnly,
140       SymbolLookupSet(Foo)));
141 
142   EXPECT_EQ(SymFlags.size(), 0U)
143       << "Symbols should have been removed from the symbol table";
144   EXPECT_TRUE(ResourceManagerGotRemove)
145       << "ResourceManager did not receive handleRemoveResources";
146   EXPECT_TRUE(MaterializationUnitDestroyed)
147       << "MaterializationUnit not destroyed in response to removal";
148 }
149 
150 TEST_F(ResourceTrackerStandardTest, BasicDefineAndRemoveAllAfterMaterializing) {
151 
152   bool ResourceManagerGotRemove = false;
153   SimpleResourceManager<> SRM(ES, [&](ResourceKey K) -> Error {
154     ResourceManagerGotRemove = true;
155     EXPECT_EQ(SRM.getRecordedResources().size(), 1U)
156         << "Unexpected number of resources recorded";
157     EXPECT_EQ(SRM.getRecordedResources().count(K), 1U)
158         << "Unexpected recorded resource";
159     SRM.removeResource(K);
160     return Error::success();
161   });
162 
163   auto MU = std::make_unique<SimpleMaterializationUnit>(
164       SymbolFlagsMap({{Foo, FooSym.getFlags()}}),
165       [&](std::unique_ptr<MaterializationResponsibility> R) {
166         cantFail(R->withResourceKeyDo(
167             [&](ResourceKey K) { SRM.recordResource(K); }));
168         cantFail(R->notifyResolved({{Foo, FooSym}}));
169         cantFail(R->notifyEmitted());
170       });
171 
172   auto RT = JD.createResourceTracker();
173   cantFail(JD.define(std::move(MU), RT));
174   cantFail(ES.lookup({&JD}, Foo));
175   cantFail(RT->remove());
176   auto SymFlags = cantFail(JD.lookupFlags(
177       LookupKind::Static, JITDylibLookupFlags::MatchExportedSymbolsOnly,
178       SymbolLookupSet(Foo)));
179 
180   EXPECT_EQ(SymFlags.size(), 0U)
181       << "Symbols should have been removed from the symbol table";
182   EXPECT_TRUE(ResourceManagerGotRemove)
183       << "ResourceManager did not receive handleRemoveResources";
184 }
185 
186 TEST_F(ResourceTrackerStandardTest, BasicDefineAndRemoveAllWhileMaterializing) {
187 
188   bool ResourceManagerGotRemove = false;
189   SimpleResourceManager<> SRM(ES, [&](ResourceKey K) -> Error {
190     ResourceManagerGotRemove = true;
191     EXPECT_EQ(SRM.getRecordedResources().size(), 0U)
192         << "Unexpected resources recorded";
193     SRM.removeResource(K);
194     return Error::success();
195   });
196 
197   std::unique_ptr<MaterializationResponsibility> MR;
198   auto MU = std::make_unique<SimpleMaterializationUnit>(
199       SymbolFlagsMap({{Foo, FooSym.getFlags()}}),
200       [&](std::unique_ptr<MaterializationResponsibility> R) {
201         MR = std::move(R);
202       });
203 
204   auto RT = JD.createResourceTracker();
205   cantFail(JD.define(std::move(MU), RT));
206 
207   ES.lookup(
208       LookupKind::Static, makeJITDylibSearchOrder(&JD), SymbolLookupSet(Foo),
209       SymbolState::Ready,
210       [](Expected<SymbolMap> Result) {
211         EXPECT_THAT_EXPECTED(Result, Failed<FailedToMaterialize>())
212             << "Lookup failed unexpectedly";
213       },
214       NoDependenciesToRegister);
215 
216   cantFail(RT->remove());
217   auto SymFlags = cantFail(JD.lookupFlags(
218       LookupKind::Static, JITDylibLookupFlags::MatchExportedSymbolsOnly,
219       SymbolLookupSet(Foo)));
220 
221   EXPECT_EQ(SymFlags.size(), 0U)
222       << "Symbols should have been removed from the symbol table";
223   EXPECT_TRUE(ResourceManagerGotRemove)
224       << "ResourceManager did not receive handleRemoveResources";
225 
226   EXPECT_THAT_ERROR(MR->withResourceKeyDo([](ResourceKey K) {
227     ADD_FAILURE() << "Should not reach withResourceKeyDo body for removed key";
228   }),
229                     Failed<ResourceTrackerDefunct>())
230       << "withResourceKeyDo on MR with removed tracker should have failed";
231   EXPECT_THAT_ERROR(MR->notifyResolved({{Foo, FooSym}}),
232                     Failed<ResourceTrackerDefunct>())
233       << "notifyResolved on MR with removed tracker should have failed";
234 
235   MR->failMaterialization();
236 }
237 
238 TEST_F(ResourceTrackerStandardTest, JITDylibClear) {
239   SimpleResourceManager<> SRM(ES);
240 
241   // Add materializer for Foo.
242   cantFail(JD.define(std::make_unique<SimpleMaterializationUnit>(
243       SymbolFlagsMap({{Foo, FooSym.getFlags()}}),
244       [&](std::unique_ptr<MaterializationResponsibility> R) {
245         cantFail(R->withResourceKeyDo(
246             [&](ResourceKey K) { ++SRM.getRecordedResources()[K]; }));
247         cantFail(R->notifyResolved({{Foo, FooSym}}));
248         cantFail(R->notifyEmitted());
249       })));
250 
251   // Add materializer for Bar.
252   cantFail(JD.define(std::make_unique<SimpleMaterializationUnit>(
253       SymbolFlagsMap({{Bar, BarSym.getFlags()}}),
254       [&](std::unique_ptr<MaterializationResponsibility> R) {
255         cantFail(R->withResourceKeyDo(
256             [&](ResourceKey K) { ++SRM.getRecordedResources()[K]; }));
257         cantFail(R->notifyResolved({{Bar, BarSym}}));
258         cantFail(R->notifyEmitted());
259       })));
260 
261   EXPECT_TRUE(SRM.getRecordedResources().empty())
262       << "Expected no resources recorded yet.";
263 
264   cantFail(
265       ES.lookup(makeJITDylibSearchOrder(&JD), SymbolLookupSet({Foo, Bar})));
266 
267   auto JDResourceKey = JD.getDefaultResourceTracker()->getKeyUnsafe();
268   EXPECT_EQ(SRM.getRecordedResources().size(), 1U)
269       << "Expected exactly one entry (for JD's ResourceKey)";
270   EXPECT_EQ(SRM.getRecordedResources().count(JDResourceKey), 1U)
271       << "Expected an entry for JD's ResourceKey";
272   EXPECT_EQ(SRM.getRecordedResources()[JDResourceKey], 2U)
273       << "Expected value of 2 for JD's ResourceKey "
274          "(+1 for each of Foo and Bar)";
275 
276   cantFail(JD.clear());
277 
278   EXPECT_TRUE(SRM.getRecordedResources().empty())
279       << "Expected no resources recorded after clear";
280 }
281 
282 TEST_F(ResourceTrackerStandardTest,
283        BasicDefineAndExplicitTransferBeforeMaterializing) {
284 
285   bool ResourceManagerGotTransfer = false;
286   SimpleResourceManager<> SRM(
287       ES,
288       [&](ResourceKey K) -> Error {
289         SRM.removeResource(K);
290         return Error::success();
291       },
292       [&](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(
342       ES,
343       [&](ResourceKey K) -> Error {
344         SRM.removeResource(K);
345         return Error::success();
346       },
347       [&](ResourceKey DstKey, ResourceKey SrcKey) {
348         ResourceManagerGotTransfer = true;
349         SRM.transferResources(DstKey, SrcKey);
350       });
351 
352   auto MakeMU = [&](SymbolStringPtr Name, JITEvaluatedSymbol Sym) {
353     return std::make_unique<SimpleMaterializationUnit>(
354         SymbolFlagsMap({{Name, Sym.getFlags()}}),
355         [=, &SRM](std::unique_ptr<MaterializationResponsibility> R) {
356           cantFail(R->withResourceKeyDo(
357               [&](ResourceKey K) { SRM.recordResource(K, 1); }));
358           cantFail(R->notifyResolved({{Name, Sym}}));
359           cantFail(R->notifyEmitted());
360         });
361   };
362 
363   auto FooRT = JD.createResourceTracker();
364   cantFail(JD.define(MakeMU(Foo, FooSym), FooRT));
365 
366   auto BarRT = JD.createResourceTracker();
367   cantFail(JD.define(MakeMU(Bar, BarSym), BarRT));
368 
369   EXPECT_EQ(SRM.getRecordedResources().size(), 0U)
370       << "Expected no recorded resources yet";
371 
372   cantFail(
373       ES.lookup(makeJITDylibSearchOrder({&JD}), SymbolLookupSet({Foo, Bar})));
374 
375   EXPECT_EQ(SRM.getRecordedResources().size(), 2U)
376       << "Expected recorded resources for both Foo and Bar";
377 
378   BarRT->transferTo(*FooRT);
379 
380   EXPECT_TRUE(ResourceManagerGotTransfer)
381       << "ResourceManager did not receive transfer";
382   EXPECT_TRUE(BarRT->isDefunct()) << "BarRT should now be defunct";
383 
384   EXPECT_EQ(SRM.getRecordedResources().size(), 1U)
385       << "Expected recorded resources for Foo only";
386   EXPECT_EQ(SRM.getRecordedResources().count(FooRT->getKeyUnsafe()), 1U)
387       << "Expected recorded resources for Foo";
388   EXPECT_EQ(SRM.getRecordedResources()[FooRT->getKeyUnsafe()], 2U)
389       << "Expected resources value for for Foo to be '2'";
390 }
391 
392 TEST_F(ResourceTrackerStandardTest,
393        BasicDefineAndExplicitTransferWhileMaterializing) {
394 
395   bool ResourceManagerGotTransfer = false;
396   SimpleResourceManager<> SRM(
397       ES,
398       [&](ResourceKey K) -> Error {
399         SRM.removeResource(K);
400         return Error::success();
401       },
402       [&](ResourceKey DstKey, ResourceKey SrcKey) {
403         ResourceManagerGotTransfer = true;
404         SRM.transferResources(DstKey, SrcKey);
405       });
406 
407   auto FooRT = JD.createResourceTracker();
408   std::unique_ptr<MaterializationResponsibility> FooMR;
409   cantFail(JD.define(std::make_unique<SimpleMaterializationUnit>(
410                          SymbolFlagsMap({{Foo, FooSym.getFlags()}}),
411                          [&](std::unique_ptr<MaterializationResponsibility> R) {
412                            FooMR = std::move(R);
413                          }),
414                      FooRT));
415 
416   auto BarRT = JD.createResourceTracker();
417 
418   ES.lookup(
419       LookupKind::Static, makeJITDylibSearchOrder(&JD), SymbolLookupSet(Foo),
420       SymbolState::Ready,
421       [](Expected<SymbolMap> Result) { cantFail(Result.takeError()); },
422       NoDependenciesToRegister);
423 
424   cantFail(FooMR->withResourceKeyDo([&](ResourceKey K) {
425     EXPECT_EQ(FooRT->getKeyUnsafe(), K)
426         << "Expected FooRT's ResourceKey for Foo here";
427     SRM.recordResource(K, 1);
428   }));
429 
430   EXPECT_EQ(SRM.getRecordedResources().size(), 1U)
431       << "Expected one recorded resource here";
432   EXPECT_EQ(SRM.getRecordedResources()[FooRT->getKeyUnsafe()], 1U)
433       << "Expected Resource value for FooRT to be '1' here";
434 
435   FooRT->transferTo(*BarRT);
436 
437   EXPECT_TRUE(ResourceManagerGotTransfer)
438       << "Expected resource manager to receive handleTransferResources call";
439 
440   cantFail(FooMR->withResourceKeyDo([&](ResourceKey K) {
441     EXPECT_EQ(BarRT->getKeyUnsafe(), K)
442         << "Expected BarRT's ResourceKey for Foo here";
443     SRM.recordResource(K, 1);
444   }));
445 
446   EXPECT_EQ(SRM.getRecordedResources().size(), 1U)
447       << "Expected one recorded resource here";
448   EXPECT_EQ(SRM.getRecordedResources().count(BarRT->getKeyUnsafe()), 1U)
449       << "Expected RecordedResources to contain an entry for BarRT";
450   EXPECT_EQ(SRM.getRecordedResources()[BarRT->getKeyUnsafe()], 2U)
451       << "Expected Resource value for BarRT to be '2' here";
452 
453   cantFail(FooMR->notifyResolved({{Foo, FooSym}}));
454   cantFail(FooMR->notifyEmitted());
455 }
456 
457 } // namespace
458