Should only decode restart points for uncompressed blocks (#3996)

Summary:
The Block object assumes contents are uncompressed. Block's constructor tries to read the number of restarts, but does not get an accurate number when its contents are compressed, which is causing issues like https://github.com/facebook/rocksdb/issues/3843.
This PR address this issue by skipping reconstruction of restart points when blocks are known to be compressed. Somehow the restart points can be read directly when Snappy is used and some tests (for example https://github.com/facebook/rocksdb/blob/master/db/db_block_cache_test.cc#L196) expects blocks to be fully constructed even when Snappy compression is used, so here we keep the restart point logic for Snappy.
Closes https://github.com/facebook/rocksdb/pull/3996

Differential Revision: D8416186

Pulled By: miasantreble

fbshipit-source-id: 002c0b62b9e5d89fb7736563d354ce0023c8cb28

db/plain_table_db_test.cc

@@ -333,21 +333,23 @@ class TestPlainTableFactory : public PlainTableFactory {
     TableProperties* props = nullptr;
     auto s =
         ReadTableProperties(file.get(), file_size, kPlainTableMagicNumber,
-                            table_reader_options.ioptions, &props);
+                            table_reader_options.ioptions, &props,
+                            true /* compression_type_missing */);
     EXPECT_TRUE(s.ok());
 
     if (store_index_in_file_) {
       BlockHandle bloom_block_handle;
       s = FindMetaBlock(file.get(), file_size, kPlainTableMagicNumber,
                         table_reader_options.ioptions,
-                        BloomBlockBuilder::kBloomBlock, &bloom_block_handle);
+                        BloomBlockBuilder::kBloomBlock, &bloom_block_handle,
+                        /* compression_type_missing */ true);
       EXPECT_TRUE(s.ok());
 
       BlockHandle index_block_handle;
       s = FindMetaBlock(file.get(), file_size, kPlainTableMagicNumber,
                         table_reader_options.ioptions,
                         PlainTableIndexBuilder::kPlainTableIndexBlock,
-                        &index_block_handle);
+                        &index_block_handle, /* compression_type_missing */ true);
       EXPECT_TRUE(s.ok());
     }
 

db/table_properties_collector_test.cc

@@ -279,7 +279,8 @@ void TestCustomizedTablePropertiesCollector(
           new test::StringSource(fwf->contents())));
   TableProperties* props;
   Status s = ReadTableProperties(fake_file_reader.get(), fwf->contents().size(),
-                                 magic_number, ioptions, &props);
+                                 magic_number, ioptions, &props,
+                                 true /* compression_type_missing */);
   std::unique_ptr<TableProperties> props_guard(props);
   ASSERT_OK(s);
 
@@ -421,7 +422,7 @@ void TestInternalKeyPropertiesCollector(
     TableProperties* props;
     Status s =
         ReadTableProperties(reader.get(), fwf->contents().size(), magic_number,
-                            ioptions, &props);
+                            ioptions, &props, true /* compression_type_missing */);
     ASSERT_OK(s);
 
     std::unique_ptr<TableProperties> props_guard(props);

db/version_set.cc

@@ -761,7 +761,8 @@ Status Version::GetTableProperties(std::shared_ptr<const TableProperties>* tp,
       new RandomAccessFileReader(std::move(file), file_name));
   s = ReadTableProperties(
       file_reader.get(), file_meta->fd.GetFileSize(),
-      Footer::kInvalidTableMagicNumber /* table's magic number */, *ioptions, &raw_table_properties);
+      Footer::kInvalidTableMagicNumber /* table's magic number */, *ioptions,
+      &raw_table_properties, false /* compression_type_missing */);
   if (!s.ok()) {
     return s;
   }

table/block.cc

