11 #include "caffe2/core/db.h" 12 #include "caffe2/core/logging.h" 17 constexpr
size_t LMDB_MAP_SIZE = 1099511627776;
19 inline void MDB_CHECK(
int mdb_status) {
20 CAFFE_ENFORCE_EQ(mdb_status, MDB_SUCCESS, mdb_strerror(mdb_status));
26 : mdb_env_(mdb_env), valid_(
false) {
27 MDB_CHECK(mdb_txn_begin(mdb_env_, NULL, MDB_RDONLY, &mdb_txn_));
28 MDB_CHECK(mdb_dbi_open(mdb_txn_, NULL, 0, &mdb_dbi_));
29 MDB_CHECK(mdb_cursor_open(mdb_txn_, mdb_dbi_, &mdb_cursor_));
33 mdb_cursor_close(mdb_cursor_);
34 mdb_dbi_close(mdb_env_, mdb_dbi_);
35 mdb_txn_abort(mdb_txn_);
39 if (key.size() == 0) {
44 mdb_key_.mv_size = key.size();
45 mdb_key_.mv_data =
const_cast<char*
>(key.c_str());
46 int mdb_status = mdb_cursor_get(
47 mdb_cursor_, &mdb_key_, &mdb_value_, MDB_SET_RANGE);
48 if (mdb_status == MDB_NOTFOUND) {
51 MDB_CHECK(mdb_status);
56 bool SupportsSeek()
override {
return true; }
60 void Next()
override { SeekLMDB(MDB_NEXT); }
62 string key()
override {
63 return string(static_cast<const char*>(mdb_key_.mv_data), mdb_key_.mv_size);
67 return string(static_cast<const char*>(mdb_value_.mv_data),
71 bool Valid()
override {
return valid_; }
74 void SeekLMDB(MDB_cursor_op op) {
75 int mdb_status = mdb_cursor_get(mdb_cursor_, &mdb_key_, &mdb_value_, op);
76 if (mdb_status == MDB_NOTFOUND) {
79 MDB_CHECK(mdb_status);
87 MDB_cursor* mdb_cursor_;
88 MDB_val mdb_key_, mdb_value_;
96 MDB_CHECK(mdb_txn_begin(mdb_env_, NULL, 0, &mdb_txn_));
97 MDB_CHECK(mdb_dbi_open(mdb_txn_, NULL, 0, &mdb_dbi_));
100 MDB_CHECK(mdb_txn_commit(mdb_txn_));
101 mdb_dbi_close(mdb_env_, mdb_dbi_);
103 void Put(
const string&
key,
const string&
value)
override;
105 MDB_CHECK(mdb_txn_commit(mdb_txn_));
106 mdb_dbi_close(mdb_env_, mdb_dbi_);
108 MDB_CHECK(mdb_txn_begin(mdb_env_, NULL, 0, &mdb_txn_));
109 MDB_CHECK(mdb_dbi_open(mdb_txn_, NULL, 0, &mdb_dbi_));
122 LMDB(
const string& source, Mode mode);
123 virtual ~
LMDB() { Close(); }
125 if (mdb_env_ != NULL) {
126 mdb_env_close(mdb_env_);
131 return make_unique<LMDBCursor>(mdb_env_);
134 return make_unique<LMDBTransaction>(mdb_env_);
141 LMDB::LMDB(
const string& source, Mode mode) :
DB(source, mode) {
142 MDB_CHECK(mdb_env_create(&mdb_env_));
143 MDB_CHECK(mdb_env_set_mapsize(mdb_env_, LMDB_MAP_SIZE));
145 #if defined(_MSC_VER) 146 CAFFE_ENFORCE_EQ(_mkdir(source.c_str()), 0,
"mkdir ", source,
" failed");
149 mkdir(source.c_str(), 0744), 0,
"mkdir ", source,
" failed");
154 flags = MDB_RDONLY | MDB_NOTLS | MDB_NOLOCK;
156 MDB_CHECK(mdb_env_open(mdb_env_, source.c_str(), flags, 0664));
157 VLOG(1) <<
"Opened lmdb " << source;
161 MDB_val mdb_key, mdb_value;
162 mdb_key.mv_data =
const_cast<char*
>(key.data());
163 mdb_key.mv_size = key.size();
164 mdb_value.mv_data =
const_cast<char*
>(value.data());
165 mdb_value.mv_size = value.size();
166 MDB_CHECK(mdb_put(mdb_txn_, mdb_dbi_, &mdb_key, &mdb_value, 0));
170 REGISTER_CAFFE2_DB(lmdb,
LMDB);
void Put(const string &key, const string &value) override
Puts the key value pair to the database.
An abstract class for the current database transaction while writing.
An abstract class for the cursor of the database while reading.
void SeekToFirst() override
Seek to the first key in the database.
void Close() override
Closes the database.
void Commit() override
Commits the current writes.
string key() override
Returns the current key.
string value() override
Returns the current value.
void Seek(const string &key) override
Seek to a specific key (or if the key does not exist, seek to the immediate next).
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.
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
An abstract class for accessing a database of key-value pairs.
void Next() override
Go to the next location in the database.
bool Valid() override
Returns whether the current location is valid - for example, if we have reached the end of the databa...