Improve distinct compression for index and data blocks (facebook#14140)

Summary:
This change enables a custom CompressionManager / Compressor to adopt custom handling for data and index blocks. In particular, index blocks for format_version >= 4 use a distinct variant of the block format. Thus, a potentially format-aware compression algorithm such as OpenZL should be told which kind of block we are compressing. (And previously I avoided passing block type in CompressBlock for efficient handling of things like dictionaries but also avoiding checks on every CompressBlock call.)

Most of the change is in BlockBasedTableBuilder to call MaybeCloneSpecialized for both kDataBlock and for kIndexBlock. But I also needed some small tweaks/additions to the public API also:
* Require a Clone() function from Compressors, to support proper implementations of MaybeCloneSpecialized() in wrapper Compressors.
* Assert that the default implementation of CompressorWrapper::MaybeCloneSpecialized() is only used in allowable cases.
* Convenience function Compressor::CloneMaybeSpecialized()

This also fixes a serious bug/oversight in ManagedPtr for (ManagedWorkingArea) that somehow wasn't showing up before. It probably doesn't need a release note because CompressionManager stuff is still considered experimental.

Pull Request resolved: facebook#14140

Test Plan: Greatly expanded DBCompressionTest.CompressionManagerWrapper to make sure the distinction between data blocks and index blocks is properly communicated to a custom CompressionManager/Compressor. The test includes processing the expected structure of data and index blocks, to serve as a tested example for structure-aware compressors.

Reviewed By: hx235

Differential Revision: D87600019

Pulled By: pdillinger

fbshipit-source-id: 252ef78910073a0e45f2c81dd45ac87ff8a41fc6

Author: pdillinger

include/rocksdb/advanced_compression.h

@@ -90,6 +90,10 @@ class Compressor {
     return CompressionType::kDisableCompressionOption;
   }
 
+  // Return a distinct but functionally equivalent Compressor. This is often
+  // needed to implement MaybeCloneSpecialized() in wrapper compressors.
+  virtual std::unique_ptr<Compressor> Clone() const = 0;
+
   // Utility struct for providing sample data for the compression dictionary.
   // Potentially extensible by callers of Compressor (but not recommended)
   struct DictSampleArgs {
@@ -131,6 +135,18 @@ class Compressor {
     return nullptr;
   }
 
+  // A convenience function when a clone is needed and may or may not be
+  // specialized.
+  std::unique_ptr<Compressor> CloneMaybeSpecialized(
+      CacheEntryRole block_type, DictSampleArgs&& dict_samples) {
+    auto clone = MaybeCloneSpecialized(block_type, std::move(dict_samples));
+    if (clone == nullptr) {
+      clone = Clone();
+      assert(clone != nullptr);
+    }
+    return clone;
+  }
+
   // A WorkingArea is an optional structure (both for callers and
   // implementations) that can enable optimizing repeated compressions by
   // reusing working space or thread-local tracking of statistics or trends.
@@ -473,9 +489,19 @@ class CompressorWrapper : public Compressor {
     return wrapped_->GetPreferredCompressionType();
   }
 
+  // NOTE: Clone() not implemented here because it needs to be in the derived
+  // class
+
+  // NOTE: MaybeCloneSpecialized() is only implemented here for convenience
+  // when the wrapped Compressor uses the default implementation of
+  // MaybeCloneSpecialized(). This needs to be overridden if not.
   std::unique_ptr<Compressor> MaybeCloneSpecialized(
       CacheEntryRole block_type, DictSampleArgs&& dict_samples) override {
-    return wrapped_->MaybeCloneSpecialized(block_type, std::move(dict_samples));
+    auto clone =
+        wrapped_->MaybeCloneSpecialized(block_type, std::move(dict_samples));
+    // Assert default no-op MaybeCloneSpecialized()
+    assert(clone == nullptr);
+    return clone;
   }
 
   ManagedWorkingArea ObtainWorkingArea() override {

include/rocksdb/data_structure.h

@@ -247,15 +247,7 @@ class ManagedPtr {
  public:
   ManagedPtr() = default;
   ManagedPtr(T* ptr, Owner* owner) : ptr_(ptr), owner_(owner) {}
-  ~ManagedPtr() {
-    if (ptr_ && owner_) {
-      if constexpr (std::is_member_function_pointer_v<decltype(Fn)>) {
-        (owner_->*Fn)(ptr_);
-      } else {
-        Fn(owner_, ptr_);
-      }
-    }
-  }
+  ~ManagedPtr() { Free(); }
   // No copies
   ManagedPtr(const ManagedPtr&) = delete;
   ManagedPtr& operator=(const ManagedPtr&) = delete;
@@ -267,6 +259,10 @@ class ManagedPtr {
     other.owner_ = nullptr;
   }
   ManagedPtr& operator=(ManagedPtr&& other) noexcept {
+    if (this == &other) {
+      return *this;
+    }
+    Free();
     ptr_ = other.ptr_;
     owner_ = other.owner_;
     other.ptr_ = nullptr;
@@ -284,6 +280,16 @@ class ManagedPtr {
  private:
   T* ptr_ = nullptr;
   Owner* owner_ = nullptr;
+
+  void Free() {
+    if (ptr_ && owner_) {
+      if constexpr (std::is_member_function_pointer_v<decltype(Fn)>) {
+        (owner_->*Fn)(ptr_);
+      } else {
+        Fn(owner_, ptr_);
+      }
+    }
+  }
 };
 
 template <typename T, typename comp>

table/block_based/block_based_table_builder.cc

@@ -110,6 +110,19 @@ FilterBlockBuilder* CreateFilterBlockBuilder(
   }
 }
 
+// A convenience function for populating the Compressor* fields; see ~Rep()
+Compressor* MaybeCloneSpecialized(
+    Compressor* compressor, CacheEntryRole block_type,
+    Compressor::DictSampleArgs&& dict_samples = {}) {
+  auto specialized =
+      compressor->MaybeCloneSpecialized(block_type, std::move(dict_samples));
+  if (specialized) {
+    // Caller is responsible for freeing when distinct
+    return specialized.release();
+  } else {
+    return compressor;
+  }
+}
 }  // namespace
 
 // kBlockBasedTableMagicNumber was picked by running
@@ -824,15 +837,17 @@ struct BlockBasedTableBuilder::Rep {
 
   // *** Compressors & decompressors - Yes, it seems like a lot here but ***
   // *** these are distinct fields to minimize extra conditionals and    ***
-  // *** field reads on hot code paths.                                  ***
+  // *** field reads on hot code paths. And to avoid interlocked         ***
+  // *** instructions associated with shared_ptr.                        ***
 
   // A compressor for blocks in general, without dictionary compression
   std::unique_ptr<Compressor> basic_compressor;
-  // A compressor using dictionary compression (when applicable)
-  std::unique_ptr<Compressor> compressor_with_dict;
-  // Once configured/determined, points to one of the above Compressors to
-  // use on data blocks.
-  Compressor* data_block_compressor = nullptr;
+  // A compressor for data blocks, which might be tuned differently and might
+  // use dictionary compression (when applicable). See ~Rep() for some details.
+  UnownedPtr<Compressor> data_block_compressor = nullptr;
+  // A compressor for index blocks, which might be tuned differently from
+  // basic_compressor. See ~Rep() for some details.
+  UnownedPtr<Compressor> index_block_compressor = nullptr;
   // A decompressor corresponding to basic_compressor (when non-nullptr).
   // Used for verification and cache warming.
   std::shared_ptr<Decompressor> basic_decompressor;
@@ -853,7 +868,7 @@ struct BlockBasedTableBuilder::Rep {
       compression_types_used;
 
   // Working area for basic_compressor when compression_parallel_threads==1
-  WorkingAreaPair basic_working_area;
+  WorkingAreaPair index_block_working_area;
   // Working area for data_block_compressor, for emit/compaction thread
   WorkingAreaPair data_block_working_area;
 
@@ -1099,7 +1114,10 @@ struct BlockBasedTableBuilder::Rep {
         filter_context, tbo.compression_opts, tbo.compression_type);
     if (basic_compressor) {
       if (table_options.enable_index_compression) {
-        basic_working_area.compress = basic_compressor->ObtainWorkingArea();
+        index_block_compressor = MaybeCloneSpecialized(
+            basic_compressor.get(), CacheEntryRole::kIndexBlock);
+        index_block_working_area.compress =
+            index_block_compressor->ObtainWorkingArea();
       }
       max_dict_sample_bytes = basic_compressor->GetMaxSampleSizeIfWantDict(
           CacheEntryRole::kDataBlock);
@@ -1114,8 +1132,10 @@ struct BlockBasedTableBuilder::Rep {
                                   tbo.compression_opts.max_dict_buffer_bytes);
         }
       } else {
-        // No distinct data block compressor using dictionary
-        data_block_compressor = basic_compressor.get();
+        // No distinct data block compressor using dictionary, but
+        // implementation might still want to specialize for data blocks
+        data_block_compressor = MaybeCloneSpecialized(
+            basic_compressor.get(), CacheEntryRole::kDataBlock);
         data_block_working_area.compress =
             data_block_compressor->ObtainWorkingArea();
       }
@@ -1129,8 +1149,9 @@ struct BlockBasedTableBuilder::Rep {
       if (table_options.verify_compression) {
         verify_decompressor = basic_decompressor.get();
         if (table_options.enable_index_compression) {
-          basic_working_area.verify = verify_decompressor->ObtainWorkingArea(
-              basic_compressor->GetPreferredCompressionType());
+          index_block_working_area.verify =
+              verify_decompressor->ObtainWorkingArea(
+                  index_block_compressor->GetPreferredCompressionType());
         }
         if (state == State::kUnbuffered) {
           assert(data_block_compressor);
@@ -1295,8 +1316,19 @@ struct BlockBasedTableBuilder::Rep {
   }
 
   ~Rep() {
+    // Delete working areas before their compressors.
+    index_block_working_area = {};
+    data_block_working_area = {};
     // Must have been cleaned up by StopParallelCompression
     assert(pc_rep == nullptr);
+    // Delete specialized compressors if they were distinct (avoiding extra
+    // fields and interlocked instructions with shared_ptr)
+    if (data_block_compressor.get() != basic_compressor.get()) {
+      delete data_block_compressor.get();
+    }
+    if (index_block_compressor.get() != basic_compressor.get()) {
+      delete index_block_compressor.get();
+    }
   }
 
   Rep(const Rep&) = delete;
@@ -1729,9 +1761,11 @@ void BlockBasedTableBuilder::WriteBlock(const Slice& uncompressed_block_data,
   assert(!r->IsParallelCompressionActive());
   CompressionType type;
   bool is_data_block = block_type == BlockType::kData;
+  // NOTE: only index and data blocks are currently compressed
+  assert(is_data_block || block_type == BlockType::kIndex);
   Status compress_status = CompressAndVerifyBlock(
       uncompressed_block_data, is_data_block,
-      is_data_block ? r->data_block_working_area : r->basic_working_area,
+      is_data_block ? r->data_block_working_area : r->index_block_working_area,
       &r->single_threaded_compressed_output, &type);
   r->SetStatus(compress_status);
   if (UNLIKELY(!ok())) {
@@ -1845,13 +1879,13 @@ Status BlockBasedTableBuilder::CompressAndVerifyBlock(
   Rep* r = rep_.get();
   Status status;
 
-  Compressor* compressor = nullptr;
+  UnownedPtr<Compressor> compressor = nullptr;
   Decompressor* verify_decomp = nullptr;
   if (is_data_block) {
     compressor = r->data_block_compressor;
     verify_decomp = r->data_block_verify_decompressor.get();
   } else {
-    compressor = r->basic_compressor.get();
+    compressor = r->index_block_compressor;
     verify_decomp = r->verify_decompressor.get();
   }
 
@@ -2116,7 +2150,7 @@ void BlockBasedTableBuilder::MaybeStartParallelCompression() {
   // that latency. So even with some optimizations, turning on the parallel
   // framework when compression is disabled just eats more CPU with little-to-no
   // improvement in throughput.
-  if (rep_->data_block_compressor == nullptr) {
+  if (!rep_->data_block_compressor) {
     // Force the generally best configuration for no compression: no parallelism
     return;
   }
@@ -2463,8 +2497,8 @@ void BlockBasedTableBuilder::WritePropertiesBlock(
 void BlockBasedTableBuilder::WriteCompressionDictBlock(
     MetaIndexBuilder* meta_index_builder) {
   Slice compression_dict;
-  if (rep_->compressor_with_dict) {
-    compression_dict = rep_->compressor_with_dict->GetSerializedDict();
+  if (rep_->data_block_compressor) {
+    compression_dict = rep_->data_block_compressor->GetSerializedDict();
   }
   if (!compression_dict.empty()) {
     BlockHandle compression_dict_block_handle;
@@ -2559,6 +2593,7 @@ void BlockBasedTableBuilder::MaybeEnterUnbuffered(
     // The below code is neither safe nor necessary for handling zero data
     // blocks.
     // For PostPopulateCompressionProperties()
+    assert(!r->data_block_compressor);
     r->data_block_compressor = r->basic_compressor.get();
     return;
   }
@@ -2600,15 +2635,12 @@ void BlockBasedTableBuilder::MaybeEnterUnbuffered(
 
   assert(samples.sample_data.size() > 0);
 
-  // final sample data block flushed, now we can generate dictionary
-  r->compressor_with_dict = r->basic_compressor->MaybeCloneSpecialized(
-      CacheEntryRole::kDataBlock, std::move(samples));
+  // final sample data block flushed, now we can generate dictionary (or it
+  // might opt not to use a dictionary and that's ok)
+  r->data_block_compressor =
+      MaybeCloneSpecialized(r->basic_compressor.get(),
+                            CacheEntryRole::kDataBlock, std::move(samples));
 
-  // The compressor might opt not to use a dictionary, in which case we
-  // can use the same compressor as for e.g. index blocks.
-  r->data_block_compressor = r->compressor_with_dict
-                                 ? r->compressor_with_dict.get()
-                                 : r->basic_compressor.get();
   Slice serialized_dict = r->data_block_compressor->GetSerializedDict();
   if (r->verify_decompressor) {
     if (serialized_dict.empty()) {
@@ -2831,8 +2863,8 @@ uint64_t BlockBasedTableBuilder::EstimatedTailSize() const {
   }
 
   // 3. Estimate compression dictionary size
-  if (rep_->compressor_with_dict) {
-    Slice dict = rep_->compressor_with_dict->GetSerializedDict();
+  if (rep_->data_block_compressor) {
+    Slice dict = rep_->data_block_compressor->GetSerializedDict();
     if (!dict.empty()) {
       estimated_tail_size += dict.size();
     }

test_util/testutil.h

@@ -766,6 +766,10 @@ struct CompressorCustomAlg : public CompressorWrapper {
     return kCompression;
   }
 
+  std::unique_ptr<Compressor> Clone() const override {
+    return std::make_unique<CompressorCustomAlg>(wrapped_->Clone());
+  }
+
   Status CompressBlock(Slice uncompressed_data, char* compressed_output,
                        size_t* compressed_output_size,
                        CompressionType* out_compression_type,
@@ -794,7 +798,7 @@ struct CompressorCustomAlg : public CompressorWrapper {
   std::unique_ptr<Compressor> MaybeCloneSpecialized(
       CacheEntryRole block_type, DictSampleArgs&& dict_samples) override {
     auto clone =
-        wrapped_->MaybeCloneSpecialized(block_type, std::move(dict_samples));
+        wrapped_->CloneMaybeSpecialized(block_type, std::move(dict_samples));
     return std::make_unique<CompressorCustomAlg>(std::move(clone));
   }
 

util/auto_tune_compressor.cc

@@ -59,6 +59,10 @@ const char* AutoSkipCompressorWrapper::Name() const {
   return "AutoSkipCompressorWrapper";
 }
 
+std::unique_ptr<Compressor> AutoSkipCompressorWrapper::Clone() const {
+  return std::make_unique<AutoSkipCompressorWrapper>(wrapped_->Clone(), opts_);
+}
+
 Status AutoSkipCompressorWrapper::CompressBlock(
     Slice uncompressed_data, char* compressed_output,
     size_t* compressed_output_size, CompressionType* out_compression_type,
@@ -174,6 +178,10 @@ CostAwareCompressor::CostAwareCompressor(const CompressionOptions& opts)
 }
 
 const char* CostAwareCompressor::Name() const { return "CostAwareCompressor"; }
+
+std::unique_ptr<Compressor> CostAwareCompressor::Clone() const {
+  return std::make_unique<CostAwareCompressor>(opts_);
+}
 size_t CostAwareCompressor::GetMaxSampleSizeIfWantDict(
     CacheEntryRole block_type) const {
   auto idx = allcompressors_index_.back();

util/auto_tune_compressor.h

@@ -64,6 +64,7 @@ class AutoSkipCompressorWrapper : public CompressorWrapper {
   explicit AutoSkipCompressorWrapper(std::unique_ptr<Compressor> compressor,
                                      const CompressionOptions& opts);
 
+  std::unique_ptr<Compressor> Clone() const override;
   Status CompressBlock(Slice uncompressed_data, char* compressed_output,
                        size_t* compressed_output_size,
                        CompressionType* out_compression_type,
@@ -149,6 +150,7 @@ class CostAwareCompressor : public Compressor {
  public:
   explicit CostAwareCompressor(const CompressionOptions& opts);
   const char* Name() const override;
+  std::unique_ptr<Compressor> Clone() const override;
   size_t GetMaxSampleSizeIfWantDict(CacheEntryRole block_type) const override;
   Slice GetSerializedDict() const override;
   CompressionType GetPreferredCompressionType() const override;