From 8fed780f14bf24954300ba37cebd2338ee7d199c Mon Sep 17 00:00:00 2001
From: fengjiayi <fengjiayi@baidu.com>
Date: Sun, 8 Apr 2018 11:26:06 +0800
Subject: [PATCH] Complete threaded reader

---
 paddle/fluid/operators/reader/CMakeLists.txt      |  1 +
 .../operators/reader/create_threaded_reader_op.cc | 15 ++++++++++++++-
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/paddle/fluid/operators/reader/CMakeLists.txt b/paddle/fluid/operators/reader/CMakeLists.txt
index 6fa0195b9ae..845528860f9 100644
--- a/paddle/fluid/operators/reader/CMakeLists.txt
+++ b/paddle/fluid/operators/reader/CMakeLists.txt
@@ -22,5 +22,6 @@ reader_library(create_batch_reader_op SRCS create_batch_reader_op.cc)
 reader_library(create_recordio_file_reader_op SRCS create_recordio_file_reader_op.cc)
 reader_library(create_double_buffer_reader_op SRCS create_double_buffer_reader_op.cc)
 reader_library(create_multi_pass_reader_op SRCS create_multi_pass_reader_op.cc)
+reader_library(create_threaded_reader_op SRCS create_threaded_reader_op.cc)
 # Export local libraries to parent
 set(READER_LIBRARY ${LOCAL_READER_LIBS} PARENT_SCOPE)
diff --git a/paddle/fluid/operators/reader/create_threaded_reader_op.cc b/paddle/fluid/operators/reader/create_threaded_reader_op.cc
index a4aebafa8bb..489866ca801 100644
--- a/paddle/fluid/operators/reader/create_threaded_reader_op.cc
+++ b/paddle/fluid/operators/reader/create_threaded_reader_op.cc
@@ -57,7 +57,15 @@ class ThreadedReader : public framework::DecoratedReader {
     return !threda_buffer.empty();
   }
 
-  void ReInit() override;
+  void ReInit() override {
+    if (!unsafe_mode_) {
+      PADDLE_THROW(
+          "ThreadedReader::ReInit() is disabled when 'unsafe_mode' is false.");
+    }
+    VLOG(5) << "ThreadedReader::ReInit() is invoked! It might be buggy in "
+               "multi-thread environment.";
+    reader_->ReInit();
+  }
 
   ~ThreadedReader() {
     for (auto& p : thread_buffers_) {
@@ -123,3 +131,8 @@ class CreateThreadedReaderOpMaker : public DecoratedReaderMakerBase {
 }  // namespace reader
 }  // namespace operators
 }  // namespace paddle
+
+namespace reader = paddle::operators::reader;
+REGISTER_FILE_READER_OPERATOR(create_threaded_reader,
+                              reader::CreateThreadedReaderOp,
+                              reader::CreateThreadedReaderOpMaker);
-- 
GitLab