1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
|
#include <atomic>
#include <condition_variable>
#include <mutex>
#include <thread> // NOLINT
#include "caffe2/core/db.h"
#include "caffe2/core/logging.h"
#include "caffe2/utils/zmq_helper.h"
namespace caffe2 {
namespace db {
class ZmqDBCursor : public Cursor {
public:
explicit ZmqDBCursor(const string& source)
: source_(source),
socket_(ZMQ_PULL),
prefetched_(false),
finalize_(false) {
socket_.Connect(source_);
// Start prefetching thread.
prefetch_thread_.reset(new std::thread([this] { this->Prefetch(); }));
// obtain the first value.
Next();
}
~ZmqDBCursor() override {
finalize_ = true;
prefetched_ = false;
producer_.notify_one();
// Wait for the prefetch thread to finish elegantly.
prefetch_thread_->join();
socket_.Disconnect(source_);
}
void Seek(const string& /*key*/) override { /* do nothing */
}
void SeekToFirst() override { /* do nothing */
}
void Next() override {
std::unique_lock<std::mutex> lock(prefetch_access_mutex_);
while (!prefetched_)
consumer_.wait(lock);
key_ = prefetch_key_;
value_ = prefetch_value_;
prefetched_ = false;
producer_.notify_one();
}
string key() override {
return key_;
}
string value() override {
return value_;
}
bool Valid() override {
return true;
}
private:
void Prefetch() {
while (!finalize_) {
std::unique_lock<std::mutex> lock(prefetch_access_mutex_);
while (prefetched_)
producer_.wait(lock);
if (finalize_) {
return;
}
ZmqMessage msg;
socket_.RecvTillSuccess(&msg);
prefetch_key_.assign(static_cast<char*>(msg.data()), msg.size());
socket_.RecvTillSuccess(&msg);
prefetch_value_.assign(static_cast<char*>(msg.data()), msg.size());
prefetched_ = true;
consumer_.notify_one();
}
}
string source_;
ZmqSocket socket_;
string key_;
string value_;
string prefetch_key_;
string prefetch_value_;
unique_ptr<std::thread> prefetch_thread_;
std::mutex prefetch_access_mutex_;
std::condition_variable producer_, consumer_;
std::atomic<bool> prefetched_;
// finalize_ is used to tell the prefetcher to quit.
std::atomic<bool> finalize_;
};
class ZmqDB : public DB {
public:
ZmqDB(const string& source, Mode mode) : DB(source, mode), source_(source) {
CAFFE_ENFORCE(mode == READ, "ZeroMQ DB only supports read mode.");
}
~ZmqDB() override {}
void Close() override {}
unique_ptr<Cursor> NewCursor() override {
return make_unique<ZmqDBCursor>(source_);
}
unique_ptr<Transaction> NewTransaction() override {
CAFFE_THROW("ZeroMQ DB does not support writing with a transaction.");
return nullptr; // dummy placeholder to suppress old compiler warnings.
}
private:
string source_;
};
REGISTER_CAFFE2_DB(ZmqDB, ZmqDB);
// For lazy-minded, one can also call with lower-case name.
REGISTER_CAFFE2_DB(zmqdb, ZmqDB);
} // namespace db
} // namespace caffe2
|
