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
SimpleResourceManager(ExecutionSession & ES)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
~SimpleResourceManager()52 ~SimpleResourceManager() { ES.deregisterResourceManager(*this); }
53
54 /// Set the HandleRemove function object.
setHandleRemove(HandleRemoveFunction HandleRemove)55 void setHandleRemove(HandleRemoveFunction HandleRemove) {
56 this->HandleRemove = std::move(HandleRemove);
57 }
58
59 /// Set the HandleTransfer function object.
setHandleTransfer(HandleTransferFunction HandleTransfer)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>>
recordResource(ResourceKey K,ResourceT Val=ResourceT (),MergeOp Merge=MergeOp ())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.
removeResource(ResourceKey K)73 void removeResource(ResourceKey K) { Resources.erase(K); }
74
75 /// Transfer resources from DstKey to SrcKey.
76 template <typename MergeOp = std::plus<ResourceT>>
transferResources(ResourceKey DstKey,ResourceKey SrcKey,MergeOp Merge=MergeOp ())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.
getRecordedResources()91 RecordedResourcesMap &getRecordedResources() { return Resources; }
getRecordedResources() const92 const RecordedResourcesMap &getRecordedResources() const { return Resources; }
93
handleRemoveResources(ResourceKey K)94 Error handleRemoveResources(ResourceKey K) override {
95 return HandleRemove(K);
96 }
97
handleTransferResources(ResourceKey DstKey,ResourceKey SrcKey)98 void handleTransferResources(ResourceKey DstKey,
99 ResourceKey SrcKey) override {
100 HandleTransfer(DstKey, SrcKey);
101 }
102
transferNotAllowed(ResourceKey DstKey,ResourceKey SrcKey)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
TEST_F(ResourceTrackerStandardTest,BasicDefineAndRemoveAllBeforeMaterializing)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
TEST_F(ResourceTrackerStandardTest,BasicDefineAndRemoveAllAfterMaterializing)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
TEST_F(ResourceTrackerStandardTest,BasicDefineAndRemoveAllWhileMaterializing)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
TEST_F(ResourceTrackerStandardTest,JITDylibClear)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
TEST_F(ResourceTrackerStandardTest,BasicDefineAndExplicitTransferBeforeMaterializing)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
TEST_F(ResourceTrackerStandardTest,BasicDefineAndExplicitTransferAfterMaterializing)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
TEST_F(ResourceTrackerStandardTest,BasicDefineAndExplicitTransferWhileMaterializing)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