diff --git a/include/caffe/util/blocking_queue.hpp b/include/caffe/util/blocking_queue.hpp new file mode 100644 index 00000000000..26f5645e32e --- /dev/null +++ b/include/caffe/util/blocking_queue.hpp @@ -0,0 +1,64 @@ +#ifndef CAFFE_UTIL_BLOCKING_QUEUE_H_ +#define CAFFE_UTIL_BLOCKING_QUEUE_H_ + +#include +#include + +namespace caffe { + +template +class blocking_queue { + public: + explicit blocking_queue() { } + virtual ~blocking_queue() { } + + void push(const T& t) { + boost::mutex::scoped_lock lock(mutex_); + queue_.push(t); + lock.unlock(); + cond_push_.notify_one(); + } + + bool empty() const { + boost::mutex::scoped_lock lock(mutex_); + 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; + } + + // Return element without removing it + T peek() { + boost::mutex::scoped_lock lock(mutex_); + while (queue_.empty()) { + cond_push_.wait(lock); + } + return queue_.front(); + } + + private: + std::queue queue_; + mutable boost::mutex mutex_; + boost::condition_variable cond_push_; + boost::condition_variable cond_empty_; + + DISABLE_COPY_AND_ASSIGN(blocking_queue); +}; + +} // namespace caffe + +#endif