database/rocksdb
1
0
Fork 0
mirror of https://github.com/facebook/rocksdb.git synced 2026-10-23 08:09:02 +03:00
Code Activity

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

Browse Source 
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
This commit is contained in:
Peter Dillinger
2025-11-21 16:34:49 -08:00
committed by meta-codesync[bot]
parent 8c7c8b8dab
commit 35148aca91
10 changed files with 611 additions and 265 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
include/rocksdb/advanced_compression.h
View File

@@ -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 {

24
include/rocksdb/data_structure.h
View File

@@ -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>

88
table/block_based/block_based_table_builder.cc
View File

@@ -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();
}

6
test_util/testutil.h
View File

@@ -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));
}

8
util/auto_tune_compressor.cc
View File

@@ -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();

2
util/auto_tune_compressor.h
View File

@@ -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;

40
util/compression.cc
View File

@@ -174,6 +174,10 @@ class BuiltinCompressorV1 : public CompressorBase {
CompressionType GetPreferredCompressionType() const override { return type_; }
std::unique_ptr<Compressor> Clone() const override {
return std::make_unique<BuiltinCompressorV1>(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<Compressor> Clone() const override {
return CloneForDict(std::string{dict_data_});
}
std::unique_ptr<Compressor> MaybeCloneSpecialized(
CacheEntryRole /*block_type*/,
DictSampleArgs&& dict_samples) final override {
@@ -238,7 +246,8 @@ class CompressorWithSimpleDictBase : public CompressorBase {
}
}
virtual std::unique_ptr<Compressor> CloneForDict(std::string&& dict_data) = 0;
virtual std::unique_ptr<Compressor> 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<Compressor> CloneForDict(std::string&& dict_data) override {
std::unique_ptr<Compressor> CloneForDict(
std::string&& dict_data) const override {
return std::make_unique<BuiltinSnappyCompressorV2>(opts_,
std::move(dict_data));
}
@@ -349,7 +359,8 @@ class BuiltinZlibCompressorV2 : public CompressorWithSimpleDictBase {
return kZlibCompression;
}
std::unique_ptr<Compressor> CloneForDict(std::string&& dict_data) override {
std::unique_ptr<Compressor> CloneForDict(
std::string&& dict_data) const override {
return std::make_unique<BuiltinZlibCompressorV2>(opts_,
std::move(dict_data));
}
@@ -447,7 +458,8 @@ class BuiltinBZip2CompressorV2 : public CompressorWithSimpleDictBase {
return kBZip2Compression;
}
std::unique_ptr<Compressor> CloneForDict(std::string&& dict_data) override {
std::unique_ptr<Compressor> CloneForDict(
std::string&& dict_data) const override {
return std::make_unique<BuiltinBZip2CompressorV2>(opts_,
std::move(dict_data));
}
@@ -526,7 +538,8 @@ class BuiltinLZ4CompressorV2WithDict : public CompressorWithSimpleDictBase {
return kLZ4Compression;
}
std::unique_ptr<Compressor> CloneForDict(std::string&& dict_data) override {
std::unique_ptr<Compressor> CloneForDict(
std::string&& dict_data) const override {
return std::make_unique<BuiltinLZ4CompressorV2WithDict>(
opts_, std::move(dict_data));
}
@@ -616,6 +629,10 @@ class BuiltinLZ4CompressorV2NoDict : public BuiltinLZ4CompressorV2WithDict {
BuiltinLZ4CompressorV2NoDict(const CompressionOptions& opts)
: BuiltinLZ4CompressorV2WithDict(opts, /*dict_data=*/{}) {}
std::unique_ptr<Compressor> Clone() const override {
return std::make_unique<BuiltinLZ4CompressorV2NoDict>(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<Compressor> CloneForDict(std::string&& dict_data) override {
std::unique_ptr<Compressor> CloneForDict(
std::string&& dict_data) const override {
return std::make_unique<BuiltinLZ4HCCompressorV2>(opts_,
std::move(dict_data));
}
@@ -782,7 +800,8 @@ class BuiltinXpressCompressorV2 : public CompressorWithSimpleDictBase {
return kXpressCompression;
}
std::unique_ptr<Compressor> CloneForDict(std::string&& dict_data) override {
std::unique_ptr<Compressor> CloneForDict(
std::string&& dict_data) const override {
return std::make_unique<BuiltinXpressCompressorV2>(opts_,
std::move(dict_data));
}
@@ -831,6 +850,13 @@ class BuiltinZSTDCompressorV2 : public CompressorBase {
CompressionType GetPreferredCompressionType() const override { return kZSTD; }
std::unique_ptr<Compressor> Clone() const override {
CompressionDict dict_copy{dict_.GetRawDict().ToString(), kZSTD,
opts_.level};
return std::make_unique<BuiltinZSTDCompressorV2>(opts_,
std::move(dict_copy));
}
size_t GetMaxSampleSizeIfWantDict(
CacheEntryRole /*block_type*/) const override {
if (opts_.max_dict_bytes == 0) {

666
util/compression_test.cc
View File

@@ -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 <bool kIndexBlockV4>
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<std::tuple<int, bool>> {
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<int> dataCheckedCount{0};
static RelaxedAtomic<int> indexCheckedCount{0};
static RelaxedAtomic<int> 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<Compressor> Clone() const override {
return std::make_unique<CheckDataBlockCompressorWrapper>(
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</*kIndexBlockV4=*/false>(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<Compressor> Clone() const override {
return std::make_unique<CheckIndexBlockCompressorWrapper>(
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</*kIndexBlockV4=*/true>(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<Compressor> Clone() const override {
return std::make_unique<MyCompressor>(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<MyWorkingArea*>(wa);
}
std::unique_ptr<Compressor> MaybeCloneSpecialized(
CacheEntryRole block_type, DictSampleArgs&& dict_samples) override {
std::unique_ptr<Compressor> result = std::make_unique<MyCompressor>(
wrapped_->CloneMaybeSpecialized(block_type, std::move(dict_samples)));
if (block_type == CacheEntryRole::kDataBlock) {
result = std::make_unique<CheckDataBlockCompressorWrapper>(
std::move(result));
} else if (block_type == CacheEntryRole::kIndexBlock) {
result = std::make_unique<CheckIndexBlockCompressorWrapper>(
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<int> 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<std::string, std::weak_ptr<CompressionManager>> 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<MyManager>("Foo");
auto mgr_bar = std::make_shared<MyManager>(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<MyManager>("Foo");
auto mgr_bar = std::make_shared<MyManager>(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<MyManager>("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<MyManager>("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<CompressionType>(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<CompressionType>(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<CompressionManager>(
mgr_bar->CompatibilityName(),
[mgr_bar](const std::string& /*uri*/,
std::unique_ptr<CompressionManager>* guard,
std::string* /*errmsg*/) {
*guard = std::make_unique<MyManager>(mgr_bar->CompatibilityName());
return guard->get();
});
Reopen(options);
// Register and re-open
auto& library = *ObjectLibrary::Default();
library.AddFactory<CompressionManager>(
mgr_bar->CompatibilityName(),
[mgr_bar](const std::string& /*uri*/,
std::unique_ptr<CompressionManager>* guard,
std::string* /*errmsg*/) {
*guard = std::make_unique<MyManager>(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) {

11
util/simple_mixed_compressor.cc
View File

@@ -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<Compressor> RandomMixedCompressor::Clone() const {
return std::make_unique<RandomMixedCompressor>(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<Compressor> RoundRobinCompressor::Clone() const {
return std::make_unique<RoundRobinCompressor>(opts_);
}
Status RoundRobinCompressor::CompressBlock(
Slice uncompressed_data, char* compressed_output,
size_t* compressed_output_size, CompressionType* out_compression_type,

3
util/simple_mixed_compressor.h
View File

@@ -28,12 +28,14 @@ class MultiCompressorWrapper : public Compressor {
CacheEntryRole block_type, DictSampleArgs&& dict_samples) override;
protected:
const CompressionOptions opts_;
std::vector<std::unique_ptr<Compressor>> compressors_;
};
struct RandomMixedCompressor : public MultiCompressorWrapper {
using MultiCompressorWrapper::MultiCompressorWrapper;
const char* Name() const override;
std::unique_ptr<Compressor> 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<Compressor> Clone() const override;
Status CompressBlock(Slice uncompressed_data, char* compressed_output,
size_t* compressed_output_size,
CompressionType* out_compression_type,