xref: /llvm-project/llvm/unittests/ExecutionEngine/Orc/ResourceTrackerTest.cpp (revision d54c4e382114650ac8084b7511083253cbe9b8d1) 
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(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(JD.lookupFlags(
140       LookupKind::Static, JITDylibLookupFlags::MatchExportedSymbolsOnly,
141       SymbolLookupSet(Foo)));
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(JD.lookupFlags(
179       LookupKind::Static, JITDylibLookupFlags::MatchExportedSymbolsOnly,
180       SymbolLookupSet(Foo)));
181 
182   EXPECT_EQ(SymFlags.size(), 0U)
183       << "Symbols should have been removed from the symbol table";
184   EXPECT_TRUE(ResourceManagerGotRemove)
185       << "ResourceManager did not receive handleRemoveResources";
186 }
187 
188 TEST_F(ResourceTrackerStandardTest, BasicDefineAndRemoveAllWhileMaterializing) {
189 
190   bool ResourceManagerGotRemove = false;
191   SimpleResourceManager<> SRM(ES);
192   SRM.setHandleRemove([&](ResourceKey K) -> Error {
193     ResourceManagerGotRemove = true;
194     EXPECT_EQ(SRM.getRecordedResources().size(), 0U)
195         << "Unexpected resources recorded";
196     SRM.removeResource(K);
197     return Error::success();
198   });
199 
200   std::unique_ptr<MaterializationResponsibility> MR;
201   auto MU = std::make_unique<SimpleMaterializationUnit>(
202       SymbolFlagsMap({{Foo, FooSym.getFlags()}}),
203       [&](std::unique_ptr<MaterializationResponsibility> R) {
204         MR = std::move(R);
205       });
206 
207   auto RT = JD.createResourceTracker();
208   cantFail(JD.define(std::move(MU), RT));
209 
210   ES.lookup(
211       LookupKind::Static, makeJITDylibSearchOrder(&JD), SymbolLookupSet(Foo),
212       SymbolState::Ready,
213       [](Expected<SymbolMap> Result) {
214         EXPECT_THAT_EXPECTED(Result, Failed<FailedToMaterialize>())
215             << "Lookup failed unexpectedly";
216       },
217       NoDependenciesToRegister);
218 
219   cantFail(RT->remove());
220   auto SymFlags = cantFail(JD.lookupFlags(
221       LookupKind::Static, JITDylibLookupFlags::MatchExportedSymbolsOnly,
222       SymbolLookupSet(Foo)));
223 
224   EXPECT_EQ(SymFlags.size(), 0U)
225       << "Symbols should have been removed from the symbol table";
226   EXPECT_TRUE(ResourceManagerGotRemove)
227       << "ResourceManager did not receive handleRemoveResources";
228 
229   EXPECT_THAT_ERROR(MR->withResourceKeyDo([](ResourceKey K) {
230     ADD_FAILURE() << "Should not reach withResourceKeyDo body for removed key";
231   }),
232                     Failed<ResourceTrackerDefunct>())
233       << "withResourceKeyDo on MR with removed tracker should have failed";
234   EXPECT_THAT_ERROR(MR->notifyResolved({{Foo, FooSym}}),
235                     Failed<ResourceTrackerDefunct>())
236       << "notifyResolved on MR with removed tracker should have failed";
237 
238   MR->failMaterialization();
239 }
240 
241 TEST_F(ResourceTrackerStandardTest, JITDylibClear) {
242   SimpleResourceManager<> SRM(ES);
243 
244   // Add materializer for Foo.
245   cantFail(JD.define(std::make_unique<SimpleMaterializationUnit>(
246       SymbolFlagsMap({{Foo, FooSym.getFlags()}}),
247       [&](std::unique_ptr<MaterializationResponsibility> R) {
248         cantFail(R->withResourceKeyDo(
249             [&](ResourceKey K) { ++SRM.getRecordedResources()[K]; }));
250         cantFail(R->notifyResolved({{Foo, FooSym}}));
251         cantFail(R->notifyEmitted());
252       })));
253 
254   // Add materializer for Bar.
255   cantFail(JD.define(std::make_unique<SimpleMaterializationUnit>(
256       SymbolFlagsMap({{Bar, BarSym.getFlags()}}),
257       [&](std::unique_ptr<MaterializationResponsibility> R) {
258         cantFail(R->withResourceKeyDo(
259             [&](ResourceKey K) { ++SRM.getRecordedResources()[K]; }));
260         cantFail(R->notifyResolved({{Bar, BarSym}}));
261         cantFail(R->notifyEmitted());
262       })));
263 
264   EXPECT_TRUE(SRM.getRecordedResources().empty())
265       << "Expected no resources recorded yet.";
266 
267   cantFail(
268       ES.lookup(makeJITDylibSearchOrder(&JD), SymbolLookupSet({Foo, Bar})));
269 
270   auto JDResourceKey = JD.getDefaultResourceTracker()->getKeyUnsafe();
271   EXPECT_EQ(SRM.getRecordedResources().size(), 1U)
272       << "Expected exactly one entry (for JD's ResourceKey)";
273   EXPECT_EQ(SRM.getRecordedResources().count(JDResourceKey), 1U)
274       << "Expected an entry for JD's ResourceKey";
275   EXPECT_EQ(SRM.getRecordedResources()[JDResourceKey], 2U)
276       << "Expected value of 2 for JD's ResourceKey "
277          "(+1 for each of Foo and Bar)";
278 
279   cantFail(JD.clear());
280 
281   EXPECT_TRUE(SRM.getRecordedResources().empty())
282       << "Expected no resources recorded after clear";
283 }
284 
285 TEST_F(ResourceTrackerStandardTest,
286        BasicDefineAndExplicitTransferBeforeMaterializing) {
287 
288   bool ResourceManagerGotTransfer = false;
289   SimpleResourceManager<> SRM(ES);
290   SRM.setHandleTransfer([&](ResourceKey DstKey, ResourceKey SrcKey) {
291     ResourceManagerGotTransfer = true;
292     auto &RR = SRM.getRecordedResources();
293     EXPECT_EQ(RR.size(), 0U) << "Expected no resources recorded yet";
294   });
295 
296   auto MakeMU = [&](SymbolStringPtr Name, JITEvaluatedSymbol Sym) {
297     return std::make_unique<SimpleMaterializationUnit>(
298         SymbolFlagsMap({{Name, Sym.getFlags()}}),
299         [=, &SRM](std::unique_ptr<MaterializationResponsibility> R) {
300           cantFail(R->withResourceKeyDo(
301               [&](ResourceKey K) { SRM.recordResource(K); }));
302           cantFail(R->notifyResolved({{Name, Sym}}));
303           cantFail(R->notifyEmitted());
304         });
305   };
306 
307   auto FooRT = JD.createResourceTracker();
308   cantFail(JD.define(MakeMU(Foo, FooSym), FooRT));
309 
310   auto BarRT = JD.createResourceTracker();
311   cantFail(JD.define(MakeMU(Bar, BarSym), BarRT));
312 
313   BarRT->transferTo(*FooRT);
314 
315   EXPECT_TRUE(ResourceManagerGotTransfer)
316       << "ResourceManager did not receive transfer";
317   EXPECT_TRUE(BarRT->isDefunct()) << "BarRT should now be defunct";
318 
319   cantFail(
320       ES.lookup(makeJITDylibSearchOrder({&JD}), SymbolLookupSet({Foo, Bar})));
321 
322   EXPECT_EQ(SRM.getRecordedResources().size(), 1U)
323       << "Expected exactly one entry (for FooRT's Key)";
324   EXPECT_EQ(SRM.getRecordedResources().count(FooRT->getKeyUnsafe()), 1U)
325       << "Expected an entry for FooRT's ResourceKey";
326   EXPECT_EQ(SRM.getRecordedResources().count(BarRT->getKeyUnsafe()), 0U)
327       << "Expected no entry for BarRT's ResourceKey";
328 
329   // We need to explicitly destroy FooRT or its resources will be implicitly
330   // transferred to the default tracker triggering a second call to our
331   // transfer function above (which expects only one call).
332   cantFail(FooRT->remove());
333 }
334 
335 TEST_F(ResourceTrackerStandardTest,
336        BasicDefineAndExplicitTransferAfterMaterializing) {
337 
338   bool ResourceManagerGotTransfer = false;
339   SimpleResourceManager<> SRM(ES);
340   SRM.setHandleTransfer([&](ResourceKey DstKey, ResourceKey SrcKey) {
341     ResourceManagerGotTransfer = true;
342     SRM.transferResources(DstKey, SrcKey);
343   });
344 
345   auto MakeMU = [&](SymbolStringPtr Name, JITEvaluatedSymbol Sym) {
346     return std::make_unique<SimpleMaterializationUnit>(
347         SymbolFlagsMap({{Name, Sym.getFlags()}}),
348         [=, &SRM](std::unique_ptr<MaterializationResponsibility> R) {
349           cantFail(R->withResourceKeyDo(
350               [&](ResourceKey K) { SRM.recordResource(K, 1); }));
351           cantFail(R->notifyResolved({{Name, Sym}}));
352           cantFail(R->notifyEmitted());
353         });
354   };
355 
356   auto FooRT = JD.createResourceTracker();
357   cantFail(JD.define(MakeMU(Foo, FooSym), FooRT));
358 
359   auto BarRT = JD.createResourceTracker();
360   cantFail(JD.define(MakeMU(Bar, BarSym), BarRT));
361 
362   EXPECT_EQ(SRM.getRecordedResources().size(), 0U)
363       << "Expected no recorded resources yet";
364 
365   cantFail(
366       ES.lookup(makeJITDylibSearchOrder({&JD}), SymbolLookupSet({Foo, Bar})));
367 
368   EXPECT_EQ(SRM.getRecordedResources().size(), 2U)
369       << "Expected recorded resources for both Foo and Bar";
370 
371   BarRT->transferTo(*FooRT);
372 
373   EXPECT_TRUE(ResourceManagerGotTransfer)
374       << "ResourceManager did not receive transfer";
375   EXPECT_TRUE(BarRT->isDefunct()) << "BarRT should now be defunct";
376 
377   EXPECT_EQ(SRM.getRecordedResources().size(), 1U)
378       << "Expected recorded resources for Foo only";
379   EXPECT_EQ(SRM.getRecordedResources().count(FooRT->getKeyUnsafe()), 1U)
380       << "Expected recorded resources for Foo";
381   EXPECT_EQ(SRM.getRecordedResources()[FooRT->getKeyUnsafe()], 2U)
382       << "Expected resources value for for Foo to be '2'";
383 }
384 
385 TEST_F(ResourceTrackerStandardTest,
386        BasicDefineAndExplicitTransferWhileMaterializing) {
387 
388   bool ResourceManagerGotTransfer = false;
389   SimpleResourceManager<> SRM(ES);
390   SRM.setHandleTransfer([&](ResourceKey DstKey, ResourceKey SrcKey) {
391     ResourceManagerGotTransfer = true;
392     SRM.transferResources(DstKey, SrcKey);
393   });
394 
395   auto FooRT = JD.createResourceTracker();
396   std::unique_ptr<MaterializationResponsibility> FooMR;
397   cantFail(JD.define(std::make_unique<SimpleMaterializationUnit>(
398                          SymbolFlagsMap({{Foo, FooSym.getFlags()}}),
399                          [&](std::unique_ptr<MaterializationResponsibility> R) {
400                            FooMR = std::move(R);
401                          }),
402                      FooRT));
403 
404   auto BarRT = JD.createResourceTracker();
405 
406   ES.lookup(
407       LookupKind::Static, makeJITDylibSearchOrder(&JD), SymbolLookupSet(Foo),
408       SymbolState::Ready,
409       [](Expected<SymbolMap> Result) { cantFail(Result.takeError()); },
410       NoDependenciesToRegister);
411 
412   cantFail(FooMR->withResourceKeyDo([&](ResourceKey K) {
413     EXPECT_EQ(FooRT->getKeyUnsafe(), K)
414         << "Expected FooRT's ResourceKey for Foo here";
415     SRM.recordResource(K, 1);
416   }));
417 
418   EXPECT_EQ(SRM.getRecordedResources().size(), 1U)
419       << "Expected one recorded resource here";
420   EXPECT_EQ(SRM.getRecordedResources()[FooRT->getKeyUnsafe()], 1U)
421       << "Expected Resource value for FooRT to be '1' here";
422 
423   FooRT->transferTo(*BarRT);
424 
425   EXPECT_TRUE(ResourceManagerGotTransfer)
426       << "Expected resource manager to receive handleTransferResources call";
427 
428   cantFail(FooMR->withResourceKeyDo([&](ResourceKey K) {
429     EXPECT_EQ(BarRT->getKeyUnsafe(), K)
430         << "Expected BarRT's ResourceKey for Foo here";
431     SRM.recordResource(K, 1);
432   }));
433 
434   EXPECT_EQ(SRM.getRecordedResources().size(), 1U)
435       << "Expected one recorded resource here";
436   EXPECT_EQ(SRM.getRecordedResources().count(BarRT->getKeyUnsafe()), 1U)
437       << "Expected RecordedResources to contain an entry for BarRT";
438   EXPECT_EQ(SRM.getRecordedResources()[BarRT->getKeyUnsafe()], 2U)
439       << "Expected Resource value for BarRT to be '2' here";
440 
441   cantFail(FooMR->notifyResolved({{Foo, FooSym}}));
442   cantFail(FooMR->notifyEmitted());
443 }
444 
445 } // namespace
446