@@ -427,13 +427,16 @@ Block::Block(BlockContents&& contents, SequenceNumber _global_seqno,
   if (size_ < sizeof(uint32_t)) {
     size_ = 0;  // Error marker
   } else {
-    num_restarts_ = NumRestarts();
-    restart_offset_ =
-        static_cast<uint32_t>(size_) - (1 + num_restarts_) * sizeof(uint32_t);
-    if (restart_offset_ > size_ - sizeof(uint32_t)) {
-      // The size is too small for NumRestarts() and therefore
-      // restart_offset_ wrapped around.
-      size_ = 0;
+    // Should only decode restart points for uncompressed blocks
+    if (compression_type() == kNoCompression) {
+      num_restarts_ = NumRestarts();
+      restart_offset_ =
+          static_cast<uint32_t>(size_) - (1 + num_restarts_) * sizeof(uint32_t);
+      if (restart_offset_ > size_ - sizeof(uint32_t)) {
+        // The size is too small for NumRestarts() and therefore
+        // restart_offset_ wrapped around.
+        size_ = 0;
+      }
     }
   }
   if (read_amp_bytes_per_bit != 0 && statistics && size_ != 0) {

table/block_based_table_reader.cc

@@ -806,7 +806,7 @@ Status BlockBasedTable::Open(const ImmutableCFOptions& ioptions,
     if (s.ok()) {
       s = ReadProperties(meta_iter->value(), rep->file.get(),
                          prefetch_buffer.get(), rep->footer, rep->ioptions,
-                         &table_properties);
+                         &table_properties, false /* compression_type_missing */);
     }
 
     if (!s.ok()) {

table/cuckoo_table_builder_test.cc

@@ -60,7 +60,7 @@ class CuckooBuilderTest : public testing::Test {
         new RandomAccessFileReader(std::move(read_file), fname));
     ASSERT_OK(ReadTableProperties(file_reader.get(), read_file_size,
                                   kCuckooTableMagicNumber, ioptions,
-                                  &props));
+                                  &props, true /* compression_type_missing */));
     // Check unused bucket.
     std::string unused_key = props->user_collected_properties[
       CuckooTablePropertyNames::kEmptyKey];

table/cuckoo_table_reader.cc

@@ -57,7 +57,7 @@ CuckooTableReader::CuckooTableReader(
   }
   TableProperties* props = nullptr;
   status_ = ReadTableProperties(file_.get(), file_size, kCuckooTableMagicNumber,
-      ioptions, &props);
+      ioptions, &props, true /* compression_type_missing */);
   if (!status_.ok()) {
     return;
   }

table/meta_blocks.cc

@@ -169,7 +169,8 @@ bool NotifyCollectTableCollectorsOnFinish(
 Status ReadProperties(const Slice& handle_value, RandomAccessFileReader* file,
                       FilePrefetchBuffer* prefetch_buffer, const Footer& footer,
                       const ImmutableCFOptions& ioptions,
-                      TableProperties** table_properties) {
+                      TableProperties** table_properties,
+                      bool compression_type_missing) {
   assert(table_properties);
 
   Slice v = handle_value;
@@ -189,6 +190,11 @@ Status ReadProperties(const Slice& handle_value, RandomAccessFileReader* file,
       file, prefetch_buffer, footer, read_options, handle, &block_contents,
       ioptions, false /* decompress */, compression_dict, cache_options);
   s = block_fetcher.ReadBlockContents();
+  // override compression_type when table file is known to contain undefined
+  // value at compression type marker
+  if (compression_type_missing) {
+    block_contents.compression_type = kNoCompression;
+  }
 
   if (!s.ok()) {
     return s;
@@ -293,7 +299,8 @@ Status ReadProperties(const Slice& handle_value, RandomAccessFileReader* file,
 Status ReadTableProperties(RandomAccessFileReader* file, uint64_t file_size,
                            uint64_t table_magic_number,
                            const ImmutableCFOptions &ioptions,
-                           TableProperties** properties) {
+                           TableProperties** properties,
+                           bool compression_type_missing) {
   // -- Read metaindex block
   Footer footer;
   auto s = ReadFooterFromFile(file, nullptr /* prefetch_buffer */, file_size,
@@ -317,6 +324,11 @@ Status ReadTableProperties(RandomAccessFileReader* file, uint64_t file_size,
   if (!s.ok()) {
     return s;
   }
+  // override compression_type when table file is known to contain undefined
+  // value at compression type marker
+  if (compression_type_missing) {
+    metaindex_contents.compression_type = kNoCompression;
+  }
   Block metaindex_block(std::move(metaindex_contents),
                         kDisableGlobalSequenceNumber);
   std::unique_ptr<InternalIterator> meta_iter(
@@ -332,7 +344,7 @@ Status ReadTableProperties(RandomAccessFileReader* file, uint64_t file_size,
   TableProperties table_properties;
   if (found_properties_block == true) {
     s = ReadProperties(meta_iter->value(), file, nullptr /* prefetch_buffer */,
-                       footer, ioptions, properties);
+                       footer, ioptions, properties, compression_type_missing);
   } else {
     s = Status::NotFound();
   }
@@ -357,7 +369,8 @@ Status FindMetaBlock(RandomAccessFileReader* file, uint64_t file_size,
                      uint64_t table_magic_number,
                      const ImmutableCFOptions &ioptions,
                      const std::string& meta_block_name,
-                     BlockHandle* block_handle) {
+                     BlockHandle* block_handle,
+                     bool compression_type_missing) {
   Footer footer;
   auto s = ReadFooterFromFile(file, nullptr /* prefetch_buffer */, file_size,
                               &footer, table_magic_number);
@@ -379,6 +392,11 @@ Status FindMetaBlock(RandomAccessFileReader* file, uint64_t file_size,
   if (!s.ok()) {
     return s;
   }
+  // override compression_type when table file is known to contain undefined
+  // value at compression type marker
+  if (compression_type_missing) {
+    metaindex_contents.compression_type = kNoCompression;
+  }
   Block metaindex_block(std::move(metaindex_contents),
                         kDisableGlobalSequenceNumber);
 
@@ -394,7 +412,7 @@ Status ReadMetaBlock(RandomAccessFileReader* file,
                      uint64_t table_magic_number,
                      const ImmutableCFOptions& ioptions,
                      const std::string& meta_block_name,
-                     BlockContents* contents) {
+                     BlockContents* contents, bool compression_type_missing) {
   Status status;
   Footer footer;
   status = ReadFooterFromFile(file, prefetch_buffer, file_size, &footer,
@@ -419,6 +437,11 @@ Status ReadMetaBlock(RandomAccessFileReader* file,
   if (!status.ok()) {
     return status;
   }
+  // override compression_type when table file is known to contain undefined
+  // value at compression type marker
+  if (compression_type_missing) {
+    metaindex_contents.compression_type = kNoCompression;
+  }
 
   // Finding metablock
   Block metaindex_block(std::move(metaindex_contents),

table/meta_blocks.h

@@ -96,16 +96,22 @@ bool NotifyCollectTableCollectorsOnFinish(
 Status ReadProperties(const Slice& handle_value, RandomAccessFileReader* file,
                       FilePrefetchBuffer* prefetch_buffer, const Footer& footer,
                       const ImmutableCFOptions& ioptions,
-                      TableProperties** table_properties);
+                      TableProperties** table_properties,
+                      bool compression_type_missing = false);
 
 // Directly read the properties from the properties block of a plain table.
 // @returns a status to indicate if the operation succeeded. On success,
 //          *table_properties will point to a heap-allocated TableProperties
 //          object, otherwise value of `table_properties` will not be modified.
+// certain tables do not have compression_type byte setup properly for
+// uncompressed blocks, caller can request to reset compression type by
+// passing compression_type_missing = true, the same applies to
+// `ReadProperties`, `FindMetaBlock`, and `ReadMetaBlock`
 Status ReadTableProperties(RandomAccessFileReader* file, uint64_t file_size,
                            uint64_t table_magic_number,
                            const ImmutableCFOptions &ioptions,
-                           TableProperties** properties);
+                           TableProperties** properties,
+                           bool compression_type_missing = false);
 
 // Find the meta block from the meta index block.
 Status FindMetaBlock(InternalIterator* meta_index_iter,
@@ -117,7 +123,8 @@ Status FindMetaBlock(RandomAccessFileReader* file, uint64_t file_size,
                      uint64_t table_magic_number,
                      const ImmutableCFOptions &ioptions,
                      const std::string& meta_block_name,
-                     BlockHandle* block_handle);
+                     BlockHandle* block_handle,
+                     bool compression_type_missing = false);
 
 // Read the specified meta block with name meta_block_name
 // from `file` and initialize `contents` with contents of this block.
@@ -127,6 +134,7 @@ Status ReadMetaBlock(RandomAccessFileReader* file,
                      uint64_t table_magic_number,
                      const ImmutableCFOptions& ioptions,
                      const std::string& meta_block_name,
-                     BlockContents* contents);
+                     BlockContents* contents,
+                     bool compression_type_missing = false);
 
 }  // namespace rocksdb

table/plain_table_reader.cc

@@ -128,7 +128,8 @@ Status PlainTableReader::Open(
 
   TableProperties* props = nullptr;
   auto s = ReadTableProperties(file.get(), file_size, kPlainTableMagicNumber,
-                               ioptions, &props);
+                               ioptions, &props,
+                               true /* compression_type_missing */);
   if (!s.ok()) {
     return s;
   }
@@ -293,7 +294,8 @@ Status PlainTableReader::PopulateIndex(TableProperties* props,
   Status s = ReadMetaBlock(file_info_.file.get(), nullptr /* prefetch_buffer */,
                            file_size_, kPlainTableMagicNumber, ioptions_,
                            PlainTableIndexBuilder::kPlainTableIndexBlock,
-                           &index_block_contents);
+                           &index_block_contents,
+                           true /* compression_type_missing */);
 
   bool index_in_file = s.ok();
 
@@ -303,7 +305,8 @@ Status PlainTableReader::PopulateIndex(TableProperties* props,
   if (index_in_file) {
     s = ReadMetaBlock(file_info_.file.get(), nullptr /* prefetch_buffer */,
                       file_size_, kPlainTableMagicNumber, ioptions_,
-                      BloomBlockBuilder::kBloomBlock, &bloom_block_contents);
+                      BloomBlockBuilder::kBloomBlock, &bloom_block_contents,
+                      true /* compression_type_missing */);
     bloom_in_file = s.ok() && bloom_block_contents.data.size() > 0;
   }
 

table/table_test.cc

@@ -2575,7 +2575,7 @@ TEST_F(PlainTableTest, BasicPlainTableProperties) {
   TableProperties* props = nullptr;
   auto s = ReadTableProperties(file_reader.get(), ss->contents().size(),
                                kPlainTableMagicNumber, ioptions,
-                               &props);
+                               &props, true /* compression_type_missing */);
   std::unique_ptr<TableProperties> props_guard(props);
   ASSERT_OK(s);
 
@@ -3169,7 +3169,7 @@ TEST_P(BlockBasedTableTest, TableWithGlobalSeqno) {
     TableProperties* props = nullptr;
     ASSERT_OK(ReadTableProperties(file_reader.get(), ss_rw.contents().size(),
                                   kBlockBasedTableMagicNumber, ioptions,
-                                  &props));
+                                  &props, true /* compression_type_missing */));
 
     UserCollectedProperties user_props = props->user_collected_properties;
     version = DecodeFixed32(
@@ -3346,7 +3346,7 @@ TEST_P(BlockBasedTableTest, BlockAlignTest) {
     TableProperties* props = nullptr;
     ASSERT_OK(ReadTableProperties(file_reader.get(), ss_rw.contents().size(),
                                   kBlockBasedTableMagicNumber, ioptions,
-                                  &props));
+                                  &props, true /* compression_type_missing */));
 
     uint64_t data_block_size = props->data_size / props->num_data_blocks;
     ASSERT_EQ(data_block_size, 4096);