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