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

Connect custom compression to crash test and ObjectLibrary (#13710)

Browse Source 
Summary:
Some pieces of follow-up to https://github.com/facebook/rocksdb/issues/13659.
_Recommend hiding whitespace for review_
* Add support for instantiating CompressionManagers through CreateFromString/ObjectLibrary.
* Pull CompressorCustomAlg and DecompressorCustomAlg out of db_test2, refactor/improvement them a bit, and put them in testutil.h for sharing with db_stress. Switched it from being built on snappy to being built on lz4 so that it can properly test dictionary compression.
* Add a custom compression manager for db_stress that uses these, and add to crash test. This depends on the ObjectLibrary stuff because some invocations of db_stress will not be configured with the custom compression manager but will need to access it to read some existing SST files.
* Remove some pieces where the concern of setting compression=kZSTD for compatibility purposes had leaked into configuring some tests and compression managers. After https://github.com/facebook/rocksdb/issues/13659 this compatibility concern is contained in the SST building code.
* Fix BuiltinDecompressorV2SnappyOnly hiding the (ignored) compression dictionary. SST read logic expects the serialized dictionary to be returned by the decompressor even if it's effectively ignored. Updated DBBlockCacheTest.CacheCompressionDict to cover this case.

For follow-up:
* Combine custom compression and mixed compression types in a file (not clean/easy without duplicating or majorly refactoring the mixed/random compressor)

Pull Request resolved: https://github.com/facebook/rocksdb/pull/13710

Test Plan: unit tests updated

Reviewed By: hx235

Differential Revision: D76928974

Pulled By: pdillinger

fbshipit-source-id: 772cf9cb048d737699b0e2887c624fb64a68aa8c
This commit is contained in:
Peter Dillinger
2025-06-19 12:54:15 -07:00
committed by Facebook GitHub Bot
parent d55655a423
commit fdc2970d37
14 changed files with 639 additions and 478 deletions
Download Patch File Download Diff File Expand all files Collapse all files

125
db/db_block_cache_test.cc
View File

@@ -829,68 +829,79 @@ TEST_F(DBBlockCacheTest, CacheCompressionDict) {
const int kNumEntriesPerFile = 128;
const int kNumBytesPerEntry = 1024;
// Try all the available libraries that support dictionary compression
std::vector<CompressionType> compression_types;
if (Zlib_Supported()) {
compression_types.push_back(kZlibCompression);
}
if (LZ4_Supported()) {
compression_types.push_back(kLZ4Compression);
compression_types.push_back(kLZ4HCCompression);
}
if (ZSTD_Supported()) {
compression_types.push_back(kZSTD);
}
std::vector<CompressionType> dict_compressions =
GetSupportedDictCompressions();
Random rnd(301);
for (auto compression_type : compression_types) {
Options options = CurrentOptions();
options.bottommost_compression = compression_type;
options.bottommost_compression_opts.max_dict_bytes = 4096;
options.bottommost_compression_opts.enabled = true;
options.create_if_missing = true;
options.num_levels = 2;
options.statistics = ROCKSDB_NAMESPACE::CreateDBStatistics();
options.target_file_size_base = kNumEntriesPerFile * kNumBytesPerEntry;
BlockBasedTableOptions table_options;
table_options.cache_index_and_filter_blocks = true;
table_options.block_cache.reset(new MockCache());
options.table_factory.reset(NewBlockBasedTableFactory(table_options));
DestroyAndReopen(options);
// Format version before and after compression handling changes
TEST_AllowUnsupportedFormatVersion() = true;
for (int format_version : {6, 7}) {
// Test all supported compression types because (at least historically)
// dictionary compression could be enabled and a dictionary block saved
// but ignored by some compression types. Ensure we at least don't crash
// or return corruption for those.
for (auto compression_type : GetSupportedCompressions()) {
// Extra handling checks only for types actually supporting dictionary
// compression.
bool dict_supported =
std::count(dict_compressions.begin(), dict_compressions.end(),
compression_type) > 0;
RecordCacheCountersForCompressionDict(options);
Options options = CurrentOptions();
options.bottommost_compression = compression_type;
options.bottommost_compression_opts.max_dict_bytes = 4096;
options.bottommost_compression_opts.enabled = true;
options.create_if_missing = true;
options.num_levels = 2;
options.statistics = ROCKSDB_NAMESPACE::CreateDBStatistics();
options.target_file_size_base = kNumEntriesPerFile * kNumBytesPerEntry;
BlockBasedTableOptions table_options;
table_options.cache_index_and_filter_blocks = true;
table_options.block_cache.reset(new MockCache());
table_options.format_version = format_version;
options.table_factory.reset(NewBlockBasedTableFactory(table_options));
DestroyAndReopen(options);
for (int i = 0; i < kNumFiles; ++i) {
ASSERT_EQ(i, NumTableFilesAtLevel(0, 0));
for (int j = 0; j < kNumEntriesPerFile; ++j) {
std::string value = rnd.RandomString(kNumBytesPerEntry);
ASSERT_OK(Put(Key(j * kNumFiles + i), value.c_str()));
RecordCacheCountersForCompressionDict(options);
for (int i = 0; i < kNumFiles; ++i) {
ASSERT_EQ(i, NumTableFilesAtLevel(0, 0));
for (int j = 0; j < kNumEntriesPerFile; ++j) {
std::string value = rnd.RandomString(kNumBytesPerEntry);
ASSERT_OK(Put(Key(j * kNumFiles + i), value.c_str()));
}
ASSERT_OK(Flush());
}
ASSERT_OK(dbfull()->TEST_WaitForCompact());
ASSERT_EQ(0, NumTableFilesAtLevel(0));
ASSERT_EQ(kNumFiles, NumTableFilesAtLevel(1));
if (dict_supported) {
// Compression dictionary blocks are preloaded.
CheckCacheCountersForCompressionDict(
options, kNumFiles /* expected_compression_dict_misses */,
0 /* expected_compression_dict_hits */,
kNumFiles /* expected_compression_dict_inserts */);
}
// Seek to a key in a file. It should cause the SST's dictionary
// meta-block to be read.
RecordCacheCounters(options);
RecordCacheCountersForCompressionDict(options);
ReadOptions read_options;
ASSERT_NE("NOT_FOUND", Get(Key(kNumFiles * kNumEntriesPerFile - 1)));
if (dict_supported) {
// Two block hits: index and dictionary since they are prefetched
// One block missed/added: data block
CheckCacheCounters(options, 1 /* expected_misses */,
2 /* expected_hits */, 1 /* expected_inserts */,
0 /* expected_failures */);
CheckCacheCountersForCompressionDict(
options, 0 /* expected_compression_dict_misses */,
1 /* expected_compression_dict_hits */,
0 /* expected_compression_dict_inserts */);
}
ASSERT_OK(Flush());
}
ASSERT_OK(dbfull()->TEST_WaitForCompact());
ASSERT_EQ(0, NumTableFilesAtLevel(0));
ASSERT_EQ(kNumFiles, NumTableFilesAtLevel(1));
// Compression dictionary blocks are preloaded.
CheckCacheCountersForCompressionDict(
options, kNumFiles /* expected_compression_dict_misses */,
0 /* expected_compression_dict_hits */,
kNumFiles /* expected_compression_dict_inserts */);
// Seek to a key in a file. It should cause the SST's dictionary meta-block
// to be read.
RecordCacheCounters(options);
RecordCacheCountersForCompressionDict(options);
ReadOptions read_options;
ASSERT_NE("NOT_FOUND", Get(Key(kNumFiles * kNumEntriesPerFile - 1)));
// Two block hits: index and dictionary since they are prefetched
// One block missed/added: data block
CheckCacheCounters(options, 1 /* expected_misses */, 2 /* expected_hits */,
1 /* expected_inserts */, 0 /* expected_failures */);
CheckCacheCountersForCompressionDict(
options, 0 /* expected_compression_dict_misses */,
1 /* expected_compression_dict_hits */,
0 /* expected_compression_dict_inserts */);
}
}

584
db/db_test2.cc
View File

@@ -1888,49 +1888,45 @@ TEST_F(DBTest2, RoundRobinManager) {
auto mgr = std::make_shared<RoundRobinManager>(
GetDefaultBuiltinCompressionManager());
for (CompressionType type : {kZSTD}) {
std::vector<std::string> values;
for (bool use_wrapper : {true}) {
SCOPED_TRACE("Compression type: " + std::to_string(type) +
(use_wrapper ? " with " : " no ") + "wrapper");
std::vector<std::string> values;
for (bool use_wrapper : {true}) {
SCOPED_TRACE((use_wrapper ? "With " : "No ") + std::string("wrapper"));
Options options = CurrentOptions();
options.compression = type;
options.statistics = ROCKSDB_NAMESPACE::CreateDBStatistics();
options.statistics->set_stats_level(StatsLevel::kExceptTimeForMutex);
BlockBasedTableOptions bbto;
bbto.enable_index_compression = false;
options.table_factory.reset(NewBlockBasedTableFactory(bbto));
options.compression_manager = use_wrapper ? mgr : nullptr;
DestroyAndReopen(options);
Options options = CurrentOptions();
options.statistics = ROCKSDB_NAMESPACE::CreateDBStatistics();
options.statistics->set_stats_level(StatsLevel::kExceptTimeForMutex);
BlockBasedTableOptions bbto;
bbto.enable_index_compression = false;
options.table_factory.reset(NewBlockBasedTableFactory(bbto));
options.compression_manager = use_wrapper ? mgr : nullptr;
DestroyAndReopen(options);
Random rnd(301);
constexpr int kCount = 13;
Random rnd(301);
constexpr int kCount = 13;
// Highly compressible blocks, except 1 non-compressible. Half of the
// compressible are morked for bypass and 1 marked for rejection. Values
// are large enough to ensure just 1 k-v per block.
for (int i = 0; i < kCount; ++i) {
std::string value;
if (i == 6) {
// One non-compressible block
value = rnd.RandomBinaryString(20000);
} else {
test::CompressibleString(&rnd, 0.1, 20000, &value);
}
values.push_back(value);
ASSERT_OK(Put(Key(i), value));
ASSERT_EQ(Get(Key(i)), value);
// Highly compressible blocks, except 1 non-compressible. Half of the
// compressible are morked for bypass and 1 marked for rejection. Values
// are large enough to ensure just 1 k-v per block.
for (int i = 0; i < kCount; ++i) {
std::string value;
if (i == 6) {
// One non-compressible block
value = rnd.RandomBinaryString(20000);
} else {
test::CompressibleString(&rnd, 0.1, 20000, &value);
}
ASSERT_OK(Flush());
// Ensure well-formed for reads
for (int i = 0; i < kCount; ++i) {
ASSERT_NE(Get(Key(i)), "NOT_FOUND");
ASSERT_EQ(Get(Key(i)), values[i]);
}
ASSERT_EQ(Get(Key(kCount)), "NOT_FOUND");
values.push_back(value);
ASSERT_OK(Put(Key(i), value));
ASSERT_EQ(Get(Key(i)), value);
}
ASSERT_OK(Flush());
// Ensure well-formed for reads
for (int i = 0; i < kCount; ++i) {
ASSERT_NE(Get(Key(i)), "NOT_FOUND");
ASSERT_EQ(Get(Key(i)), values[i]);
}
ASSERT_EQ(Get(Key(kCount)), "NOT_FOUND");
}
}
}
@@ -1939,54 +1935,49 @@ TEST_F(DBTest2, RandomMixedCompressionManager) {
if (ZSTD_Supported()) {
auto mgr = std::make_shared<RandomMixedCompressionManager>(
GetDefaultBuiltinCompressionManager());
// Currently mixedmanager only supports with preffered compression manager
// zstd
for (CompressionType type : {kZSTD}) {
std::vector<std::string> values;
for (bool use_wrapper : {true}) {
SCOPED_TRACE("Compression type: " + std::to_string(type) +
(use_wrapper ? " with " : " no ") + "wrapper");
std::vector<std::string> values;
for (bool use_wrapper : {true}) {
SCOPED_TRACE((use_wrapper ? "With " : "No ") + std::string("wrapper"));
Options options = CurrentOptions();
options.compression = type;
options.statistics = ROCKSDB_NAMESPACE::CreateDBStatistics();
options.statistics->set_stats_level(StatsLevel::kExceptTimeForMutex);
BlockBasedTableOptions bbto;
bbto.enable_index_compression = false;
options.table_factory.reset(NewBlockBasedTableFactory(bbto));
options.compression_manager = use_wrapper ? mgr : nullptr;
DestroyAndReopen(options);
Options options = CurrentOptions();
options.statistics = ROCKSDB_NAMESPACE::CreateDBStatistics();
options.statistics->set_stats_level(StatsLevel::kExceptTimeForMutex);
BlockBasedTableOptions bbto;
bbto.enable_index_compression = false;
options.table_factory.reset(NewBlockBasedTableFactory(bbto));
options.compression_manager = use_wrapper ? mgr : nullptr;
DestroyAndReopen(options);
Random rnd(301);
constexpr int kCount = 13;
Random rnd(301);
constexpr int kCount = 13;
// Highly compressible blocks, except 1 non-compressible. Half of the
// compressible are morked for bypass and 1 marked for rejection. Values
// are large enough to ensure just 1 k-v per block.
for (int i = 0; i < kCount; ++i) {
std::string value;
if (i == 6) {
// One non-compressible block
value = rnd.RandomBinaryString(20000);
} else {
test::CompressibleString(&rnd, 0.1, 20000, &value);
}
values.push_back(value);
ASSERT_OK(Put(Key(i), value));
ASSERT_EQ(Get(Key(i)), value);
// Highly compressible blocks, except 1 non-compressible. Half of the
// compressible are morked for bypass and 1 marked for rejection. Values
// are large enough to ensure just 1 k-v per block.
for (int i = 0; i < kCount; ++i) {
std::string value;
if (i == 6) {
// One non-compressible block
value = rnd.RandomBinaryString(20000);
} else {
test::CompressibleString(&rnd, 0.1, 20000, &value);
}
ASSERT_OK(Flush());
// Ensure well-formed for reads
for (int i = 0; i < kCount; ++i) {
ASSERT_NE(Get(Key(i)), "NOT_FOUND");
ASSERT_EQ(Get(Key(i)), values[i]);
}
ASSERT_EQ(Get(Key(kCount)), "NOT_FOUND");
values.push_back(value);
ASSERT_OK(Put(Key(i), value));
ASSERT_EQ(Get(Key(i)), value);
}
ASSERT_OK(Flush());
// Ensure well-formed for reads
for (int i = 0; i < kCount; ++i) {
ASSERT_NE(Get(Key(i)), "NOT_FOUND");
ASSERT_EQ(Get(Key(i)), values[i]);
}
ASSERT_EQ(Get(Key(kCount)), "NOT_FOUND");
}
}
}
TEST_F(DBTest2, CompressionManagerWrapper) {
// Test that we can use a custom CompressionManager to wrap the built-in
// CompressionManager, thus adopting a custom *strategy* based on existing
@@ -2103,110 +2094,19 @@ TEST_F(DBTest2, CompressionManagerWrapper) {
}
}
namespace {
template <CompressionType kCompression>
struct CompressorCustomAlg : public CompressorWrapper {
explicit CompressorCustomAlg(const CompressionOptions& opts)
: CompressorWrapper(GetDefaultBuiltinCompressionManager()->GetCompressor(
opts, kSnappyCompression)) {}
explicit CompressorCustomAlg(std::unique_ptr<Compressor> compressor)
: CompressorWrapper(std::move(compressor)) {}
const char* Name() const override { return "CompressorCustomAlg"; }
Status CompressBlock(Slice uncompressed_data, std::string* compressed_output,
CompressionType* out_compression_type,
ManagedWorkingArea* working_area) override {
Status s = wrapped_->CompressBlock(uncompressed_data, compressed_output,
out_compression_type, working_area);
if (*out_compression_type != kNoCompression) {
assert(*out_compression_type == kSnappyCompression);
compressed_output->insert(/*pos=*/0, /*count=*/1,
lossless_cast<char>(kCompression));
*out_compression_type = kCompression;
}
return s;
}
std::unique_ptr<Compressor> MaybeCloneSpecialized(
CacheEntryRole block_type, DictSampleArgs&& dict_samples) override {
std::unique_ptr<Compressor> rv =
wrapped_->MaybeCloneSpecialized(block_type, std::move(dict_samples));
if (rv) {
rv = std::make_unique<CompressorCustomAlg>(std::move(rv));
}
return rv;
}
};
struct DecompressorCustomAlg : public DecompressorWrapper {
DecompressorCustomAlg()
: DecompressorWrapper(
GetDefaultBuiltinCompressionManager()->GetDecompressor()) {}
explicit DecompressorCustomAlg(std::shared_ptr<Decompressor> decompressor)
: DecompressorWrapper(std::move(decompressor)) {}
const char* Name() const override { return "DecompressorCustomAlg"; }
Status MaybeCloneForDict(const Slice& serialized_dict,
std::unique_ptr<Decompressor>* out) override {
Status s = wrapped_->MaybeCloneForDict(serialized_dict, out);
if (s.ok()) {
*out = std::make_unique<DecompressorCustomAlg>(std::move(*out));
}
return s;
}
Status ExtractUncompressedSize(Args& args) override {
if (args.compression_type > kLastBuiltinCompression) {
assert(args.compressed_data.size() > 0);
assert(args.compressed_data[0] ==
lossless_cast<char>(args.compression_type));
// It's ok to modify args if we restore to original
SaveAndRestore<Slice> save_compressed_slice(&args.compressed_data);
args.compressed_data.remove_prefix(1);
SaveAndRestore<CompressionType> save_compression_type(
&args.compression_type);
args.compression_type = kSnappyCompression;
return wrapped_->ExtractUncompressedSize(args);
} else {
// Also support built-in compressions
return wrapped_->ExtractUncompressedSize(args);
}
}
Status DecompressBlock(const Args& args, char* uncompressed_output) override {
if (args.compression_type > kLastBuiltinCompression) {
assert(args.compressed_data.size() > 0);
assert(args.compressed_data[0] ==
lossless_cast<char>(args.compression_type));
// Or we can copy args and modify
Args modified_args = args;
modified_args.compressed_data.remove_prefix(1);
modified_args.compression_type = kSnappyCompression;
return wrapped_->DecompressBlock(modified_args, uncompressed_output);
} else {
// Also support built-in compressions
return wrapped_->DecompressBlock(args, uncompressed_output);
}
}
};
} // anonymous namespace
TEST_F(DBTest2, CompressionManagerCustomCompression) {
if (!Snappy_Supported()) {
fprintf(stderr, "snappy compression not supported, skip this test\n");
return;
}
// 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.
using Compressor8A = CompressorCustomAlg<kCustomCompression8A>;
using Compressor8B = CompressorCustomAlg<kCustomCompression8B>;
using Compressor8C = CompressorCustomAlg<kCustomCompression8C>;
using Compressor8A = test::CompressorCustomAlg<kCustomCompression8A>;
using Compressor8B = test::CompressorCustomAlg<kCustomCompression8B>;
using Compressor8C = test::CompressorCustomAlg<kCustomCompression8C>;
if (!Compressor8A::Supported() || !LZ4_Supported()) {
fprintf(stderr,
"Prerequisite compression library not supported. Skipping\n");
return;
}
class MyManager : public CompressionManager {
public:
@@ -2230,13 +2130,13 @@ TEST_F(DBTest2, CompressionManagerCustomCompression) {
switch (static_cast<unsigned char>(type)) {
case kCustomCompression8A:
used_compressor8A_count_++;
return std::make_unique<Compressor8A>(opts);
return std::make_unique<Compressor8A>();
case kCustomCompression8B:
used_compressor8B_count_++;
return std::make_unique<Compressor8B>(opts);
return std::make_unique<Compressor8B>();
case kCustomCompression8C:
used_compressor8C_count_++;
return std::make_unique<Compressor8C>(opts);
return std::make_unique<Compressor8C>();
// Also support built-in compression algorithms
default:
return GetDefaultBuiltinCompressionManager()->GetCompressor(opts,
@@ -2244,9 +2144,8 @@ TEST_F(DBTest2, CompressionManagerCustomCompression) {
}
}
// TODO: test limited-scope decompressors
std::shared_ptr<Decompressor> GetDecompressor() override {
return std::make_shared<DecompressorCustomAlg>();
return std::make_shared<test::DecompressorCustomAlg>();
}
CompressionType last_specific_decompressor_type_ = kNoCompression;
@@ -2256,7 +2155,9 @@ TEST_F(DBTest2, CompressionManagerCustomCompression) {
const CompressionType* types_end) override {
assert(types_end > types_begin);
last_specific_decompressor_type_ = *types_begin;
return std::make_shared<DecompressorCustomAlg>();
auto decomp = std::make_shared<test::DecompressorCustomAlg>();
decomp->SetAllowedTypes(types_begin, types_end);
return decomp;
}
void AddFriend(const std::shared_ptr<CompressionManager>& mgr) {
@@ -2283,181 +2184,192 @@ TEST_F(DBTest2, CompressionManagerCustomCompression) {
std::map<std::string, std::weak_ptr<CompressionManager>> friends_;
};
// 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.
auto mgr_foo = std::make_shared<MyManager>("Foo");
auto mgr_bar = std::make_shared<MyManager>("Bar");
for (bool use_dict : {false, true}) {
SCOPED_TRACE(use_dict ? "With dict" : "No dict");
// 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");
// 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.
auto mgr_foo = std::make_shared<MyManager>("Foo");
auto mgr_bar = std::make_shared<MyManager>("Bar");
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
options.compression = kSnappyCompression;
DestroyAndReopen(options);
// 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;
Random rnd(404);
std::string value;
ASSERT_OK(Put("a", test::CompressibleString(&rnd, 0.1, kValueSize, &value)));
ASSERT_OK(Flush());
constexpr uint16_t kValueSize = 10000;
// 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);
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);
// 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, "Snappy");
Random rnd(404);
std::string value;
ASSERT_OK(
Put("a", test::CompressibleString(&rnd, 0.1, kValueSize, &value)));
ASSERT_OK(Flush());
// 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);
// 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);
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);
// 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");
// Custom compression schema (different CompatibilityName) not supported
// before format_version=7
options.compression_manager = mgr_foo;
options.compression = kSnappyCompression;
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);
// TODO: eliminate this hack when format_version=7 is published
SaveAndRestore guard(&TEST_AllowUnsupportedFormatVersion(), true);
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);
// Set new format version
bbto.format_version = 7;
options.table_factory.reset(NewBlockBasedTableFactory(bbto));
// 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 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);
// TODO: eliminate this hack when format_version=7 is published
SaveAndRestore guard(&TEST_AllowUnsupportedFormatVersion(), true);
// Using a built-in compression type with fv=7 but named custom schema
options.compression_manager = mgr_foo;
options.compression = kSnappyCompression;
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);
// Set new format version
bbto.format_version = 7;
options.table_factory.reset(NewBlockBasedTableFactory(bbto));
// Verify it was compressed with snappy
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;01;");
EXPECT_EQ(mgr_foo->last_specific_decompressor_type_, kSnappyCompression);
// 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
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);
// Using a built-in compression type with fv=7 but named custom schema
options.compression_manager = mgr_foo;
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 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_, kCustomCompression8A);
// Verify it was compressed with snappy
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_, kLZ4Compression);
// Also dynamically changeable, because the compression manager will respect
// the current setting as reported under the legacy logic
ASSERT_OK(dbfull()->SetOptions({{"compression", "kSnappyCompression"}}));
ASSERT_OK(Put("d", test::CompressibleString(&rnd, 0.1, kValueSize, &value)));
ASSERT_OK(Flush());
ASSERT_EQ(NumTableFilesAtLevel(0), 4);
ASSERT_EQ(Get("d"), value);
// 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 snappy
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;01;");
EXPECT_EQ(mgr_foo->last_specific_decompressor_type_, kSnappyCompression);
// 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_, kCustomCompression8A);
// 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);
// 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 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_, kCustomCompression8B);
// Verify it was compressed with snappy
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_, kLZ4Compression);
// 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 = kSnappyCompression;
ASSERT_EQ(TryReopen(options).code(), Status::Code::kNotFound);
// 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);
// But should re-open if we make Bar aware of the Foo compression manager
mgr_bar->AddFriend(mgr_foo);
Reopen(options);
// 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_, kCustomCompression8B);
// 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);
// 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);
// TODO: mix of compatibility names in same DB
// TODO: test old version of a compression manager unable to read a
// compression type
// TODO: test getting compression manager from object registry
// 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);
// TODO: mix of compatibility names in same DB
// TODO: test old version of a compression manager unable to read a
// compression type
// TODO: test getting compression manager from object registry
}
}
class CompactionStallTestListener : public EventListener {

65
db_stress_tool/db_stress_compression_manager.h Normal file
View File

@@ -0,0 +1,65 @@
// Copyright (c) Meta Platforms, Inc. and affiliates.
// This source code is licensed under both the GPLv2 (found in the
// COPYING file in the root directory) and Apache 2.0 License
// (found in the LICENSE.Apache file in the root directory).
#pragma once
#include "test_util/testutil.h"
namespace ROCKSDB_NAMESPACE {
class DbStressCustomCompressionManager : public CompressionManager {
public:
const char* Name() const override {
return "DbStressCustomCompressionManager";
}
const char* CompatibilityName() const override { return "DbStressCustom1"; }
bool SupportsCompressionType(CompressionType type) const override {
return default_->SupportsCompressionType(type) ||
type == kCustomCompressionAA || type == kCustomCompressionAB ||
type == kCustomCompressionAC;
}
std::unique_ptr<Compressor> GetCompressor(const CompressionOptions& opts,
CompressionType type) override {
// db_stress never specifies a custom type, so we randomly use them anyway
// when this compression manager is used.
std::array<CompressionType, 4> choices = {
type, kCustomCompressionAA, kCustomCompressionAB, kCustomCompressionAC};
type = choices[Random::GetTLSInstance()->Uniform(4)];
switch (static_cast<unsigned char>(type)) {
case kCustomCompressionAA:
return std::make_unique<
test::CompressorCustomAlg<kCustomCompressionAA>>();
case kCustomCompressionAB:
return std::make_unique<
test::CompressorCustomAlg<kCustomCompressionAB>>();
case kCustomCompressionAC:
return std::make_unique<
test::CompressorCustomAlg<kCustomCompressionAC>>();
// Also support built-in compression algorithms
default:
return GetDefaultBuiltinCompressionManager()->GetCompressor(opts, type);
}
}
std::shared_ptr<Decompressor> GetDecompressor() override {
return std::make_shared<test::DecompressorCustomAlg>();
}
std::shared_ptr<Decompressor> GetDecompressorForTypes(
const CompressionType* types_begin,
const CompressionType* types_end) override {
auto decomp = std::make_shared<test::DecompressorCustomAlg>();
decomp->SetAllowedTypes(types_begin, types_end);
return decomp;
}
protected:
std::shared_ptr<CompressionManager> default_ =
GetDefaultBuiltinCompressionManager();
};
} // namespace ROCKSDB_NAMESPACE

52
db_stress_tool/db_stress_test_base.cc
View File

@@ -11,6 +11,7 @@
#include <ios>
#include <thread>
#include "db_stress_tool/db_stress_compression_manager.h"
#include "db_stress_tool/db_stress_listener.h"
#include "rocksdb/io_status.h"
#include "rocksdb/options.h"
@@ -3412,29 +3413,34 @@ void StressTest::Open(SharedState* shared, bool reopen) {
InitializeOptionsFromFlags(cache_, filter_policy_, options_);
}
InitializeOptionsGeneral(cache_, filter_policy_, sqfc_factory_, options_);
if (!strcasecmp(FLAGS_compression_manager.c_str(), "mixed") ||
!strcasecmp(FLAGS_compression_manager.c_str(), "randommixed")) {
// Currently limited to ZSTD compression. Table property compression_name
// needs to set to zstd for now even when there can be more than one
// algorithm in the table under your compressor.
if (!ZSTD_Supported()) {
fprintf(stderr,
"ZSTD compression not supported thus mixed compression cannot be "
"used\n");
exit(1);
}
if (!strcasecmp(FLAGS_compression_manager.c_str(), "mixed")) {
auto mgr = std::make_shared<RoundRobinManager>(
GetDefaultBuiltinCompressionManager());
options_.compression_manager = mgr;
} else if (!strcasecmp(FLAGS_compression_manager.c_str(), "randommixed")) {
auto mgr = std::make_shared<RandomMixedCompressionManager>(
GetDefaultBuiltinCompressionManager());
options_.compression_manager = mgr;
}
options_.compression = kZSTD;
options_.bottommost_compression = kZSTD;
{
// We must register any compression managers with a custom
// CompatibilityName() so that if it was used in a past invocation but not
// the current invocation, we can still read the SST files requiring it.
static std::once_flag loaded;
std::call_once(loaded, [&]() {
TEST_AllowUnsupportedFormatVersion() = true;
auto& library = *ObjectLibrary::Default();
library.AddFactory<CompressionManager>(
DbStressCustomCompressionManager().CompatibilityName(),
[](const std::string& /*uri*/,
std::unique_ptr<CompressionManager>* guard,
std::string* /*errmsg*/) {
*guard = std::make_unique<DbStressCustomCompressionManager>();
return guard->get();
});
});
}
if (!strcasecmp(FLAGS_compression_manager.c_str(), "custom")) {
options_.compression_manager =
std::make_shared<DbStressCustomCompressionManager>();
} else if (!strcasecmp(FLAGS_compression_manager.c_str(), "mixed")) {
options_.compression_manager = std::make_shared<RoundRobinManager>(
GetDefaultBuiltinCompressionManager());
} else if (!strcasecmp(FLAGS_compression_manager.c_str(), "randommixed")) {
options_.compression_manager =
std::make_shared<RandomMixedCompressionManager>(
GetDefaultBuiltinCompressionManager());
} else if (!strcasecmp(FLAGS_compression_manager.c_str(), "autoskip")) {
options_.compression_manager =
CreateAutoSkipCompressionManager(GetDefaultBuiltinCompressionManager());

10
include/rocksdb/advanced_compression.h
View File

@@ -273,11 +273,11 @@ class Decompressor {
// dictionary is processed into a form reusable by repeated compressions in
// many threads, that happens within this call.
//
// Must return OK if storing a result in `out`. Otherwise, could return values
// like NotSupported - dictionary compression is not (yet) supported for this
// kind of Decompressor.
// Corruption - dictionary is malformed (though many implementations will
// accept any data as a dictionary)
// Must return OK if and only if storing a result in `out`. Otherwise, could
// return values like NotSupported - dictionary compression is not (yet)
// supported for this kind of Decompressor. Corruption - dictionary is
// malformed (though many implementations will accept any data as a
// dictionary)
virtual Status MaybeCloneForDict(const Slice& /*serialized_dict*/,
std::unique_ptr<Decompressor>* /*out*/) {
return Status::NotSupported(

8
options/options_settable_test.cc
View File

@@ -703,8 +703,12 @@ TEST_F(OptionsSettableTest, ColumnFamilyOptionsAllFieldsSettable) {
new_options->compaction_options_fifo.file_temperature_age_thresholds[0]
.age,
12345);
ASSERT_EQ(new_options->compression_manager,
GetBuiltinCompressionManager(/*compression_format_version*/ 2));
// TODO: try to enhance ObjectLibrary to support singletons
// ASSERT_EQ(new_options->compression_manager,
// GetBuiltinCompressionManager(/*compression_format_version*/ 2));
ASSERT_STREQ(
new_options->compression_manager->Name(),
GetBuiltinCompressionManager(/*compression_format_version*/ 2)->Name());
ColumnFamilyOptions rnd_filled_options = *new_options;

12
table/block_based/block_based_table_builder.cc
View File

@@ -832,15 +832,17 @@ 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(tbo.compression_type);
basic_working_area.verify = verify_decompressor->ObtainWorkingArea(
basic_compressor->GetPreferredCompressionType());
}
if (state == State::kUnbuffered) {
assert(data_block_compressor);
data_block_verify_decompressor = verify_decompressor.get();
for (uint32_t i = 0; i < compression_parallel_threads; i++) {
data_block_working_areas[i].verify =
verify_decompressor->ObtainWorkingArea(tbo.compression_type);
data_block_verify_decompressor->ObtainWorkingArea(
data_block_compressor->GetPreferredCompressionType());
}
data_block_verify_decompressor = verify_decompressor.get();
}
}
}
@@ -2075,6 +2077,8 @@ void BlockBasedTableBuilder::EnterUnbuffered() {
}
}
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));

126
test_util/testutil.h
View File

@@ -23,6 +23,7 @@
#include "rocksdb/slice.h"
#include "rocksdb/table.h"
#include "table/internal_iterator.h"
#include "util/defer.h"
#include "util/mutexlock.h"
#ifdef ROCKSDB_UNITTESTS_WITH_CUSTOM_OBJECTS_FROM_STATIC_LIBS
@@ -731,6 +732,131 @@ class StringFS : public FileSystemWrapper {
std::unordered_map<std::string, std::string> files_;
};
// A compressor that essentially implements a custom compression algorithm
// by leveraging an existing compression algorithm and putting a custom header
// on it to detect any attempts to decompress it with the wrong compression
// type or dictionary.
template <CompressionType kCompression>
struct CompressorCustomAlg : public CompressorWrapper {
static bool Supported() { return LZ4_Supported(); }
explicit CompressorCustomAlg(std::unique_ptr<Compressor> wrapped =
GetDefaultBuiltinCompressionManager()
->GetCompressor({}, kLZ4Compression))
: CompressorWrapper(std::move(wrapped)),
dictionary_hash_(GetSliceHash(wrapped_->GetSerializedDict())) {
static_assert(kCompression > kLastBuiltinCompression);
}
const char* Name() const override { return "CompressorCustomAlg"; }
Status CompressBlock(Slice uncompressed_data, std::string* compressed_output,
CompressionType* out_compression_type,
ManagedWorkingArea* working_area) override {
Status s = wrapped_->CompressBlock(uncompressed_data, compressed_output,
out_compression_type, working_area);
if (*out_compression_type != kNoCompression) {
assert(*out_compression_type == kLZ4Compression);
std::string header(/*size=*/5, 0);
header[0] = lossless_cast<char>(kCompression);
EncodeFixed32(&header[1], dictionary_hash_);
compressed_output->insert(0, header);
*out_compression_type = kCompression;
}
return s;
}
std::unique_ptr<Compressor> MaybeCloneSpecialized(
CacheEntryRole block_type, DictSampleArgs&& dict_samples) override {
auto clone =
wrapped_->MaybeCloneSpecialized(block_type, std::move(dict_samples));
return std::make_unique<CompressorCustomAlg>(std::move(clone));
}
protected:
uint32_t dictionary_hash_;
};
// A decompressor suitable for all the instantiable CompressorCustomAlg
// implementations. Can be configured to check that it is only used to
// decompress certain types using SetAllowedTypes().
struct DecompressorCustomAlg : public DecompressorWrapper {
using TypeSet = SmallEnumSet<CompressionType, kDisableCompressionOption>;
DecompressorCustomAlg(
std::shared_ptr<Decompressor> wrapped =
GetDefaultBuiltinCompressionManager()->GetDecompressor())
: DecompressorWrapper(std::move(wrapped)),
dictionary_hash_(GetSliceHash(wrapped_->GetSerializedDict())),
allowed_types_(TypeSet::All()) {}
const char* Name() const override { return "DecompressorCustomAlg"; }
Status MaybeCloneForDict(const Slice& serialized_dict,
std::unique_ptr<Decompressor>* out) override {
Status s = wrapped_->MaybeCloneForDict(serialized_dict, out);
if (s.ok()) {
assert(*out != nullptr);
auto clone = std::make_unique<DecompressorCustomAlg>(std::move(*out));
clone->SetAllowedTypes(allowed_types_);
*out = std::move(clone);
assert(out->get()->GetSerializedDict() == serialized_dict);
} else {
assert(*out == nullptr);
}
return s;
}
Status ExtractUncompressedSize(Args& args) override {
if (args.compression_type > kLastBuiltinCompression) {
assert(args.compressed_data.size() > 0);
assert(args.compressed_data[0] ==
lossless_cast<char>(args.compression_type));
assert(DecodeFixed32(args.compressed_data.data() + 1) ==
dictionary_hash_);
// Strip off our header because ExtractUncompressedSize() is also going
// to strip off the uncompressed size data.
args.compressed_data.remove_prefix(5);
// It's ok to modify other parts of args if we restore to original
SaveAndRestore<CompressionType> save_compression_type(
&args.compression_type, kLZ4Compression);
return wrapped_->ExtractUncompressedSize(args);
} else {
// Also support built-in compressions
return wrapped_->ExtractUncompressedSize(args);
}
}
Status DecompressBlock(const Args& args, char* uncompressed_output) override {
if (args.compression_type > kLastBuiltinCompression) {
// Also allowed to copy args and modify
Args modified_args = args;
modified_args.compression_type = kLZ4Compression;
return wrapped_->DecompressBlock(modified_args, uncompressed_output);
} else {
// Also support built-in compressions
return wrapped_->DecompressBlock(args, uncompressed_output);
}
}
void SetAllowedTypes(const CompressionType* types_begin,
const CompressionType* types_end) {
TypeSet allowed_types;
for (auto type = types_begin; type != types_end; ++type) {
allowed_types.Add(*type);
}
allowed_types_ = std::move(allowed_types);
}
void SetAllowedTypes(TypeSet allowed_types) {
allowed_types_ = std::move(allowed_types);
}
protected:
uint32_t dictionary_hash_;
SmallEnumSet<CompressionType, kDisableCompressionOption> allowed_types_;
};
// Randomly initialize the given DBOptions
void RandomInitDBOptions(DBOptions* db_opt, Random* rnd);

22
tools/db_crashtest.py
View File

@@ -178,7 +178,7 @@ default_params = {
"verify_checksum": 1,
"write_buffer_size": lambda: random.choice([1024 * 1024, 4 * 1024 * 1024]),
"writepercent": 35,
"format_version": lambda: random.choice([2, 3, 4, 5, 6, 6]),
"format_version": lambda: random.choice([2, 3, 4, 5, 6, 7, 7]),
"index_block_restart_interval": lambda: random.choice(range(1, 16)),
"use_multiget": lambda: random.randint(0, 1),
"use_get_entity": lambda: random.choice([0] * 7 + [1]),
@@ -350,7 +350,11 @@ default_params = {
"ingest_wbwi_one_in": lambda: random.choice([0, 0, 100, 500]),
"universal_reduce_file_locking": lambda: random.randint(0, 1),
"compression_manager": lambda: random.choice(
["mixed"] * 1 + ["none"] * 2 + ["autoskip"] * 2 + ["randommixed"] * 2
["mixed"] * 1
+ ["none"] * 2
+ ["autoskip"] * 2
+ ["randommixed"] * 2
+ ["custom"] * 3
),
}
@@ -1004,15 +1008,17 @@ def finalize_and_sanitize(src_params):
# have to disable metadata write fault injection to other file
dest_params["exclude_wal_from_write_fault_injection"] = 1
dest_params["metadata_write_fault_one_in"] = 0
# Disabling block align if mixed manager is neing used
if (
# Disabling block align if mixed manager is being used
if dest_params.get("compression_manager") == "custom":
if dest_params.get("block_align") == 1:
dest_params["block_align"] = 0
if dest_params["format_version"] < 7:
dest_params["format_version"] = 7
elif (
dest_params.get("compression_manager") == "mixed"
or dest_params.get("compression_manager") == "randommixed"
):
if dest_params.get("block_align") == 1:
dest_params["block_align"] = 0
dest_params["compression_type"] = "zstd"
dest_params["bottommost_compression_type"] = "zstd"
dest_params["block_align"] = 0
elif dest_params.get("compression_manager") == "autoskip":
# ensuring the compression is being used
if dest_params.get("compression_type") == "none":

2
tools/sst_dump_tool.cc
View File

@@ -197,6 +197,8 @@ int SSTDumpTool::Run(int argc, char const* const* argv, Options options) {
int64_t tmp_val;
TEST_AllowUnsupportedFormatVersion() = true;
for (int i = 1; i < argc; i++) {
if (strncmp(argv[i], "--env_uri=", 10) == 0) {
env_uri = argv[i] + 10;

79
util/compression.cc
View File

@@ -7,6 +7,7 @@
#include "options/options_helper.h"
#include "rocksdb/convenience.h"
#include "rocksdb/utilities/object_registry.h"
namespace ROCKSDB_NAMESPACE {
@@ -700,13 +701,6 @@ class BuiltinDecompressorV2SnappyOnly : public BuiltinDecompressorV2 {
assert(args.compression_type == kSnappyCompression);
return Snappy_DecompressBlock(args, uncompressed_output);
}
Status MaybeCloneForDict(const Slice&,
std::unique_ptr<Decompressor>* out) override {
// NOTE: quietly ignores the dictionary (for compatibility)
*out = std::make_unique<BuiltinDecompressorV2SnappyOnly>();
return Status::OK();
}
};
class BuiltinDecompressorV2WithDict : public BuiltinDecompressorV2 {
@@ -752,6 +746,17 @@ class BuiltinDecompressorV2WithDict : public BuiltinDecompressorV2 {
Status BuiltinDecompressorV2::MaybeCloneForDict(
const Slice& dict, std::unique_ptr<Decompressor>* out) {
// Because of unfortunate decisions in handling built-in compression types,
// all the compression types before ZSTD that do not actually support
// dictionary compression pretend to support it. Specifically, we have to be
// able to read files with a compression dictionary block using those
// compression types even though the compression dictionary is ignored by
// the compression algorithm. And the Decompressor has to return the
// configured dictionary from GetSerializedDict() even if it is ignored. This
// unfortunately means that a new schema version (BuiltinV3?) would be needed
// toactually support dictionary compression in the future for these
// algorithms (if the libraries add support).
// TODO: can we make this a better/cleaner experience?
*out = std::make_unique<BuiltinDecompressorV2WithDict>(dict);
return Status::OK();
}
@@ -955,27 +960,53 @@ const std::shared_ptr<BuiltinCompressionManagerV2>
} // namespace
Status CompressionManager::CreateFromString(
const ConfigOptions& config_options, const std::string& id,
const ConfigOptions& config_options, const std::string& value,
std::shared_ptr<CompressionManager>* result) {
if (id == kNullptrString || id.empty()) {
if (value == kNullptrString || value.empty()) {
result->reset();
return Status::OK();
} else if (id.compare(kBuiltinCompressionManagerV1->CompatibilityName()) ==
0 ||
id.compare(kBuiltinCompressionManagerV1->Name()) == 0) {
*result = kBuiltinCompressionManagerV1;
return Status::OK();
} else if (id.compare(kBuiltinCompressionManagerV2->CompatibilityName()) ==
0 ||
id.compare(kBuiltinCompressionManagerV2->Name()) == 0) {
*result = kBuiltinCompressionManagerV2;
return Status::OK();
} else if (config_options.ignore_unsupported_options) {
return Status::OK();
} else {
return Status::NotFound("Compatible compression manager for \"" + id +
"\"");
}
static std::once_flag loaded;
std::call_once(loaded, [&]() {
auto& library = *ObjectLibrary::Default();
// TODO: try to enhance ObjectLibrary to support singletons
library.AddFactory<CompressionManager>(
kBuiltinCompressionManagerV1->CompatibilityName(),
[](const std::string& /*uri*/,
std::unique_ptr<CompressionManager>* guard,
std::string* /*errmsg*/) {
*guard = std::make_unique<BuiltinCompressionManagerV1>();
return guard->get();
});
library.AddFactory<CompressionManager>(
kBuiltinCompressionManagerV2->CompatibilityName(),
[](const std::string& /*uri*/,
std::unique_ptr<CompressionManager>* guard,
std::string* /*errmsg*/) {
*guard = std::make_unique<BuiltinCompressionManagerV2>();
return guard->get();
});
});
std::string id;
std::unordered_map<std::string, std::string> opt_map;
Status status = Customizable::GetOptionsMap(config_options, result->get(),
value, &id, &opt_map);
if (!status.ok()) { // GetOptionsMap failed
return status;
} else if (id.empty()) { // We have no Id but have options. Not good
return Status::NotSupported("Cannot reset object ", id);
} else {
status = config_options.registry->NewSharedObject(id, result);
}
if (config_options.ignore_unsupported_options && status.IsNotSupported()) {
return Status::OK();
} else if (status.ok()) {
status = Customizable::ConfigureNewObject(config_options, result->get(),
opt_map);
}
return status;
}
std::shared_ptr<CompressionManager>

1
util/compression.h
View File

@@ -248,6 +248,7 @@ struct DecompressorDict {
decompressor_ = std::make_unique<FailureDecompressor>(std::move(s));
} else {
assert(s.ok());
assert(decompressor_->GetSerializedDict() == dict);
}
memory_usage_ = sizeof(struct DecompressorDict);

30
util/simple_mixed_compressor.cc
View File

@@ -16,20 +16,16 @@ namespace ROCKSDB_NAMESPACE {
// MultiCompressorWrapper implementation
MultiCompressorWrapper::MultiCompressorWrapper(const CompressionOptions& opts,
CompressionType type,
CompressionDict&& dict) {
assert(type != kNoCompression);
assert(type == kZSTD);
auto builtInManager = GetDefaultBuiltinCompressionManager();
const auto& compressions = GetSupportedCompressions();
for (auto type_ : compressions) {
if (type_ == kNoCompression) {
for (auto type : compressions) {
if (type == kNoCompression) {
continue;
}
compressors_.push_back(builtInManager->GetCompressor(opts, type_));
compressors_.push_back(builtInManager->GetCompressor(opts, type));
}
(void)dict;
(void)type;
}
size_t MultiCompressorWrapper::GetMaxSampleSizeIfWantDict(
@@ -42,7 +38,7 @@ Slice MultiCompressorWrapper::GetSerializedDict() const {
}
CompressionType MultiCompressorWrapper::GetPreferredCompressionType() const {
return kZSTD;
return compressors_.back()->GetPreferredCompressionType();
}
Compressor::ManagedWorkingArea MultiCompressorWrapper::ObtainWorkingArea() {
@@ -51,6 +47,8 @@ Compressor::ManagedWorkingArea MultiCompressorWrapper::ObtainWorkingArea() {
std::unique_ptr<Compressor> MultiCompressorWrapper::MaybeCloneSpecialized(
CacheEntryRole block_type, DictSampleArgs&& dict_samples) {
// TODO: full dictionary compression support. Currently this just falls
// back on a non-multi compressor when asked to use a dictionary.
return compressors_.back()->MaybeCloneSpecialized(block_type,
std::move(dict_samples));
}
@@ -75,11 +73,9 @@ const char* RandomMixedCompressionManager::Name() const {
}
std::unique_ptr<Compressor> RandomMixedCompressionManager::GetCompressorForSST(
const FilterBuildingContext& context, const CompressionOptions& opts,
CompressionType preferred) {
assert(preferred == kZSTD);
(void)context;
return std::make_unique<RandomMixedCompressor>(opts, preferred);
const FilterBuildingContext& /*context*/, const CompressionOptions& opts,
CompressionType /*preferred*/) {
return std::make_unique<RandomMixedCompressor>(opts);
}
// RoundRobinCompressor implementation
@@ -103,11 +99,9 @@ RelaxedAtomic<uint64_t> RoundRobinCompressor::block_counter{0};
const char* RoundRobinManager::Name() const { return "RoundRobinManager"; }
std::unique_ptr<Compressor> RoundRobinManager::GetCompressorForSST(
const FilterBuildingContext& context, const CompressionOptions& opts,
CompressionType preferred) {
assert(preferred == kZSTD);
(void)context;
return std::make_unique<RoundRobinCompressor>(opts, preferred);
const FilterBuildingContext& /*context*/, const CompressionOptions& opts,
CompressionType /*preferred*/) {
return std::make_unique<RoundRobinCompressor>(opts);
}
} // namespace ROCKSDB_NAMESPACE

1
util/simple_mixed_compressor.h
View File

@@ -18,7 +18,6 @@ namespace ROCKSDB_NAMESPACE {
class MultiCompressorWrapper : public Compressor {
public:
explicit MultiCompressorWrapper(const CompressionOptions& opts,
CompressionType type,
CompressionDict&& dict = {});
size_t GetMaxSampleSizeIfWantDict(CacheEntryRole block_type) const override;