diff --git a/Common/Thread/ThreadManager.cpp b/Common/Thread/ThreadManager.cpp index bf3ed3cf98..bc1cf37797 100644 --- a/Common/Thread/ThreadManager.cpp +++ b/Common/Thread/ThreadManager.cpp @@ -26,9 +26,10 @@ const int EXTRA_THREADS = 4; // For I/O limited tasks struct GlobalThreadContext { std::mutex mutex; // associated with each respective condition variable std::deque queue; + std::atomic queue_size; std::vector threads_; - int roundRobin = 0; + std::atomic roundRobin; }; struct ThreadContext { @@ -43,7 +44,8 @@ struct ThreadContext { }; ThreadManager::ThreadManager() : global_(new GlobalThreadContext()) { - + global_->queue_size = 0; + global_->roundRobin = 0; } ThreadManager::~ThreadManager() { @@ -76,6 +78,7 @@ static void WorkerThreadFunc(GlobalThreadContext *global, ThreadContext *thread) if (!global->queue.empty()) { task = global->queue.front(); global->queue.pop_front(); + global->queue_size--; // We are processing one now, so mark that. thread->queue_size++; @@ -88,7 +91,7 @@ static void WorkerThreadFunc(GlobalThreadContext *global, ThreadContext *thread) if (!thread->private_queue.empty()) { task = thread->private_queue.front(); thread->private_queue.pop_front(); - } else if (thread->private_single.load() == nullptr) { + } else if (thread->private_single.load() == nullptr && global->queue_size.load() == 0) { thread->cond.wait(lock); } } @@ -162,9 +165,14 @@ void ThreadManager::EnqueueTask(Task *task, TaskType taskType) { { std::unique_lock lock(global_->mutex); global_->queue.push_back(task); - global_->threads_[global_->roundRobin % maxThread]->cond.notify_one(); - global_->roundRobin++; + global_->queue_size++; } + + // Lock the thread to ensure it gets the message. + int chosenIndex = global_->roundRobin++; + ThreadContext *&chosenThread = global_->threads_[chosenIndex % maxThread]; + std::unique_lock lock(chosenThread->mutex); + chosenThread->cond.notify_one(); } void ThreadManager::EnqueueTaskOnThread(int threadNum, Task *task, TaskType taskType) {