Make failure to load UDI when opening an SST a soft failure (facebook#13921)

Summary:
If user_defined_index_factory in BlockBasedTableOptions is configured and we try to open an SST file without the corresponding UDI (either during DB open or file ingestion), ignore a failure to load the UDI by default. If fail_if_no_udi_on_open in BlockBasedTableOptions is true, then treat it as a fatal error.

Pull Request resolved: facebook#13921

Test Plan: Update unit tests

Reviewed By: xingbowang

Differential Revision: D81826054

Pulled By: anand1976

fbshipit-source-id: f4fe0b13ccb02b9448622af487680131e349c52b

Author: anand1976

include/rocksdb/statistics.h

@@ -542,6 +542,9 @@ enum Tickers : uint32_t {
   // TransactionOptions::large_txn_commit_optimize_threshold.
   NUMBER_WBWI_INGEST,
 
+  // Failure to load the UDI during SST table open
+  SST_USER_DEFINED_INDEX_LOAD_FAIL_COUNT,
+
   TICKER_ENUM_MAX
 };
 

include/rocksdb/table.h

@@ -502,6 +502,12 @@ struct BlockBasedTableOptions {
   // during table building.
   std::shared_ptr<UserDefinedIndexFactory> user_defined_index_factory = nullptr;
 
+  // EXPERIMENTAL
+  //
+  // Return an error Status if a user_defined_index_factory is configured,
+  // but there's no corresponding UDI block in the SST file being opened.
+  bool fail_if_no_udi_on_open = false;
+
   // If true, place whole keys in the filter (not just prefixes).
   // This must generally be true for gets to be efficient.
   bool whole_key_filtering = true;

monitoring/statistics.cc

@@ -273,6 +273,8 @@ const std::vector<std::pair<Tickers, std::string>> TickersNameMap = {
     {FILE_READ_CORRUPTION_RETRY_SUCCESS_COUNT,
      "rocksdb.file.read.corruption.retry.success.count"},
     {NUMBER_WBWI_INGEST, "rocksdb.number.wbwi.ingest"},
+    {SST_USER_DEFINED_INDEX_LOAD_FAIL_COUNT,
+     "rocksdb.sst.user.defined.index.load.fail.count"},
 };
 
 const std::vector<std::pair<Histograms, std::string>> HistogramsNameMap = {

options/options_settable_test.cc

@@ -203,7 +203,8 @@ TEST_F(OptionsSettableTest, BlockBasedTableOptionsAllFieldsSettable) {
       "max_auto_readahead_size=0;"
       "prepopulate_block_cache=kDisable;"
       "initial_auto_readahead_size=0;"
-      "num_file_reads_for_auto_readahead=0",
+      "num_file_reads_for_auto_readahead=0;"
+      "fail_if_no_udi_on_open=true",
       new_bbto));
 
   ASSERT_EQ(unset_bytes_base,

table/block_based/block_based_table_factory.cc

@@ -399,6 +399,9 @@ static struct BlockBasedTableTypeInfo {
          {offsetof(struct BlockBasedTableOptions,
                    num_file_reads_for_auto_readahead),
           OptionType::kUInt64T, OptionVerificationType::kNormal}},
+        {"fail_if_no_udi_on_open",
+         {offsetof(struct BlockBasedTableOptions, fail_if_no_udi_on_open),
+          OptionType::kBoolean, OptionVerificationType::kNormal}},
     };
   }
 } block_based_table_type_info;
@@ -874,6 +877,14 @@ std::string BlockBasedTableFactory::GetPrintableOptions() const {
                ? "nullptr"
                : table_options_.filter_policy->Name());
   ret.append(buffer);
+  snprintf(buffer, kBufferSize, "  user_defined_index_factory: %s\n",
+           table_options_.user_defined_index_factory == nullptr
+               ? "nullptr"
+               : table_options_.user_defined_index_factory->Name());
+  ret.append(buffer);
+  snprintf(buffer, kBufferSize, "  fail_if_no_udi_on_open: %d\n",
+           table_options_.fail_if_no_udi_on_open);
+  ret.append(buffer);
   snprintf(buffer, kBufferSize, "  whole_key_filtering: %d\n",
            table_options_.whole_key_filtering);
   ret.append(buffer);

table/block_based/block_based_table_reader.cc

