Support Super Block Alignment

Summary:

Add restart point in index builder when super block alignment padded the
data block

Test Plan:
Unit Test

Reviewers:

Subscribers:

Tasks:
T227973776

Tags:

Author: Xingbo Wang

db/db_flush_test.cc

@@ -3561,6 +3561,69 @@ TEST_F(DBFlushTest, VerifyOutputRecordCount) {
     ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->DisableProcessing();
   }
 }
+
+std::string formatKey(int i) {
+  int desired_length = 10;
+  char buffer[64];
+  sprintf(buffer, "%0*d", desired_length, i);
+  return buffer;
+}
+
+std::string formatValue(int i) { return std::string(1234, 'a' + (i % 26)); }
+
+TEST_F(DBFlushTest, SuperBlock) {
+  constexpr int key_count = 12345;
+  Options options;
+  BlockBasedTableOptions block_options;
+  block_options.block_align = false;
+  block_options.super_block_align = true;
+  block_options.super_block_alignment_size = 64 * 1024;
+  block_options.super_block_alignment_max_padding_size = 4 * 1024;
+  options.table_factory.reset(NewBlockBasedTableFactory(block_options));
+
+  Reopen(options);
+
+  int super_block_pad_count = 0;
+  SyncPoint::GetInstance()->SetCallBack(
+      "BlockBasedTableBuilder::WriteMaybeCompressedBlock:"
+      "SuperBlockAlignment",
+      [&super_block_pad_count](void* /*arg*/) { super_block_pad_count++; });
+  SyncPoint::GetInstance()->EnableProcessing();
+
+  // Add lots of keys
+  for (int i = 0; i < key_count; ++i) {
+    Put(formatKey(i), formatValue(i));
+  }
+
+  // flush the data in memory to disk to verify with super block alignment, the
+  // data could be read back properly
+  Reopen(options);
+
+  SyncPoint::GetInstance()->DisableProcessing();
+  SyncPoint::GetInstance()->ClearAllCallBacks();
+
+  ASSERT_GT(super_block_pad_count, 0);
+
+  // verify the values are correct
+  // Test GET
+  for (int i = 0; i < key_count; ++i) {
+    PinnableSlice value;
+    ASSERT_OK(Get(formatKey(i), &value));
+    ASSERT_EQ(value.ToString(), formatValue(i));
+  }
+  // Test iterator
+  {
+    std::unique_ptr<Iterator> it(db_->NewIterator(ReadOptions()));
+    int i = 0;
+    for (it->SeekToFirst(); it->Valid(); it->Next()) {
+      ASSERT_OK(it->status());
+      ASSERT_EQ((it->key()).ToString(), formatKey(i));
+      ASSERT_EQ((it->value()).ToString(), formatValue(i));
+      i++;
+    }
+    ASSERT_EQ(i, key_count);
+  }
+}
 }  // namespace ROCKSDB_NAMESPACE
 
 int main(int argc, char** argv) {

include/rocksdb/table.h

@@ -597,6 +597,17 @@ struct BlockBasedTableOptions {
   // Align data blocks on lesser of page size and block size
   bool block_align = false;
 
+  // Align data blocks on super block alignment. Avoid a data block split across
+  // super block boundaries.
+  bool super_block_align = false;
+
+  // Super block alignment size. Default to 512 KB. It has to be a power of 2
+  // and higher than block size.
+  size_t super_block_alignment_size = 512 * 1024;
+
+  // Maximum number of bytes allowed to be padded for super block alignment.
+  size_t super_block_alignment_max_padding_size = 32 * 1024;
+
   // This enum allows trading off increased index size for improved iterator
   // seek performance in some situations, particularly when block cache is
   // disabled (ReadOptions::fill_cache = false) and direct IO is

table/block_based/block_based_table_builder.cc

@@ -1319,7 +1319,9 @@ void BlockBasedTableBuilder::EmitBlock(std::string& uncompressed,
                                              r->get_offset());
     r->pc_rep->EmitBlock(block_rep);
   } else {
-    WriteBlock(uncompressed, &r->pending_handle, BlockType::kData);
+    bool should_add_restart_point_on_index_block = false;
+    WriteBlock(uncompressed, &r->pending_handle, BlockType::kData,
+               &should_add_restart_point_on_index_block);
     if (ok()) {
       // We do not emit the index entry for a block until we have seen the
       // first key for the next data block.  This allows us to use shorter
@@ -1330,14 +1332,14 @@ void BlockBasedTableBuilder::EmitBlock(std::string& uncompressed,
       // blocks.
       r->index_builder->AddIndexEntry(
           last_key_in_current_block, first_key_in_next_block, r->pending_handle,
-          &r->index_separator_scratch);
+          &r->index_separator_scratch, should_add_restart_point_on_index_block);
     }
   }
 }
 
