From 35148aca91cda84d6fa9b295eb5500d6d965dca6 Mon Sep 17 00:00:00 2001 From: Peter Dillinger Date: Fri, 21 Nov 2025 16:34:49 -0800 Subject: [PATCH] Improve distinct compression for index and data blocks (#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: https://github.com/facebook/rocksdb/pull/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 --- include/rocksdb/advanced_compression.h | 28 +- include/rocksdb/data_structure.h | 24 +- .../block_based/block_based_table_builder.cc | 88 ++- test_util/testutil.h | 6 +- util/auto_tune_compressor.cc | 8 + util/auto_tune_compressor.h | 2 + util/compression.cc | 40 +- util/compression_test.cc | 666 ++++++++++++------ util/simple_mixed_compressor.cc | 11 +- util/simple_mixed_compressor.h | 3 + 10 files changed, 611 insertions(+), 265 deletions(-) diff --git a/include/rocksdb/advanced_compression.h b/include/rocksdb/advanced_compression.h index bd02949498..be3da68f51 100644 --- a/include/rocksdb/advanced_compression.h +++ b/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 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 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 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 { diff --git a/include/rocksdb/data_structure.h b/include/rocksdb/data_structure.h index bf0144cd29..7563a83abf 100644 --- a/include/rocksdb/data_structure.h +++ b/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) { - (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) { + (owner_->*Fn)(ptr_); + } else { + Fn(owner_, ptr_); + } + } + } }; template diff --git a/table/block_based/block_based_table_builder.cc b/table/block_based/block_based_table_builder.cc index 3b1befdfbc..74c90edea0 100644 --- a/table/block_based/block_based_table_builder.cc +++ b/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 basic_compressor; - // A compressor using dictionary compression (when applicable) - std::unique_ptr 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 data_block_compressor = nullptr; + // A compressor for index blocks, which might be tuned differently from + // basic_compressor. See ~Rep() for some details. + UnownedPtr index_block_compressor = nullptr; // A decompressor corresponding to basic_compressor (when non-nullptr). // Used for verification and cache warming. std::shared_ptr 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 = 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(); } diff --git a/test_util/testutil.h b/test_util/testutil.h index fc172b8e4b..aa837a972b 100644 --- a/test_util/testutil.h +++ b/test_util/testutil.h @@ -766,6 +766,10 @@ struct CompressorCustomAlg : public CompressorWrapper { return kCompression; } + std::unique_ptr Clone() const override { + return std::make_unique(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 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(std::move(clone)); } diff --git a/util/auto_tune_compressor.cc b/util/auto_tune_compressor.cc index ed3bff8127..221ebe3007 100644 --- a/util/auto_tune_compressor.cc +++ b/util/auto_tune_compressor.cc @@ -59,6 +59,10 @@ const char* AutoSkipCompressorWrapper::Name() const { return "AutoSkipCompressorWrapper"; } +std::unique_ptr AutoSkipCompressorWrapper::Clone() const { + return std::make_unique(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 CostAwareCompressor::Clone() const { + return std::make_unique(opts_); +} size_t CostAwareCompressor::GetMaxSampleSizeIfWantDict( CacheEntryRole block_type) const { auto idx = allcompressors_index_.back(); diff --git a/util/auto_tune_compressor.h b/util/auto_tune_compressor.h index 818d8c43e7..ebd09f1c1e 100644 --- a/util/auto_tune_compressor.h +++ b/util/auto_tune_compressor.h @@ -64,6 +64,7 @@ class AutoSkipCompressorWrapper : public CompressorWrapper { explicit AutoSkipCompressorWrapper(std::unique_ptr compressor, const CompressionOptions& opts); + std::unique_ptr 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 Clone() const override; size_t GetMaxSampleSizeIfWantDict(CacheEntryRole block_type) const override; Slice GetSerializedDict() const override; CompressionType GetPreferredCompressionType() const override; diff --git a/util/compression.cc b/util/compression.cc index 5831643d46..afa1e3d535 100644 --- a/util/compression.cc +++ b/util/compression.cc @@ -174,6 +174,10 @@ class BuiltinCompressorV1 : public CompressorBase { CompressionType GetPreferredCompressionType() const override { return type_; } + std::unique_ptr Clone() const override { + return std::make_unique(opts_, type_); + } + Status CompressBlock(Slice uncompressed_data, char* compressed_output, size_t* compressed_output_size, CompressionType* out_compression_type, @@ -226,6 +230,10 @@ class CompressorWithSimpleDictBase : public CompressorBase { // NOTE: empty dict is equivalent to no dict Slice GetSerializedDict() const override { return dict_data_; } + std::unique_ptr Clone() const override { + return CloneForDict(std::string{dict_data_}); + } + std::unique_ptr MaybeCloneSpecialized( CacheEntryRole /*block_type*/, DictSampleArgs&& dict_samples) final override { @@ -238,7 +246,8 @@ class CompressorWithSimpleDictBase : public CompressorBase { } } - virtual std::unique_ptr CloneForDict(std::string&& dict_data) = 0; + virtual std::unique_ptr CloneForDict( + std::string&& dict_data) const = 0; protected: const std::string dict_data_; @@ -257,7 +266,8 @@ class BuiltinSnappyCompressorV2 : public CompressorWithSimpleDictBase { return kSnappyCompression; } - std::unique_ptr CloneForDict(std::string&& dict_data) override { + std::unique_ptr CloneForDict( + std::string&& dict_data) const override { return std::make_unique(opts_, std::move(dict_data)); } @@ -349,7 +359,8 @@ class BuiltinZlibCompressorV2 : public CompressorWithSimpleDictBase { return kZlibCompression; } - std::unique_ptr CloneForDict(std::string&& dict_data) override { + std::unique_ptr CloneForDict( + std::string&& dict_data) const override { return std::make_unique(opts_, std::move(dict_data)); } @@ -447,7 +458,8 @@ class BuiltinBZip2CompressorV2 : public CompressorWithSimpleDictBase { return kBZip2Compression; } - std::unique_ptr CloneForDict(std::string&& dict_data) override { + std::unique_ptr CloneForDict( + std::string&& dict_data) const override { return std::make_unique(opts_, std::move(dict_data)); } @@ -526,7 +538,8 @@ class BuiltinLZ4CompressorV2WithDict : public CompressorWithSimpleDictBase { return kLZ4Compression; } - std::unique_ptr CloneForDict(std::string&& dict_data) override { + std::unique_ptr CloneForDict( + std::string&& dict_data) const override { return std::make_unique( opts_, std::move(dict_data)); } @@ -616,6 +629,10 @@ class BuiltinLZ4CompressorV2NoDict : public BuiltinLZ4CompressorV2WithDict { BuiltinLZ4CompressorV2NoDict(const CompressionOptions& opts) : BuiltinLZ4CompressorV2WithDict(opts, /*dict_data=*/{}) {} + std::unique_ptr Clone() const override { + return std::make_unique(opts_); + } + ManagedWorkingArea ObtainWorkingArea() override { // Using an LZ4_stream_t between compressions and resetting with // LZ4_resetStream_fast is actually slower than using a fresh LZ4_stream_t @@ -687,7 +704,8 @@ class BuiltinLZ4HCCompressorV2 : public CompressorWithSimpleDictBase { return kLZ4HCCompression; } - std::unique_ptr CloneForDict(std::string&& dict_data) override { + std::unique_ptr CloneForDict( + std::string&& dict_data) const override { return std::make_unique(opts_, std::move(dict_data)); } @@ -782,7 +800,8 @@ class BuiltinXpressCompressorV2 : public CompressorWithSimpleDictBase { return kXpressCompression; } - std::unique_ptr CloneForDict(std::string&& dict_data) override { + std::unique_ptr CloneForDict( + std::string&& dict_data) const override { return std::make_unique(opts_, std::move(dict_data)); } @@ -831,6 +850,13 @@ class BuiltinZSTDCompressorV2 : public CompressorBase { CompressionType GetPreferredCompressionType() const override { return kZSTD; } + std::unique_ptr Clone() const override { + CompressionDict dict_copy{dict_.GetRawDict().ToString(), kZSTD, + opts_.level}; + return std::make_unique(opts_, + std::move(dict_copy)); + } + size_t GetMaxSampleSizeIfWantDict( CacheEntryRole /*block_type*/) const override { if (opts_.max_dict_bytes == 0) { diff --git a/util/compression_test.cc b/util/compression_test.cc index da95a91af2..c4a3baa7a4 100644 --- a/util/compression_test.cc +++ b/util/compression_test.cc @@ -15,10 +15,12 @@ #include "table/block_based/block_builder.h" #include "test_util/testutil.h" #include "util/auto_tune_compressor.h" +#include "util/coding.h" #include "util/random.h" #include "util/simple_mixed_compressor.h" namespace ROCKSDB_NAMESPACE { + class DBCompressionTest : public DBTestBase { public: DBCompressionTest() : DBTestBase("compression_test", /*env_do_fsync=*/true) {} @@ -1108,7 +1110,131 @@ TEST_F(DBCompressionTest, RandomMixedCompressionManager) { } } -TEST_F(DBCompressionTest, CompressionManagerWrapper) { +namespace { +// Template parameter to distinguish data blocks vs. v4+ index blocks +template +static Status ValidateRocksBlock(Slice data) { + const char* src = data.data(); + size_t srcSize = data.size(); + const char* const block_type_str = + kIndexBlockV4 ? "Index block" : "Data block"; + + // Minimum RocksDB block content size: at least 1 entry + restarts + if (srcSize < 8) { + return Status::Corruption(std::string(block_type_str) + " too small"); + } + + uint32_t numRestarts = DecodeFixed32(src + srcSize - sizeof(uint32_t)); + + // Sanity check: num_restarts should be reasonable + // TODO: also support data block hash index + if (numRestarts > srcSize / 4 || numRestarts == 0) { + return Status::Corruption(std::string("Invalid num_restarts in ") + + block_type_str); + } + + size_t restartsSize = numRestarts * sizeof(uint32_t) + sizeof(uint32_t); + if (srcSize < restartsSize) { + return Status::Corruption(std::string(block_type_str) + + " too small for restarts array"); + } + + size_t entriesSize = srcSize - restartsSize; + const char* entriesEnd = src + entriesSize; + + // Parse entries + const char* p = src; + while (p < entriesEnd) { + // Parse shared_bytes varint + uint32_t shared; + const char* next = GetVarint32Ptr(p, entriesEnd, &shared); + if (next == nullptr) { + return Status::Corruption(std::string("Invalid shared_bytes varint in ") + + block_type_str); + } + p = next; + + // Parse unshared_bytes varint + uint32_t unshared; + next = GetVarint32Ptr(p, entriesEnd, &unshared); + if (next == nullptr) { + return Status::Corruption( + std::string("Invalid unshared_bytes varint in ") + block_type_str); + } + p = next; + + uint32_t valueLen = 0; + if constexpr (!kIndexBlockV4) { + // For data blocks, parse value_length varint + next = GetVarint32Ptr(p, entriesEnd, &valueLen); + if (next == nullptr) { + return Status::Corruption( + std::string("Invalid value_length varint in ") + block_type_str); + } + p = next; + } + + // Validate key delta + if (p + unshared > entriesEnd) { + return Status::Corruption( + std::string("Key delta exceeds end of entries in ") + block_type_str); + } + p += unshared; + + if constexpr (kIndexBlockV4) { + // For v4 index blocks, value is self-describing (varints) + // Parse first varint (always present) + uint32_t v1; + next = GetVarint32Ptr(p, entriesEnd, &v1); + if (next == nullptr) { + return Status::Corruption(std::string("Invalid value varint in ") + + block_type_str); + } + p = next; + + // If shared_bytes == 0, there's a second varint + if (shared == 0) { + uint32_t v2; + next = GetVarint32Ptr(p, entriesEnd, &v2); + if (next == nullptr) { + return Status::Corruption( + std::string("Invalid second value varint in ") + block_type_str); + } + p = next; + } + } else { + // For data blocks, validate value + if (p + valueLen > entriesEnd) { + return Status::Corruption( + std::string("Value exceeds end of entries in ") + block_type_str); + } + p += valueLen; + } + } + + return Status::OK(); +} +} // anonymous namespace + +class DBCompressionTestMaybeParallel + : public DBCompressionTest, + public testing::WithParamInterface> { + public: + DBCompressionTestMaybeParallel() + : DBCompressionTest(), + parallel_threads_(std::get<0>(GetParam())), + use_dict_(std::get<1>(GetParam())) {} + + protected: + int parallel_threads_; + bool use_dict_; +}; + +INSTANTIATE_TEST_CASE_P(DBCompressionTest, DBCompressionTestMaybeParallel, + ::testing::Combine(::testing::Values(1, 4), + ::testing::Values(false, true))); + +TEST_P(DBCompressionTestMaybeParallel, CompressionManagerWrapper) { // Test that we can use a custom CompressionManager to wrap the built-in // CompressionManager, thus adopting a custom *strategy* based on existing // algorithms. This will "mark" some blocks (in their contents) as "do not @@ -1119,14 +1245,81 @@ TEST_F(DBCompressionTest, CompressionManagerWrapper) { static std::string kDoNotCompress = "do_not_compress"; static std::string kRejectCompression = "reject_compression"; - struct MyCompressor : public CompressorWrapper { + static RelaxedAtomic dataCheckedCount{0}; + static RelaxedAtomic indexCheckedCount{0}; + static RelaxedAtomic compressCalledCount{0}; + + // We also have wrappers here to help verify that when RocksDB asks to + // specialize the Compressor for a particular kind of block, it only passes in + // that kind of block to ensure proper grouping of related data for + // compression. We check this by parsing the subtly distinct schemas of data + // blocks vs. v4+ index blocks. This also ensures that structure-aware + // compressions like OpenZL can parse the data block and index block formats. + struct CheckDataBlockCompressorWrapper : public CompressorWrapper { using CompressorWrapper::CompressorWrapper; - const char* Name() const override { return "MyCompressor"; } + const char* Name() const override { return "CheckDataBlockCompressor"; } + + std::unique_ptr Clone() const override { + return std::make_unique( + wrapped_->Clone()); + } Status CompressBlock(Slice uncompressed_data, char* compressed_output, size_t* compressed_output_size, CompressionType* out_compression_type, ManagedWorkingArea* working_area) override { + dataCheckedCount.FetchAddRelaxed(1); + // Parse and validate data block format before compressing + Status s = ValidateRocksBlock(uncompressed_data); + if (!s.ok()) { + return s; + } + // Delegate to wrapped compressor on success + return wrapped_->CompressBlock(uncompressed_data, compressed_output, + compressed_output_size, + out_compression_type, working_area); + } + }; + + struct CheckIndexBlockCompressorWrapper : public CompressorWrapper { + using CompressorWrapper::CompressorWrapper; + const char* Name() const override { return "CheckIndexBlockCompressor"; } + + std::unique_ptr Clone() const override { + return std::make_unique( + wrapped_->Clone()); + } + + Status CompressBlock(Slice uncompressed_data, char* compressed_output, + size_t* compressed_output_size, + CompressionType* out_compression_type, + ManagedWorkingArea* working_area) override { + indexCheckedCount.FetchAddRelaxed(1); + // Parse and validate index block v4 format before compressing + Status s = ValidateRocksBlock(uncompressed_data); + if (!s.ok()) { + return s; + } + // Delegate to wrapped compressor on success + return wrapped_->CompressBlock(uncompressed_data, compressed_output, + compressed_output_size, + out_compression_type, working_area); + } + }; + + struct MyCompressor : public CompressorWrapper { + using CompressorWrapper::CompressorWrapper; + const char* Name() const override { return "MyCompressor"; } + + std::unique_ptr Clone() const override { + return std::make_unique(wrapped_->Clone()); + } + + Status CompressBlock(Slice uncompressed_data, char* compressed_output, + size_t* compressed_output_size, + CompressionType* out_compression_type, + ManagedWorkingArea* working_area) override { + compressCalledCount.FetchAddRelaxed(1); auto begin = uncompressed_data.data(); auto end = uncompressed_data.data() + uncompressed_data.size(); if (std::search(begin, end, kDoNotCompress.begin(), @@ -1154,6 +1347,7 @@ TEST_F(DBCompressionTest, CompressionManagerWrapper) { : wrapped_(std::move(wrapped)) {} ManagedWorkingArea wrapped_; }; + ManagedWorkingArea ObtainWorkingArea() override { ManagedWorkingArea rv{ new MyWorkingArea{CompressorWrapper::ObtainWorkingArea()}, this}; @@ -1168,6 +1362,20 @@ TEST_F(DBCompressionTest, CompressionManagerWrapper) { void ReleaseWorkingArea(WorkingArea* wa) override { delete static_cast(wa); } + + std::unique_ptr MaybeCloneSpecialized( + CacheEntryRole block_type, DictSampleArgs&& dict_samples) override { + std::unique_ptr result = std::make_unique( + wrapped_->CloneMaybeSpecialized(block_type, std::move(dict_samples))); + if (block_type == CacheEntryRole::kDataBlock) { + result = std::make_unique( + std::move(result)); + } else if (block_type == CacheEntryRole::kIndexBlock) { + result = std::make_unique( + std::move(result)); + } + return result; + } }; struct MyManager : public CompressionManagerWrapper { using CompressionManagerWrapper::CompressionManagerWrapper; @@ -1191,10 +1399,15 @@ TEST_F(DBCompressionTest, CompressionManagerWrapper) { Options options = CurrentOptions(); options.compression = type; + options.compression_opts.parallel_threads = parallel_threads_; + options.compression_opts.max_dict_bytes = use_dict_ ? 4096 : 0; options.statistics = ROCKSDB_NAMESPACE::CreateDBStatistics(); options.statistics->set_stats_level(StatsLevel::kExceptTimeForMutex); BlockBasedTableOptions bbto; - bbto.enable_index_compression = false; + bbto.enable_index_compression = true; + bbto.index_type = BlockBasedTableOptions::kTwoLevelIndexSearch; + bbto.partition_filters = true; + bbto.filter_policy.reset(NewBloomFilterPolicy(5)); options.table_factory.reset(NewBlockBasedTableFactory(bbto)); options.compression_manager = use_wrapper ? mgr : nullptr; DestroyAndReopen(options); @@ -1228,14 +1441,21 @@ TEST_F(DBCompressionTest, CompressionManagerWrapper) { } ASSERT_OK(Flush()); - if (use_wrapper) { - EXPECT_EQ(kCount / 2 - 1, PopStat(NUMBER_BLOCK_COMPRESSED)); + // Index partition is compressed + constexpr int kIdxComp = 1; + // Top level index block is rejected for compression + constexpr int kIdxRej = 1; + + if (use_dict_) { + // FIXME: why don't the stats match? (for now, checking for crashes) + } else if (use_wrapper) { + EXPECT_EQ(kCount / 2 - 1 + kIdxComp, PopStat(NUMBER_BLOCK_COMPRESSED)); EXPECT_EQ(kCount / 2, PopStat(NUMBER_BLOCK_COMPRESSION_BYPASSED)); - EXPECT_EQ(1 + 1, PopStat(NUMBER_BLOCK_COMPRESSION_REJECTED)); + EXPECT_EQ(1 + 1 + kIdxRej, PopStat(NUMBER_BLOCK_COMPRESSION_REJECTED)); } else { - EXPECT_EQ(kCount - 1, PopStat(NUMBER_BLOCK_COMPRESSED)); + EXPECT_EQ(kCount - 1 + kIdxComp, PopStat(NUMBER_BLOCK_COMPRESSED)); EXPECT_EQ(0, PopStat(NUMBER_BLOCK_COMPRESSION_BYPASSED)); - EXPECT_EQ(1, PopStat(NUMBER_BLOCK_COMPRESSION_REJECTED)); + EXPECT_EQ(1 + kIdxRej, PopStat(NUMBER_BLOCK_COMPRESSION_REJECTED)); } // Ensure well-formed for reads @@ -1243,11 +1463,27 @@ TEST_F(DBCompressionTest, CompressionManagerWrapper) { ASSERT_NE(Get(Key(i)), "NOT_FOUND"); } ASSERT_EQ(Get(Key(kCount)), "NOT_FOUND"); + + // Ensure expected checks were performed + EXPECT_EQ(indexCheckedCount.ExchangeRelaxed(0), + use_wrapper ? kIdxComp + kIdxRej : 0); + EXPECT_EQ(dataCheckedCount.ExchangeRelaxed(0), use_wrapper ? kCount : 0); + // And every use of MyCompressor went through either the data block + // checker or index block checker + EXPECT_EQ(compressCalledCount.ExchangeRelaxed(0), + use_wrapper ? kIdxComp + kIdxRej + kCount : 0); } } } -TEST_F(DBCompressionTest, CompressionManagerCustomCompression) { +namespace { +std::string UniqueName(const std::string& base) { + static RelaxedAtomic counter{0}; + return base + std::to_string(counter.FetchAddRelaxed(1)); +} +} // anonymous namespace + +TEST_P(DBCompressionTestMaybeParallel, CompressionManagerCustomCompression) { // Test that we can use a custom CompressionManager to implement custom // compression algorithms, and that there are appropriate schema guard rails // to ensure data is not processed by the wrong algorithm. @@ -1263,9 +1499,12 @@ TEST_F(DBCompressionTest, CompressionManagerCustomCompression) { class MyManager : public CompressionManager { public: - explicit MyManager(const char* compat_name) : compat_name_(compat_name) {} + explicit MyManager(const std::string& compat_name) + : compat_name_(compat_name), name_("MyManager:" + compat_name_) {} const char* Name() const override { return name_.c_str(); } - const char* CompatibilityName() const override { return compat_name_; } + const char* CompatibilityName() const override { + return compat_name_.c_str(); + } bool SupportsCompressionType(CompressionType type) const override { return type == kCustomCompression8A || type == kCustomCompression8B || @@ -1330,246 +1569,237 @@ TEST_F(DBCompressionTest, CompressionManagerCustomCompression) { } private: - const char* compat_name_; + std::string compat_name_; std::string name_; // weak_ptr to avoid cycles std::map> friends_; }; - for (bool use_dict : {false, true}) { - SCOPED_TRACE(use_dict ? "With dict" : "No dict"); + // Although these compression managers are actually compatible, we must + // respect their distinct compatibility names and treat them as incompatible + // (or else risk processing data incorrectly) + // NOTE: these are not registered in ObjectRegistry to test what happens + // when the original CompressionManager might not be available, but + // mgr_bar will be registered during the test, with different names to + // prevent interference between iterations. + auto mgr_foo = std::make_shared("Foo"); + auto mgr_bar = std::make_shared(UniqueName("Bar")); - // Although these compression managers are actually compatible, we must - // respect their distinct compatibility names and treat them as incompatible - // (or else risk processing data incorrectly) - // NOTE: these are not registered in ObjectRegistry to test what happens - // when the original CompressionManager might not be available, but - // mgr_bar will be registered during the test, with different names to - // prevent interference between iterations. - auto mgr_foo = std::make_shared("Foo"); - auto mgr_bar = std::make_shared(use_dict ? "Bar1" : "Bar2"); + // And this one claims to be fully compatible with the built-in compression + // manager when it's not fully compatible (for custom CompressionTypes) + auto mgr_claim_compatible = std::make_shared("BuiltinV2"); - // And this one claims to be fully compatible with the built-in compression - // manager when it's not fully compatible (for custom CompressionTypes) - auto mgr_claim_compatible = std::make_shared("BuiltinV2"); + constexpr uint16_t kValueSize = 10000; - constexpr uint16_t kValueSize = 10000; + Options options = CurrentOptions(); + options.level0_file_num_compaction_trigger = 20; + BlockBasedTableOptions bbto; + bbto.enable_index_compression = false; + bbto.format_version = 6; // Before custom compression alg support + options.table_factory.reset(NewBlockBasedTableFactory(bbto)); + // Claims not to use custom compression (and doesn't unless setting a custom + // CompressionType) + options.compression_manager = mgr_claim_compatible; + // Use a built-in compression type with dictionary support + options.compression = kLZ4Compression; + options.compression_opts.max_dict_bytes = use_dict_ ? kValueSize / 2 : 0; + options.compression_opts.parallel_threads = parallel_threads_; + DestroyAndReopen(options); - Options options = CurrentOptions(); - options.level0_file_num_compaction_trigger = 20; - BlockBasedTableOptions bbto; - bbto.enable_index_compression = false; - bbto.format_version = 6; // Before custom compression alg support - options.table_factory.reset(NewBlockBasedTableFactory(bbto)); - // Claims not to use custom compression (and doesn't unless setting a custom - // CompressionType) - options.compression_manager = mgr_claim_compatible; - // Use a built-in compression type with dictionary support - options.compression = kLZ4Compression; - options.compression_opts.max_dict_bytes = kValueSize / 2; - DestroyAndReopen(options); + Random rnd(404); + std::string value; + ASSERT_OK(Put("a", test::CompressibleString(&rnd, 0.1, kValueSize, &value))); + ASSERT_OK(Flush()); - Random rnd(404); - std::string value; - ASSERT_OK( - Put("a", test::CompressibleString(&rnd, 0.1, kValueSize, &value))); - ASSERT_OK(Flush()); + // That data should be readable without access to the original compression + // manager, because it used the built-in CompatibilityName and a built-in + // CompressionType + options.compression_manager = nullptr; + Reopen(options); + ASSERT_EQ(Get("a"), value); - // That data should be readable without access to the original compression - // manager, because it used the built-in CompatibilityName and a built-in - // CompressionType - options.compression_manager = nullptr; - Reopen(options); - ASSERT_EQ(Get("a"), value); + // Verify it was compressed + Range r = {"a", "a0"}; + TablePropertiesCollection tables_properties; + ASSERT_OK(db_->GetPropertiesOfTablesInRange(db_->DefaultColumnFamily(), &r, 1, + &tables_properties)); + ASSERT_EQ(tables_properties.size(), 1U); + EXPECT_LT(tables_properties.begin()->second->data_size, kValueSize / 2); + EXPECT_EQ(tables_properties.begin()->second->compression_name, "LZ4"); - // Verify it was compressed - Range r = {"a", "a0"}; - TablePropertiesCollection tables_properties; - ASSERT_OK(db_->GetPropertiesOfTablesInRange(db_->DefaultColumnFamily(), &r, - 1, &tables_properties)); - ASSERT_EQ(tables_properties.size(), 1U); - EXPECT_LT(tables_properties.begin()->second->data_size, kValueSize / 2); - EXPECT_EQ(tables_properties.begin()->second->compression_name, "LZ4"); + // Disallow setting a custom CompressionType with a CompressionManager + // claiming to be built-in compatible. + options.compression_manager = mgr_claim_compatible; + options.compression = kCustomCompression8A; + ASSERT_EQ(TryReopen(options).code(), Status::Code::kInvalidArgument); - // Disallow setting a custom CompressionType with a CompressionManager - // claiming to be built-in compatible. - options.compression_manager = mgr_claim_compatible; - options.compression = kCustomCompression8A; - ASSERT_EQ(TryReopen(options).code(), Status::Code::kInvalidArgument); + options.compression_manager = nullptr; + options.compression = kCustomCompressionFE; + ASSERT_EQ(TryReopen(options).code(), Status::Code::kInvalidArgument); + options.compression = + static_cast(kLastBuiltinCompression + 1); + ASSERT_EQ(TryReopen(options).code(), Status::Code::kInvalidArgument); - options.compression_manager = nullptr; - options.compression = kCustomCompressionFE; - ASSERT_EQ(TryReopen(options).code(), Status::Code::kInvalidArgument); - options.compression = - static_cast(kLastBuiltinCompression + 1); - ASSERT_EQ(TryReopen(options).code(), Status::Code::kInvalidArgument); + // Custom compression schema (different CompatibilityName) not supported + // before format_version=7 + options.compression_manager = mgr_foo; + options.compression = kLZ4Compression; + ASSERT_EQ(TryReopen(options).code(), Status::Code::kInvalidArgument); - // Custom compression schema (different CompatibilityName) not supported - // before format_version=7 - options.compression_manager = mgr_foo; - options.compression = kLZ4Compression; - ASSERT_EQ(TryReopen(options).code(), Status::Code::kInvalidArgument); + // Set format version supporting custom compression + bbto.format_version = 7; + options.table_factory.reset(NewBlockBasedTableFactory(bbto)); - // Set format version supporting custom compression - bbto.format_version = 7; - options.table_factory.reset(NewBlockBasedTableFactory(bbto)); + // Custom compression type not supported with built-in schema name, even + // with format_version=7 + options.compression_manager = mgr_claim_compatible; + options.compression = kCustomCompression8B; + ASSERT_EQ(TryReopen(options).code(), Status::Code::kInvalidArgument); - // Custom compression type not supported with built-in schema name, even - // with format_version=7 - options.compression_manager = mgr_claim_compatible; - options.compression = kCustomCompression8B; - ASSERT_EQ(TryReopen(options).code(), Status::Code::kInvalidArgument); + // Custom compression schema, but specifying a custom compression type it + // doesn't support. + options.compression_manager = mgr_foo; + options.compression = kCustomCompressionF0; + ASSERT_EQ(TryReopen(options).code(), Status::Code::kNotSupported); - // Custom compression schema, but specifying a custom compression type it - // doesn't support. - options.compression_manager = mgr_foo; - options.compression = kCustomCompressionF0; - ASSERT_EQ(TryReopen(options).code(), Status::Code::kNotSupported); + // Using a built-in compression type with fv=7 but named custom schema + options.compression = kLZ4Compression; + Reopen(options); + ASSERT_OK(Put("b", test::CompressibleString(&rnd, 0.1, kValueSize, &value))); + ASSERT_OK(Flush()); + ASSERT_EQ(NumTableFilesAtLevel(0), 2); + ASSERT_EQ(Get("b"), value); - // Using a built-in compression type with fv=7 but named custom schema - options.compression = kLZ4Compression; - Reopen(options); - ASSERT_OK( - Put("b", test::CompressibleString(&rnd, 0.1, kValueSize, &value))); - ASSERT_OK(Flush()); - ASSERT_EQ(NumTableFilesAtLevel(0), 2); - ASSERT_EQ(Get("b"), value); + // Verify it was compressed with LZ4 + r = {"b", "b0"}; + tables_properties.clear(); + ASSERT_OK(db_->GetPropertiesOfTablesInRange(db_->DefaultColumnFamily(), &r, 1, + &tables_properties)); + ASSERT_EQ(tables_properties.size(), 1U); + EXPECT_LT(tables_properties.begin()->second->data_size, kValueSize / 2); + // Uses new format for "compression_name" property + EXPECT_EQ(tables_properties.begin()->second->compression_name, "Foo;04;"); + EXPECT_EQ(mgr_foo->last_specific_decompressor_type_.LoadRelaxed(), + kLZ4Compression); - // Verify it was compressed with LZ4 - r = {"b", "b0"}; - tables_properties.clear(); - ASSERT_OK(db_->GetPropertiesOfTablesInRange(db_->DefaultColumnFamily(), &r, - 1, &tables_properties)); - ASSERT_EQ(tables_properties.size(), 1U); - EXPECT_LT(tables_properties.begin()->second->data_size, kValueSize / 2); - // Uses new format for "compression_name" property - EXPECT_EQ(tables_properties.begin()->second->compression_name, "Foo;04;"); - EXPECT_EQ(mgr_foo->last_specific_decompressor_type_.LoadRelaxed(), - kLZ4Compression); + // Custom compression type + options.compression = kCustomCompression8A; + Reopen(options); + ASSERT_OK(Put("c", test::CompressibleString(&rnd, 0.1, kValueSize, &value))); + EXPECT_EQ(mgr_foo->used_compressor8A_count_, 0); + ASSERT_OK(Flush()); + ASSERT_EQ(NumTableFilesAtLevel(0), 3); + ASSERT_EQ(Get("c"), value); + EXPECT_EQ(mgr_foo->used_compressor8A_count_, 1); - // Custom compression type - options.compression = kCustomCompression8A; - Reopen(options); - ASSERT_OK( - Put("c", test::CompressibleString(&rnd, 0.1, kValueSize, &value))); - EXPECT_EQ(mgr_foo->used_compressor8A_count_, 0); - ASSERT_OK(Flush()); - ASSERT_EQ(NumTableFilesAtLevel(0), 3); - ASSERT_EQ(Get("c"), value); - EXPECT_EQ(mgr_foo->used_compressor8A_count_, 1); + // Verify it was compressed with custom format + r = {"c", "c0"}; + tables_properties.clear(); + ASSERT_OK(db_->GetPropertiesOfTablesInRange(db_->DefaultColumnFamily(), &r, 1, + &tables_properties)); + ASSERT_EQ(tables_properties.size(), 1U); + EXPECT_LT(tables_properties.begin()->second->data_size, kValueSize / 2); + EXPECT_EQ(tables_properties.begin()->second->compression_name, "Foo;8A;"); + EXPECT_EQ(mgr_foo->last_specific_decompressor_type_.LoadRelaxed(), + kCustomCompression8A); - // Verify it was compressed with custom format - r = {"c", "c0"}; - tables_properties.clear(); - ASSERT_OK(db_->GetPropertiesOfTablesInRange(db_->DefaultColumnFamily(), &r, - 1, &tables_properties)); - ASSERT_EQ(tables_properties.size(), 1U); - EXPECT_LT(tables_properties.begin()->second->data_size, kValueSize / 2); - EXPECT_EQ(tables_properties.begin()->second->compression_name, "Foo;8A;"); - EXPECT_EQ(mgr_foo->last_specific_decompressor_type_.LoadRelaxed(), - kCustomCompression8A); + // Also dynamically changeable, because the compression manager will respect + // the current setting as reported under the legacy logic + ASSERT_OK(dbfull()->SetOptions({{"compression", "kLZ4Compression"}})); + ASSERT_OK(Put("d", test::CompressibleString(&rnd, 0.1, kValueSize, &value))); + ASSERT_OK(Flush()); + ASSERT_EQ(NumTableFilesAtLevel(0), 4); + ASSERT_EQ(Get("d"), value); - // Also dynamically changeable, because the compression manager will respect - // the current setting as reported under the legacy logic - ASSERT_OK(dbfull()->SetOptions({{"compression", "kLZ4Compression"}})); - ASSERT_OK( - Put("d", test::CompressibleString(&rnd, 0.1, kValueSize, &value))); - ASSERT_OK(Flush()); - ASSERT_EQ(NumTableFilesAtLevel(0), 4); - ASSERT_EQ(Get("d"), value); + // Verify it was compressed with LZ4 + r = {"d", "d0"}; + tables_properties.clear(); + ASSERT_OK(db_->GetPropertiesOfTablesInRange(db_->DefaultColumnFamily(), &r, 1, + &tables_properties)); + ASSERT_EQ(tables_properties.size(), 1U); + EXPECT_LT(tables_properties.begin()->second->data_size, kValueSize / 2); + EXPECT_EQ(tables_properties.begin()->second->compression_name, "Foo;04;"); + EXPECT_EQ(mgr_foo->last_specific_decompressor_type_.LoadRelaxed(), + kLZ4Compression); - // Verify it was compressed with LZ4 - r = {"d", "d0"}; - tables_properties.clear(); - ASSERT_OK(db_->GetPropertiesOfTablesInRange(db_->DefaultColumnFamily(), &r, - 1, &tables_properties)); - ASSERT_EQ(tables_properties.size(), 1U); - EXPECT_LT(tables_properties.begin()->second->data_size, kValueSize / 2); - EXPECT_EQ(tables_properties.begin()->second->compression_name, "Foo;04;"); - EXPECT_EQ(mgr_foo->last_specific_decompressor_type_.LoadRelaxed(), - kLZ4Compression); + // Dynamically changeable to custom compressions also + ASSERT_OK(dbfull()->SetOptions({{"compression", "kCustomCompression8B"}})); + ASSERT_OK(Put("e", test::CompressibleString(&rnd, 0.1, kValueSize, &value))); + ASSERT_OK(Flush()); + ASSERT_EQ(NumTableFilesAtLevel(0), 5); + ASSERT_EQ(Get("e"), value); - // Dynamically changeable to custom compressions also - ASSERT_OK(dbfull()->SetOptions({{"compression", "kCustomCompression8B"}})); - ASSERT_OK( - Put("e", test::CompressibleString(&rnd, 0.1, kValueSize, &value))); - ASSERT_OK(Flush()); - ASSERT_EQ(NumTableFilesAtLevel(0), 5); - ASSERT_EQ(Get("e"), value); + // Verify it was compressed with custom format + r = {"e", "e0"}; + tables_properties.clear(); + ASSERT_OK(db_->GetPropertiesOfTablesInRange(db_->DefaultColumnFamily(), &r, 1, + &tables_properties)); + ASSERT_EQ(tables_properties.size(), 1U); + EXPECT_LT(tables_properties.begin()->second->data_size, kValueSize / 2); + EXPECT_EQ(tables_properties.begin()->second->compression_name, "Foo;8B;"); + EXPECT_EQ(mgr_foo->last_specific_decompressor_type_.LoadRelaxed(), + kCustomCompression8B); - // Verify it was compressed with custom format - r = {"e", "e0"}; - tables_properties.clear(); - ASSERT_OK(db_->GetPropertiesOfTablesInRange(db_->DefaultColumnFamily(), &r, - 1, &tables_properties)); - ASSERT_EQ(tables_properties.size(), 1U); - EXPECT_LT(tables_properties.begin()->second->data_size, kValueSize / 2); - EXPECT_EQ(tables_properties.begin()->second->compression_name, "Foo;8B;"); - EXPECT_EQ(mgr_foo->last_specific_decompressor_type_.LoadRelaxed(), - kCustomCompression8B); + // Fails to re-open with incompatible compression manager (can't find + // compression manager Foo because it's not registered nor known by Bar) + options.compression_manager = mgr_bar; + options.compression = kLZ4Compression; + ASSERT_EQ(TryReopen(options).code(), Status::Code::kNotSupported); - // Fails to re-open with incompatible compression manager (can't find - // compression manager Foo because it's not registered nor known by Bar) - options.compression_manager = mgr_bar; - options.compression = kLZ4Compression; - ASSERT_EQ(TryReopen(options).code(), Status::Code::kNotSupported); + // But should re-open if we make Bar aware of the Foo compression manager + mgr_bar->AddFriend(mgr_foo); + Reopen(options); - // But should re-open if we make Bar aware of the Foo compression manager - mgr_bar->AddFriend(mgr_foo); - Reopen(options); + // Can still read everything + ASSERT_EQ(Get("a").size(), kValueSize); + ASSERT_EQ(Get("b").size(), kValueSize); + ASSERT_EQ(Get("c").size(), kValueSize); + ASSERT_EQ(Get("d").size(), kValueSize); + ASSERT_EQ(Get("e").size(), kValueSize); - // Can still read everything - ASSERT_EQ(Get("a").size(), kValueSize); - ASSERT_EQ(Get("b").size(), kValueSize); - ASSERT_EQ(Get("c").size(), kValueSize); - ASSERT_EQ(Get("d").size(), kValueSize); - ASSERT_EQ(Get("e").size(), kValueSize); + // Add a file using mgr_bar + ASSERT_OK(Put("f", test::CompressibleString(&rnd, 0.1, kValueSize, &value))); + ASSERT_OK(Flush()); + ASSERT_EQ(NumTableFilesAtLevel(0), 6); + ASSERT_EQ(Get("f"), value); - // Add a file using mgr_bar - ASSERT_OK( - Put("f", test::CompressibleString(&rnd, 0.1, kValueSize, &value))); - ASSERT_OK(Flush()); - ASSERT_EQ(NumTableFilesAtLevel(0), 6); - ASSERT_EQ(Get("f"), value); + // Verify it was compressed appropriately + r = {"f", "f0"}; + tables_properties.clear(); + ASSERT_OK(db_->GetPropertiesOfTablesInRange(db_->DefaultColumnFamily(), &r, 1, + &tables_properties)); + ASSERT_EQ(tables_properties.size(), 1U); + EXPECT_LT(tables_properties.begin()->second->data_size, kValueSize / 2); + EXPECT_EQ(mgr_bar->last_specific_decompressor_type_.LoadRelaxed(), + kLZ4Compression); - // Verify it was compressed appropriately - r = {"f", "f0"}; - tables_properties.clear(); - ASSERT_OK(db_->GetPropertiesOfTablesInRange(db_->DefaultColumnFamily(), &r, - 1, &tables_properties)); - ASSERT_EQ(tables_properties.size(), 1U); - EXPECT_LT(tables_properties.begin()->second->data_size, kValueSize / 2); - EXPECT_EQ(mgr_bar->last_specific_decompressor_type_.LoadRelaxed(), - kLZ4Compression); + // Fails to re-open with incompatible compression manager (can't find + // compression manager Bar because it's not registered nor known by Foo) + options.compression_manager = mgr_foo; + ASSERT_EQ(TryReopen(options).code(), Status::Code::kNotSupported); - // Fails to re-open with incompatible compression manager (can't find - // compression manager Bar because it's not registered nor known by Foo) - options.compression_manager = mgr_foo; - ASSERT_EQ(TryReopen(options).code(), Status::Code::kNotSupported); + // Register and re-open + auto& library = *ObjectLibrary::Default(); + library.AddFactory( + mgr_bar->CompatibilityName(), + [mgr_bar](const std::string& /*uri*/, + std::unique_ptr* guard, + std::string* /*errmsg*/) { + *guard = std::make_unique(mgr_bar->CompatibilityName()); + return guard->get(); + }); + Reopen(options); - // Register and re-open - auto& library = *ObjectLibrary::Default(); - library.AddFactory( - mgr_bar->CompatibilityName(), - [mgr_bar](const std::string& /*uri*/, - std::unique_ptr* guard, - std::string* /*errmsg*/) { - *guard = std::make_unique(mgr_bar->CompatibilityName()); - return guard->get(); - }); - Reopen(options); + // Can still read everything + ASSERT_EQ(Get("a").size(), kValueSize); + ASSERT_EQ(Get("b").size(), kValueSize); + ASSERT_EQ(Get("c").size(), kValueSize); + ASSERT_EQ(Get("d").size(), kValueSize); + ASSERT_EQ(Get("e").size(), kValueSize); + ASSERT_EQ(Get("f").size(), kValueSize); - // Can still read everything - ASSERT_EQ(Get("a").size(), kValueSize); - ASSERT_EQ(Get("b").size(), kValueSize); - ASSERT_EQ(Get("c").size(), kValueSize); - ASSERT_EQ(Get("d").size(), kValueSize); - ASSERT_EQ(Get("e").size(), kValueSize); - ASSERT_EQ(Get("f").size(), kValueSize); - - // TODO: test old version of a compression manager unable to read a - // compression type - } + // TODO: test old version of a compression manager unable to read a + // compression type } TEST_F(DBCompressionTest, FailWhenCompressionNotSupportedTest) { diff --git a/util/simple_mixed_compressor.cc b/util/simple_mixed_compressor.cc index a1ee40481d..282fcd97a5 100644 --- a/util/simple_mixed_compressor.cc +++ b/util/simple_mixed_compressor.cc @@ -16,7 +16,8 @@ namespace ROCKSDB_NAMESPACE { // MultiCompressorWrapper implementation MultiCompressorWrapper::MultiCompressorWrapper(const CompressionOptions& opts, - CompressionDict&& dict) { + CompressionDict&& dict) + : opts_(opts) { // TODO: make the compression manager a field auto builtInManager = GetBuiltinV2CompressionManager(); const auto& compressions = GetSupportedCompressions(); @@ -59,6 +60,10 @@ const char* RandomMixedCompressor::Name() const { return "RandomMixedCompressor"; } +std::unique_ptr RandomMixedCompressor::Clone() const { + return std::make_unique(opts_); +} + Status RandomMixedCompressor::CompressBlock( Slice uncompressed_data, char* compressed_output, size_t* compressed_output_size, CompressionType* out_compression_type, @@ -86,6 +91,10 @@ const char* RoundRobinCompressor::Name() const { return "RoundRobinCompressor"; } +std::unique_ptr RoundRobinCompressor::Clone() const { + return std::make_unique(opts_); +} + Status RoundRobinCompressor::CompressBlock( Slice uncompressed_data, char* compressed_output, size_t* compressed_output_size, CompressionType* out_compression_type, diff --git a/util/simple_mixed_compressor.h b/util/simple_mixed_compressor.h index 0c12d88a0a..e55b724976 100644 --- a/util/simple_mixed_compressor.h +++ b/util/simple_mixed_compressor.h @@ -28,12 +28,14 @@ class MultiCompressorWrapper : public Compressor { CacheEntryRole block_type, DictSampleArgs&& dict_samples) override; protected: + const CompressionOptions opts_; std::vector> compressors_; }; struct RandomMixedCompressor : public MultiCompressorWrapper { using MultiCompressorWrapper::MultiCompressorWrapper; const char* Name() const override; + std::unique_ptr Clone() const override; Status CompressBlock(Slice uncompressed_data, char* compressed_output, size_t* compressed_output_size, CompressionType* out_compression_type, @@ -51,6 +53,7 @@ class RandomMixedCompressionManager : public CompressionManagerWrapper { struct RoundRobinCompressor : public MultiCompressorWrapper { using MultiCompressorWrapper::MultiCompressorWrapper; const char* Name() const override; + std::unique_ptr Clone() const override; Status CompressBlock(Slice uncompressed_data, char* compressed_output, size_t* compressed_output_size, CompressionType* out_compression_type,