Allow for Customizable DB Open Hooks for DB Bench

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

BUCK

@@ -405,6 +405,7 @@ cpp_library_wrapper(name="rocksdb_tools_lib", srcs=[
         "test_util/testutil.cc",
         "tools/block_cache_analyzer/block_cache_trace_analyzer.cc",
         "tools/db_bench_tool.cc",
+        "tools/tool_hooks.cc",
         "tools/simulated_hybrid_file_system.cc",
         "tools/trace_analyzer_tool.cc",
     ], deps=[":rocksdb_lib"], headers=[], link_whole=False, extra_test_libs=False)

include/rocksdb/db_bench_tool.h

@@ -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

include/rocksdb/tool_hooks.h

@@ -0,0 +1,112 @@
+// Copyright (c) 2011-present, Facebook, Inc.  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).
+// Copyright (c) 2011 The LevelDB Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file. See the AUTHORS file for names of contributors.
+
+#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

src.mk

@@ -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                                     \
@@ -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                                                \
@@ -376,12 +376,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                           \
@@ -406,7 +407,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 =                                                 \
@@ -443,14 +444,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                                     \
@@ -484,7 +484,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                                       \
@@ -651,15 +651,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                            \

tools/db_bench_tool.cc

@@ -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"
@@ -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;
@@ -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) {
@@ -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
@@ -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);
@@ -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;
     }
@@ -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;
@@ -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();
@@ -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();
       }
@@ -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;
@@ -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;
       }
@@ -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) {
@@ -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:     "
@@ -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;
@@ -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;