github.com/igggame/nebulas-go@v2.1.0+incompatible/nbre/fs/rocksdb_storage.cpp (about)

     1  // Copyright (C) 2018 go-nebulas authors
     2  //
     3  // This file is part of the go-nebulas library.
     4  //
     5  // the go-nebulas library is free software: you can redistribute it and/or
     6  // modify
     7  // it under the terms of the GNU General Public License as published by
     8  // the Free Software Foundation, either version 3 of the License, or
     9  // (at your option) any later version.
    10  //
    11  // the go-nebulas library is distributed in the hope that it will be useful,
    12  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    13  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    14  // GNU General Public License for more details.
    15  //
    16  // You should have received a copy of the GNU General Public License
    17  // along with the go-nebulas library.  If not, see
    18  // <http://www.gnu.org/licenses/>.
    19  //
    20  
    21  #include "rocksdb_storage.h"
    22  #include <rocksdb/advanced_options.h>
    23  #include <rocksdb/cache.h>
    24  #include <rocksdb/filter_policy.h>
    25  #include <rocksdb/options.h>
    26  #include <rocksdb/slice.h>
    27  #include <rocksdb/table.h>
    28  
    29  namespace neb{
    30  namespace fs {
    31  
    32  rocksdb_storage::rocksdb_storage() : m_db(nullptr), m_enable_batch(false) {}
    33  
    34  rocksdb_storage::~rocksdb_storage() { close_database(); }
    35  
    36  void rocksdb_storage::open_database(const std::string &db_name,
    37                                      storage_open_flag flag) {
    38    rocksdb::DB *db = nullptr;
    39    rocksdb::Status status;
    40  
    41    if (nullptr == m_db) {
    42      if (flag == storage_open_for_readonly) {
    43        rocksdb::Options options;
    44        options.keep_log_file_num = 1;
    45        options.max_open_files = 500;
    46        status = rocksdb::DB::OpenForReadOnly(options, db_name, &db, false);
    47        m_enable_batch = true;
    48      } else {
    49        rocksdb::Options options;
    50        options.keep_log_file_num = 1;
    51  
    52        //! TODO setup bloomfilter, LRUCache, writer buffer size
    53        rocksdb::BlockBasedTableOptions table_options;
    54        table_options.filter_policy.reset(rocksdb::NewBloomFilterPolicy(10));
    55        table_options.block_cache = rocksdb::NewLRUCache(512 << 20);
    56        options.table_factory.reset(
    57            rocksdb::NewBlockBasedTableFactory(table_options));
    58  
    59        options.create_if_missing = true;
    60        options.max_open_files = 500;
    61        options.write_buffer_size = 64 * 1024 * 1024;
    62        options.IncreaseParallelism(4);
    63        status = rocksdb::DB::Open(options, db_name, &db);
    64        m_enable_batch = false;
    65      }
    66  
    67      if (status.ok()) {
    68        m_db = std::unique_ptr<rocksdb::DB>(db);
    69      } else {
    70        LOG(ERROR) << "open db error: " << status.ToString();
    71        throw storage_general_failure(status.ToString());
    72      }
    73    } else {
    74      throw std::runtime_error("database already open");
    75    }
    76  }
    77  
    78  void rocksdb_storage::close_database() {
    79    if (!m_db) {
    80      return;
    81    }
    82    rocksdb::Status status = m_db->Close();
    83    if (!status.ok()) {
    84      throw std::runtime_error("close database failed");
    85    }
    86    m_db.reset(nullptr);
    87  }
    88  
    89  bytes rocksdb_storage::get_bytes(const bytes &key) {
    90    if (!m_db) {
    91      throw storage_exception_no_init();
    92    }
    93    rocksdb::Slice s(reinterpret_cast<const char *>(key.value()), key.size());
    94    std::string value;
    95    auto status = m_db->Get(rocksdb::ReadOptions(), s, &value);
    96    if (!status.ok()) {
    97      throw storage_general_failure(status.ToString());
    98    }
    99    return string_to_byte(value);
   100  }
   101  
   102  void rocksdb_storage::put_bytes(const bytes &key, const bytes &val) {
   103    if (!m_db) {
   104      throw storage_exception_no_init();
   105    }
   106    if (m_enable_batch) {
   107      m_batched_ops.emplace_back([key, val](rocksdb::WriteBatch &wb) {
   108        std::string str_value = byte_to_string(val);
   109        rocksdb::Slice s(reinterpret_cast<const char *>(key.value()), key.size());
   110        auto status = wb.Put(s, str_value);
   111        if (!status.ok()) {
   112          throw storage_general_failure(status.ToString());
   113        }
   114      });
   115      return;
   116    }
   117    std::string str_value = byte_to_string(val);
   118    rocksdb::Slice s(reinterpret_cast<const char *>(key.value()), key.size());
   119    auto status = m_db->Put(rocksdb::WriteOptions(), s, str_value);
   120    if (!status.ok()) {
   121      throw storage_general_failure(status.ToString());
   122    }
   123  }
   124  
   125  void rocksdb_storage::del_by_bytes(const bytes &key) {
   126    if (!m_db) {
   127      throw storage_exception_no_init();
   128    }
   129    if (m_enable_batch) {
   130      m_batched_ops.emplace_back([key](rocksdb::WriteBatch &wb) {
   131        rocksdb::Slice s(reinterpret_cast<const char *>(key.value()), key.size());
   132        auto status = wb.Delete(s);
   133        if (!status.ok()) {
   134          throw storage_general_failure(status.ToString());
   135        }
   136      });
   137      return;
   138    }
   139    rocksdb::Slice s(reinterpret_cast<const char *>(key.value()), key.size());
   140    auto status = m_db->Delete(rocksdb::WriteOptions(), s);
   141    if (!status.ok()) {
   142      throw storage_general_failure(status.ToString());
   143    }
   144  }
   145  
   146  void rocksdb_storage::enable_batch() { m_enable_batch = true; }
   147  void rocksdb_storage::disable_batch() {
   148    if (m_enable_batch) {
   149      flush();
   150    }
   151    m_enable_batch = false;
   152  }
   153  void rocksdb_storage::flush() {
   154    if (!m_enable_batch) {
   155      return;
   156    }
   157    if (!m_db) {
   158      return;
   159    }
   160    rocksdb::WriteBatch wb(0, 0);
   161    std::for_each(m_batched_ops.begin(), m_batched_ops.end(),
   162                  [&wb](const batch_operation_t &f) { f(wb); });
   163    auto status = m_db->Write(rocksdb::WriteOptions(), &wb);
   164    if (!status.ok()) {
   165      throw storage_general_failure(status.ToString());
   166    }
   167  }
   168  
   169  void rocksdb_storage::display(
   170      const std::function<void(rocksdb::Iterator *)> &cb) {
   171    rocksdb::Iterator *it = m_db->NewIterator(rocksdb::ReadOptions());
   172    cb(it);
   173  }
   174  
   175  } // end namespace fs
   176  } // end namespace neb
   177