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_test.cc

@@ -14,13 +14,64 @@ 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);
+  }
+  {
+    // 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);
+  }
+}
+
+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);
+}
+
+TEST(Channel, ConcurrentSendNonConcurrentReceiveWithSufficientBufferSize) {
+  const size_t buffer_size = 10;
+  auto ch = MakeChannel<size_t>(buffer_size);
 
-  Channel<int>* ch = MakeChannel<int>(10);
+  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);
   CloseChannel(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 {
@@ -40,7 +41,7 @@ class Buffered : public paddle::framework::Channel<T> {
   std::condition_variable full_cond_var_;
   std::deque<T> channel_;
 
-  Buffered(size_t cap) : cap_(cap) {}
+ Buffered(size_t cap) : cap_(cap) { PADDLE_ENFORCE_GT(cap, 0); }
   virtual ~Buffered();
 
   void NotifyAllSenders(std::unique_lock<std::mutex>*);

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 {