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
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
|
#include <cstdio>
#include <iomanip>
#include <sstream>
#include <thread>
#include <gtest/gtest.h>
#include "caffe2/core/blob_serialization.h"
#include "caffe2/core/db.h"
#include "caffe2/core/logging.h"
#include "caffe2/proto/caffe2_pb.h"
#include "common/gtest/gtest_extensions.h"
namespace caffe2 {
namespace db {
constexpr int kMaxItems = 10;
static bool CreateAndFill(const string& db_type, const string& name) {
VLOG(1) << "Creating db: " << name;
std::unique_ptr<DB> db(CreateDB(db_type, name, NEW));
if (!db.get()) {
LOG(ERROR) << "Cannot create db of type " << db_type;
return false;
}
std::unique_ptr<Transaction> trans(db->NewTransaction());
for (int i = 0; i < kMaxItems; ++i) {
std::stringstream ss;
ss << std::setw(2) << std::setfill('0') << i;
trans->Put(ss.str(), std::string{ss.str()});
}
trans->Commit();
trans.reset();
db.reset();
return true;
}
static void TestCursor(Cursor* cursor) {
// Test the first key.
cursor->SeekToFirst();
EXPECT_EQ(cursor->key(), "00");
EXPECT_EQ(cursor->value(), "00");
// Test if Next() works.
cursor->Next();
EXPECT_EQ(cursor->key(), "01");
cursor->Next();
EXPECT_EQ(cursor->key(), "02");
// Test if we can return to the first key.
cursor->SeekToFirst();
EXPECT_EQ(cursor->key(), "00");
// Test seeking to a key that exists.
cursor->Seek("05");
EXPECT_EQ(cursor->key(), "05");
// Test seeking to a key that does not exist - that should give us the
// immediate next key.
cursor->Seek("07.5");
EXPECT_EQ(cursor->key(), "08");
// Test seeking over the end of the db - that should make the current
// cursor invalid.
cursor->Seek("11");
EXPECT_FALSE(cursor->Valid());
// Test seeking to empty string, aka the beginning
cursor->Seek("");
EXPECT_EQ(cursor->key(), "00");
}
static void DBSeekTestWrapper(const string& db_type) {
std::string name = std::tmpnam(nullptr);
if (!CreateAndFill(db_type, name)) {
// Manually fail the test, and not do anything onwards.
EXPECT_TRUE(0);
} else {
std::unique_ptr<DB> db(CreateDB(db_type, name, READ));
std::unique_ptr<Cursor> cursor(db->NewCursor());
TestCursor(cursor.get());
}
}
TEST(DBSeekTest, RocksDB) {
SKIP() << "The test is broken. So skip.";
DBSeekTestWrapper("rocksdb");
}
TEST(DBSeekTest, LevelDB) {
DBSeekTestWrapper("leveldb");
}
TEST(DBSeekTest, LMDB) {
SKIP() << "The test is broken. So skip.";
DBSeekTestWrapper("lmdb");
}
TEST(DBReaderTest, Reader) {
std::string name = std::tmpnam(nullptr);
CreateAndFill("leveldb", name);
std::unique_ptr<DBReader> reader(new DBReader("leveldb", name));
EXPECT_TRUE(reader->cursor() != nullptr);
// DBReader should have a full-fledged cursor.
TestCursor(reader->cursor());
// Test the Read() functionality.
reader->cursor()->Seek("05");
EXPECT_EQ(reader->cursor()->key(), "05");
string key;
string value;
reader->Read(&key, &value);
EXPECT_EQ(key, "05");
EXPECT_EQ(value, "05");
reader->Read(&key, &value);
EXPECT_EQ(key, "06");
EXPECT_EQ(value, "06");
// Test if we are able to serialize it using the blob serialization
// interface.
reader->cursor()->Seek("05");
EXPECT_EQ(reader->cursor()->key(), "05");
Blob reader_blob;
reader_blob.Reset(reader.release());
std::string str = SerializeBlob(reader_blob, "saved_reader");
// Release to close the old reader.
reader_blob.Reset();
BlobProto blob_proto;
CHECK(blob_proto.ParseFromString(str));
EXPECT_EQ(blob_proto.name(), "saved_reader");
EXPECT_EQ(blob_proto.type(), "DBReader");
DBReaderProto proto;
CHECK(proto.ParseFromString(blob_proto.content()));
EXPECT_EQ(proto.source(), name);
EXPECT_EQ(proto.db_type(), "leveldb");
EXPECT_EQ(proto.key(), "05");
// Test restoring the reader from the serialized proto.
EXPECT_NO_THROW(DeserializeBlob(str, &reader_blob));
EXPECT_TRUE(reader_blob.IsType<DBReader>());
const DBReader& new_reader = reader_blob.Get<DBReader>();
EXPECT_TRUE(new_reader.cursor() != nullptr);
EXPECT_EQ(new_reader.cursor()->key(), "05");
// Test Reader's multi-threading capability.
vector<unique_ptr<std::thread>> threads(kMaxItems);
vector<string> keys(kMaxItems);
vector<string> values(kMaxItems);
for (int i = 0; i < kMaxItems; ++i) {
threads[i].reset(new std::thread(
[&new_reader](string* key, string* value) {
new_reader.Read(key, value);
},
&keys[i],
&values[i]));
}
for (int i = 0; i < kMaxItems; ++i) {
threads[i]->join();
EXPECT_TRUE(keys[i].size() > 0);
}
// Check if the names are all unique by putting them into a set and
// checking the size.
std::set<string> keys_set(keys.begin(), keys.end());
EXPECT_EQ(keys_set.size(), kMaxItems);
}
TEST(DBReaderShardedTest, Reader) {
std::string name = std::tmpnam(nullptr);
CreateAndFill("leveldb", name);
std::unique_ptr<DBReader> reader0(new DBReader("leveldb", name, 3, 0));
string key;
string value;
reader0->Read(&key, &value);
EXPECT_EQ(key, "00");
EXPECT_EQ(value, "00");
reader0->Read(&key, &value);
EXPECT_EQ(key, "03");
EXPECT_EQ(value, "03");
reader0->Read(&key, &value);
EXPECT_EQ(key, "06");
EXPECT_EQ(value, "06");
reader0->Read(&key, &value);
EXPECT_EQ(key, "09");
EXPECT_EQ(value, "09");
reader0->Read(&key, &value);
EXPECT_EQ(key, "00");
EXPECT_EQ(value, "00");
reader0->Read(&key, &value);
EXPECT_EQ(key, "03");
EXPECT_EQ(value, "03");
CreateAndFill("leveldb", name + "1");
std::unique_ptr<DBReader> reader1(new DBReader("leveldb", name + "1", 3, 1));
reader1->Read(&key, &value);
EXPECT_EQ(key, "01");
EXPECT_EQ(value, "01");
reader1->Read(&key, &value);
EXPECT_EQ(key, "04");
EXPECT_EQ(value, "04");
CreateAndFill("leveldb", name + "2");
std::unique_ptr<DBReader> reader2(new DBReader("leveldb", name + "2", 3, 2));
reader2->Read(&key, &value);
EXPECT_EQ(key, "02");
EXPECT_EQ(value, "02");
reader2->Read(&key, &value);
EXPECT_EQ(key, "05");
EXPECT_EQ(value, "05");
}
} // namespace db
} // namespace caffe2
|
