Refine channel test

paddle/framework/CMakeLists.txt

@@ -26,7 +26,7 @@ nv_test(lod_tensor_gpu_test SRCS lod_tensor_test.cu DEPS lod_tensor)
 
 cc_test(variable_test SRCS variable_test.cc)
 
-cc_library(threadpool SRCS threadpool.cc)
+cc_library(threadpool SRCS threadpool.cc DEPS enforce)
 cc_test(threadpool_test SRCS threadpool_test.cc DEPS threadpool)
 
 cc_library(scope SRCS scope.cc DEPS glog threadpool)

paddle/framework/channel.h

@@ -26,9 +26,7 @@ class Channel {
  virtual void Send(T*) = 0;
  virtual void Receive(T*) = 0;
  virtual size_t Cap() = 0;
-
- // Don't delete channels; instead, call Channel::Close.
- protected:
+ virtual void Close() = 0;
  virtual ~Channel() {}
 };
 
@@ -50,11 +48,7 @@ Channel<T>* MakeChannel(size_t buffer_size) {
 
 template <typename T>
 void CloseChannel(Channel<T>* ch) {
- if (ch->Cap() > 0) {
- delete dynamic_cast<details::Buffered<T>*>(ch);
- } else {
- delete dynamic_cast<details::UnBuffered<T>*>(ch);
- }
+ ch->Close();
 }
 
 } // namespace framework

paddle/framework/channel_test.cc

@@ -14,13 +14,67 @@ limitations under the License. */
 
 #include "paddle/framework/channel.h"
 
+#include <chrono>
+#include <thread>
+
 #include "gtest/gtest.h"
 
+using paddle::framework::Channel;
+using paddle::framework::MakeChannel;
+using paddle::framework::CloseChannel;
+
 TEST(Channel, MakeAndClose) {
- using paddle::framework::Channel;
- using paddle::framework::MakeChannel;
- using paddle::framework::CloseChannel;
+ using paddle::framework::details::Buffered;
+ using paddle::framework::details::UnBuffered;
+ {
+ // MakeChannel should return a buffered channel is buffer_size > 0.
+ auto ch = MakeChannel<int>(10);
+ EXPECT_NE(dynamic_cast<Buffered<int>*>(ch), nullptr);
+ EXPECT_EQ(dynamic_cast<UnBuffered<int>*>(ch), nullptr);
+ CloseChannel(ch);
+ delete ch;
+ }
+ {
+ // MakeChannel should return an un-buffered channel is buffer_size = 0.
+ auto ch = MakeChannel<int>(0);
+ EXPECT_EQ(dynamic_cast<Buffered<int>*>(ch), nullptr);
+ EXPECT_NE(dynamic_cast<UnBuffered<int>*>(ch), nullptr);
+ CloseChannel(ch);
+ delete ch;
+ }
+}
+
+TEST(Channel, SufficientBufferSizeDoesntBlock) {
+ const size_t buffer_size = 10;
+ auto ch = MakeChannel<size_t>(buffer_size);
+ for (size_t i = 0; i < buffer_size; ++i) {
+ ch->Send(&i); // should not block
+ }
+
+ size_t out;
+ for (size_t i = 0; i < buffer_size; ++i) {
+ ch->Receive(&out); // should not block
+ EXPECT_EQ(out, i);
+ }
+ CloseChannel(ch);
+ delete ch;
+}
+
+TEST(Channel, ConcurrentSendNonConcurrentReceiveWithSufficientBufferSize) {
+ const size_t buffer_size = 10;
+ auto ch = MakeChannel<size_t>(buffer_size);
+ size_t sum = 0;
+ std::thread t([&]() {
+ // Try to write more than buffer size.
+ for (size_t i = 0; i < 2 * buffer_size; ++i) {
+ ch->Send(&i); // should not block
+ sum += i;
+ }
+ });
+ std::this_thread::sleep_for(std::chrono::milliseconds(100)); // wait 0.5 sec
+ EXPECT_EQ(sum, 45U);
 
- Channel<int>* ch = MakeChannel<int>(10);
  CloseChannel(ch);
+ t.join();
+ delete ch;
 }

paddle/framework/details/buffered_channel.h

@@ -18,6 +18,7 @@ limitations under the License. */
 #include <mutex>
 
 #include "paddle/framework/channel.h"
+#include "paddle/platform/enforce.h"
 
 namespace paddle {
 namespace framework {
@@ -32,49 +33,68 @@ class Buffered : public paddle::framework::Channel<T> {
  virtual void Send(T*);
  virtual void Receive(T*);
  virtual size_t Cap() { return cap_; }
+ virtual void Close();
+ virtual ~Buffered();
 
  private:
  size_t cap_;
  std::mutex mu_;
  std::condition_variable empty_cond_var_;
  std::condition_variable full_cond_var_;
  std::deque<T> channel_;
+ bool closed_;
 
- Buffered(size_t cap) : cap_(cap) {}
- virtual ~Buffered();
+ Buffered(size_t cap) : cap_(cap), closed_(false) {
+ PADDLE_ENFORCE_GT(cap, 0);
+ }
 
  void NotifyAllSenders(std::unique_lock<std::mutex>*);
 };
 
 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_ || closed_; });
+ if (!closed_) {
+ channel_.push_back(std::move(*item));
+ lock.unlock();
+ empty_cond_var_.notify_one();
+ }
 }
 
 template <typename T>
 void Buffered<T>::Receive(T* item) {
  std::unique_lock<std::mutex> lock(mu_);
- empty_cond_var_.wait(lock, [this]() { return !channel_.empty(); });
- *item = std::move(channel_.front());
- channel_.pop_front();
+ empty_cond_var_.wait(lock, [this]() { return !channel_.empty() || closed_; });
+ if (!closed_) {
+ *item = std::move(channel_.front());
+ channel_.pop_front();
+ NotifyAllSenders(&lock);
+ } else {
+ item = nullptr;
+ }
+}
+
+template <typename T>
+void Buffered<T>::Close() {
+ std::unique_lock<std::mutex> lock(mu_);
+ closed_ = true;
  NotifyAllSenders(&lock);
 }
 
 template <typename T>
 Buffered<T>::~Buffered() {
  std::unique_lock<std::mutex> lock(mu_);
+ closed_ = 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

paddle/framework/details/unbuffered_channel.h

@@ -32,10 +32,11 @@ class UnBuffered : public paddle::framework::Channel<T> {
  virtual void Send(T*);
  virtual void Receive(T*);
  virtual size_t Cap() { return 0; }
+ virtual void Close();
+ virtual ~UnBuffered();
 
  private:
  UnBuffered() {}
- virtual ~UnBuffered();
 };
 
 template <typename T>
@@ -44,6 +45,9 @@ void UnBuffered<T>::Send(T* channel_element) {}
 template <typename T>
 void UnBuffered<T>::Receive(T*) {}
 
+template <typename T>
+void UnBuffered<T>::Close() {}
+
 template <typename T>
 UnBuffered<T>::~UnBuffered() {}
 

paddle/framework/threadpool.cc

@@ -14,6 +14,8 @@
 
 #include "paddle/framework/threadpool.h"
 
+#include "paddle/platform/enforce.h"
+
 namespace paddle {
 namespace framework {
 

paddle/framework/threadpool.h

@@ -22,7 +22,7 @@ limitations under the License. */
 #include <thread>
 #include <vector>
 
-#include "paddle/platform/enforce.h"
+#include "paddle/platform/macros.h" // for DISABLE_COPY_AND_ASSIGN
 
 namespace paddle {
 namespace framework {