Skip to content

Commit

Permalink
Allow for Customizable DB Open Hooks for DB Bench (#13326)
Browse files Browse the repository at this point in the history 
Summary:

This diff introduces ToolHooks, a class which allows for users to interpose their own set of logic for various functionality with db_bench_tool (i.e., various OpenDB implementations).

Differential Revision: D67868126
  • Loading branch information
@krhancoc @facebook-github-bot
krhancoc authored and facebook-github-bot committed Jan 24, 2025
1 parent a8bd6a3 commit 5287196
Show file tree
Hide file tree
Showing 6 changed files with 233 additions and 34 deletions.
1 change: 1 addition & 0 deletions CMakeLists.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1572,6 +1572,7 @@ if(WITH_BENCHMARK_TOOLS)
add_executable(db_bench${ARTIFACT_SUFFIX}
tools/simulated_hybrid_file_system.cc
tools/db_bench.cc
tools/tool_hooks.cc
tools/db_bench_tool.cc)
target_link_libraries(db_bench${ARTIFACT_SUFFIX}
${ROCKSDB_LIB} ${THIRDPARTY_LIBS})
Expand Down
3 changes: 2 additions & 1 deletion include/rocksdb/db_bench_tool.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,8 @@
#pragma once

#include "rocksdb/rocksdb_namespace.h"
#include "rocksdb/tool_hooks.h"

namespace ROCKSDB_NAMESPACE {
int db_bench_tool(int argc, char** argv);
int db_bench_tool(int argc, char** argv, ToolHooks& hooks = defaultHooks);
} // namespace ROCKSDB_NAMESPACE
110 changes: 110 additions & 0 deletions include/rocksdb/tool_hooks.h
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,110 @@
// Copyright (c) Meta Platforms, Inc. and affiliates. All rights reserved.
// This source code is licensed under both the GPLv2 (found in the
// COPYING file in the root directory) and Apache 2.0 License
// (found in the LICENSE.Apache file in the root directory).


#pragma once

#include "rocksdb/db.h"


namespace ROCKSDB_NAMESPACE {

struct TransactionDBOptions;
class TransactionDB;
class OptimisticTransactionDB;

namespace blob_db {
struct BlobDBOptions;
class BlobDB;
}

/*
* ToolHooks is currently a WORK IN PROGRESS, API is subject to change.
* ToolHooks is a class that allows users to override the default behavior of the various OpenDB calls used by
* db_bench_tool. This allows users to easily extend the functionality of db_bench_tool to support their own
* open implementations.
*/
class ToolHooks {
public:
ToolHooks() = default;
virtual ~ToolHooks() = default;
virtual Status Open(const Options& db_options, const std::string& name, DB** dbptr) = 0;
virtual Status Open(const DBOptions& db_options, const std::string& name,
const std::vector<ColumnFamilyDescriptor>& column_families,
std::vector<ColumnFamilyHandle*>* handles, DB** dbptr) = 0;
virtual Status OpenForReadOnly(const Options& options, const std::string& name,
DB** dbptr,
bool error_if_wal_file_exists) = 0;
virtual Status OpenForReadOnly(const Options& options, const std::string& name,
const std::vector<ColumnFamilyDescriptor>& column_families,
std::vector<ColumnFamilyHandle*>* handles, DB** dbptr) = 0;
virtual Status OpenTransactionDB(const Options& db_options,
const TransactionDBOptions& txn_db_options,
const std::string& dbname,
TransactionDB** dbptr) = 0;
virtual Status OpenTransactionDB(const DBOptions& db_options,
const TransactionDBOptions& txn_db_options,
const std::string& dbname,
const std::vector<ColumnFamilyDescriptor>& column_families,
std::vector<ColumnFamilyHandle*>* handles,
TransactionDB** dbptr) = 0;
virtual Status OpenOptimisticTransactionDB(const Options& options, const std::string& dbname,
OptimisticTransactionDB** dbptr) = 0;
virtual Status OpenOptimisticTransactionDB(const DBOptions& db_options, const std::string& dbname,
const std::vector<ColumnFamilyDescriptor>& column_families,
std::vector<ColumnFamilyHandle*>* handles,
OptimisticTransactionDB** dbptr) = 0;
virtual Status OpenAsSecondary(const Options& options, const std::string& name,
const std::string& secondary_path, DB** dbptr) = 0;
virtual Status OpenAsFollower(const Options& options, const std::string& name,
const std::string& leader_path,
std::unique_ptr<DB>* dbptr) = 0;
virtual Status Open(const Options& options, const blob_db::BlobDBOptions& bdb_options,
const std::string& dbname, blob_db::BlobDB** blob_db) = 0;

};

class DefaultHooks : public ToolHooks {
public:
DefaultHooks() = default;
~DefaultHooks() override = default;
virtual Status Open(const Options& db_options, const std::string& name, DB** dbptr) override;
virtual Status Open(const DBOptions& db_options, const std::string& name,
const std::vector<ColumnFamilyDescriptor>& column_families,
std::vector<ColumnFamilyHandle*>* handles, DB** dbptr) override;
virtual Status OpenForReadOnly(const Options& options, const std::string& name,
DB** dbptr,
bool error_if_wal_file_exists) override;
virtual Status OpenForReadOnly(const Options& options, const std::string& name,
const std::vector<ColumnFamilyDescriptor>& column_families,
std::vector<ColumnFamilyHandle*>* handles, DB** dbptr) override;
virtual Status OpenTransactionDB(const Options& db_options,
const TransactionDBOptions& txn_db_options,
const std::string& dbname,
TransactionDB** dbptr) override;
virtual Status OpenTransactionDB(const DBOptions& db_options,
const TransactionDBOptions& txn_db_options,
const std::string& dbname,
const std::vector<ColumnFamilyDescriptor>& column_families,
std::vector<ColumnFamilyHandle*>* handles,
TransactionDB** dbptr) override;
virtual Status OpenOptimisticTransactionDB(const Options& options, const std::string& dbname,
OptimisticTransactionDB** dbptr) override;
virtual Status OpenOptimisticTransactionDB(const DBOptions& db_options, const std::string& dbname,
const std::vector<ColumnFamilyDescriptor>& column_families,
std::vector<ColumnFamilyHandle*>* handles,
OptimisticTransactionDB** dbptr) override;
virtual Status OpenAsSecondary(const Options& options, const std::string& name,
const std::string& secondary_path, DB** dbptr) override;
virtual Status OpenAsFollower(const Options& options, const std::string& name,
const std::string& leader_path,
std::unique_ptr<DB>* dbptr) override;
virtual Status Open(const Options& options, const blob_db::BlobDBOptions& bdb_options,
const std::string& dbname, blob_db::BlobDB** blob_db) override;
};

extern DefaultHooks defaultHooks;

} // namespace ROCKSDB_NAMESPACE
22 changes: 11 additions & 11 deletions src.mk
View file
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,7 @@ LIB_SOURCES = \
env/env_encryption.cc \
env/env_posix.cc \
env/file_system.cc \
env/fs_on_demand.cc \
env/fs_on_demand.cc \
env/fs_posix.cc \
env/fs_remap.cc \
env/file_system_tracer.cc \
Expand Down Expand Up @@ -235,7 +235,7 @@ LIB_SOURCES = \
trace_replay/trace_replay.cc \
trace_replay/block_cache_tracer.cc \
trace_replay/io_tracer.cc \
util/async_file_reader.cc \
util/async_file_reader.cc \
util/build_version.cc \
util/cleanable.cc \
util/coding.cc \
Expand Down Expand Up @@ -377,12 +377,13 @@ MOCK_LIB_SOURCES = \

BENCH_LIB_SOURCES = \
tools/db_bench_tool.cc \
tools/tool_hooks.cc \
tools/simulated_hybrid_file_system.cc \

CACHE_BENCH_LIB_SOURCES = \
CACHE_BENCH_LIB_SOURCES = \
cache/cache_bench_tool.cc \

STRESS_LIB_SOURCES = \
STRESS_LIB_SOURCES = \
db_stress_tool/batched_ops_stress.cc \
db_stress_tool/cf_consistency_stress.cc \
db_stress_tool/db_stress_common.cc \
Expand All @@ -407,7 +408,7 @@ TEST_LIB_SOURCES = \
test_util/secondary_cache_test_util.cc \
test_util/testharness.cc \
test_util/testutil.cc \
utilities/agg_merge/test_agg_merge.cc \
utilities/agg_merge/test_agg_merge.cc \
utilities/cassandra/test_utils.cc \

FOLLY_SOURCES = \
Expand Down Expand Up @@ -444,14 +445,13 @@ BENCH_MAIN_SOURCES = \
tools/db_bench.cc \
util/filter_bench.cc \
utilities/persistent_cache/persistent_cache_bench.cc \
#util/log_write_bench.cc \

TEST_MAIN_SOURCES = \
cache/cache_test.cc \
cache/cache_reservation_manager_test.cc \
cache/compressed_secondary_cache_test.cc \
cache/lru_cache_test.cc \
cache/tiered_secondary_cache_test.cc \
cache/tiered_secondary_cache_test.cc \
db/blob/blob_counting_iterator_test.cc \
db/blob/blob_file_addition_test.cc \
db/blob/blob_file_builder_test.cc \
Expand Down Expand Up @@ -485,7 +485,7 @@ TEST_MAIN_SOURCES = \
db/db_dynamic_level_test.cc \
db/db_encryption_test.cc \
db/db_flush_test.cc \
db/db_follower_test.cc \
db/db_follower_test.cc \
db/db_readonly_with_timestamp_test.cc \
db/db_with_timestamp_basic_test.cc \
db/import_column_family_test.cc \
Expand Down Expand Up @@ -652,15 +652,15 @@ TEST_MAIN_SOURCES = \
utilities/util_merge_operators_test.cc \
utilities/write_batch_with_index/write_batch_with_index_test.cc \

TEST_MAIN_SOURCES_C = \
TEST_MAIN_SOURCES_C = \
db/c_test.c \

WITH_FAISS_TEST_MAIN_SOURCES = \
WITH_FAISS_TEST_MAIN_SOURCES = \
utilities/secondary_index/faiss_ivf_index_test.cc \

MICROBENCH_SOURCES = \
microbench/ribbon_bench.cc \
microbench/db_basic_bench.cc \
microbench/db_basic_bench.cc \

JNI_NATIVE_SOURCES = \
java/rocksjni/backupenginejni.cc \
Expand Down
46 changes: 24 additions & 22 deletions tools/db_bench_tool.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@
#include "rocksdb/slice_transform.h"
#include "rocksdb/stats_history.h"
#include "rocksdb/table.h"
#include "rocksdb/tool_hooks.h"
#include "rocksdb/utilities/backup_engine.h"
#include "rocksdb/utilities/object_registry.h"
#include "rocksdb/utilities/optimistic_transaction_db.h"
Expand Down Expand Up @@ -3448,11 +3449,11 @@ class Benchmark {
exit(1);
}

void Run() {
void Run(ToolHooks &hooks) {
if (!SanityCheck()) {
ErrorExit();
}
Open(&open_options_);
Open(&open_options_, hooks);
PrintHeader(open_options_);
std::stringstream benchmark_stream(FLAGS_benchmarks);
std::string name;
Expand Down Expand Up @@ -3794,7 +3795,7 @@ class Benchmark {
}
multi_dbs_.clear();
}
Open(&open_options_); // use open_options for the last accessed
Open(&open_options_, hooks); // use open_options for the last accessed
}

if (method != nullptr) {
Expand Down Expand Up @@ -4753,7 +4754,7 @@ class Benchmark {
options.paranoid_memory_checks = FLAGS_paranoid_memory_checks;
}

void InitializeOptionsGeneral(Options* opts) {
void InitializeOptionsGeneral(Options* opts, ToolHooks& hooks) {
// Be careful about what is set here to avoid accidentally overwriting
// settings already configured by OPTIONS file. Only configure settings that
// are needed for the benchmark to run, settings for shared objects that
Expand Down Expand Up @@ -4846,7 +4847,7 @@ class Benchmark {
}

if (FLAGS_num_multi_db <= 1) {
OpenDb(options, FLAGS_db, &db_);
OpenDb(options, hooks, FLAGS_db, &db_);
} else {
multi_dbs_.clear();
multi_dbs_.resize(FLAGS_num_multi_db);
Expand All @@ -4855,7 +4856,7 @@ class Benchmark {
if (!wal_dir.empty()) {
options.wal_dir = GetPathForMultiple(wal_dir, i);
}
OpenDb(options, GetPathForMultiple(FLAGS_db, i), &multi_dbs_[i]);
OpenDb(options, hooks, GetPathForMultiple(FLAGS_db, i), &multi_dbs_[i]);
}
options.wal_dir = wal_dir;
}
Expand All @@ -4882,15 +4883,15 @@ class Benchmark {
}
}

void Open(Options* opts) {
void Open(Options* opts, ToolHooks& hooks) {
if (!InitializeOptionsFromFile(opts)) {
InitializeOptionsFromFlags(opts);
}

InitializeOptionsGeneral(opts);
InitializeOptionsGeneral(opts, hooks);
}

void OpenDb(Options options, const std::string& db_name,
void OpenDb(Options options, ToolHooks& hooks, const std::string& db_name,
DBWithColumnFamilies* db) {
uint64_t open_start = FLAGS_report_open_timing ? FLAGS_env->NowNanos() : 0;
Status s;
Expand Down Expand Up @@ -4931,10 +4932,10 @@ class Benchmark {
}
}
if (FLAGS_readonly) {
s = DB::OpenForReadOnly(options, db_name, column_families, &db->cfh,
s = hooks.OpenForReadOnly(options, db_name, column_families, &db->cfh,
&db->db);
} else if (FLAGS_optimistic_transaction_db) {
s = OptimisticTransactionDB::Open(options, db_name, column_families,
s = hooks.OpenOptimisticTransactionDB(options, db_name, column_families,
&db->cfh, &db->opt_txn_db);
if (s.ok()) {
db->db = db->opt_txn_db->GetBaseDB();
Expand All @@ -4947,22 +4948,22 @@ class Benchmark {
txn_db_options.skip_concurrency_control = true;
txn_db_options.write_policy = WRITE_PREPARED;
}
s = TransactionDB::Open(options, txn_db_options, db_name,
s = hooks.OpenTransactionDB(options, txn_db_options, db_name,
column_families, &db->cfh, &ptr);
if (s.ok()) {
db->db = ptr;
}
} else {
s = DB::Open(options, db_name, column_families, &db->cfh, &db->db);
s = hooks.Open(options, db_name, column_families, &db->cfh, &db->db);
}
db->cfh.resize(FLAGS_num_column_families);
db->num_created = num_hot;
db->num_hot = num_hot;
db->cfh_idx_to_prob = std::move(cfh_idx_to_prob);
} else if (FLAGS_readonly) {
s = DB::OpenForReadOnly(options, db_name, &db->db);
s = hooks.OpenForReadOnly(options, db_name, &db->db, false);
} else if (FLAGS_optimistic_transaction_db) {
s = OptimisticTransactionDB::Open(options, db_name, &db->opt_txn_db);
s = hooks.OpenOptimisticTransactionDB(options, db_name, &db->opt_txn_db);
if (s.ok()) {
db->db = db->opt_txn_db->GetBaseDB();
}
Expand All @@ -4976,7 +4977,7 @@ class Benchmark {
}
s = CreateLoggerFromOptions(db_name, options, &options.info_log);
if (s.ok()) {
s = TransactionDB::Open(options, txn_db_options, db_name, &ptr);
s = hooks.OpenTransactionDB(options, txn_db_options, db_name, &ptr);
}
if (s.ok()) {
db->db = ptr;
Expand All @@ -4994,7 +4995,7 @@ class Benchmark {
blob_db_options.blob_file_size = FLAGS_blob_db_file_size;
blob_db_options.compression = FLAGS_blob_db_compression_type_e;
blob_db::BlobDB* ptr = nullptr;
s = blob_db::BlobDB::Open(options, blob_db_options, db_name, &ptr);
s = hooks.Open(options, blob_db_options, db_name, &ptr);
if (s.ok()) {
db->db = ptr;
}
Expand All @@ -5005,7 +5006,7 @@ class Benchmark {
default_secondary_path += "/dbbench_secondary";
FLAGS_secondary_path = default_secondary_path;
}
s = DB::OpenAsSecondary(options, db_name, FLAGS_secondary_path, &db->db);
s = hooks.OpenAsSecondary(options, db_name, FLAGS_secondary_path, &db->db);
if (s.ok() && FLAGS_secondary_update_interval > 0) {
secondary_update_thread_.reset(new port::Thread(
[this](int interval, DBWithColumnFamilies* _db) {
Expand All @@ -5026,12 +5027,12 @@ class Benchmark {
}
} else if (FLAGS_open_as_follower) {
std::unique_ptr<DB> dbptr;
s = DB::OpenAsFollower(options, db_name, FLAGS_leader_path, &dbptr);
s = hooks.OpenAsFollower(options, db_name, FLAGS_leader_path, &dbptr);
if (s.ok()) {
db->db = dbptr.release();
}
} else {
s = DB::Open(options, db_name, &db->db);
s = hooks.Open(options, db_name, &db->db);
}
if (FLAGS_report_open_timing) {
std::cout << "OpenDb: "
Expand Down Expand Up @@ -8623,7 +8624,8 @@ class Benchmark {
}
};

int db_bench_tool(int argc, char** argv) {
int db_bench_tool(int argc, char** argv, ToolHooks& hooks) {

ROCKSDB_NAMESPACE::port::InstallStackTraceHandler();
ConfigOptions config_options;
static bool initialized = false;
Expand Down Expand Up @@ -8775,7 +8777,7 @@ int db_bench_tool(int argc, char** argv) {
}

ROCKSDB_NAMESPACE::Benchmark benchmark;
benchmark.Run();
benchmark.Run(hooks);

if (FLAGS_print_malloc_stats) {
std::string stats_string;
Expand Down
Loading

0 comments on commit 5287196

Please sign in to comment.