Caffe2 - C++ API
A deep learning, cross platform ML framework
lmdb.cc
1 
17 #include "lmdb.h" // NOLINT
18 
19 #if defined(_MSC_VER)
20 #include <direct.h>
21 #endif
22 
23 #include <sys/stat.h>
24 
25 #include <string>
26 
27 #include "caffe2/core/db.h"
28 #include "caffe2/core/logging.h"
29 
30 namespace caffe2 {
31 namespace db {
32 
33 constexpr size_t LMDB_MAP_SIZE = 1099511627776; // 1 TB
34 
35 inline void MDB_CHECK(int mdb_status) {
36  CAFFE_ENFORCE_EQ(mdb_status, MDB_SUCCESS, mdb_strerror(mdb_status));
37 }
38 
39 class LMDBCursor : public Cursor {
40  public:
41  explicit LMDBCursor(MDB_env* mdb_env)
42  : mdb_env_(mdb_env), valid_(false) {
43  MDB_CHECK(mdb_txn_begin(mdb_env_, NULL, MDB_RDONLY, &mdb_txn_));
44  MDB_CHECK(mdb_dbi_open(mdb_txn_, NULL, 0, &mdb_dbi_));
45  MDB_CHECK(mdb_cursor_open(mdb_txn_, mdb_dbi_, &mdb_cursor_));
46  SeekToFirst();
47  }
48  virtual ~LMDBCursor() {
49  mdb_cursor_close(mdb_cursor_);
50  mdb_dbi_close(mdb_env_, mdb_dbi_);
51  mdb_txn_abort(mdb_txn_);
52  }
53 
54  void Seek(const string& key) override {
55  if (key.size() == 0) {
56  SeekToFirst();
57  return;
58  }
59  // a key of 16k size should be enough? I am not sure though.
60  mdb_key_.mv_size = key.size();
61  mdb_key_.mv_data = const_cast<char*>(key.c_str());
62  int mdb_status = mdb_cursor_get(
63  mdb_cursor_, &mdb_key_, &mdb_value_, MDB_SET_RANGE);
64  if (mdb_status == MDB_NOTFOUND) {
65  valid_ = false;
66  } else {
67  MDB_CHECK(mdb_status);
68  valid_ = true;
69  }
70  }
71 
72  bool SupportsSeek() override { return true; }
73 
74  void SeekToFirst() override { SeekLMDB(MDB_FIRST); }
75 
76  void Next() override { SeekLMDB(MDB_NEXT); }
77 
78  string key() override {
79  return string(static_cast<const char*>(mdb_key_.mv_data), mdb_key_.mv_size);
80  }
81 
82  string value() override {
83  return string(static_cast<const char*>(mdb_value_.mv_data),
84  mdb_value_.mv_size);
85  }
86 
87  bool Valid() override { return valid_; }
88 
89  private:
90  void SeekLMDB(MDB_cursor_op op) {
91  int mdb_status = mdb_cursor_get(mdb_cursor_, &mdb_key_, &mdb_value_, op);
92  if (mdb_status == MDB_NOTFOUND) {
93  valid_ = false;
94  } else {
95  MDB_CHECK(mdb_status);
96  valid_ = true;
97  }
98  }
99 
100  MDB_env* mdb_env_;
101  MDB_txn* mdb_txn_;
102  MDB_dbi mdb_dbi_;
103  MDB_cursor* mdb_cursor_;
104  MDB_val mdb_key_, mdb_value_;
105  bool valid_;
106 };
107 
108 class LMDBTransaction final : public Transaction {
109  public:
110  explicit LMDBTransaction(MDB_env* mdb_env)
111  : mdb_env_(mdb_env) {
112  MDB_CHECK(mdb_txn_begin(mdb_env_, NULL, 0, &mdb_txn_));
113  MDB_CHECK(mdb_dbi_open(mdb_txn_, NULL, 0, &mdb_dbi_));
114  }
115  ~LMDBTransaction() {
116  MDB_CHECK(mdb_txn_commit(mdb_txn_));
117  mdb_dbi_close(mdb_env_, mdb_dbi_);
118  }
119  void Put(const string& key, const string& value) override;
120  void Commit() override {
121  MDB_CHECK(mdb_txn_commit(mdb_txn_));
122  mdb_dbi_close(mdb_env_, mdb_dbi_);
123  // Begin a new transaction.
124  MDB_CHECK(mdb_txn_begin(mdb_env_, NULL, 0, &mdb_txn_));
125  MDB_CHECK(mdb_dbi_open(mdb_txn_, NULL, 0, &mdb_dbi_));
126  }
127 
128  private:
129  MDB_env* mdb_env_;
130  MDB_dbi mdb_dbi_;
131  MDB_txn* mdb_txn_;
132 
133  DISABLE_COPY_AND_ASSIGN(LMDBTransaction);
134 };
135 
136 class LMDB : public DB {
137  public:
138  LMDB(const string& source, Mode mode);
139  virtual ~LMDB() { Close(); }
140  void Close() override {
141  if (mdb_env_ != NULL) {
142  mdb_env_close(mdb_env_);
143  mdb_env_ = NULL;
144  }
145  }
146  unique_ptr<Cursor> NewCursor() override {
147  return make_unique<LMDBCursor>(mdb_env_);
148  }
149  unique_ptr<Transaction> NewTransaction() override {
150  return make_unique<LMDBTransaction>(mdb_env_);
151  }
152 
153  private:
154  MDB_env* mdb_env_;
155 };
156 
157 LMDB::LMDB(const string& source, Mode mode) : DB(source, mode) {
158  MDB_CHECK(mdb_env_create(&mdb_env_));
159  MDB_CHECK(mdb_env_set_mapsize(mdb_env_, LMDB_MAP_SIZE));
160  if (mode == NEW) {
161 #if defined(_MSC_VER)
162  CAFFE_ENFORCE_EQ(_mkdir(source.c_str()), 0, "mkdir ", source, " failed");
163 #else
164  CAFFE_ENFORCE_EQ(
165  mkdir(source.c_str(), 0744), 0, "mkdir ", source, " failed");
166 #endif
167  }
168  int flags = 0;
169  if (mode == READ) {
170  flags = MDB_RDONLY | MDB_NOTLS | MDB_NOLOCK;
171  }
172  MDB_CHECK(mdb_env_open(mdb_env_, source.c_str(), flags, 0664));
173  VLOG(1) << "Opened lmdb " << source;
174 }
175 
176 void LMDBTransaction::Put(const string& key, const string& value) {
177  MDB_val mdb_key, mdb_value;
178  mdb_key.mv_data = const_cast<char*>(key.data());
179  mdb_key.mv_size = key.size();
180  mdb_value.mv_data = const_cast<char*>(value.data());
181  mdb_value.mv_size = value.size();
182  MDB_CHECK(mdb_put(mdb_txn_, mdb_dbi_, &mdb_key, &mdb_value, 0));
183 }
184 
185 REGISTER_CAFFE2_DB(LMDB, LMDB);
186 REGISTER_CAFFE2_DB(lmdb, LMDB);
187 
188 } // namespace db
189 } // namespace caffe2
void Put(const string &key, const string &value) override
Puts the key value pair to the database.
Definition: lmdb.cc:176
An abstract class for the current database transaction while writing.
Definition: db.h:77
An abstract class for the cursor of the database while reading.
Definition: db.h:38
void SeekToFirst() override
Seek to the first key in the database.
Definition: lmdb.cc:74
void Close() override
Closes the database.
Definition: lmdb.cc:140
void Commit() override
Commits the current writes.
Definition: lmdb.cc:120
string key() override
Returns the current key.
Definition: lmdb.cc:78
string value() override
Returns the current value.
Definition: lmdb.cc:82
void Seek(const string &key) override
Seek to a specific key (or if the key does not exist, seek to the immediate next).
Definition: lmdb.cc:54
unique_ptr< Transaction > NewTransaction() override
Returns a transaction to write data to the database.
Definition: lmdb.cc:149
unique_ptr< Cursor > NewCursor() override
Returns a cursor to read the database.
Definition: lmdb.cc:146
Copyright (c) 2016-present, Facebook, Inc.
An abstract class for accessing a database of key-value pairs.
Definition: db.h:96
void Next() override
Go to the next location in the database.
Definition: lmdb.cc:76
bool Valid() override
Returns whether the current location is valid - for example, if we have reached the end of the databa...
Definition: lmdb.cc:87