From f4a76078d033320576490c436e9d7f5796dade90 Mon Sep 17 00:00:00 2001
From: Qiao Longfei <qiaolongfei@baidu.com>
Date: Tue, 6 Nov 2018 11:07:51 +0800
Subject: [PATCH] optimize thread pool

---
 paddle/fluid/framework/threadpool.cc |  8 ++++++--
 paddle/fluid/framework/threadpool.h  | 11 ++++++++---
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/paddle/fluid/framework/threadpool.cc b/paddle/fluid/framework/threadpool.cc
index a588cb417ae..a471c83115f 100644
--- a/paddle/fluid/framework/threadpool.cc
+++ b/paddle/fluid/framework/threadpool.cc
@@ -59,8 +59,8 @@ ThreadPool::~ThreadPool() {
     // notify all threads to stop running
     std::lock_guard<std::mutex> l(mutex_);
     running_ = false;
-    scheduled_.notify_all();
   }
+  scheduled_.notify_all();
 
   for (auto& t : threads_) {
     t->join();
@@ -75,10 +75,14 @@ void ThreadPool::TaskLoop() {
     scheduled_.wait(
         lock, [this] { return !this->tasks_.empty() || !this->running_; });
 
-    if (!running_ || tasks_.empty()) {
+    if (!running_ && tasks_.empty()) {
       return;
     }
 
+    if (tasks_.empty()) {
+      PADDLE_THROW("This thread has no task to Run");
+    }
+
     // pop a task from the task queue
     auto task = std::move(tasks_.front());
     tasks_.pop();
diff --git a/paddle/fluid/framework/threadpool.h b/paddle/fluid/framework/threadpool.h
index 0687e628aaa..7a51d18fbbf 100644
--- a/paddle/fluid/framework/threadpool.h
+++ b/paddle/fluid/framework/threadpool.h
@@ -58,7 +58,7 @@ class ThreadPool {
   ~ThreadPool();
 
   // Run pushes a function to the task queue and returns a std::future
-  // object.  To wait for the completion of the task, call
+  // object. To wait for the completion of the task, call
   // std::future::wait().
   template <typename Callback>
   std::future<void> Run(Callback fn) {
@@ -69,7 +69,6 @@ class ThreadPool {
   template <typename Callback>
   std::future<std::unique_ptr<platform::EnforceNotMet>> RunAndGetException(
       Callback fn) {
-    std::unique_lock<std::mutex> lock(mutex_);
     Task task([fn]() -> std::unique_ptr<platform::EnforceNotMet> {
       try {
         fn();
@@ -84,7 +83,13 @@ class ThreadPool {
       return nullptr;
     });
     std::future<std::unique_ptr<platform::EnforceNotMet>> f = task.get_future();
-    tasks_.push(std::move(task));
+    {
+      std::unique_lock<std::mutex> lock(mutex_);
+      if (!running_) {
+        PADDLE_THROW("enqueue on stopped ThreadPool");
+      }
+      tasks_.push(std::move(task));
+    }
     scheduled_.notify_one();
     return f;
   }
-- 
GitLab