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