Fix multiscan crash when fill_cache=false (facebook#13889)

Summary:
When fill_cache is ReadOptions is false, multi scan Prepare crashes with the following assertion failure. In this case, CreateAndPibBlockInCache needs to directly create a block with full ownership.

facebook#9  0x00007f2fc003bc93 in __GI___assert_fail (assertion=0x7f2fc2147361 "pinned_data_blocks_guard[block_idx].GetValue()", file=0x7f2fc2146e08 "table/block_based/block_based_table_iterator.cc", line=1178, function=0x7f2fc2147262 "virtual void rocksdb::BlockBasedTableIterator::Prepare(const rocksdb::MultiScanArgs *)") at assert.c:101
101 in assert.c
facebook#10 0x00007f2fc1d73088 in rocksdb::BlockBasedTableIterator::Prepare(rocksdb::MultiScanArgs const*) () from /data/users/anand76/rocksdb_anand76/librocksdb.so.10.6

Pull Request resolved: facebook#13889

Test Plan: Parameterize the DBMultiScanIteratorTest tests with fill_cache

Reviewed By: cbi42

Differential Revision: D80552069

Pulled By: anand1976

fbshipit-source-id: 1a0b64af1e14c63d826add1f994a832ebff12757

Author: anand1976

db/db_iterator_test.cc

@@ -4142,13 +4142,17 @@ TEST_P(DBIteratorTest, AverageMemtableOpsScanFlushTriggerByOverwrites) {
   ASSERT_EQ(1, NumTableFilesAtLevel(0));
 }
 
-class DBMultiScanIteratorTest : public DBTestBase {
+class DBMultiScanIteratorTest : public DBTestBase,
+                                public ::testing::WithParamInterface<bool> {
  public:
   DBMultiScanIteratorTest()
       : DBTestBase("db_multi_scan_iterator_test", /*env_do_fsync=*/true) {}
 };
 
-TEST_F(DBMultiScanIteratorTest, BasicTest) {
+INSTANTIATE_TEST_CASE_P(DBMultiScanIteratorTest, DBMultiScanIteratorTest,
+                        ::testing::Bool());
+
+TEST_P(DBMultiScanIteratorTest, BasicTest) {
   // Create a file
   for (int i = 0; i < 100; ++i) {
     std::stringstream ss;
@@ -4159,6 +4163,7 @@ TEST_F(DBMultiScanIteratorTest, BasicTest) {
 
   std::vector<std::string> key_ranges({"k03", "k10", "k25", "k50"});
   ReadOptions ro;
+  ro.fill_cache = GetParam();
   MultiScanArgs scan_options(BytewiseComparator());
   scan_options.insert(key_ranges[0], key_ranges[1]);
   scan_options.insert(key_ranges[2], key_ranges[3]);
@@ -4249,7 +4254,7 @@ TEST_F(DBMultiScanIteratorTest, BasicTest) {
   iter.reset();
 }
 
-TEST_F(DBMultiScanIteratorTest, MixedBoundsTest) {
+TEST_P(DBMultiScanIteratorTest, MixedBoundsTest) {
   // Create a file
   for (int i = 0; i < 100; ++i) {
     std::stringstream ss;
@@ -4261,6 +4266,7 @@ TEST_F(DBMultiScanIteratorTest, MixedBoundsTest) {
   std::vector<std::string> key_ranges(
       {"k03", "k10", "k25", "k50", "k75", "k90"});
   ReadOptions ro;
+  ro.fill_cache = GetParam();
   MultiScanArgs scan_options(BytewiseComparator());
   scan_options.insert(key_ranges[0], key_ranges[1]);
   scan_options.insert(key_ranges[2]);
@@ -4335,7 +4341,7 @@ TEST_F(DBMultiScanIteratorTest, MixedBoundsTest) {
   iter.reset();
 }
 
-TEST_F(DBMultiScanIteratorTest, RangeAcrossFiles) {
+TEST_P(DBMultiScanIteratorTest, RangeAcrossFiles) {
   auto options = CurrentOptions();
   options.target_file_size_base = 100 << 10;  // 20KB
   options.compaction_style = kCompactionStyleUniversal;
@@ -4354,6 +4360,7 @@ TEST_F(DBMultiScanIteratorTest, RangeAcrossFiles) {
   ASSERT_EQ(2, NumTableFilesAtLevel(49));
   std::vector<std::string> key_ranges({Key(10), Key(90)});
   ReadOptions ro;
+  ro.fill_cache = GetParam();
   MultiScanArgs scan_options(BytewiseComparator());
   scan_options.insert(key_ranges[0], key_ranges[1]);
   ColumnFamilyHandle* cfh = dbfull()->DefaultColumnFamily();

table/block_based/block_based_table_iterator.cc

@@ -1199,7 +1199,8 @@ void BlockBasedTableIterator::Prepare(const MultiScanArgs* multiscan_opts) {
 #endif
         assert(pinned_data_blocks_guard[block_idx].IsEmpty());
         s = table_->CreateAndPinBlockInCache<Block_kData>(
-            read_options_, block, &tmp_contents,
+            read_options_, block, table_->get_rep()->decompressor.get(),
+            &tmp_contents,
             &(pinned_data_blocks_guard[block_idx].As<Block_kData>()));
         if (!s.ok()) {
           // Abort: failed to create and pin block in cache

table/block_based/block_based_table_reader.cc

@@ -109,8 +109,8 @@ CacheAllocationPtr CopyBufferToHeap(MemoryAllocator* allocator, Slice& buf) {
       CachableEntry<T>* out_parsed_block) const;                               \
   template Status BlockBasedTable::CreateAndPinBlockInCache<T>(                \
       const ReadOptions& ro, const BlockHandle& handle,                        \
-      BlockContents* block_contents, CachableEntry<T>* out_parsed_block)       \
-      const;
+      UnownedPtr<Decompressor> decomp, BlockContents* block_contents,          \
+      CachableEntry<T>* out_parsed_block) const;
 
 INSTANTIATE_BLOCKLIKE_TEMPLATES(ParsedFullFilterBlock);
 INSTANTIATE_BLOCKLIKE_TEMPLATES(DecompressorDict);
@@ -1741,13 +1741,55 @@ Status BlockBasedTable::LookupAndPinBlocksInCache(
 
 template <typename TBlocklike>
 Status BlockBasedTable::CreateAndPinBlockInCache(
-    const ReadOptions& ro, const BlockHandle& handle, BlockContents* contents,
+    const ReadOptions& ro, const BlockHandle& handle,
+    UnownedPtr<Decompressor> decomp, BlockContents* contents,
     CachableEntry<TBlocklike>* out_parsed_block) const {
-  return MaybeReadBlockAndLoadToCache(
-      nullptr, ro, handle, rep_->decompressor.get(),
-      /*for_compaction=*/false, out_parsed_block, nullptr, nullptr, contents,
-      /*async_read=*/false,
-      /*use_block_cache_for_lookup=*/true);
+  CompressionType compression_type = GetBlockCompressionType(*contents);
+  // If we don't own the contents and we don't need to decompress, copy
+  // the block to heap in order to have ownership. If decompression is
+  // needed, then the decompressor will allocate a buffer.
+  if (!contents->own_bytes() && compression_type == kNoCompression) {
+    Slice src = Slice(contents->data.data(), BlockSizeWithTrailer(handle));
+    *contents = BlockContents(
+        CopyBufferToHeap(GetMemoryAllocator(rep_->table_options), src),
+        handle.size());
+#ifndef NDEBUG
+    contents->has_trailer = true;
+#endif
+  }
+
+  Status s;
+  if (ro.fill_cache) {
+    s = MaybeReadBlockAndLoadToCache(nullptr, ro, handle, decomp,
+                                     /*for_compaction=*/false, out_parsed_block,
+                                     nullptr, nullptr, contents,
+                                     /*async_read=*/false,
+                                     /*use_block_cache_for_lookup=*/true);
+  }
+
+  if (!s.ok()) {
+    return s;
+  }
+
+  // fill_cache could be false, or no block cache is configured. In that
+  // case, decompress if necessary and take ownership of the block
+  if (out_parsed_block->GetValue() == nullptr && contents != nullptr) {
+    BlockContents tmp_contents;
+    if (compression_type != kNoCompression) {
+      s = DecompressSerializedBlock(contents->data.data(), handle.size(),
+                                    compression_type, *decomp, &tmp_contents,
+                                    rep_->ioptions,
+                                    GetMemoryAllocator(rep_->table_options));
+    } else {
+      tmp_contents = std::move(*contents);
+    }
+    if (s.ok()) {
+      std::unique_ptr<TBlocklike> block_holder;
+      rep_->create_context.Create(&block_holder, std::move(tmp_contents));
+      out_parsed_block->SetOwnedValue(std::move(block_holder));
+    }
+  }
+  return s;
 }
 
 // If contents is nullptr, this function looks up the block caches for the

table/block_based/block_based_table_reader.h

@@ -311,7 +311,7 @@ class BlockBasedTable : public TableReader {
   template <typename TBlocklike>
   Status CreateAndPinBlockInCache(
       const ReadOptions& ro, const BlockHandle& handle,
-      BlockContents* block_contents,
+      UnownedPtr<Decompressor> decomp, BlockContents* block_contents,
       CachableEntry<TBlocklike>* out_parsed_block) const;
 
   struct Rep;

table/block_based/block_based_table_reader_sync_and_async.h

@@ -37,8 +37,6 @@ DEFINE_SYNC_AND_ASYNC(void, BlockBasedTable::RetrieveMultipleBlocks)
   RandomAccessFileReader* file = rep_->file.get();
   const Footer& footer = rep_->footer;
   const ImmutableOptions& ioptions = rep_->ioptions;
-  size_t read_amp_bytes_per_bit = rep_->table_options.read_amp_bytes_per_bit;
-  MemoryAllocator* memory_allocator = GetMemoryAllocator(rep_->table_options);
 
   if (ioptions.allow_mmap_reads) {
     size_t idx_in_batch = 0;
@@ -266,79 +264,8 @@ DEFINE_SYNC_AND_ASYNC(void, BlockBasedTable::RetrieveMultipleBlocks)
     }
 
     if (s.ok()) {
-      // When the blocks share the same underlying buffer (scratch or direct io
-      // buffer), we may need to manually copy the block into heap if the
-      // serialized block has to be inserted into a cache. That falls into the
-      // following cases -
-      // 1. serialized block is not compressed, it needs to be inserted into
-      //    the uncompressed block cache if there is one
-      // 2. If the serialized block is compressed, it needs to be inserted
-      //    into the compressed block cache if there is one
-      //
-      // In all other cases, the serialized block is either uncompressed into a
-      // heap buffer or there is no cache at all.
-      CompressionType compression_type =
-          GetBlockCompressionType(serialized_block);
-      if ((use_fs_scratch || use_shared_buffer) &&
-          compression_type == kNoCompression) {
-        Slice serialized =
-            Slice(req.result.data() + req_offset, BlockSizeWithTrailer(handle));
-        serialized_block = BlockContents(
-            CopyBufferToHeap(GetMemoryAllocator(rep_->table_options),
-                             serialized),
-            handle.size());
-#ifndef NDEBUG
-        serialized_block.has_trailer = true;
-#endif
-      }
-    }
-
-    if (s.ok()) {
-      if (options.fill_cache) {
-        CachableEntry<Block_kData>* block_entry = &results[idx_in_batch];
-        // MaybeReadBlockAndLoadToCache will insert into the block caches if
-        // necessary. Since we're passing the serialized block contents, it
-        // will avoid looking up the block cache
-        s = MaybeReadBlockAndLoadToCache(
-            nullptr, options, handle, decomp,
-            /*for_compaction=*/false, block_entry, mget_iter->get_context,
-            /*lookup_context=*/nullptr, &serialized_block,
-            /*async_read=*/false, /*use_block_cache_for_lookup=*/true);
-
-        if (!s.ok()) {
-          statuses[idx_in_batch] = s;
-          continue;
-        }
-        // block_entry value could be null if no block cache is present, i.e
-        // BlockBasedTableOptions::no_block_cache is true and no compressed
-        // block cache is configured. In that case, fall
-        // through and set up the block explicitly
-        if (block_entry->GetValue() != nullptr) {
-          continue;
-        }
-      }
-
-      CompressionType compression_type =
-          GetBlockCompressionType(serialized_block);
-      BlockContents contents;
-      if (compression_type != kNoCompression) {
-        s = DecompressSerializedBlock(
-            req.result.data() + req_offset, handle.size(), compression_type,
-            *decomp, &contents, rep_->ioptions, memory_allocator);
-      } else {
-        // There are two cases here:
-        // 1) caller uses the shared buffer (scratch or direct io buffer);
-        // 2) we use the requst buffer.
-        // If scratch buffer or direct io buffer is used, we ensure that
-        // all serialized blocks are copyed to the heap as single blocks. If
-        // scratch buffer is not used, we also have no combined read, so the
-        // serialized block can be used directly.
-        contents = std::move(serialized_block);
-      }
-      if (s.ok()) {
-        results[idx_in_batch].SetOwnedValue(std::make_unique<Block_kData>(
-            std::move(contents), read_amp_bytes_per_bit, ioptions.stats));
-      }
+      s = CreateAndPinBlockInCache(options, handle, decomp, &serialized_block,
+                                   &results[idx_in_batch]);
     }
     statuses[idx_in_batch] = s;
   }

unreleased_history/bug_fixes/multiscan_fill_cache.md

@@ -0,0 +1 @@
+Fix a crash in iterator Prepare() when fill_cache=false