16
16
17
17
#include < sys/time.h>
18
18
19
+ #include < chrono> // NOLINT
19
20
#include < cstdlib>
20
21
#include < fstream>
21
22
#include < iostream>
@@ -39,6 +40,11 @@ void ReadThread(const std::vector<std::string>& file_list,
39
40
int thread_id, std::vector<ReaderThreadStatus>* thread_status,
40
41
std::shared_ptr<LoDTensorBlockingQueue> queue);
41
42
43
+ // monitor all running thread, if they are all stopped,
44
+ // then push an empty data into LoDTensorBlockingQueue
45
+ void MonitorThread (std::vector<ReaderThreadStatus>* thread_status,
46
+ std::shared_ptr<LoDTensorBlockingQueue> queue);
47
+
42
48
class CTRReader : public framework ::FileReader {
43
49
public:
44
50
explicit CTRReader (const std::shared_ptr<LoDTensorBlockingQueue>& queue,
@@ -58,7 +64,7 @@ class CTRReader : public framework::FileReader {
58
64
}
59
65
}
60
66
61
- ~CTRReader () { Shutdown (); }
67
+ ~CTRReader () {}
62
68
63
69
void ReadNext (std::vector<framework::LoDTensor>* out) override {
64
70
bool success;
@@ -68,12 +74,19 @@ class CTRReader : public framework::FileReader {
68
74
69
75
void Shutdown () override {
70
76
VLOG (3 ) << " Shutdown reader" ;
77
+ if (status_ == ReaderStatus::kStopped ) {
78
+ return ;
79
+ }
71
80
// shutdown should stop all the reader thread
72
81
for (auto & read_thread : read_threads_) {
73
82
read_thread->join ();
74
83
}
84
+ monitor_thread_->join ();
85
+
75
86
read_threads_.clear ();
87
+ monitor_thread_.reset (nullptr );
76
88
queue_->Close ();
89
+ status_ = ReaderStatus::kStopped ;
77
90
}
78
91
79
92
void Start () override {
@@ -87,6 +100,9 @@ class CTRReader : public framework::FileReader {
87
100
std::bind (&ReadThread, file_groups_[thread_id], slots_, batch_size_,
88
101
thread_id, &read_thread_status_, queue_)));
89
102
}
103
+ monitor_thread_.reset (new std::thread (
104
+ std::bind (&MonitorThread, &read_thread_status_, queue_)));
105
+ status_ = ReaderStatus::kRunning ;
90
106
}
91
107
92
108
private:
@@ -107,6 +123,7 @@ class CTRReader : public framework::FileReader {
107
123
const std::vector<std::string> file_list_;
108
124
std::shared_ptr<LoDTensorBlockingQueue> queue_;
109
125
std::vector<std::unique_ptr<std::thread>> read_threads_;
126
+ std::unique_ptr<std::thread> monitor_thread_;
110
127
std::vector<ReaderThreadStatus> read_thread_status_;
111
128
std::vector<std::vector<std::string>> file_groups_;
112
129
};
0 commit comments