refine close channel

PaddlePaddle · Jan 30, 2018 · 3e7ca5c · 3e7ca5c
1 parent be0525f
commit 3e7ca5c
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 6 deletions.
diff --git a/paddle/framework/channel.h b/paddle/framework/channel.h
@@ -26,6 +26,7 @@ class Channel {
   virtual void Send(T*) = 0;
   virtual void Receive(T*) = 0;
   virtual size_t Cap() = 0;
+  virtual void Close() = 0;
 
   // Don't delete channels; instead, call Channel::Close.
  protected:
@@ -50,6 +51,11 @@ Channel<T>* MakeChannel(size_t buffer_size) {
 
 template <typename T>
 void CloseChannel(Channel<T>* ch) {
+  ch->Close();
+}
+
+template <typename T>
+void DeleteChannel(Channel<T>* ch) {
   if (ch->Cap() > 0) {
     delete dynamic_cast<details::Buffered<T>*>(ch);
   } else {

diff --git a/paddle/framework/channel_test.cc b/paddle/framework/channel_test.cc
@@ -72,4 +72,5 @@ TEST(Channel, ConcurrentSendNonConcurrentReceiveWithSufficientBufferSize) {
   std::this_thread::sleep_for(std::chrono::milliseconds(100));  // wait 0.5 sec
   EXPECT_EQ(sum, 45U);
   CloseChannel(ch);
+  t.join();
 }
diff --git a/paddle/framework/details/buffered_channel.h b/paddle/framework/details/buffered_channel.h
@@ -33,15 +33,17 @@ class Buffered : public paddle::framework::Channel<T> {
   virtual void Send(T*);
   virtual void Receive(T*);
   virtual size_t Cap() { return cap_; }
+  virtual void Close();
 
  private:
   size_t cap_;
   std::mutex mu_;
   std::condition_variable empty_cond_var_;
   std::condition_variable full_cond_var_;
   std::deque<T> channel_;
+  bool close;
 
-  Buffered(size_t cap) : cap_(cap) { PADDLE_ENFORCE_GT(cap, 0); }
+  Buffered(size_t cap) : cap_(cap), close(false) { PADDLE_ENFORCE_GT(cap, 0); }
   virtual ~Buffered();
 
   void NotifyAllSenders(std::unique_lock<std::mutex>*);
@@ -50,10 +52,13 @@ class Buffered : public paddle::framework::Channel<T> {
 template <typename T>
 void Buffered<T>::Send(T* item) {
   std::unique_lock<std::mutex> lock(mu_);
-  full_cond_var_.wait(lock, [this]() { return channel_.size() < cap_; });
-  channel_.push_back(std::move(*item));
-  lock.unlock();
-  empty_cond_var_.notify_one();
+  full_cond_var_.wait(lock,
+                      [this]() { return channel_.size() < cap_ || close; });
+  if (!close) {
+    channel_.push_back(std::move(*item));
+    lock.unlock();
+    empty_cond_var_.notify_one();
+  }
 }
 
 template <typename T>
@@ -65,17 +70,25 @@ void Buffered<T>::Receive(T* item) {
   NotifyAllSenders(&lock);
 }
 
+template <typename T>
+void Buffered<T>::Close() {
+  std::unique_lock<std::mutex> lock(mu_);
+  close = true;
+  NotifyAllSenders(&lock);
+}
+
 template <typename T>
 Buffered<T>::~Buffered() {
   std::unique_lock<std::mutex> lock(mu_);
+  close = true;
   channel_.clear();
   NotifyAllSenders(&lock);
 }
 
 template <typename T>
 void Buffered<T>::NotifyAllSenders(std::unique_lock<std::mutex>* lock) {
   lock->unlock();
-  full_cond_var_.notify_one();
+  full_cond_var_.notify_all();
 }
 
 }  // namespace details

diff --git a/paddle/framework/details/unbuffered_channel.h b/paddle/framework/details/unbuffered_channel.h
@@ -32,6 +32,7 @@ class UnBuffered : public paddle::framework::Channel<T> {
   virtual void Send(T*);
   virtual void Receive(T*);
   virtual size_t Cap() { return 0; }
+  virtual void Close();
 
  private:
   UnBuffered() {}
@@ -41,6 +42,9 @@ class UnBuffered : public paddle::framework::Channel<T> {
 template <typename T>
 void UnBuffered<T>::Send(T* channel_element) {}
 
+template <typename T>
+void UnBuffered<T>::Close() {}
+
 template <typename T>
 void UnBuffered<T>::Receive(T*) {}