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