TitanDBImpl::Open only open base DB once (#114)

Summary:
Refactor `TitanDBImpl::Open` to open base DB only once. Previously we open the DB twice - the first time only to obtain `cf_id` for each CF, and then initialize internals before open base DB for real. The patch change to open base DB only once, initialize internals after that. Before `TitanDBImpl` is initialized, `BaseDbListener` will be no-op and `TableBuilder` will not output blob files.

Test Plan:
Added new tests

include/titan/options.h

@@ -158,9 +158,8 @@ struct TitanCFOptions : public ColumnFamilyOptions {
   TitanCFOptions() = default;
   explicit TitanCFOptions(const ColumnFamilyOptions& options)
       : ColumnFamilyOptions(options) {}
-  explicit TitanCFOptions(const ColumnFamilyOptions&,
-                          const ImmutableTitanCFOptions&,
-                          const MutableTitanCFOptions&);
+  TitanCFOptions(const ColumnFamilyOptions&, const ImmutableTitanCFOptions&,
+                 const MutableTitanCFOptions&);
 
   TitanCFOptions& operator=(const ColumnFamilyOptions& options) {
     *dynamic_cast<ColumnFamilyOptions*>(this) = options;

src/base_db_listener.cc

 #include "base_db_listener.h"
 
+#include "db_impl.h"
+
 namespace rocksdb {
 namespace titandb {
 
-BaseDbListener::BaseDbListener(TitanDBImpl* db) : db_impl_(db) {}
+BaseDbListener::BaseDbListener(TitanDBImpl* db) : db_impl_(db) {
+  assert(db_impl_ != nullptr);
+}
 
 BaseDbListener::~BaseDbListener() {}
 
 void BaseDbListener::OnFlushCompleted(DB* /*db*/,
                                       const FlushJobInfo& flush_job_info) {
-  db_impl_->OnFlushCompleted(flush_job_info);
+  if (db_impl_->initialized()) {
+    db_impl_->OnFlushCompleted(flush_job_info);
+  }
 }
 
 void BaseDbListener::OnCompactionCompleted(
     DB* /* db */, const CompactionJobInfo& compaction_job_info) {
-  db_impl_->OnCompactionCompleted(compaction_job_info);
+  if (db_impl_->initialized()) {
+    db_impl_->OnCompactionCompleted(compaction_job_info);
+  }
 }
 
 }  // namespace titandb

src/base_db_listener.h

 #pragma once
 
-#include "db_impl.h"
 #include "rocksdb/listener.h"
 
 namespace rocksdb {
-
 namespace titandb {
 
+class TitanDBImpl;
+
 class BaseDbListener final : public EventListener {
  public:
   BaseDbListener(TitanDBImpl* db);
@@ -22,5 +22,4 @@ class BaseDbListener final : public EventListener {
 };
 
 }  // namespace titandb
-
 }  // namespace rocksdb

src/db.cc

@@ -29,7 +29,6 @@ Status TitanDB::Open(const TitanDBOptions& db_options,
   auto s = impl->Open(descs, handles);
   if (s.ok()) {
     *db = impl;
-    impl->StartBackgroundTasks();
   } else {
     *db = nullptr;
     delete impl;

src/db_impl.h

@@ -128,12 +128,16 @@ class TitanDBImpl : public TitanDB {
   bool GetIntProperty(ColumnFamilyHandle* column_family, const Slice& property,
                       uint64_t* value) override;
 
+  bool initialized() const { return initialized_; }
+
   void OnFlushCompleted(const FlushJobInfo& flush_job_info);
 
   void OnCompactionCompleted(const CompactionJobInfo& compaction_job_info);
 
   void StartBackgroundTasks();
 
+  void TEST_set_initialized(bool initialized) { initialized_ = initialized; }
+
   Status TEST_StartGC(uint32_t column_family_id);
   Status TEST_PurgeObsoleteFiles();
 
@@ -150,6 +154,9 @@ class TitanDBImpl : public TitanDB {
   friend class TitanDBTest;
   friend class TitanThreadSafetyTest;
 
+  Status OpenImpl(const std::vector<TitanCFDescriptor>& descs,
+                  std::vector<ColumnFamilyHandle*>* handles);
+
   Status ValidateOptions(
       const TitanDBOptions& options,
       const std::vector<TitanCFDescriptor>& column_families) const;
@@ -242,6 +249,9 @@ class TitanDBImpl : public TitanDB {
   EnvOptions env_options_;
   DBImpl* db_impl_;
   TitanDBOptions db_options_;
+
+  std::atomic<bool> initialized_{false};
+
   // Turn DB into read-only if background error happened
   Status bg_error_;
   std::atomic_bool has_bg_error_{false};

src/table_builder_test.cc

@@ -6,6 +6,7 @@
 #include "blob_file_manager.h"
 #include "blob_file_reader.h"
 #include "blob_file_set.h"
+#include "db_impl.h"
 #include "table_builder.h"
 #include "table_factory.h"
 
@@ -98,6 +99,60 @@ class FileManager : public BlobFileManager {
   BlobFileSet* blob_file_set_;
 };
 
+class TestTableFactory : public TableFactory {
+ private:
+  std::shared_ptr<TableFactory> base_factory_;
+  mutable TableBuilder* latest_table_builder_ = nullptr;
+
+ public:
+  TestTableFactory(std::shared_ptr<TableFactory>& base_factory)
+      : base_factory_(base_factory) {}
+
+  const char* Name() const override { return "TestTableFactory"; }
+
+  Status NewTableReader(
+      const TableReaderOptions& options,
+      std::unique_ptr<RandomAccessFileReader>&& file, uint64_t file_size,
+      std::unique_ptr<TableReader>* result,
+      bool prefetch_index_and_filter_in_cache) const override {
+    return base_factory_->NewTableReader(options, std::move(file), file_size,
+                                         result,
+                                         prefetch_index_and_filter_in_cache);
+  }
+
+  TableBuilder* NewTableBuilder(const TableBuilderOptions& options,
+                                uint32_t column_family_id,
+                                WritableFileWriter* file) const override {
+    latest_table_builder_ =
+        base_factory_->NewTableBuilder(options, column_family_id, file);
+    return latest_table_builder_;
+  }
+
+  std::string GetPrintableTableOptions() const override {
+    return base_factory_->GetPrintableTableOptions();
+  }
+
+  Status SanitizeOptions(const DBOptions& db_options,
+                         const ColumnFamilyOptions& cf_options) const override {
+    // Override this when we need to validate our options.
+    return base_factory_->SanitizeOptions(db_options, cf_options);
+  }
+
+  Status GetOptionString(std::string* opt_string,
+                         const std::string& delimiter) const override {
+    // Override this when we need to persist our options.
+    return base_factory_->GetOptionString(opt_string, delimiter);
+  }
+
+  void* GetOptions() override { return base_factory_->GetOptions(); }
+
+  bool IsDeleteRangeSupported() const override {
+    return base_factory_->IsDeleteRangeSupported();
+  }
+
+  TableBuilder* latest_table_builder() const { return latest_table_builder_; }
+};
+
 class TableBuilderTest : public testing::Test {
  public:
   TableBuilderTest()
@@ -110,11 +165,18 @@ class TableBuilderTest : public testing::Test {
     cf_options_.min_blob_size = kMinBlobSize;
     blob_file_set_.reset(new BlobFileSet(db_options_, nullptr));
     std::map<uint32_t, TitanCFOptions> cfs{{0, cf_options_}};
+    db_impl_.reset(new TitanDBImpl(db_options_, tmpdir_));
+    db_impl_->TEST_set_initialized(true);
     blob_file_set_->AddColumnFamilies(cfs);
     blob_manager_.reset(new FileManager(db_options_, blob_file_set_.get()));
-    table_factory_.reset(new TitanTableFactory(db_options_, cf_options_,
-                                               blob_manager_, &mutex_,
-                                               blob_file_set_.get(), nullptr));
+    // Replace base table facotry.
+    base_table_factory_ =
+        std::make_shared<TestTableFactory>(cf_options_.table_factory);
+    cf_options_.table_factory = base_table_factory_;
+    cf_ioptions_.table_factory = base_table_factory_.get();
+    table_factory_.reset(new TitanTableFactory(
+        db_options_, cf_options_, db_impl_.get(), blob_manager_, &mutex_,
+        blob_file_set_.get(), nullptr));
   }
 
   ~TableBuilderTest() {
@@ -205,11 +267,40 @@ class TableBuilderTest : public testing::Test {
   std::string tmpdir_;
   std::string base_name_;
   std::string blob_name_;
+  std::unique_ptr<TitanDBImpl> db_impl_;
+  std::shared_ptr<TestTableFactory> base_table_factory_;
   std::unique_ptr<TableFactory> table_factory_;
   std::shared_ptr<BlobFileManager> blob_manager_;
   std::unique_ptr<BlobFileSet> blob_file_set_;
 };
 
+// Before TitanDBImpl initialized, table factory should return base table
+// builder.
+TEST_F(TableBuilderTest, BeforeDBInitialized) {
+  CompressionOptions compression_opts;
+  TableBuilderOptions opts(cf_ioptions_, cf_moptions_,
+                           cf_ioptions_.internal_comparator, &collectors_,
+                           kNoCompression, 0 /*sample_for_compression*/,
+                           compression_opts, false /*skip_filters*/,
+                           kDefaultColumnFamilyName, 0 /*target_level*/);
+
+  db_impl_->TEST_set_initialized(false);
+  std::unique_ptr<WritableFileWriter> file1;
+  NewBaseFileWriter(&file1);
+  std::unique_ptr<TableBuilder> builder1(
+      table_factory_->NewTableBuilder(opts, 0 /*cf_id*/, file1.get()));
+  ASSERT_EQ(builder1.get(), base_table_factory_->latest_table_builder());
+  builder1->Abandon();
+
+  db_impl_->TEST_set_initialized(true);
+  std::unique_ptr<WritableFileWriter> file2;
+  NewBaseFileWriter(&file2);
+  std::unique_ptr<TableBuilder> builder2(
+      table_factory_->NewTableBuilder(opts, 0 /*cf_id*/, file2.get()));
+  ASSERT_NE(builder2.get(), base_table_factory_->latest_table_builder());
+  builder2->Abandon();
+}
+
 TEST_F(TableBuilderTest, Basic) {
   std::unique_ptr<WritableFileWriter> base_file;
   NewBaseFileWriter(&base_file);
@@ -356,9 +447,9 @@ TEST_F(TableBuilderTest, NumEntries) {
 // To test size of each blob file is around blob_file_target_size after building
 TEST_F(TableBuilderTest, TargetSize) {
   cf_options_.blob_file_target_size = kTargetBlobFileSize;
-  table_factory_.reset(new TitanTableFactory(db_options_, cf_options_,
-                                             blob_manager_, &mutex_,
-                                             blob_file_set_.get(), nullptr));
+  table_factory_.reset(new TitanTableFactory(
+      db_options_, cf_options_, db_impl_.get(), blob_manager_, &mutex_,
+      blob_file_set_.get(), nullptr));
   std::unique_ptr<WritableFileWriter> base_file;
   NewBaseFileWriter(&base_file);
   std::unique_ptr<TableBuilder> table_builder;
@@ -387,9 +478,9 @@ TEST_F(TableBuilderTest, TargetSize) {
 // correct
 TEST_F(TableBuilderTest, LevelMerge) {
   cf_options_.level_merge = true;
-  table_factory_.reset(new TitanTableFactory(db_options_, cf_options_,
-                                             blob_manager_, &mutex_,
-                                             blob_file_set_.get(), nullptr));
+  table_factory_.reset(new TitanTableFactory(
+      db_options_, cf_options_, db_impl_.get(), blob_manager_, &mutex_,
+      blob_file_set_.get(), nullptr));
   std::unique_ptr<WritableFileWriter> base_file;
   NewBaseFileWriter(&base_file);
   std::unique_ptr<TableBuilder> table_builder;

src/table_factory.cc

 #include "table_factory.h"
 
+#include "db_impl.h"
 #include "table_builder.h"
 
 namespace rocksdb {
@@ -20,6 +21,9 @@ TableBuilder* TitanTableFactory::NewTableBuilder(
     WritableFileWriter* file) const {
   std::unique_ptr<TableBuilder> base_builder(
       base_factory_->NewTableBuilder(options, column_family_id, file));
+  if (!db_impl_->initialized()) {
+    return base_builder.release();
+  }
   TitanCFOptions cf_options = cf_options_;
   cf_options.blob_run_mode = blob_run_mode_.load();
   std::weak_ptr<BlobStorage> blob_storage;

src/table_factory.h

@@ -11,10 +11,12 @@
 namespace rocksdb {
 namespace titandb {
 
+class TitanDBImpl;
+
 class TitanTableFactory : public TableFactory {
  public:
   TitanTableFactory(const TitanDBOptions& db_options,
-                    const TitanCFOptions& cf_options,
+                    const TitanCFOptions& cf_options, TitanDBImpl* db_impl,
                     std::shared_ptr<BlobFileManager> blob_manager,
                     port::Mutex* db_mutex, BlobFileSet* blob_file_set,
                     TitanStats* stats)
@@ -22,6 +24,7 @@ class TitanTableFactory : public TableFactory {
         cf_options_(cf_options),
         blob_run_mode_(cf_options.blob_run_mode),
         base_factory_(cf_options.table_factory),
+        db_impl_(db_impl),
         blob_manager_(blob_manager),
         db_mutex_(db_mutex),
         blob_file_set_(blob_file_set),
@@ -66,6 +69,7 @@ class TitanTableFactory : public TableFactory {
   const TitanCFOptions cf_options_;
   std::atomic<TitanBlobRunMode> blob_run_mode_;
   std::shared_ptr<TableFactory> base_factory_;
+  TitanDBImpl* db_impl_;
   std::shared_ptr<BlobFileManager> blob_manager_;
   port::Mutex* db_mutex_;
   BlobFileSet* blob_file_set_;

src/titan_db_test.cc

@@ -310,6 +310,33 @@ class TitanDBTest : public testing::Test {
   std::vector<ColumnFamilyHandle*> cf_handles_;
 };
 
+TEST_F(TitanDBTest, Open) {
+  std::atomic<bool> checked_before_initialized{false};
+  std::atomic<bool> background_job_started{false};
+  SyncPoint::GetInstance()->SetCallBack(
+      "TitanDBImpl::OnFlushCompleted:Begin",
+      [&](void*) { background_job_started = true; });
+  SyncPoint::GetInstance()->SetCallBack(
+      "TitanDBImpl::OnCompactionCompleted:Begin",
+      [&](void*) { background_job_started = true; });
+  SyncPoint::GetInstance()->SetCallBack(
+      "TitanDBImpl::OpenImpl:BeforeInitialized", [&](void* arg) {
+        checked_before_initialized = true;
+        TitanDBImpl* db = reinterpret_cast<TitanDBImpl*>(arg);
+        // Try to trigger flush and compaction. Listeners should not be call.
+        ASSERT_OK(db->Put(WriteOptions(), "k1", "v1"));
+        ASSERT_OK(db->Flush(FlushOptions()));
+        ASSERT_OK(db->Put(WriteOptions(), "k1", "v2"));
+        ASSERT_OK(db->Put(WriteOptions(), "k2", "v3"));
+        ASSERT_OK(db->Flush(FlushOptions()));
+        ASSERT_OK(db->CompactRange(CompactRangeOptions(), nullptr, nullptr));
+      });
+  SyncPoint::GetInstance()->EnableProcessing();
+  Open();
+  ASSERT_TRUE(checked_before_initialized.load());
+  ASSERT_FALSE(background_job_started.load());
+}
+
 TEST_F(TitanDBTest, Basic) {
   const uint64_t kNumKeys = 100;
   std::map<std::string, std::string> data;

src/titan_stats.h

@@ -168,7 +168,7 @@ class TitanStats : public Statistics {
   // created after DB open.
   Status Initialize(std::map<uint32_t, TitanCFOptions> cf_options) {
     for (auto& opts : cf_options) {
-      internal_stats_[opts.first] = NewTitanInternalStats(opts.second);
+      internal_stats_[opts.first] = std::make_shared<TitanInternalStats>();
     }
     return Status::OK();
   }
@@ -279,11 +279,6 @@ class TitanStats : public Statistics {
       tickers_;
   std::array<HistogramImpl, INTERNAL_HISTOGRAM_ENUM_MAX - HISTOGRAM_ENUM_MAX>
       histograms_;
-
-  std::shared_ptr<TitanInternalStats> NewTitanInternalStats(
-      TitanCFOptions& opts) {
-    return std::make_shared<TitanInternalStats>();
-  }
 };
 
 // Utility functions for Titan ticker and histogram stats types