[llvm] [OrcJIT] distinguish normal and abrupt TaskDispatcher shutdown (PR #154356)
Jameson Nash via llvm-commits
llvm-commits at lists.llvm.org
Tue Aug 19 07:56:14 PDT 2025
https://github.com/vtjnash created https://github.com/llvm/llvm-project/pull/154356
Add a `run_to_complete` complement to `shutdown`, which does nearly the same operations, but does not cancel future / pending work. The distinction is often mainly academic, but split for clarity of intent of each caller.
Annotate each and every tests that assumes that any previously scheduled tasks have finished running with a call of `run_to_completion` before reading the variables written by those tasks. This more clearly emphasizes the expectations of each of these tests and makes them robust against data-races or changes to the default scheduler.
These should have no impact currently, but have been tested and shown to pass with a scheduler which defers all work (as if being run on a thread) until an explicit synchronization point "acquires" that state, demonstrating that this is each and every test that would be affected by changes to the default scheduler.
>From 1cd472372b6c653fb002d1bb2977c4a385f6a426 Mon Sep 17 00:00:00 2001
From: Jameson Nash <vtjnash at gmail.com>
Date: Tue, 19 Aug 2025 14:32:33 +0000
Subject: [PATCH] [OrcJIT] distinguish normal and abrupt TaskDispatcher
shutdown
Add a `run_to_complete` complement to `shutdown`, which does nearly the
same operations, but does not cancel future / pending work. The
distinction is often mainly academic, but split for clarity of intent of
each caller.
Annotate each and every tests that assumes that any previously scheduled
tasks have finished running with a call of `run_to_completion` before
reading the variables written by those tasks. This more clearly
emphasizes the expectations of each of these tests and makes them robust
against data-races or changes to the default scheduler.
These should have no impact currently, but have been tested and shown to
pass with a scheduler which defers all work (as if being run on a
thread) until an explicit synchronization point "acquires" that state,
demonstrating that this is each and every test that would be affected by
changes to the default scheduler.
---
.../llvm/ExecutionEngine/Orc/TaskDispatch.h | 19 +++++-
llvm/lib/ExecutionEngine/Orc/TaskDispatch.cpp | 7 ++-
.../JITLink/JITLinkTestUtils.h | 2 +-
.../ExecutionEngine/Orc/CoreAPIsTest.cpp | 58 ++++++++++++++++++-
.../ExecutionEngine/Orc/OrcTestCommon.h | 8 +++
.../Orc/ResourceTrackerTest.cpp | 9 +++
.../ExecutionEngine/Orc/TaskDispatchTest.cpp | 2 +
7 files changed, 96 insertions(+), 9 deletions(-)
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/TaskDispatch.h b/llvm/include/llvm/ExecutionEngine/Orc/TaskDispatch.h
index 9cf6e00ad7131..075044a06cb15 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/TaskDispatch.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/TaskDispatch.h
@@ -112,14 +112,26 @@ class LLVM_ABI TaskDispatcher {
virtual void dispatch(std::unique_ptr<Task> T) = 0;
/// Called by ExecutionSession. Waits until all tasks have completed.
- virtual void shutdown() = 0;
+ /// The TaskDispatcher may be reused immediatly afterwards.
+ void run_to_complete() { run(false); }
+
+ /// Called by ExecutionSession. Halts all in-progress work as soon as
+ /// possible. May cause deadlocks since promises will not be set, so this
+ /// should only be used immediately before exiting.
+ /// The TaskDispatcher should not be reused afterwards.
+ void shutdown() { run(true); }
+
+private:
+ virtual void run(bool cancel) = 0;
};
/// Runs all tasks on the current thread.
class LLVM_ABI InPlaceTaskDispatcher : public TaskDispatcher {
public:
void dispatch(std::unique_ptr<Task> T) override;
- void shutdown() override;
+
+private:
+ void run(bool cancel) override;
};
#if LLVM_ENABLE_THREADS
@@ -131,8 +143,9 @@ class LLVM_ABI DynamicThreadPoolTaskDispatcher : public TaskDispatcher {
: MaxMaterializationThreads(MaxMaterializationThreads) {}
void dispatch(std::unique_ptr<Task> T) override;
- void shutdown() override;
+
private:
+ void run(bool cancel) override;
bool canRunMaterializationTaskNow();
bool canRunIdleTaskNow();
diff --git a/llvm/lib/ExecutionEngine/Orc/TaskDispatch.cpp b/llvm/lib/ExecutionEngine/Orc/TaskDispatch.cpp
index e87a14f3ea7c4..e4d7588fc28ed 100644
--- a/llvm/lib/ExecutionEngine/Orc/TaskDispatch.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/TaskDispatch.cpp
@@ -26,7 +26,7 @@ TaskDispatcher::~TaskDispatcher() = default;
void InPlaceTaskDispatcher::dispatch(std::unique_ptr<Task> T) { T->run(); }
-void InPlaceTaskDispatcher::shutdown() {}
+void InPlaceTaskDispatcher::run(bool cancel) {}
#if LLVM_ENABLE_THREADS
void DynamicThreadPoolTaskDispatcher::dispatch(std::unique_ptr<Task> T) {
@@ -105,9 +105,10 @@ void DynamicThreadPoolTaskDispatcher::dispatch(std::unique_ptr<Task> T) {
}).detach();
}
-void DynamicThreadPoolTaskDispatcher::shutdown() {
+void DynamicThreadPoolTaskDispatcher::run(bool cancel) {
std::unique_lock<std::mutex> Lock(DispatchMutex);
- Shutdown = true;
+ if (cancel)
+ Shutdown = true;
OutstandingCV.wait(Lock, [this]() { return Outstanding == 0; });
}
diff --git a/llvm/unittests/ExecutionEngine/JITLink/JITLinkTestUtils.h b/llvm/unittests/ExecutionEngine/JITLink/JITLinkTestUtils.h
index dc077f900d195..6bd299bbd37a4 100644
--- a/llvm/unittests/ExecutionEngine/JITLink/JITLinkTestUtils.h
+++ b/llvm/unittests/ExecutionEngine/JITLink/JITLinkTestUtils.h
@@ -133,7 +133,7 @@ class MockJITLinkContext : public llvm::jitlink::JITLinkContext {
HandleFailed(std::move(HandleFailed)) {}
~MockJITLinkContext() {
- if (auto Err = MJMM->deallocate(std::move(FinalizedAllocs)))
+ if (auto Err = MJMM->Deallocate(std::move(FinalizedAllocs)))
notifyFailed(std::move(Err));
}
diff --git a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp b/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
index ec94083859bc5..a66007704997f 100644
--- a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
@@ -98,14 +98,17 @@ TEST_F(CoreAPIsStandardTest, BasicSuccessfulLookup) {
ES.lookup(LookupKind::Static, makeJITDylibSearchOrder(&JD),
SymbolLookupSet(Foo), SymbolState::Ready, OnCompletion,
NoDependenciesToRegister);
+ getDispatcher().run_to_complete();
EXPECT_FALSE(OnCompletionRun) << "Should not have been resolved yet";
cantFail(FooMR->notifyResolved({{Foo, FooSym}}));
EXPECT_FALSE(OnCompletionRun) << "Should not be ready yet";
+ getDispatcher().run_to_complete();
cantFail(FooMR->notifyEmitted({}));
+ getDispatcher().run_to_complete();
EXPECT_TRUE(OnCompletionRun) << "Should have been marked ready";
}
@@ -120,6 +123,7 @@ TEST_F(CoreAPIsStandardTest, EmptyLookup) {
ES.lookup(LookupKind::Static, makeJITDylibSearchOrder(&JD), SymbolLookupSet(),
SymbolState::Ready, OnCompletion, NoDependenciesToRegister);
+ getDispatcher().run_to_complete();
EXPECT_TRUE(OnCompletionRun) << "OnCompletion was not run for empty query";
}
@@ -173,6 +177,7 @@ TEST_F(CoreAPIsStandardTest, MaterializationSideEffctsOnlyBasic) {
},
NoDependenciesToRegister);
+ getDispatcher().run_to_complete();
EXPECT_FALSE(Result) << "Lookup returned unexpectedly";
EXPECT_TRUE(FooR) << "Lookup failed to trigger materialization";
EXPECT_THAT_ERROR(FooR->notifyEmitted({}), Succeeded())
@@ -254,6 +259,7 @@ TEST_F(CoreAPIsStandardTest, RemoveSymbolsTest) {
OnCompletionRun = true;
},
NoDependenciesToRegister);
+ getDispatcher().run_to_complete();
{
// Attempt 1: Search for a missing symbol, Qux.
@@ -281,6 +287,7 @@ TEST_F(CoreAPIsStandardTest, RemoveSymbolsTest) {
auto Err = JD.remove({Foo, Bar, Baz});
EXPECT_FALSE(!!Err) << "Expected success";
}
+ getDispatcher().run_to_complete();
EXPECT_TRUE(BarDiscarded) << "\"Bar\" should have been discarded";
EXPECT_TRUE(BarMaterializerDestructed)
@@ -507,12 +514,14 @@ TEST_F(CoreAPIsStandardTest, TestTrivialCircularDependency) {
ES.lookup(LookupKind::Static, makeJITDylibSearchOrder(&JD),
SymbolLookupSet({Foo}), SymbolState::Ready, OnCompletion,
NoDependenciesToRegister);
+ getDispatcher().run_to_complete();
EXPECT_THAT_ERROR(FooR->notifyResolved({{Foo, FooSym}}), Succeeded())
<< "No symbols marked failed, but Foo failed to resolve";
SymbolDependenceGroup SDG({{Foo}, {{&JD, SymbolNameSet({Foo})}}});
EXPECT_THAT_ERROR(FooR->notifyEmitted(SDG), Succeeded())
<< "No symbols marked failed, but Foo failed to emit";
+ getDispatcher().run_to_complete();
EXPECT_TRUE(FooReady)
<< "Self-dependency prevented symbol from being marked ready";
@@ -554,6 +563,7 @@ TEST_F(CoreAPIsStandardTest, TestBasicQueryDependenciesReporting) {
EXPECT_THAT_EXPECTED(std::move(Result), Succeeded());
},
NoDependenciesToRegister);
+ getDispatcher().run_to_complete();
cantFail(FooR->notifyResolved({{Foo, FooSym}}));
cantFail(FooR->notifyEmitted({}));
@@ -577,6 +587,7 @@ TEST_F(CoreAPIsStandardTest, TestBasicQueryDependenciesReporting) {
EXPECT_TRUE(Deps.count(Baz));
DependenciesCallbackRan = true;
});
+ getDispatcher().run_to_complete();
cantFail(BarR->notifyEmitted({}));
@@ -640,6 +651,7 @@ TEST_F(CoreAPIsStandardTest, TestCircularDependenceInOneJITDylib) {
ES.lookup(LookupKind::Static, makeJITDylibSearchOrder(&JD),
SymbolLookupSet(Foo), SymbolState::Ready, std::move(OnFooReady),
NoDependenciesToRegister);
+ getDispatcher().run_to_complete();
bool BarResolved = false;
bool BarReady = false;
@@ -660,6 +672,7 @@ TEST_F(CoreAPIsStandardTest, TestCircularDependenceInOneJITDylib) {
ES.lookup(LookupKind::Static, makeJITDylibSearchOrder(&JD),
SymbolLookupSet(Bar), SymbolState::Ready, std::move(OnBarReady),
NoDependenciesToRegister);
+ getDispatcher().run_to_complete();
bool BazResolved = false;
bool BazReady = false;
@@ -681,6 +694,7 @@ TEST_F(CoreAPIsStandardTest, TestCircularDependenceInOneJITDylib) {
ES.lookup(LookupKind::Static, makeJITDylibSearchOrder(&JD),
SymbolLookupSet(Baz), SymbolState::Ready, std::move(OnBazReady),
NoDependenciesToRegister);
+ getDispatcher().run_to_complete();
// Check that nothing has been resolved yet.
EXPECT_FALSE(FooResolved) << "\"Foo\" should not be resolved yet";
@@ -694,6 +708,7 @@ TEST_F(CoreAPIsStandardTest, TestCircularDependenceInOneJITDylib) {
<< "No symbols failed, but Bar failed to resolve";
EXPECT_THAT_ERROR(BazR->notifyResolved({{Baz, BazSym}}), Succeeded())
<< "No symbols failed, but Baz failed to resolve";
+ getDispatcher().run_to_complete();
// Verify that the symbols have been resolved, but are not ready yet.
EXPECT_TRUE(FooResolved) << "\"Foo\" should be resolved now";
@@ -714,6 +729,7 @@ TEST_F(CoreAPIsStandardTest, TestCircularDependenceInOneJITDylib) {
EXPECT_THAT_ERROR(BarR->notifyEmitted(BarDeps), Succeeded())
<< "No symbols failed, but Bar failed to emit";
}
+ getDispatcher().run_to_complete();
// Verify that nothing is ready until the circular dependence is resolved.
EXPECT_FALSE(FooReady) << "\"Foo\" still should not be ready";
@@ -726,6 +742,7 @@ TEST_F(CoreAPIsStandardTest, TestCircularDependenceInOneJITDylib) {
EXPECT_THAT_ERROR(BazR->notifyEmitted(BazDeps), Succeeded())
<< "No symbols failed, but Baz failed to emit";
}
+ getDispatcher().run_to_complete();
// Verify that everything becomes ready once the circular dependence resolved.
EXPECT_TRUE(FooReady) << "\"Foo\" should be ready now";
@@ -764,6 +781,7 @@ TEST_F(CoreAPIsStandardTest, FailureInDependency) {
ES.lookup(LookupKind::Static, makeJITDylibSearchOrder(&JD),
SymbolLookupSet(Foo), SymbolState::Ready, std::move(OnFooReady),
NoDependenciesToRegister);
+ getDispatcher().run_to_complete();
bool OnBarReadyRun = false;
auto OnBarReady = [&](Expected<SymbolMap> Result) {
@@ -774,6 +792,7 @@ TEST_F(CoreAPIsStandardTest, FailureInDependency) {
ES.lookup(LookupKind::Static, makeJITDylibSearchOrder(&JD),
SymbolLookupSet(Bar), SymbolState::Ready, std::move(OnBarReady),
NoDependenciesToRegister);
+ getDispatcher().run_to_complete();
// Fail bar.
BarR->failMaterialization();
@@ -842,6 +861,7 @@ TEST_F(CoreAPIsStandardTest, AddDependencyOnFailedSymbol) {
NoDependenciesToRegister);
// Fail bar.
+ getDispatcher().run_to_complete();
BarR->failMaterialization();
// We expect Bar's query to fail immediately, but Foo's query not to have run
@@ -859,6 +879,7 @@ TEST_F(CoreAPIsStandardTest, AddDependencyOnFailedSymbol) {
}
FooR->failMaterialization();
+ getDispatcher().run_to_complete();
// Foo's query should have failed before we return from addDependencies.
EXPECT_TRUE(OnFooReadyRun) << "Query for \"Foo\" was not run";
@@ -1028,6 +1049,7 @@ TEST_F(CoreAPIsStandardTest, AddAndMaterializeLazySymbol) {
ES.lookup(LookupKind::Static, makeJITDylibSearchOrder(&JD),
SymbolLookupSet(Foo), SymbolState::Ready, std::move(OnCompletion),
NoDependenciesToRegister);
+ getDispatcher().run_to_complete();
EXPECT_TRUE(FooMaterialized) << "Foo was not materialized";
EXPECT_TRUE(BarDiscarded) << "Bar was not discarded";
@@ -1073,6 +1095,7 @@ TEST_F(CoreAPIsStandardTest, TestBasicWeakSymbolMaterialization) {
ES.lookup(LookupKind::Static, makeJITDylibSearchOrder(&JD),
SymbolLookupSet(Bar), SymbolState::Ready, std::move(OnCompletion),
NoDependenciesToRegister);
+ getDispatcher().run_to_complete();
EXPECT_TRUE(OnCompletionRun) << "OnCompletion not run";
EXPECT_TRUE(BarMaterialized) << "Bar was not materialized at all";
@@ -1097,10 +1120,11 @@ TEST_F(CoreAPIsStandardTest, RedefineBoundWeakSymbol) {
TEST_F(CoreAPIsStandardTest, DefineMaterializingSymbol) {
bool ExpectNoMoreMaterialization = false;
+ auto &D = getDispatcher();
DispatchOverride = [&](std::unique_ptr<Task> T) {
if (ExpectNoMoreMaterialization && isa<MaterializationTask>(*T))
ADD_FAILURE() << "Unexpected materialization";
- T->run();
+ D.dispatch_super(std::move(T));
};
auto MU = std::make_unique<SimpleMaterializationUnit>(
@@ -1116,12 +1140,14 @@ TEST_F(CoreAPIsStandardTest, DefineMaterializingSymbol) {
cantFail(ES.lookup(makeJITDylibSearchOrder(&JD), Foo));
// Assert that materialization is complete by now.
+ D.run_to_complete();
ExpectNoMoreMaterialization = true;
// Look up bar to verify that no further materialization happens.
auto BarResult = cantFail(ES.lookup(makeJITDylibSearchOrder(&JD), Bar));
EXPECT_EQ(BarResult.getAddress(), BarSym.getAddress())
<< "Expected Bar == BarSym";
+ D.run_to_complete();
}
TEST_F(CoreAPIsStandardTest, GeneratorTest) {
@@ -1215,11 +1241,13 @@ TEST_F(CoreAPIsStandardTest, SimpleAsynchronousGeneratorTest) {
}
},
NoDependenciesToRegister);
+ getDispatcher().run_to_complete();
EXPECT_FALSE(LookupCompleted);
cantFail(JD.define(absoluteSymbols({{Foo, FooSym}})));
G.takeLookup().LS.continueLookup(Error::success());
+ getDispatcher().run_to_complete();
EXPECT_TRUE(LookupCompleted);
}
@@ -1238,11 +1266,13 @@ TEST_F(CoreAPIsStandardTest, ErrorFromSuspendedAsynchronousGeneratorTest) {
EXPECT_THAT_EXPECTED(Result, Failed());
},
NoDependenciesToRegister);
+ getDispatcher().run_to_complete();
EXPECT_FALSE(LookupCompleted);
G.takeLookup().LS.continueLookup(
make_error<StringError>("boom", inconvertibleErrorCode()));
+ getDispatcher().run_to_complete();
EXPECT_TRUE(LookupCompleted);
}
@@ -1261,6 +1291,7 @@ TEST_F(CoreAPIsStandardTest, ErrorFromAutoSuspendedAsynchronousGeneratorTest) {
EXPECT_THAT_EXPECTED(Result, Failed());
},
NoDependenciesToRegister);
+ getDispatcher().run_to_complete();
EXPECT_EQ(LookupsCompleted, 0U);
@@ -1276,15 +1307,18 @@ TEST_F(CoreAPIsStandardTest, ErrorFromAutoSuspendedAsynchronousGeneratorTest) {
EXPECT_THAT_EXPECTED(Result, Failed());
},
NoDependenciesToRegister);
+ getDispatcher().run_to_complete();
EXPECT_EQ(LookupsCompleted, 0U);
// Unsuspend the first lookup.
LS1.continueLookup(make_error<StringError>("boom", inconvertibleErrorCode()));
+ getDispatcher().run_to_complete();
// Unsuspend the second.
G.takeLookup().LS.continueLookup(
make_error<StringError>("boom", inconvertibleErrorCode()));
+ getDispatcher().run_to_complete();
EXPECT_EQ(LookupsCompleted, 2U);
}
@@ -1316,6 +1350,7 @@ TEST_F(CoreAPIsStandardTest, BlockedGeneratorAutoSuspensionTest) {
}
},
NoDependenciesToRegister);
+ getDispatcher().run_to_complete();
// The generator should immediately see the first lookup.
EXPECT_NE(G.Lookup, std::nullopt);
@@ -1345,6 +1380,7 @@ TEST_F(CoreAPIsStandardTest, BlockedGeneratorAutoSuspensionTest) {
}
},
NoDependenciesToRegister);
+ getDispatcher().run_to_complete();
// Add lookup 3.
//
@@ -1363,6 +1399,7 @@ TEST_F(CoreAPIsStandardTest, BlockedGeneratorAutoSuspensionTest) {
}
},
NoDependenciesToRegister);
+ getDispatcher().run_to_complete();
// Add lookup 4.
//
@@ -1379,6 +1416,7 @@ TEST_F(CoreAPIsStandardTest, BlockedGeneratorAutoSuspensionTest) {
}
},
NoDependenciesToRegister);
+ getDispatcher().run_to_complete();
// All lookups have been started, but none should have been completed yet.
EXPECT_FALSE(Lookup1Completed);
@@ -1393,6 +1431,7 @@ TEST_F(CoreAPIsStandardTest, BlockedGeneratorAutoSuspensionTest) {
// allow both 2 and 3 to complete.
cantFail(JD.define(absoluteSymbols({{Foo, FooSym}})));
G.takeLookup().LS.continueLookup(Error::success());
+ getDispatcher().run_to_complete();
EXPECT_TRUE(Lookup1Completed);
EXPECT_TRUE(Lookup2Completed);
@@ -1407,6 +1446,7 @@ TEST_F(CoreAPIsStandardTest, BlockedGeneratorAutoSuspensionTest) {
cantFail(JD.define(absoluteSymbols({{Baz, BazSym}})));
G.takeLookup().LS.continueLookup(Error::success());
+ getDispatcher().run_to_complete();
EXPECT_TRUE(Lookup4Completed);
}
@@ -1456,7 +1496,7 @@ TEST_F(CoreAPIsStandardTest, FailEmissionAfterResolution) {
ES.lookup(
LookupKind::Static, makeJITDylibSearchOrder(&JD),
SymbolLookupSet({Baz}), SymbolState::Resolved,
- [&](Expected<SymbolMap> Result) {
+ [R = std::move(R)](Expected<SymbolMap> Result) {
// Called when "baz" is resolved. We don't actually depend
// on or care about baz, but use it to trigger failure of
// this materialization before Baz has been finalized in
@@ -1472,6 +1512,7 @@ TEST_F(CoreAPIsStandardTest, FailEmissionAfterResolution) {
auto Result =
ES.lookup(makeJITDylibSearchOrder(&JD), SymbolLookupSet({Foo, Bar}));
+ getDispatcher().run_to_complete();
EXPECT_THAT_EXPECTED(std::move(Result), Failed())
<< "Unexpected success while trying to test error propagation";
@@ -1500,6 +1541,7 @@ TEST_F(CoreAPIsStandardTest, FailAfterPartialResolution) {
QueryHandlerRun = true;
},
NoDependenciesToRegister);
+ getDispatcher().run_to_complete();
EXPECT_TRUE(QueryHandlerRun) << "Query handler never ran";
}
@@ -1526,6 +1568,7 @@ TEST_F(CoreAPIsStandardTest, FailDefineMaterializingDueToDefunctTracker) {
ES.lookup(LookupKind::Static, makeJITDylibSearchOrder(&JD),
SymbolLookupSet(Foo), SymbolState::Ready, OnCompletion,
NoDependenciesToRegister);
+ getDispatcher().run_to_complete();
cantFail(RT->remove());
@@ -1533,6 +1576,7 @@ TEST_F(CoreAPIsStandardTest, FailDefineMaterializingDueToDefunctTracker) {
<< "defineMaterializing should have failed due to a defunct tracker";
FooMR->failMaterialization();
+ getDispatcher().run_to_complete();
EXPECT_TRUE(OnCompletionRan) << "OnCompletion handler did not run.";
}
@@ -1560,15 +1604,22 @@ TEST_F(CoreAPIsStandardTest, TestLookupWithThreadedMaterialization) {
std::mutex WorkThreadsMutex;
SmallVector<std::thread, 0> WorkThreads;
+ // auto &D = getDispatcher();
DispatchOverride = [&](std::unique_ptr<Task> T) {
std::lock_guard<std::mutex> Lock(WorkThreadsMutex);
WorkThreads.push_back(
std::thread([T = std::move(T)]() mutable { T->run(); }));
+ // WorkThreads.push_back(std::thread(
+ // [T = std::move(T), &D]() mutable {
+ // D.dispatch_super(std::move(T));
+ // D.run_to_complete();
+ // }));
};
cantFail(JD.define(absoluteSymbols({{Foo, FooSym}})));
auto FooLookupResult = cantFail(ES.lookup(makeJITDylibSearchOrder(&JD), Foo));
+ getDispatcher().run_to_complete();
EXPECT_EQ(FooLookupResult.getAddress(), FooSym.getAddress())
<< "lookup returned an incorrect address";
@@ -1683,6 +1734,7 @@ TEST_F(CoreAPIsStandardTest, TestMaterializeWeakSymbol) {
ES.lookup(LookupKind::Static, makeJITDylibSearchOrder(&JD),
SymbolLookupSet({Foo}), SymbolState::Ready, std::move(OnCompletion),
NoDependenciesToRegister);
+ getDispatcher().run_to_complete();
auto MU2 = std::make_unique<SimpleMaterializationUnit>(
SymbolFlagsMap({{Foo, JITSymbolFlags::Exported}}),
@@ -1699,6 +1751,7 @@ TEST_F(CoreAPIsStandardTest, TestMaterializeWeakSymbol) {
// No dependencies registered, can't fail:
cantFail(FooR->notifyResolved(SymbolMap({{Foo, FooSym}})));
cantFail(FooR->notifyEmitted({}));
+ getDispatcher().run_to_complete();
}
static bool linkOrdersEqual(const std::vector<JITDylibSP> &LHS,
@@ -1850,6 +1903,7 @@ TEST_F(CoreAPIsStandardTest, RemoveJITDylibs) {
}
},
NoDependenciesToRegister);
+ getDispatcher().run_to_complete();
// Remove the JITDylib.
auto Err = ES.removeJITDylib(JD);
diff --git a/llvm/unittests/ExecutionEngine/Orc/OrcTestCommon.h b/llvm/unittests/ExecutionEngine/Orc/OrcTestCommon.h
index 469de2a3665a0..77e1c2baa23f9 100644
--- a/llvm/unittests/ExecutionEngine/Orc/OrcTestCommon.h
+++ b/llvm/unittests/ExecutionEngine/Orc/OrcTestCommon.h
@@ -57,6 +57,9 @@ class CoreAPIsBasedStandardTest : public testing::Test {
public:
OverridableDispatcher(CoreAPIsBasedStandardTest &Parent) : Parent(Parent) {}
void dispatch(std::unique_ptr<Task> T) override;
+ void dispatch_super(std::unique_ptr<Task> T) {
+ InPlaceTaskDispatcher::dispatch(std::move(T));
+ }
private:
CoreAPIsBasedStandardTest &Parent;
@@ -65,6 +68,11 @@ class CoreAPIsBasedStandardTest : public testing::Test {
std::unique_ptr<llvm::orc::ExecutorProcessControl>
makeEPC(std::shared_ptr<SymbolStringPool> SSP);
+ OverridableDispatcher &getDispatcher() {
+ return static_cast<OverridableDispatcher &>(
+ ES.getExecutorProcessControl().getDispatcher());
+ }
+
std::shared_ptr<SymbolStringPool> SSP = std::make_shared<SymbolStringPool>();
ExecutionSession ES{makeEPC(SSP)};
JITDylib &JD = ES.createBareJITDylib("JD");
diff --git a/llvm/unittests/ExecutionEngine/Orc/ResourceTrackerTest.cpp b/llvm/unittests/ExecutionEngine/Orc/ResourceTrackerTest.cpp
index 0e5d151714fcc..0b3c35c5ae8a0 100644
--- a/llvm/unittests/ExecutionEngine/Orc/ResourceTrackerTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/ResourceTrackerTest.cpp
@@ -141,6 +141,7 @@ TEST_F(ResourceTrackerStandardTest,
LookupKind::Static,
{{&JD, JITDylibLookupFlags::MatchExportedSymbolsOnly}},
SymbolLookupSet(Foo, SymbolLookupFlags::WeaklyReferencedSymbol)));
+ getDispatcher().run_to_complete();
EXPECT_EQ(SymFlags.size(), 0U)
<< "Symbols should have been removed from the symbol table";
@@ -181,6 +182,7 @@ TEST_F(ResourceTrackerStandardTest, BasicDefineAndRemoveAllAfterMaterializing) {
LookupKind::Static,
{{&JD, JITDylibLookupFlags::MatchExportedSymbolsOnly}},
SymbolLookupSet(Foo, SymbolLookupFlags::WeaklyReferencedSymbol)));
+ getDispatcher().run_to_complete();
EXPECT_EQ(SymFlags.size(), 0U)
<< "Symbols should have been removed from the symbol table";
@@ -218,6 +220,7 @@ TEST_F(ResourceTrackerStandardTest, BasicDefineAndRemoveAllWhileMaterializing) {
<< "Lookup failed unexpectedly";
},
NoDependenciesToRegister);
+ getDispatcher().run_to_complete();
cantFail(RT->remove());
auto SymFlags = cantFail(ES.lookupFlags(
@@ -240,6 +243,7 @@ TEST_F(ResourceTrackerStandardTest, BasicDefineAndRemoveAllWhileMaterializing) {
<< "notifyResolved on MR with removed tracker should have failed";
MR->failMaterialization();
+ getDispatcher().run_to_complete();
}
TEST_F(ResourceTrackerStandardTest, JITDylibClear) {
@@ -270,6 +274,7 @@ TEST_F(ResourceTrackerStandardTest, JITDylibClear) {
cantFail(
ES.lookup(makeJITDylibSearchOrder(&JD), SymbolLookupSet({Foo, Bar})));
+ getDispatcher().run_to_complete();
auto JDResourceKey = JD.getDefaultResourceTracker()->getKeyUnsafe();
EXPECT_EQ(SRM.getRecordedResources().size(), 1U)
@@ -323,6 +328,7 @@ TEST_F(ResourceTrackerStandardTest,
cantFail(
ES.lookup(makeJITDylibSearchOrder({&JD}), SymbolLookupSet({Foo, Bar})));
+ getDispatcher().run_to_complete();
EXPECT_EQ(SRM.getRecordedResources().size(), 1U)
<< "Expected exactly one entry (for FooRT's Key)";
@@ -370,6 +376,7 @@ TEST_F(ResourceTrackerStandardTest,
cantFail(
ES.lookup(makeJITDylibSearchOrder({&JD}), SymbolLookupSet({Foo, Bar})));
+ getDispatcher().run_to_complete();
EXPECT_EQ(SRM.getRecordedResources().size(), 2U)
<< "Expected recorded resources for both Foo and Bar";
@@ -415,6 +422,7 @@ TEST_F(ResourceTrackerStandardTest,
SymbolState::Ready,
[](Expected<SymbolMap> Result) { cantFail(Result.takeError()); },
NoDependenciesToRegister);
+ getDispatcher().run_to_complete();
cantFail(FooMR->withResourceKeyDo([&](ResourceKey K) {
EXPECT_EQ(FooRT->getKeyUnsafe(), K)
@@ -447,6 +455,7 @@ TEST_F(ResourceTrackerStandardTest,
cantFail(FooMR->notifyResolved({{Foo, FooSym}}));
cantFail(FooMR->notifyEmitted({}));
+ getDispatcher().run_to_complete();
}
} // namespace
diff --git a/llvm/unittests/ExecutionEngine/Orc/TaskDispatchTest.cpp b/llvm/unittests/ExecutionEngine/Orc/TaskDispatchTest.cpp
index 4931fc26a417e..d4a552775f492 100644
--- a/llvm/unittests/ExecutionEngine/Orc/TaskDispatchTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/TaskDispatchTest.cpp
@@ -19,6 +19,7 @@ TEST(InPlaceTaskDispatchTest, GenericNamedTask) {
auto D = std::make_unique<InPlaceTaskDispatcher>();
bool B = false;
D->dispatch(makeGenericNamedTask([&]() { B = true; }));
+ D->run_to_complete();
EXPECT_TRUE(B);
D->shutdown();
}
@@ -31,6 +32,7 @@ TEST(DynamicThreadPoolDispatchTest, GenericNamedTask) {
D->dispatch(makeGenericNamedTask(
[P = std::move(P)]() mutable { P.set_value(true); }));
EXPECT_TRUE(F.get());
+ D->run_to_complete();
D->shutdown();
}
#endif
More information about the llvm-commits
mailing list