@@ -1333,33 +1333,62 @@ Status BlockBasedTable::PrefetchIndexAndFilterBlocks(
     s = FindMetaBlock(meta_iter, kUserDefinedIndexPrefix + udi_name,
                       &udi_block_handle);
     if (!s.ok()) {
-      return s;
-    }
-    // Read the block, and allocate on heap or pin in cache. The UDI block is
-    // not compressed. RetrieveBlock will verify the checksum.
-    s = RetrieveBlock(prefetch_buffer, ro, udi_block_handle,
-                      rep_->decompressor.get(), &rep_->udi_block,
-                      /*get_context=*/nullptr, lookup_context,
-                      /*for_compaction=*/false, use_cache, /*async_read=*/false,
-                      /*use_block_cache_for_lookup=*/false);
-    if (!s.ok()) {
-      return s;
+      RecordTick(rep_->ioptions.statistics.get(),
+                 SST_USER_DEFINED_INDEX_LOAD_FAIL_COUNT);
+      if (table_options.fail_if_no_udi_on_open) {
+        ROCKS_LOG_ERROR(rep_->ioptions.logger,
+                        "Failed to find the the UDI block %s in file %s; %s",
+                        udi_name.c_str(), rep_->file->file_name().c_str(),
+                        s.ToString().c_str());
+        // MAke the status more informative
+        s = Status::Corruption(s.ToString(), rep_->file->file_name());
+        return s;
+      } else {
+        // Emit a warning, but ignore the error status
+        ROCKS_LOG_WARN(rep_->ioptions.logger,
+                       "Failed to find the the UDI block %s in file %s; %s",
+                       udi_name.c_str(), rep_->file->file_name().c_str(),
+                       s.ToString().c_str());
+        s = Status::OK();
+      }
     }
-    assert(!rep_->udi_block.IsEmpty());
 
-    std::unique_ptr<UserDefinedIndexReader> udi_reader =
-        table_options.user_defined_index_factory->NewReader(
-            rep_->udi_block.GetValue()->data);
-    index_reader = std::make_unique<UserDefinedIndexReaderWrapper>(
-        udi_name, std::move(index_reader), std::move(udi_reader));
+    // If the UDI block size is 0, that means there's effectively no user
+    // defined index. In that case, skip setting up the reader.
+    if (udi_block_handle.size() > 0) {
+      // Read the block, and allocate on heap or pin in cache. The UDI block is
+      // not compressed. RetrieveBlock will verify the checksum.
+      if (s.ok()) {
+        s = RetrieveBlock(prefetch_buffer, ro, udi_block_handle,
+                          rep_->decompressor.get(), &rep_->udi_block,
+                          /*get_context=*/nullptr, lookup_context,
+                          /*for_compaction=*/false, use_cache,
+                          /*async_read=*/false,
+                          /*use_block_cache_for_lookup=*/false);
+      }
+      if (s.ok()) {
+        assert(!rep_->udi_block.IsEmpty());
+
+        std::unique_ptr<UserDefinedIndexReader> udi_reader =
+            table_options.user_defined_index_factory->NewReader(
+                rep_->udi_block.GetValue()->data);
+        if (udi_reader) {
+          index_reader = std::make_unique<UserDefinedIndexReaderWrapper>(
+              udi_name, std::move(index_reader), std::move(udi_reader));
+        } else {
+          s = Status::Corruption("Failed to create UDI reader for " + udi_name +
+                                 " in file " + rep_->file->file_name());
+        }
+      }
+    }
   }
 
   rep_->index_reader = std::move(index_reader);
 
   // The partitions of partitioned index are always stored in cache. They
   // are hence follow the configuration for pin and prefetch regardless of
   // the value of cache_index_and_filter_blocks
