add blocking_queue::wait_for_empty for blocking until done

BVLC · Mar 20, 2015 · cb80038 · cb80038
1 parent ede1f2a
commit cb80038
Showing 1 changed file with 12 additions and 1 deletion.
diff --git a/include/caffe/util/blocking_queue.hpp b/include/caffe/util/blocking_queue.hpp
@@ -1,8 +1,8 @@
 #ifndef CAFFE_UTIL_BLOCKING_QUEUE_H_
 #define CAFFE_UTIL_BLOCKING_QUEUE_H_
 
-#include <queue>
 #include <boost/thread.hpp>
+#include <queue>
 
 namespace caffe {
 
@@ -24,10 +24,20 @@ class blocking_queue {
     return queue_.empty();
   }
 
+  void wait_for_empty() {
+    boost::mutex::scoped_lock lock(mutex_);
+    while (!queue_.empty()) {
+      cond_empty_.wait(lock);
+    }
+  }
+
   T pop() {
     T t = peek();
     boost::mutex::scoped_lock lock(mutex_);
     queue_.pop();
+    if (queue_.empty()) {
+      cond_empty_.notify_all();
+    }
     return t;
   }
 
@@ -44,6 +54,7 @@ class blocking_queue {
   std::queue<T> queue_;
   mutable boost::mutex mutex_;
   boost::condition_variable cond_push_;
+  boost::condition_variable cond_empty_;
 
   DISABLE_COPY_AND_ASSIGN(blocking_queue);
 };