7 #include "caffe2/core/db.h" 8 #include "caffe2/core/logging.h" 13 constexpr
size_t LMDB_MAP_SIZE = 1099511627776;
15 inline void MDB_CHECK(
int mdb_status) {
16 CAFFE_ENFORCE_EQ(mdb_status, MDB_SUCCESS, mdb_strerror(mdb_status));
22 : mdb_env_(mdb_env), valid_(
false) {
23 MDB_CHECK(mdb_txn_begin(mdb_env_, NULL, MDB_RDONLY, &mdb_txn_));
24 MDB_CHECK(mdb_dbi_open(mdb_txn_, NULL, 0, &mdb_dbi_));
25 MDB_CHECK(mdb_cursor_open(mdb_txn_, mdb_dbi_, &mdb_cursor_));
29 mdb_cursor_close(mdb_cursor_);
30 mdb_dbi_close(mdb_env_, mdb_dbi_);
31 mdb_txn_abort(mdb_txn_);
35 if (key.size() == 0) {
40 mdb_key_.mv_size = key.size();
41 mdb_key_.mv_data =
const_cast<char*
>(key.c_str());
42 int mdb_status = mdb_cursor_get(
43 mdb_cursor_, &mdb_key_, &mdb_value_, MDB_SET_RANGE);
44 if (mdb_status == MDB_NOTFOUND) {
47 MDB_CHECK(mdb_status);
52 bool SupportsSeek()
override {
return true; }
56 void Next()
override { SeekLMDB(MDB_NEXT); }
58 string key()
override {
59 return string(static_cast<const char*>(mdb_key_.mv_data), mdb_key_.mv_size);
63 return string(static_cast<const char*>(mdb_value_.mv_data),
67 bool Valid()
override {
return valid_; }
70 void SeekLMDB(MDB_cursor_op op) {
71 int mdb_status = mdb_cursor_get(mdb_cursor_, &mdb_key_, &mdb_value_, op);
72 if (mdb_status == MDB_NOTFOUND) {
75 MDB_CHECK(mdb_status);
83 MDB_cursor* mdb_cursor_;
84 MDB_val mdb_key_, mdb_value_;
92 MDB_CHECK(mdb_txn_begin(mdb_env_, NULL, 0, &mdb_txn_));
93 MDB_CHECK(mdb_dbi_open(mdb_txn_, NULL, 0, &mdb_dbi_));
96 MDB_CHECK(mdb_txn_commit(mdb_txn_));
97 mdb_dbi_close(mdb_env_, mdb_dbi_);
99 void Put(
const string&
key,
const string&
value)
override;
101 MDB_CHECK(mdb_txn_commit(mdb_txn_));
102 mdb_dbi_close(mdb_env_, mdb_dbi_);
104 MDB_CHECK(mdb_txn_begin(mdb_env_, NULL, 0, &mdb_txn_));
105 MDB_CHECK(mdb_dbi_open(mdb_txn_, NULL, 0, &mdb_dbi_));
118 LMDB(
const string& source, Mode mode);
119 virtual ~
LMDB() { Close(); }
121 if (mdb_env_ != NULL) {
122 mdb_env_close(mdb_env_);
127 return make_unique<LMDBCursor>(mdb_env_);
130 return make_unique<LMDBTransaction>(mdb_env_);
137 LMDB::LMDB(
const string& source, Mode mode) :
DB(source, mode) {
138 MDB_CHECK(mdb_env_create(&mdb_env_));
139 MDB_CHECK(mdb_env_set_mapsize(mdb_env_, LMDB_MAP_SIZE));
142 mkdir(source.c_str(), 0744), 0,
"mkdir ", source,
" failed");
146 flags = MDB_RDONLY | MDB_NOTLS | MDB_NOLOCK;
148 MDB_CHECK(mdb_env_open(mdb_env_, source.c_str(), flags, 0664));
149 VLOG(1) <<
"Opened lmdb " << source;
153 MDB_val mdb_key, mdb_value;
154 mdb_key.mv_data =
const_cast<char*
>(key.data());
155 mdb_key.mv_size = key.size();
156 mdb_value.mv_data =
const_cast<char*
>(value.data());
157 mdb_value.mv_size = value.size();
158 MDB_CHECK(mdb_put(mdb_txn_, mdb_dbi_, &mdb_key, &mdb_value, 0));
162 REGISTER_CAFFE2_DB(lmdb,
LMDB);
An abstract class for accessing a database of key-value pairs.
void Commit() override
Commits the current writes.
string value() override
Returns the current value.
bool Valid() override
Returns whether the current location is valid - for example, if we have reached the end of the databa...
An abstract class for the cursor of the database while reading.
string key() override
Returns the current key.
void Next() override
Go to the next location in the database.
void Put(const string &key, const string &value) override
Puts the key value pair to the database.
void Close() override
Closes the database.
Simple registry implementation in Caffe2 that uses static variables to register object creators durin...
void SeekToFirst() override
Seek to the first key in the database.
void Seek(const string &key) override
Seek to a specific key (or if the key does not exist, seek to the immediate next).
An abstract class for the current database transaction while writing.
unique_ptr< Transaction > NewTransaction() override
Returns a transaction to write data to the database.
unique_ptr< Cursor > NewCursor() override
Returns a cursor to read the database.