-void BlockBasedTableBuilder::WriteBlock(const Slice& uncompressed_block_data,
-                                        BlockHandle* handle,
-                                        BlockType block_type) {
+void BlockBasedTableBuilder::WriteBlock(
+    const Slice& uncompressed_block_data, BlockHandle* handle,
+    BlockType block_type, bool* should_add_restart_point_on_index_block) {
   Rep* r = rep_;
   assert(r->state == Rep::State::kUnbuffered);
   CompressionType type;
@@ -1358,7 +1360,8 @@ void BlockBasedTableBuilder::WriteBlock(const Slice& uncompressed_block_data,
   WriteMaybeCompressedBlock(type == kNoCompression
                                 ? uncompressed_block_data
                                 : Slice(r->single_threaded_compressed_output),
-                            type, handle, block_type, &uncompressed_block_data);
+                            type, handle, block_type, &uncompressed_block_data,
+                            should_add_restart_point_on_index_block);
   r->single_threaded_compressed_output.Reset();
   if (is_data_block) {
     r->props.data_size = r->get_offset();
@@ -1489,13 +1492,17 @@ void BlockBasedTableBuilder::CompressAndVerifyBlock(
 
 void BlockBasedTableBuilder::WriteMaybeCompressedBlock(
     const Slice& block_contents, CompressionType comp_type, BlockHandle* handle,
-    BlockType block_type, const Slice* uncompressed_block_data) {
+    BlockType block_type, const Slice* uncompressed_block_data,
+    bool* should_add_restart_point_on_index_block) {
   // File format contains a sequence of blocks where each block has:
   //    block_data: uint8[n]
   //    compression_type: uint8
   //    checksum: uint32
   Rep* r = rep_;
   bool is_data_block = block_type == BlockType::kData;
+  if (should_add_restart_point_on_index_block != nullptr) {
+    *should_add_restart_point_on_index_block = false;
+  }
   IOOptions io_options;
   IOStatus io_s =
       WritableFileWriter::PrepareIOOptions(r->write_options, io_options);
@@ -1505,7 +1512,37 @@ void BlockBasedTableBuilder::WriteMaybeCompressedBlock(
   }
   // Old, misleading name of this function: WriteRawBlock
   StopWatch sw(r->ioptions.clock, r->ioptions.stats, WRITE_RAW_BLOCK_MICROS);
-  const uint64_t offset = r->get_offset();
+
+  uint64_t offset = r->get_offset();
+  // try to align the page to the super alignment size, if enabled
+  if (r->table_options.super_block_align && is_data_block) {
+    auto super_block_alignment_mask =
+        r->table_options.super_block_alignment_size - 1;
+    if ((offset & (~super_block_alignment_mask)) !=
+        ((offset + block_contents.size()) & (~super_block_alignment_mask))) {
+      // new block would cross the super block boundary
+      auto pad_bytes = r->table_options.super_block_alignment_size -
+                       (offset & super_block_alignment_mask);
+      if (pad_bytes <=
+          r->table_options.super_block_alignment_max_padding_size) {
+        io_s = r->file->Pad(io_options, pad_bytes);
+        if (!io_s.ok()) {
+          r->SetIOStatus(io_s);
+          return;
+        }
+        r->pre_compression_size += pad_bytes;
+        offset += pad_bytes;
+        r->set_offset(offset);
+        if (should_add_restart_point_on_index_block != nullptr) {
+          *should_add_restart_point_on_index_block = true;
+        }
+        TEST_SYNC_POINT(
+            "BlockBasedTableBuilder::WriteMaybeCompressedBlock:"
+            "SuperBlockAlignment");
+      }
+    }
+  }
+
   handle->set_offset(offset);
   handle->set_size(block_contents.size());
   assert(status().ok());

table/block_based/block_based_table_builder.h

@@ -126,11 +126,13 @@ class BlockBasedTableBuilder : public TableBuilder {
 
   // Compress and write block content to the file.
   void WriteBlock(const Slice& block_contents, BlockHandle* handle,
-                  BlockType block_type);
+                  BlockType block_type,
+                  bool* should_add_restart_point_on_index_block = nullptr);
   // Directly write data to the file.
   void WriteMaybeCompressedBlock(
       const Slice& block_contents, CompressionType, BlockHandle* handle,
-      BlockType block_type, const Slice* uncompressed_block_data = nullptr);
+      BlockType block_type, const Slice* uncompressed_block_data = nullptr,
+      bool* should_add_restart_point_on_index_block = nullptr);
 
   void SetupCacheKeyPrefix(const TableBuilderOptions& tbo);
 

table/block_based/block_based_table_factory.cc

@@ -690,6 +690,24 @@ Status BlockBasedTableFactory::ValidateOptions(
     return Status::InvalidArgument(
         "block size exceeds maximum number (4GiB) allowed");
   }
+  if (table_options_.super_block_align &&
+      (table_options_.super_block_alignment_size &
+       (table_options_.super_block_alignment_size - 1))) {
+    return Status::InvalidArgument(
+        "Super Block alignment requested but super block alignment size is not "
+        "a power of 2");
+  }
+  if (table_options_.super_block_alignment_size >
+      std::numeric_limits<uint32_t>::max()) {
+    return Status::InvalidArgument(
+        "Super block alignment size exceeds maximum number (4GiB) allowed");
+  }
+  if (table_options_.super_block_alignment_max_padding_size >
+      table_options_.super_block_alignment_size) {
+    return Status::InvalidArgument(
+        "Super block alignment max padding size exceeds super block alignment "
+        "size");
+  }
   if (table_options_.data_block_index_type ==
           BlockBasedTableOptions::kDataBlockBinaryAndHash &&
       table_options_.data_block_hash_table_util_ratio <= 0) {

table/block_based/block_builder.h

@@ -74,6 +74,15 @@ class BlockBuilder {
                             : 0);
   }
 
+  // Add a restart point, if it is not at a restart point.
+  void Restart() {
+    if (counter_ != 0) {
+      restarts_.push_back(static_cast<uint32_t>(buffer_.size()));
+      estimate_ += sizeof(uint32_t);
+      counter_ = 0;
+    }
+  }
+
   // Returns an estimated block size after appending key and value.
   size_t EstimateSizeAfterKV(const Slice& key, const Slice& value) const;
 

table/block_based/index_builder.cc

@@ -248,16 +248,16 @@ void PartitionedIndexBuilder::FinishIndexEntry(const BlockHandle& block_handle,
 Slice PartitionedIndexBuilder::AddIndexEntry(
     const Slice& last_key_in_current_block,
     const Slice* first_key_in_next_block, const BlockHandle& block_handle,
-    std::string* separator_scratch) {
+    std::string* separator_scratch, bool add_restart_point) {
   // At least when running without parallel compression, maintain behavior of
   // avoiding a last index partition with just one entry
   if (first_key_in_next_block) {
     MaybeFlush(last_key_in_current_block, block_handle);
   }
 
-  auto sep = sub_index_builder_->AddIndexEntry(last_key_in_current_block,
-                                               first_key_in_next_block,
-                                               block_handle, separator_scratch);
+  auto sep = sub_index_builder_->AddIndexEntry(
+      last_key_in_current_block, first_key_in_next_block, block_handle,
+      separator_scratch, add_restart_point);
   entries_.back().key.assign(sep.data(), sep.size());
 
   if (!must_use_separator_with_seq_ &&

table/block_based/index_builder.h

@@ -75,7 +75,8 @@ class IndexBuilder {
   virtual Slice AddIndexEntry(const Slice& last_key_in_current_block,
                               const Slice* first_key_in_next_block,
                               const BlockHandle& block_handle,
-                              std::string* separator_scratch) = 0;
+                              std::string* separator_scratch,
+                              bool add_restart_point) = 0;
 
   // An abstract (extensible) holder for passing data from PrepareIndexEntry to
   // FinishIndexEntry (see below).
@@ -286,12 +287,14 @@ class ShortenedIndexBuilder : public IndexBuilder {
   void AddIndexEntryImpl(const Slice& separator_with_seq,
                          const Slice& first_internal_key,
                          const BlockHandle& block_handle,
-                         bool must_use_separator_with_seq) {
+                         bool must_use_separator_with_seq,
+                         bool add_restart_point) {
     IndexValue entry(block_handle, first_internal_key);
     std::string encoded_entry;
     std::string delta_encoded_entry;
     entry.EncodeTo(&encoded_entry, include_first_key_, nullptr);
-    if (use_value_delta_encoding_ && !last_encoded_handle_.IsNull()) {
+    if (use_value_delta_encoding_ && !last_encoded_handle_.IsNull() &&
+        !add_restart_point) {
       entry.EncodeTo(&delta_encoded_entry, include_first_key_,
                      &last_encoded_handle_);
     } else {
@@ -301,6 +304,10 @@ class ShortenedIndexBuilder : public IndexBuilder {
     last_encoded_handle_ = block_handle;
     const Slice delta_encoded_entry_slice(delta_encoded_entry);
 
+    if (add_restart_point) {
+      index_block_builder_.Restart();
+    }
+
     // TODO(yuzhangyu): fix this when "FindShortInternalKeySuccessor"
     //  optimization is available.
     // Timestamp aware comparator currently doesn't provide override for
@@ -322,15 +329,16 @@ class ShortenedIndexBuilder : public IndexBuilder {
   Slice AddIndexEntry(const Slice& last_key_in_current_block,
                       const Slice* first_key_in_next_block,
                       const BlockHandle& block_handle,
-                      std::string* separator_scratch) override {
+                      std::string* separator_scratch,
+                      bool add_restart_point) override {
     Slice separator_with_seq = GetSeparatorWithSeq(
         last_key_in_current_block, first_key_in_next_block, separator_scratch);
 
     std::string first_internal_key_buf;
     Slice first_internal_key = GetFirstInternalKey(&first_internal_key_buf);
 
     AddIndexEntryImpl(separator_with_seq, first_internal_key, block_handle,
-                      must_use_separator_with_seq_);
+                      must_use_separator_with_seq_, add_restart_point);
     current_block_first_internal_key_.clear();
     return separator_with_seq;
   }
@@ -388,7 +396,7 @@ class ShortenedIndexBuilder : public IndexBuilder {
     ShortenedPreparedIndexEntry* entry =
         static_cast<ShortenedPreparedIndexEntry*>(base_entry);
     AddIndexEntryImpl(entry->separator_with_seq, entry->first_internal_key,
-                      block_handle, entry->must_use_separator_with_seq);
+                      block_handle, entry->must_use_separator_with_seq, false);
   }
 
   using IndexBuilder::Finish;
@@ -483,11 +491,12 @@ class HashIndexBuilder : public IndexBuilder {
   Slice AddIndexEntry(const Slice& last_key_in_current_block,
                       const Slice* first_key_in_next_block,
                       const BlockHandle& block_handle,
-                      std::string* separator_scratch) override {
+                      std::string* separator_scratch,
+                      bool add_restart_point) override {
     ++current_restart_index_;
     return primary_index_builder_.AddIndexEntry(
         last_key_in_current_block, first_key_in_next_block, block_handle,
-        separator_scratch);
+        separator_scratch, add_restart_point);
   }
 
   std::unique_ptr<PreparedIndexEntry> CreatePreparedIndexEntry() override {
@@ -611,7 +620,8 @@ class PartitionedIndexBuilder : public IndexBuilder {
   Slice AddIndexEntry(const Slice& last_key_in_current_block,
                       const Slice* first_key_in_next_block,
                       const BlockHandle& block_handle,
-                      std::string* separator_scratch) override;
+                      std::string* separator_scratch,
+                      bool add_restart_point) override;
 
   std::unique_ptr<PreparedIndexEntry> CreatePreparedIndexEntry() override;
   void PrepareIndexEntry(const Slice& last_key_in_current_block,

table/block_based/partitioned_filter_block_test.cc

@@ -315,7 +315,8 @@ class PartitionedFilterBlockTest
         std::string(*InternalKey(user_key, 0, ValueType::kTypeValue).rep());
     BlockHandle dont_care_block_handle(1, 1);
     std::string scratch;
-    builder->AddIndexEntry(key, nullptr, dont_care_block_handle, &scratch);
+    builder->AddIndexEntry(key, nullptr, dont_care_block_handle, &scratch,
+                           false);
   }
 
   void CutABlock(PartitionedIndexBuilder* builder, const std::string& user_key,
@@ -327,7 +328,8 @@ class PartitionedFilterBlockTest
     BlockHandle dont_care_block_handle(1, 1);
     Slice slice = Slice(next_key.data(), next_key.size());
     std::string scratch;
-    builder->AddIndexEntry(key, &slice, dont_care_block_handle, &scratch);
+    builder->AddIndexEntry(key, &slice, dont_care_block_handle, &scratch,
+                           false);
   }
 
   int CountNumOfIndexPartitions(PartitionedIndexBuilder* builder) {

table/block_based/user_defined_index_wrapper.h

@@ -41,7 +41,8 @@ class UserDefinedIndexBuilderWrapper : public IndexBuilder {
   Slice AddIndexEntry(const Slice& last_key_in_current_block,
                       const Slice* first_key_in_next_block,
                       const BlockHandle& block_handle,
-                      std::string* separator_scratch) override {
+                      std::string* separator_scratch,
+                      bool add_restart_point) override {
     UserDefinedIndexBuilder::BlockHandle handle;
     handle.offset = block_handle.offset();
     handle.size = block_handle.size();
@@ -66,7 +67,7 @@ class UserDefinedIndexBuilderWrapper : public IndexBuilder {
     }
     return internal_index_builder_->AddIndexEntry(
         last_key_in_current_block, first_key_in_next_block, block_handle,
-        separator_scratch);
+        separator_scratch, add_restart_point);
   }
 
   // Not supported with parallel compression