Caffe2 - C++ API
A deep learning, cross platform ML framework
leveldb.cc
1 #include "caffe2/core/db.h"
2 #include "caffe2/core/logging.h"
3 #include "caffe2/core/flags.h"
4 #include "leveldb/db.h"
5 #include "leveldb/write_batch.h"
6 
7 C10_DEFINE_int(
8  caffe2_leveldb_block_size,
9  65536,
10  "The caffe2 leveldb block size when writing a leveldb.");
11 
12 namespace caffe2 {
13 namespace db {
14 
15 class LevelDBCursor : public Cursor {
16  public:
17  explicit LevelDBCursor(leveldb::DB* db)
18  : iter_(db->NewIterator(leveldb::ReadOptions())) {
19  SeekToFirst();
20  }
21  ~LevelDBCursor() override {}
22  void Seek(const string& key) override { iter_->Seek(key); }
23  bool SupportsSeek() override { return true; }
24  void SeekToFirst() override { iter_->SeekToFirst(); }
25  void Next() override { iter_->Next(); }
26  string key() override { return iter_->key().ToString(); }
27  string value() override { return iter_->value().ToString(); }
28  bool Valid() override { return iter_->Valid(); }
29 
30  private:
31  std::unique_ptr<leveldb::Iterator> iter_;
32 };
33 
35  public:
36  explicit LevelDBTransaction(leveldb::DB* db) : db_(db) {
37  CAFFE_ENFORCE(db_);
38  batch_.reset(new leveldb::WriteBatch());
39  }
40  ~LevelDBTransaction() override {
41  Commit();
42  }
43  void Put(const string& key, const string& value) override {
44  batch_->Put(key, value);
45  }
46  void Commit() override {
47  leveldb::Status status = db_->Write(leveldb::WriteOptions(), batch_.get());
48  batch_.reset(new leveldb::WriteBatch());
49  CAFFE_ENFORCE(
50  status.ok(),
51  "Failed to write batch to leveldb. ", status.ToString());
52  }
53 
54  private:
55  leveldb::DB* db_;
56  std::unique_ptr<leveldb::WriteBatch> batch_;
57 
58  C10_DISABLE_COPY_AND_ASSIGN(LevelDBTransaction);
59 };
60 
61 class LevelDB : public DB {
62  public:
63  LevelDB(const string& source, Mode mode) : DB(source, mode) {
64  leveldb::Options options;
65  options.block_size = FLAGS_caffe2_leveldb_block_size;
66  options.write_buffer_size = 268435456;
67  options.max_open_files = 100;
68  options.error_if_exists = mode == NEW;
69  options.create_if_missing = mode != READ;
70  leveldb::DB* db_temp;
71  leveldb::Status status = leveldb::DB::Open(options, source, &db_temp);
72  CAFFE_ENFORCE(
73  status.ok(),
74  "Failed to open leveldb ", source, ". ", status.ToString());
75  db_.reset(db_temp);
76  VLOG(1) << "Opened leveldb " << source;
77  }
78 
79  void Close() override { db_.reset(); }
80  unique_ptr<Cursor> NewCursor() override {
81  return make_unique<LevelDBCursor>(db_.get());
82  }
83  unique_ptr<Transaction> NewTransaction() override {
84  return make_unique<LevelDBTransaction>(db_.get());
85  }
86 
87  private:
88  std::unique_ptr<leveldb::DB> db_;
89 };
90 
91 REGISTER_CAFFE2_DB(LevelDB, LevelDB);
92 // For lazy-minded, one can also call with lower-case name.
93 REGISTER_CAFFE2_DB(leveldb, LevelDB);
94 
95 } // namespace db
96 } // namespace caffe2
An abstract class for the current database transaction while writing.
Definition: db.h:61
An abstract class for the cursor of the database while reading.
Definition: db.h:22
void Put(const string &key, const string &value) override
Puts the key value pair to the database.
Definition: leveldb.cc:43
bool Valid() override
Returns whether the current location is valid - for example, if we have reached the end of the databa...
Definition: leveldb.cc:28
void SeekToFirst() override
Seek to the first key in the database.
Definition: leveldb.cc:24
void Next() override
Go to the next location in the database.
Definition: leveldb.cc:25
unique_ptr< Cursor > NewCursor() override
Returns a cursor to read the database.
Definition: leveldb.cc:80
void Commit() override
Commits the current writes.
Definition: leveldb.cc:46
void Seek(const string &key) override
Seek to a specific key (or if the key does not exist, seek to the immediate next).
Definition: leveldb.cc:22
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
Definition: blob.h:13
An abstract class for accessing a database of key-value pairs.
Definition: db.h:80
string value() override
Returns the current value.
Definition: leveldb.cc:27
void Close() override
Closes the database.
Definition: leveldb.cc:79
unique_ptr< Transaction > NewTransaction() override
Returns a transaction to write data to the database.
Definition: leveldb.cc:83
string key() override
Returns the current key.
Definition: leveldb.cc:26