Skip to content

Commit

Permalink
[ORC] Fix Task cleanup during DynamicThreadPoolTaskDispatcher::shutdown.
Browse files Browse the repository at this point in the history
Threads created by DynamicThreadPoolTaskDispatcher::dispatch had been holding a
unique_ptr to the most recent Task, meaning that the Task would be destroyed
when the thread object was destroyed, but this would happen *after* the thread
signaled the Dispatcher that it was finished. This could cause
DynamicThreadPoolTaskDispatcher::shutdown to return (and consequently
ExecutionSession to be destroyed) before all Tasks were destroyed, with Task
destructors accessing ExecutionSession and related objects after they were
freed.

The fix is to reset the Task pointer immediately after it is run to trigger
cleanup, *then* (if there are no other tasks to run) signal the Dispatcher that
the thread is finished.

This patch also updates DynamicThreadPoolTaskDispatcher::dispatch to reject any
new Tasks dispatched after DynamicThreadPoolTaskDispatcher::shutdown is called.
  • Loading branch information
lhames committed Jan 9, 2025
1 parent 3cb9648 commit 8312876
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 3 deletions.
2 changes: 1 addition & 1 deletion llvm/include/llvm/ExecutionEngine/Orc/TaskDispatch.h
Original file line number Diff line number Diff line change
Expand Up @@ -122,7 +122,7 @@ class DynamicThreadPoolTaskDispatcher : public TaskDispatcher {
void shutdown() override;
private:
std::mutex DispatchMutex;
bool Running = true;
bool Shutdown = false;
size_t Outstanding = 0;
std::condition_variable OutstandingCV;

Expand Down
15 changes: 13 additions & 2 deletions llvm/lib/ExecutionEngine/Orc/TaskDispatch.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,10 @@ void DynamicThreadPoolTaskDispatcher::dispatch(std::unique_ptr<Task> T) {
{
std::lock_guard<std::mutex> Lock(DispatchMutex);

// Reject new tasks if they're dispatched after a call to shutdown.
if (Shutdown)
return;

if (IsMaterializationTask) {

// If this is a materialization task and there are too many running
Expand All @@ -54,6 +58,14 @@ void DynamicThreadPoolTaskDispatcher::dispatch(std::unique_ptr<Task> T) {
// Run the task.
T->run();

// Reset the task to free any resources. We need this to happen *before*
// we notify anyone (via Outstanding) that this thread is done to ensure
// that we don't proceed with JIT shutdown while still holding resources.
// (E.g. this was causing "Dangling SymbolStringPtr" assertions).
T.reset();

// Check the work queue state and either proceed with the next task or
// end this thread.
std::lock_guard<std::mutex> Lock(DispatchMutex);
if (!MaterializationTaskQueue.empty()) {
// If there are any materialization tasks running then steal that work.
Expand All @@ -64,7 +76,6 @@ void DynamicThreadPoolTaskDispatcher::dispatch(std::unique_ptr<Task> T) {
IsMaterializationTask = true;
}
} else {
// Otherwise decrement work counters.
if (IsMaterializationTask)
--NumMaterializationThreads;
--Outstanding;
Expand All @@ -78,7 +89,7 @@ void DynamicThreadPoolTaskDispatcher::dispatch(std::unique_ptr<Task> T) {

void DynamicThreadPoolTaskDispatcher::shutdown() {
std::unique_lock<std::mutex> Lock(DispatchMutex);
Running = false;
Shutdown = true;
OutstandingCV.wait(Lock, [this]() { return Outstanding == 0; });
}
#endif
Expand Down

0 comments on commit 8312876

Please sign in to comment.