Skip to content

Commit

Permalink
ENH: fix Python multi-processing hang on unix
Browse files Browse the repository at this point in the history
  • Loading branch information
dzenanz committed Jul 16, 2021
1 parent bdad19f commit 1716c96
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 3 deletions.
14 changes: 12 additions & 2 deletions Modules/Core/Common/include/itkThreadPool.h
Original file line number Diff line number Diff line change
Expand Up @@ -124,13 +124,23 @@ auto result = pool->AddWork([](int param) { return param; }, 7);
SetDoNotWaitForThreads(bool doNotWaitForThreads);

protected:
/* We need access to the mutex in AddWork, and the variable is only
/** We need access to the mutex in AddWork, and the variable is only
* visible in .cxx file, so this method returns it. */
std::mutex &
GetMutex();

ThreadPool();
~ThreadPool() override;

/** Stop the pool and release threads. To be called by the destructor and atfork. */
void
CleanUp();

~ThreadPool() override { this->CleanUp(); }

static void
PrepareForFork();
static void
ResumeFromFork();

private:
/** Only used to synchronize the global variable across static libraries.*/
Expand Down
21 changes: 20 additions & 1 deletion Modules/Core/Common/src/itkThreadPool.cxx
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,9 @@ ThreadPool::GetInstance()
{
new ThreadPool(); // constructor sets m_PimplGlobals->m_ThreadPoolInstance
}
#if defined(ITK_USE_PTHREADS)
pthread_atfork(ThreadPool::PrepareForFork, ThreadPool::ResumeFromFork, ThreadPool::ResumeFromFork);
#endif
}
}
return m_PimplGlobals->m_ThreadPoolInstance;
Expand Down Expand Up @@ -132,7 +135,8 @@ ThreadPool::GetNumberOfCurrentlyIdleThreads() const
return int(m_Threads.size()) - int(m_WorkQueue.size()); // lousy approximation
}

ThreadPool::~ThreadPool()
void
ThreadPool::CleanUp()
{
{
std::unique_lock<std::mutex> mutexHolder(m_PimplGlobals->m_Mutex);
Expand All @@ -154,6 +158,21 @@ ThreadPool::~ThreadPool()
}
}

void
ThreadPool::PrepareForFork()
{
m_PimplGlobals->m_ThreadPoolInstance->CleanUp();
}

void
ThreadPool::ResumeFromFork()
{
ThreadPool * instance = m_PimplGlobals->m_ThreadPoolInstance.GetPointer();
ThreadIdType threadCount = instance->m_Threads.size();
instance->m_Threads.clear();
instance->m_Stopping = false;
instance->AddThreads(threadCount);
}

void
ThreadPool::ThreadExecute()
Expand Down

0 comments on commit 1716c96

Please sign in to comment.