-  if (prefetch_all || pin_partition) {
+  if (s.ok() && (prefetch_all || pin_partition)) {
     s = rep_->index_reader->CacheDependencies(ro, pin_partition,
                                               prefetch_buffer);
   }

table/table_test.cc

@@ -7480,6 +7480,9 @@ class UserDefinedIndexTest : public BlockBasedTableTestBase {
                           const Slice* first_key_in_next_block,
                           const BlockHandle& block_handle,
                           std::string* separator_scratch) override {
+        if (keys_added_ == 0) {
+          return last_key_in_current_block;
+        }
         EXPECT_EQ(last_key_in_current_block.size(), 5);
         if (first_key_in_next_block) {
           EXPECT_EQ(first_key_in_next_block->size(), 5);
@@ -7500,12 +7503,19 @@ class UserDefinedIndexTest : public BlockBasedTableTestBase {
 
       void OnKeyAdded(const Slice& key, ValueType /*value*/,
                       const Slice& /*value*/) override {
+        if (key.starts_with("dummy")) {
+          return;
+        }
         EXPECT_EQ(key.size(), 5);
         // Track keys added to the index
         keys_added_++;
       }
 
       Status Finish(Slice* index_contents) override {
+        if (entries_added_ == 0) {
+          *index_contents = Slice();
+          return Status::OK();
+        }
         // Serialize the index data
         std::string result;
         for (const auto& entry : index_data_) {
@@ -8020,6 +8030,7 @@ TEST_F(UserDefinedIndexTest, IngestTest) {
 
 // Verify that external file ingestion fails if we try to ingest an SST file
 // without the UDI and a UDI factory is configured in BlockBasedTableOptions
+// and fail_if_no_udi_on_open is true in BlockBasedTableOptions.
 TEST_F(UserDefinedIndexTest, IngestFailTest) {
   Options options;
   BlockBasedTableOptions table_options;
@@ -8051,6 +8062,7 @@ TEST_F(UserDefinedIndexTest, IngestFailTest) {
   auto user_defined_index_factory =
       std::make_shared<TestUserDefinedIndexFactory>();
   table_options.user_defined_index_factory = user_defined_index_factory;
+  table_options.fail_if_no_udi_on_open = true;
   options.table_factory.reset(NewBlockBasedTableFactory(table_options));
 
   std::unique_ptr<DB> db;
@@ -8065,6 +8077,72 @@ TEST_F(UserDefinedIndexTest, IngestFailTest) {
   s = db->IngestExternalFile(cfh, {ingest_file}, ifo);
   ASSERT_NOK(s);
 
+  ASSERT_OK(db->SetOptions(
+      cfh, {{"block_based_table_factory", "{fail_if_no_udi_on_open=false;}"}}));
+  s = db->IngestExternalFile(cfh, {ingest_file}, ifo);
+  ASSERT_OK(s);
+
+  ASSERT_OK(db->DestroyColumnFamilyHandle(cfh));
+  ASSERT_OK(db->Close());
+  ASSERT_OK(DestroyDB(dbname, options));
+}
+
+TEST_F(UserDefinedIndexTest, IngestEmptyUDI) {
+  Options options;
+  BlockBasedTableOptions table_options;
+  std::string dbname = test::PerThreadDBPath("user_defined_index_test");
+  std::string ingest_file = dbname + "test.sst";
+  std::string ingest_file2 = dbname + "dummy.sst";
+
+  // Set up the user-defined index factory
+  auto user_defined_index_factory =
+      std::make_shared<TestUserDefinedIndexFactory>();
+  table_options.user_defined_index_factory = user_defined_index_factory;
+  // Set up custom flush block policy that flushes every 3 keys
+  table_options.flush_block_policy_factory =
+      std::make_shared<CustomFlushBlockPolicyFactory>();
+
+  options.table_factory.reset(NewBlockBasedTableFactory(table_options));
+
+  std::unique_ptr<SstFileWriter> writer;
+  writer.reset(new SstFileWriter(EnvOptions(), options));
+  ASSERT_OK(writer->Open(ingest_file));
+
+  // Add 100 keys instead of just 5
+  for (int i = 0; i < 100; i++) {
+    std::stringstream ss;
+    ss << std::setw(2) << std::setfill('0') << i;
+    std::string key = "key" + ss.str();
+    std::string value = "value" + ss.str();
+    ASSERT_OK(writer->Put(key, value));
+  }
+  ASSERT_OK(writer->Finish());
+  writer.reset();
+  writer.reset(new SstFileWriter(EnvOptions(), options));
+  ASSERT_OK(writer->Open(ingest_file2));
+  ASSERT_OK(writer->Put("dummy", "val"));
+  ASSERT_OK(writer->Finish());
+  writer.reset();
+
+  table_options.fail_if_no_udi_on_open = true;
+  options.table_factory.reset(NewBlockBasedTableFactory(table_options));
+
+  std::unique_ptr<DB> db;
+  options.create_if_missing = true;
+  Status s = DB::Open(options, dbname, &db);
+  ASSERT_OK(s);
+  ASSERT_TRUE(db != nullptr);
+  ColumnFamilyHandle* cfh = nullptr;
+  ASSERT_OK(db->CreateColumnFamily(options, "new_cf", &cfh));
+
+  std::vector<IngestExternalFileArg> ifa;
+  ifa.emplace_back();
+  ifa[0].column_family = cfh;
+  ifa[0].external_files.emplace_back(ingest_file);
+  ifa[0].external_files.emplace_back(ingest_file2);
+  s = db->IngestExternalFiles(ifa);
+  ASSERT_OK(s);
+
   ASSERT_OK(db->DestroyColumnFamilyHandle(cfh));
   ASSERT_OK(db->Close());
   ASSERT_OK(DestroyDB(dbname, options));

unreleased_history/new_features/fail_if_no_udi_on_open.md

@@ -0,0 +1 @@
+Add the fail_if_no_udi_on_open flag in BlockBasedTableOption to control whether a missing user defined index block in a SST is a hard error or not.