Split db.hpp into leveldb_db.hpp and lmdb_db.hpp

3 files changed, 74 insertions, 55 deletions

diff --git a/src/caffe/util/db.cpp b/src/caffe/util/db.cpp
index 7f701810..f55420e9 100644
--- a/src/caffe/util/db.cpp
+++ b/src/caffe/util/db.cpp
@@ -1,64 +1,11 @@
 #include "caffe/util/db.hpp"
+#include "caffe/util/db_leveldb.hpp"
+#include "caffe/util/db_lmdb.hpp"
 
-#include <sys/stat.h>
 #include <string>
 
 namespace caffe { namespace db {
 
-const size_t LMDB_MAP_SIZE = 1099511627776;  // 1 TB
-
-void LevelDB::Open(const string& source, Mode mode) {
-  leveldb::Options options;
-  options.block_size = 65536;
-  options.write_buffer_size = 268435456;
-  options.max_open_files = 100;
-  options.error_if_exists = mode == NEW;
-  options.create_if_missing = mode != READ;
-  leveldb::Status status = leveldb::DB::Open(options, source, &db_);
-  CHECK(status.ok()) << "Failed to open leveldb " << source
-                     << std::endl << status.ToString();
-  LOG(INFO) << "Opened leveldb " << source;
-}
-
-void LMDB::Open(const string& source, Mode mode) {
-  MDB_CHECK(mdb_env_create(&mdb_env_));
-  MDB_CHECK(mdb_env_set_mapsize(mdb_env_, LMDB_MAP_SIZE));
-  if (mode == NEW) {
-    CHECK_EQ(mkdir(source.c_str(), 0744), 0) << "mkdir " << source << "failed";
-  }
-  int flags = 0;
-  if (mode == READ) {
-    flags = MDB_RDONLY | MDB_NOTLS;
-  }
-  MDB_CHECK(mdb_env_open(mdb_env_, source.c_str(), flags, 0664));
-  LOG(INFO) << "Opened lmdb " << source;
-}
-
-LMDBCursor* LMDB::NewCursor() {
-  MDB_txn* mdb_txn;
-  MDB_cursor* mdb_cursor;
-  MDB_CHECK(mdb_txn_begin(mdb_env_, NULL, MDB_RDONLY, &mdb_txn));
-  MDB_CHECK(mdb_dbi_open(mdb_txn, NULL, 0, &mdb_dbi_));
-  MDB_CHECK(mdb_cursor_open(mdb_txn, mdb_dbi_, &mdb_cursor));
-  return new LMDBCursor(mdb_txn, mdb_cursor);
-}
-
-LMDBTransaction* LMDB::NewTransaction() {
-  MDB_txn* mdb_txn;
-  MDB_CHECK(mdb_txn_begin(mdb_env_, NULL, 0, &mdb_txn));
-  MDB_CHECK(mdb_dbi_open(mdb_txn, NULL, 0, &mdb_dbi_));
-  return new LMDBTransaction(&mdb_dbi_, mdb_txn);
-}
-
-void LMDBTransaction::Put(const string& key, const string& value) {
-  MDB_val mdb_key, mdb_value;
-  mdb_key.mv_data = const_cast<char*>(key.data());
-  mdb_key.mv_size = key.size();
-  mdb_value.mv_data = const_cast<char*>(value.data());
-  mdb_value.mv_size = value.size();
-  MDB_CHECK(mdb_put(mdb_txn_, *mdb_dbi_, &mdb_key, &mdb_value, 0));
-}
-
 DB* GetDB(DataParameter::DB backend) {
   switch (backend) {
   case DataParameter_DB_LEVELDB:
diff --git a/src/caffe/util/db_leveldb.cpp b/src/caffe/util/db_leveldb.cpp
new file mode 100644
index 00000000..06c46627
--- /dev/null
+++ b/src/caffe/util/db_leveldb.cpp
@@ -0,0 +1,21 @@
+#include "caffe/util/db_leveldb.hpp"
+
+#include <string>
+
+namespace caffe { namespace db {
+
+void LevelDB::Open(const string& source, Mode mode) {
+  leveldb::Options options;
+  options.block_size = 65536;
+  options.write_buffer_size = 268435456;
+  options.max_open_files = 100;
+  options.error_if_exists = mode == NEW;
+  options.create_if_missing = mode != READ;
+  leveldb::Status status = leveldb::DB::Open(options, source, &db_);
+  CHECK(status.ok()) << "Failed to open leveldb " << source
+                     << std::endl << status.ToString();
+  LOG(INFO) << "Opened leveldb " << source;
+}
+
+}  // namespace db
+}  // namespace caffe
diff --git a/src/caffe/util/db_lmdb.cpp b/src/caffe/util/db_lmdb.cpp
new file mode 100644
index 00000000..a054b796
--- /dev/null
+++ b/src/caffe/util/db_lmdb.cpp
@@ -0,0 +1,51 @@
+#include "caffe/util/db_lmdb.hpp"
+
+#include <sys/stat.h>
+
+#include <string>
+
+namespace caffe { namespace db {
+
+const size_t LMDB_MAP_SIZE = 1099511627776;  // 1 TB
+
+void LMDB::Open(const string& source, Mode mode) {
+  MDB_CHECK(mdb_env_create(&mdb_env_));
+  MDB_CHECK(mdb_env_set_mapsize(mdb_env_, LMDB_MAP_SIZE));
+  if (mode == NEW) {
+    CHECK_EQ(mkdir(source.c_str(), 0744), 0) << "mkdir " << source << "failed";
+  }
+  int flags = 0;
+  if (mode == READ) {
+    flags = MDB_RDONLY | MDB_NOTLS;
+  }
+  MDB_CHECK(mdb_env_open(mdb_env_, source.c_str(), flags, 0664));
+  LOG(INFO) << "Opened lmdb " << source;
+}
+
+LMDBCursor* LMDB::NewCursor() {
+  MDB_txn* mdb_txn;
+  MDB_cursor* mdb_cursor;
+  MDB_CHECK(mdb_txn_begin(mdb_env_, NULL, MDB_RDONLY, &mdb_txn));
+  MDB_CHECK(mdb_dbi_open(mdb_txn, NULL, 0, &mdb_dbi_));
+  MDB_CHECK(mdb_cursor_open(mdb_txn, mdb_dbi_, &mdb_cursor));
+  return new LMDBCursor(mdb_txn, mdb_cursor);
+}
+
+LMDBTransaction* LMDB::NewTransaction() {
+  MDB_txn* mdb_txn;
+  MDB_CHECK(mdb_txn_begin(mdb_env_, NULL, 0, &mdb_txn));
+  MDB_CHECK(mdb_dbi_open(mdb_txn, NULL, 0, &mdb_dbi_));
+  return new LMDBTransaction(&mdb_dbi_, mdb_txn);
+}
+
+void LMDBTransaction::Put(const string& key, const string& value) {
+  MDB_val mdb_key, mdb_value;
+  mdb_key.mv_data = const_cast<char*>(key.data());
+  mdb_key.mv_size = key.size();
+  mdb_value.mv_data = const_cast<char*>(value.data());
+  mdb_value.mv_size = value.size();
+  MDB_CHECK(mdb_put(mdb_txn_, *mdb_dbi_, &mdb_key, &mdb_value, 0));
+}
+
+}  // namespace db
+}  // namespace caffe