Fix a nullptr access bug in MultiScan (facebook#14062)

Summary:
Fixing a nullptr access in multiscan, under following situation.

```
Block Based Table: blk1:[k1,k2], blk2:[k3,                k8], blk3:[k9]
Scan ranges:            [k1,             k4), [k5,k6), [k7,            k10)
Prepared block ranges:  [0,2],                [2,2],   [1,3]
```

1. Seek key k1 on the first range, read key k1, k2.
2. Seek key k4 on the 2nd range, blocks 0,1 would be unpinned.
3. Seek key k9, block 1 would be accessed, but it is unpinned, which trigger assert failure in debug mode and nullptr access on release build.

This fix changes how blocks are unpinned. It is now only unpinning the block, when the cur_data_block_idx has passed it.

Pull Request resolved: facebook#14062

Test Plan:
Unit Test
rand_seed 304010984 on UserDefinedIndexStressTest

Reviewed By: cbi42

Differential Revision: D84976410

Pulled By: xingbowang

fbshipit-source-id: 6b99bf85fc9d4108c5267ae77be77ccfe08923cd

Author: Xingbo Wang

table/block_based/block_based_table_iterator.cc

@@ -1157,10 +1157,6 @@ bool BlockBasedTableIterator::SeekMultiScanImpl(const Slice* seek_target) {
       // Seeking a range that is out side of prepared ranges.
       return out_of_bound;
     }
-    // unpin block, then do a seek.
-    if (multi_scan_->next_scan_idx > 0) {
-      UnpinPreviousScanBlocks(multi_scan_->next_scan_idx);
-    }
 
     auto [cur_scan_start_idx, cur_scan_end_idx] =
         multi_scan_->block_index_ranges_per_scan[multi_scan_->next_scan_idx];
@@ -1231,6 +1227,8 @@ void BlockBasedTableIterator::MultiScanUnexpectedSeekTarget(
 
 void BlockBasedTableIterator::MultiScanSeekTargetFromBlock(
     const Slice* seek_target, size_t block_idx) {
+  assert(multi_scan_->cur_data_block_idx <= block_idx);
+
   if (!block_iter_points_to_real_block_ ||
       multi_scan_->cur_data_block_idx != block_idx) {
     if (block_iter_points_to_real_block_) {
@@ -1245,32 +1243,20 @@ void BlockBasedTableIterator::MultiScanSeekTargetFromBlock(
       return;
     }
   }
-  multi_scan_->cur_data_block_idx = block_idx;
-  block_iter_points_to_real_block_ = true;
-  block_iter_.Seek(*seek_target);
-  FindKeyForward();
-}
 
-void BlockBasedTableIterator::UnpinPreviousScanBlocks(size_t current_scan_idx) {
-  // TODO: support aborting and clearn up async IO requests, currently
-  // only unpins already initialized blocks
-  assert(multi_scan_);
-  assert(current_scan_idx < multi_scan_->block_index_ranges_per_scan.size());
-  if (current_scan_idx == 0) return;
-
-  auto prev_start_block_idx = std::get<0>(
-      multi_scan_->block_index_ranges_per_scan[current_scan_idx - 1]);
-  // Since a block can be shared between consecutive scans, we need
-  // curr_start_block_idx here instead of just release blocks
-  // up to the end of previous range block index.
-  auto curr_start_block_idx =
-      std::get<0>(multi_scan_->block_index_ranges_per_scan[current_scan_idx]);
-  for (size_t block_idx = prev_start_block_idx;
-       block_idx < curr_start_block_idx; ++block_idx) {
-    if (!multi_scan_->pinned_data_blocks[block_idx].IsEmpty()) {
-      multi_scan_->pinned_data_blocks[block_idx].Reset();
+  // Move current data block index forward until block_idx, meantime, unpin all
+  // the blocks in between
+  while (multi_scan_->cur_data_block_idx < block_idx) {
+    // unpin block
+    if (!multi_scan_->pinned_data_blocks[multi_scan_->cur_data_block_idx]
+             .IsEmpty()) {
+      multi_scan_->pinned_data_blocks[multi_scan_->cur_data_block_idx].Reset();
     }
+    multi_scan_->cur_data_block_idx++;
   }
+  block_iter_points_to_real_block_ = true;
+  block_iter_.Seek(*seek_target);
+  FindKeyForward();
 }
 
 void BlockBasedTableIterator::FindBlockForwardInMultiScan() {
@@ -1307,6 +1293,11 @@ void BlockBasedTableIterator::FindBlockForwardInMultiScan() {
     }
     // Move to the next pinned data block
     ResetDataIter();
+    // Unpin previous block if it is not reset by data iterator
+    if (!multi_scan_->pinned_data_blocks[multi_scan_->cur_data_block_idx]
+             .IsEmpty()) {
+      multi_scan_->pinned_data_blocks[multi_scan_->cur_data_block_idx].Reset();
+    }
     ++multi_scan_->cur_data_block_idx;
 
     if (MultiScanLoadDataBlock(multi_scan_->cur_data_block_idx)) {

table/block_based/block_based_table_iterator.h

@@ -620,9 +620,6 @@ class BlockBasedTableIterator : public InternalIteratorBase<Slice> {
 
   void FindBlockForwardInMultiScan();
 
-  // Unpins blocks from the immediately previous scan range.
-  void UnpinPreviousScanBlocks(size_t current_scan_idx);
-
   void PrepareReadAsyncCallBack(FSReadRequest& req, void* cb_arg) {
     // Record status, result and sanity check offset from `req`.
     AsyncReadState* async_state = static_cast<AsyncReadState*>(cb_arg);

table/table_test.cc

@@ -87,6 +87,7 @@ namespace ROCKSDB_NAMESPACE {
 namespace {
 
 const std::string kDummyValue(10000, 'o');
+constexpr auto kVerbose = false;
 
 // DummyPropertiesCollector used to test BlockBasedTableProperties
 class DummyPropertiesCollector : public TablePropertiesCollector {
@@ -934,7 +935,6 @@ class HarnessTest : public testing::Test {
 
   void TestRandomAccess(Random* rnd, const std::vector<std::string>& keys,
                         const stl_wrappers::KVMap& data) {
-    static const bool kVerbose = false;
     InternalIterator* iter = constructor_->NewIterator();
     ASSERT_TRUE(!iter->Valid());
     stl_wrappers::KVMap::const_iterator model_iter = data.begin();
@@ -7827,7 +7827,6 @@ class UserDefinedIndexTestBase : public BlockBasedTableTestBase {
     read_opts.iterate_upper_bound = &ub;
     std::unique_ptr<Iterator> iter(db->NewIterator(read_opts, cfh));
     iter->Prepare(scan_opts);
-    static const bool kVerbose = false;
     for (auto opt : opts) {
       ub = opt.range.limit.value();
       iter->Seek(opt.range.start.value());
@@ -8979,8 +8978,6 @@ std::ostream& operator<<(std::ostream& os,
             << param.enable_compaction_with_sst_partitioner << "}";
 }
 
-constexpr auto kVerbose = false;
-
 struct DataRange {
   size_t start;  // inclusive
   size_t end;    // exclusive
@@ -9131,23 +9128,47 @@ class UserDefinedIndexStressTest
   }
 
   void CreateSstFileWithRanges(const std::string& ingest_file,
-                               const DataRange& range) {
-    std::unique_ptr<SstFileWriter> writer =
-        std::make_unique<SstFileWriter>(EnvOptions(), options_);
-    ASSERT_OK(writer->Open(ingest_file));
+                               const std::vector<DataRange>& ranges,
+                               bool& data_added) {
+    std::unique_ptr<SstFileWriter> writer;
+    data_added = false;
 
-    assert(range.start != range.end);
+    std::vector<DataRange> ranges_in_file;
 
-    if (range.is_range_delete) {
-      ASSERT_OK(writer->DeleteRange(range.start_key, range.end_key));
-    } else {
-      for (size_t i = range.start; i != range.end;) {
-        auto key = FormatKey(i);
-        range.start < range.end ? i++ : i--;
-        ASSERT_OK(writer->Put(key, range.value));
+    for (auto const& range : ranges) {
+      assert(range.start != range.end);
+      if (range.skipped) {
+        continue;
+      }
+
+      if (writer == nullptr) {
+        writer = std::make_unique<SstFileWriter>(EnvOptions(), options_);
+        ASSERT_OK(writer->Open(ingest_file));
+      }
+      ranges_in_file.push_back(range);
+
+      data_added = true;
+
+      if (range.is_range_delete) {
+        ASSERT_OK(writer->DeleteRange(range.start_key, range.end_key));
+      } else {
+        for (size_t i = range.start; i != range.end;) {
+          auto key = FormatKey(i);
+          range.start < range.end ? i++ : i--;
+          ASSERT_OK(writer->Put(key, range.value));
+        }
+      }
+    }
+    if (kVerbose) {
+      std::cout << "Ingested file: " + ingest_file + "; Range: {" << std::endl;
+      for (const auto& range : ranges_in_file) {
+        std::cout << "    " << range.ToString() << "," << std::endl;
       }
+      std::cout << "}" << std::endl;
+    }
+    if (data_added) {
+      ASSERT_OK(writer->Finish());
     }
-    ASSERT_OK(writer->Finish()) << range.ToString();
   }
 
   void RangeScan(std::unique_ptr<Iterator>& iter,
@@ -9267,17 +9288,42 @@ class UserDefinedIndexStressTest
   void IngestFilesInOneLevel(const std::vector<DataRange>& ranges_in_level,
                              const std::string& ingest_file_name_prefix,
                              size_t& ingest_file_count,
-                             const IngestExternalFileOptions& ifo) {
+                             const IngestExternalFileOptions& ifo,
+                             bool combine_ranges = false) {
     std::vector<std::string> ingest_files;
     // Generate SST file and bulk load them one level at a time
-    for (auto const& range : ranges_in_level) {
-      if (!range.skipped) {
+    if (combine_ranges) {
+      size_t i = 0;
+      while (i < ranges_in_level.size()) {
+        // if combine ranges, generate 1 SST file that combines muliple ranges
+        // together
+        size_t batch_end_idx =
+            std::min(i + rnd.Uniform(3) + 2, ranges_in_level.size());
+        bool data_added = false;
         ASSERT_NO_FATAL_FAILURE(CreateSstFileWithRanges(
             ingest_file_name_prefix + std::to_string(ingest_file_count),
-            range));
-        ingest_files.push_back(ingest_file_name_prefix +
-                               std::to_string(ingest_file_count));
-        ingest_file_count++;
+            {ranges_in_level.begin() + i,
+             ranges_in_level.begin() + batch_end_idx},
+            data_added));
+        if (data_added) {
+          ingest_files.push_back(ingest_file_name_prefix +
+                                 std::to_string(ingest_file_count));
+          ingest_file_count++;
+        }
+        i = batch_end_idx;
+      }
+    } else {
+      for (auto const& range : ranges_in_level) {
+        if (!range.skipped) {
+          bool data_added = false;
+          ASSERT_NO_FATAL_FAILURE(CreateSstFileWithRanges(
+              ingest_file_name_prefix + std::to_string(ingest_file_count),
+              {range}, data_added));
+          ASSERT_TRUE(data_added);
+          ingest_files.push_back(ingest_file_name_prefix +
+                                 std::to_string(ingest_file_count));
+          ingest_file_count++;
+        }
       }
     }
 
@@ -9314,9 +9360,9 @@ class UserDefinedIndexStressTest
 // Becuase query count == 2, level n+1 would only prepare 3-5. but since 4-6
 // got deleted in the upper level, they are not returned, so only 3 is
 // returned. Meantime the query should have return [3, 6]
-// One way to fix this is by preparing more data blocks once prepared blocks are
-// exhausted, but upper bound is not reached yet.
-// This requires following changes:
+// One way to fix this is by preparing more data blocks once prepared blocks
+// are exhausted, but upper bound is not reached yet. This requires following
+// changes:
 // 1. Fix out of bound flag in block table iterator. Only set it if the key is
 // larger than the upper bound.
 // 2. Refactor the prepared block single dimension vector into 2 dimension of
@@ -9349,12 +9395,126 @@ TEST_P(UserDefinedIndexStressTest, DISABLED_PartialDeleteRange) {
   ASSERT_NO_FATAL_FAILURE(ValidateQueryResult());
 }
 
+TEST_P(UserDefinedIndexStressTest, DeleteRangeMixedWithDataFile) {
+  // Create 2 column families. One use normal put/del, the other uses sst
+  // ingest.
+  // Test the case where there are 3 levels, the middle level is a delete
+  // range file that span across the entire key space. The top level file have
+  // multiple files and each one has both data and delete range Scan same
+  // range between the 2 CF and validate the result is same
+  SCOPED_TRACE("Start with random seed: " + std::to_string(rand_seed_));
+  dbname_ = test::PerThreadDBPath(
+      "UserDefinedIndexStressTest_DeleteRangeMixedWithDataFile");
+  SCOPED_TRACE("dbname: " + dbname_);
+  ASSERT_NO_FATAL_FAILURE(SetupDB(dbname_));
+
+  // Test 3 levels.
+  // bottom level is normal data files.
+  ranges_in_levels_.push_back(GenerateKeyRanges(rnd.Uniform(3) + 4, 2, "L6"));
+  // middle level delete range between each level
+  if (is_reverse_comparator_) {
+    ranges_in_levels_.push_back({{.start = 100,
+                                  .end = 0,
+                                  .is_range_delete = true,
+                                  .skipped = false,
+                                  .start_key = "keyz",
+                                  .end_key = "key"}});
+  } else {
+    ranges_in_levels_.push_back({{.start = 0,
+                                  .end = 100,
+                                  .is_range_delete = true,
+                                  .skipped = false,
+                                  .start_key = "key",
+                                  .end_key = "keyz"}});
+  }
+
+  // Top level is normal data files
+  ranges_in_levels_.push_back(GenerateKeyRanges(rnd.Uniform(3) + 4, 2, "L4"));
+
+  IngestExternalFileOptions ifo;
+  ifo.snapshot_consistency = false;
+  auto ingest_file_name_prefix = dbname_ + "ingest_file_";
+  size_t ingest_file_count = 0;
+  auto first_level = true;
+  for (auto const& ranges_in_level : ranges_in_levels_) {
+    ASSERT_NO_FATAL_FAILURE(
+        IngestFilesInOneLevel(ranges_in_level, ingest_file_name_prefix,
+                              ingest_file_count, ifo, true));
+    if (first_level) {
+      first_level = false;
+      if (enable_compaction_with_sst_partitioner_) {
+        // When compaction is enabled, do a compaction at the first level
+        ASSERT_NO_FATAL_FAILURE(CompactIngestedCF());
+      }
+    }
+  }
+
+  ASSERT_NO_FATAL_FAILURE(AddDataToRegularCF());
+
+  ASSERT_NO_FATAL_FAILURE(ValidateQueryResult());
+}
+
 TEST_P(UserDefinedIndexStressTest, DeleteRange) {
   // Create 2 column families. One use normal put/del, the other uses sst
   // ingest.
-  // Test the case where there are 3 levels, the middle level is a delete range
-  // file that span across the entire key space.
-  // Range scan same range between the 2 CF and validate the result is same
+  // Test the case where there are 3 levels, the middle level is a delete
+  // range file that span across the entire key space. Range scan same range
+  // between the 2 CF and validate the result is same
+  SCOPED_TRACE("Start with random seed: " + std::to_string(rand_seed_));
+  dbname_ = test::PerThreadDBPath("UserDefinedIndexStressTest_DeleteRange");
+  SCOPED_TRACE("dbname: " + dbname_);
+  ASSERT_NO_FATAL_FAILURE(SetupDB(dbname_));
+
+  // Test 3 levels.
+  // bottom level is normal data files.
+  ranges_in_levels_.push_back(GenerateKeyRanges(rnd.Uniform(3) + 4, 2, "L6"));
+  // middle level delete range between each level
+  if (is_reverse_comparator_) {
+    ranges_in_levels_.push_back({{.start = 100,
+                                  .end = 0,
+                                  .is_range_delete = true,
+                                  .skipped = false,
+                                  .start_key = "keyz",
+                                  .end_key = "key"}});
+  } else {
+    ranges_in_levels_.push_back({{.start = 0,
+                                  .end = 100,
+                                  .is_range_delete = true,
+                                  .skipped = false,
+                                  .start_key = "key",
+                                  .end_key = "keyz"}});
+  }
+  // Top level is normal data files
+  ranges_in_levels_.push_back(GenerateKeyRanges(rnd.Uniform(3) + 4, 2, "L4"));
+
+  IngestExternalFileOptions ifo;
+  ifo.snapshot_consistency = false;
+  auto ingest_file_name_prefix = dbname_ + "ingest_file_";
+  size_t ingest_file_count = 0;
+  auto first_level = true;
+  for (auto const& ranges_in_level : ranges_in_levels_) {
+    ASSERT_NO_FATAL_FAILURE(IngestFilesInOneLevel(
+        ranges_in_level, ingest_file_name_prefix, ingest_file_count, ifo));
+    if (first_level) {
+      first_level = false;
+      if (enable_compaction_with_sst_partitioner_) {
+        // When compaction is enabled, do a compaction at the first level
+        ASSERT_NO_FATAL_FAILURE(CompactIngestedCF());
+      }
+    }
+  }
+
+  ASSERT_NO_FATAL_FAILURE(AddDataToRegularCF());
+
+  ASSERT_NO_FATAL_FAILURE(ValidateQueryResult());
+}
+
+TEST_P(UserDefinedIndexStressTest, AtomicReplaceBulkLoad) {
+  // Create 2 column families. One use normal put/del, the other uses sst
+  // ingest.
+  // Test the case where there are 3 levels, the middle level is a delete
+  // range file that span across the entire key space. Range scan same range
+  // between the 2 CF and validate the result is same
   SCOPED_TRACE("Start with random seed: " + std::to_string(rand_seed_));
   dbname_ = test::PerThreadDBPath("UserDefinedIndexStressTest_DeleteRange");
   SCOPED_TRACE("dbname: " + dbname_);
@@ -9399,6 +9559,23 @@ TEST_P(UserDefinedIndexStressTest, DeleteRange) {
     }
   }
 
+  // Ingest the a new file with atomic replace with full key space, this layer
+  // is exactly same as the one at Level 4
+  bool data_added;
+  ASSERT_NO_FATAL_FAILURE(CreateSstFileWithRanges(
+      ingest_file_name_prefix + std::to_string(++ingest_file_count),
+      ranges_in_levels_[2], data_added));
+
+  IngestExternalFileArg ingest_arg;
+  ingest_arg.column_family = ingest_cfh_;
+  ingest_arg.options = ifo;
+  ingest_arg.external_files.push_back(ingest_file_name_prefix +
+                                      std::to_string(ingest_file_count));
+  ingest_arg.atomic_replace_range = RangeOpt(nullptr, nullptr);
+
+  ASSERT_OK(db_->IngestExternalFiles(
+      std::vector<IngestExternalFileArg>({ingest_arg})));
+
   ASSERT_NO_FATAL_FAILURE(AddDataToRegularCF());
 
   ASSERT_NO_FATAL_FAILURE(ValidateQueryResult());

unreleased_history/bug_fixes/multi_scan_page_unpin_bug_fix.md

@@ -0,0 +1 @@
+Fix a bug in Page unpinning in MultiScan