Implement MixedCompressor that Round robins on compression algorithm (#13647)

Summary:
**Summary**
This pull request introduces a mixed compressor, RoundRobinManager and RoundRobinCompressor, which selects algorithms in a loop. This implementation replaces the current hacky approach to round-robin compression in BuiltInCompressorV2. Additionally, it configures RocksDB to optionally utilize this customized compressor in the db stress test.

**Testing**
Testing was performed by verifying the stdout output from both RoundRobinCompressor and BuiltInCompressorV2.

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

Reviewed By: pdillinger

Differential Revision: D75921997

Pulled By: shubhajeet

fbshipit-source-id: 8f42ac46f08ba982b2cd70241bd7dc13ff5a1225

db/db_test2.cc

@@ -30,6 +30,7 @@
 #include "test_util/testutil.h"
 #include "util/defer.h"
 #include "util/random.h"
+#include "util/simple_mixed_compressor.h"
 #include "utilities/fault_injection_env.h"
 
 namespace ROCKSDB_NAMESPACE {
@@ -1883,6 +1884,110 @@ TEST_F(DBTest2, CompressionOptions) {
   }
 }
 
+TEST_F(DBTest2, RoundRobinManager) {
+  if (ZSTD_Supported()) {
+    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");
+
+        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);
+
+        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);
+        }
+        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, SimpleMixedCompressionManager) {
+  if (ZSTD_Supported()) {
+    auto mgr = std::make_shared<SimpleMixedCompressionManager>(
+        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");
+
+        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);
+
+        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);
+        }
+        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

db_stress_tool/db_stress_common.h

@@ -398,6 +398,7 @@ DECLARE_bool(use_adaptive_mutex_lru);
 DECLARE_uint32(compress_format_version);
 DECLARE_uint64(manifest_preallocation_size);
 DECLARE_bool(enable_checksum_handoff);
+DECLARE_string(compression_manager);
 DECLARE_uint64(max_total_wal_size);
 DECLARE_double(high_pri_pool_ratio);
 DECLARE_double(low_pri_pool_ratio);

db_stress_tool/db_stress_gflags.cc

@@ -467,7 +467,9 @@ DEFINE_uint64(blob_file_size,
 DEFINE_string(blob_compression_type, "none",
               "[Integrated BlobDB] The compression algorithm to use for large "
               "values stored in blob files.");
-
+DEFINE_string(compression_manager, "mixed",
+              "Ability to change compression manager specified in "
+              "simple_mixed_manager.h (mixed -> roundRobin)");
 DEFINE_bool(enable_blob_garbage_collection,
             ROCKSDB_NAMESPACE::AdvancedColumnFamilyOptions()
                 .enable_blob_garbage_collection,

db_stress_tool/db_stress_test_base.cc

@@ -34,6 +34,7 @@
 #include "rocksdb/utilities/write_batch_with_index.h"
 #include "test_util/testutil.h"
 #include "util/cast_util.h"
+#include "util/simple_mixed_compressor.h"
 #include "utilities/backup/backup_engine_impl.h"
 #include "utilities/fault_injection_fs.h"
 #include "utilities/fault_injection_secondary_cache.h"
@@ -3411,7 +3412,28 @@ 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")) {
+    // 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.
+    options_.compression = kZSTD;
+    options_.bottommost_compression = kZSTD;
+    if (!ZSTD_Supported()) {
+      fprintf(stderr,
+              "ZSTD compression not supported thus mixed compression cannot be "
+              "used\n");
+      exit(1);
+    }
+    auto mgr = std::make_shared<RoundRobinManager>(
+        GetDefaultBuiltinCompressionManager());
+    options_.compression_manager = mgr;
+  } else if (!strcasecmp(FLAGS_compression_manager.c_str(), "none")) {
+    // Nothing to do using default compression manager
+  } else {
+    fprintf(stderr, "Unknown compression manager: %s\n",
+            FLAGS_compression_manager.c_str());
+    exit(1);
+  }
   if (FLAGS_prefix_size == 0 && FLAGS_rep_factory == kHashSkipList) {
     fprintf(stderr,
             "prefeix_size cannot be zero if memtablerep == prefix_hash\n");

tools/db_crashtest.py

@@ -347,7 +347,9 @@ default_params = {
     "memtable_op_scan_flush_trigger": lambda: random.choice([0, 10, 100, 1000]),
     "memtable_avg_op_scan_flush_trigger": lambda: random.choice([0, 2, 20, 200]),
     "ingest_wbwi_one_in": lambda: random.choice([0, 0, 100, 500]),
+    "compression_manager": lambda: random.choice(["mixed", "none"]),
 }
+
 _TEST_DIR_ENV_VAR = "TEST_TMPDIR"
 # If TEST_TMPDIR_EXPECTED is not specified, default value will be TEST_TMPDIR
 _TEST_EXPECTED_DIR_ENV_VAR = "TEST_TMPDIR_EXPECTED"
@@ -995,10 +997,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
-    # Enabling block_align with compression is not supported
-    if dest_params.get("block_align") == 1:
-        dest_params["compression_type"] = "none"
+    # Disabling block align if mixed manager is neing used
+    if dest_params.get("compression_manager") == "mixed":
+        if dest_params.get("block_align") == 1:
+            dest_params["block_align"] = 0
+        dest_params["compression_type"] = "zstd"
         dest_params["bottommost_compression_type"] = "none"
+    else:
+        # Enabling block_align with compression is not supported
+        if dest_params.get("block_align") == 1:
+            dest_params["compression_type"] = "none"
+            dest_params["bottommost_compression_type"] = "none"
     # If periodic_compaction_seconds is not set, daily_offpeak_time_utc doesn't do anything
     if dest_params.get("periodic_compaction_seconds") == 0:
         dest_params["daily_offpeak_time_utc"] = ""

tools/ldb_cmd.cc

@@ -44,6 +44,7 @@
 #include "util/cast_util.h"
 #include "util/coding.h"
 #include "util/file_checksum_helper.h"
+#include "util/simple_mixed_compressor.h"
 #include "util/stderr_logger.h"
 #include "util/string_util.h"
 #include "util/write_batch_util.h"
@@ -867,11 +868,12 @@ bool LDBCommand::ParseCompressionTypeOption(
             "No compressions are supported in this build for \"mixed\".");
         return false;
       }
-      // A temporary hack to generate an SST file with a mix of compression
-      // types, as this has been *de facto* supported for a long time on the
-      // read side with no code to generate them on the write side. We can test
-      // that functionality, e.g. in check_format_compatible.sh, with this hack
-      g_hack_mixed_compression.StoreRelaxed(1);
+      options_.compression = kZSTD;
+      options_.bottommost_compression = kZSTD;
+      auto mgr = std::make_shared<RoundRobinManager>(
+          GetDefaultBuiltinCompressionManager());
+      options_.compression_manager = mgr;
+
       // Need to list zstd in the compression_name table property if it's
       // potentially used by being in the mix (i.e., potentially at least one
       // data block in the table is compressed by zstd). This ensures proper

util/compression.cc

@@ -254,7 +254,6 @@ class BuiltinCompressorV2 : public Compressor {
   void ReleaseWorkingArea(WorkingArea* wa) override {
     delete static_cast<CompressionContext*>(wa);
   }
-
   Status CompressBlock(Slice uncompressed_data, std::string* compressed_output,
                        CompressionType* out_compression_type,
                        ManagedWorkingArea* wa) override {
@@ -264,17 +263,6 @@ class BuiltinCompressorV2 : public Compressor {
       ctx = static_cast<CompressionContext*>(wa->get());
     }
     CompressionType type = type_;
-#ifndef NDEBUG
-    if (type != kNoCompression && g_hack_mixed_compression.LoadRelaxed() > 0U) {
-      // To assert that if zstd is in the mix, the compression_name table
-      // property (which comes from `type_`) needs to be set to kZSTD, for
-      // proper handling of context and dictionaries.
-      assert(!ZSTD_Supported() || type == kZSTD);
-      const auto& compressions = GetSupportedCompressions();
-      auto counter = g_hack_mixed_compression.FetchAddRelaxed(1);
-      type = compressions[counter % compressions.size()];
-    }
-#endif  // !NDEBUG
     if (ctx == nullptr) {
       tmp_ctx.emplace(type, opts_);
       ctx = &*tmp_ctx;
@@ -821,7 +809,6 @@ Status BuiltinDecompressorV2OptimizeZstd::MaybeCloneForDict(
       serialized_dict);
   return Status::OK();
 }
-
 class BuiltinCompressionManagerV2 : public CompressionManager {
  public:
   BuiltinCompressionManagerV2() = default;
@@ -954,7 +941,4 @@ GetDefaultBuiltinCompressionManager() {
 // END built-in implementation of customization interface
 // ***********************************************************************
 
-#ifndef NDEBUG
-RelaxedAtomic<uint64_t> g_hack_mixed_compression{0};
-#endif  // !NDEBUG
 }  // namespace ROCKSDB_NAMESPACE

util/compression.h

@@ -1988,10 +1988,4 @@ class ZSTDStreamingUncompress final : public StreamingUncompress {
 #endif
 };
 
-#ifndef NDEBUG
-// 0 == disable the hack
-// > 0 => counter for rotating through compression types
-extern RelaxedAtomic<uint64_t> g_hack_mixed_compression;
-#endif
-
 }  // namespace ROCKSDB_NAMESPACE

util/simple_mixed_compressor.h

@@ -0,0 +1,130 @@
+// Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
+//  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).
+//
+// Copyright (c) 2011 The LevelDB Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file. See the AUTHORS file for names of contributors.
+//
+#pragma once
+#include <random>
+
+#include "compression.h"
+#include "options/options_helper.h"
+#include "rocksdb/advanced_compression.h"
+
+namespace ROCKSDB_NAMESPACE {
+
+class MultiCompressorWrapper : public Compressor {
+ public:
+  explicit 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) {  // Avoid no compression
+        continue;
+      }
+      compressors_.push_back(builtInManager->GetCompressor(opts, type_));
+    }
+    (void)dict;
+    (void)type;
+  }
+  size_t GetMaxSampleSizeIfWantDict(CacheEntryRole block_type) const override {
+    return compressors_.back()->GetMaxSampleSizeIfWantDict(block_type);
+  }
+
+  Slice GetSerializedDict() const override {
+    return compressors_.back()->GetSerializedDict();
+  }
+
+  CompressionType GetPreferredCompressionType() const override { return kZSTD; }
+
+  ManagedWorkingArea ObtainWorkingArea() override {
+    return compressors_.back()->ObtainWorkingArea();
+  }
+  virtual std::unique_ptr<Compressor> MaybeCloneSpecialized(
+      CacheEntryRole block_type, DictSampleArgs&& dict_samples) override {
+    return compressors_.back()->MaybeCloneSpecialized(block_type,
+                                                      std::move(dict_samples));
+  }
+
+ protected:
+  std::vector<std::unique_ptr<Compressor>> compressors_;
+
+ private:
+  mutable std::mutex mutex_;  // Protects access to current_index_
+};
+struct SimpleMixedCompressor : public MultiCompressorWrapper {
+  using MultiCompressorWrapper::MultiCompressorWrapper;
+  Status CompressBlock(Slice uncompressed_data, std::string* compressed_output,
+                       CompressionType* out_compression_type,
+                       ManagedWorkingArea* wa) override {
+    const auto& compressions = GetSupportedCompressions();
+    std::random_device rd;
+    std::mt19937 gen(rd());
+    std::uniform_int_distribution<> dis(
+        1, (int)compressions.size() - 2);  // avoiding no compression and zstd
+    auto selected = dis(gen);
+    auto& compressor = compressors_[selected % compressors_.size()];
+    // fprintf(stdout, "[MultiCompressorWrapper] selected compressor
+    // typeint:%d\n",
+    //         selected);
+    Status status = compressor->CompressBlock(
+        uncompressed_data, compressed_output, out_compression_type, wa);
+    return status;
+  }
+};
+
+class SimpleMixedCompressionManager : public CompressionManagerWrapper {
+  using CompressionManagerWrapper::CompressionManagerWrapper;
+  const char* Name() const override { return wrapped_->Name(); }
+  std::unique_ptr<Compressor> GetCompressorForSST(
+      const FilterBuildingContext& context, const CompressionOptions& opts,
+      CompressionType preferred) override {
+    assert(preferred == kZSTD);
+    (void)context;
+    return std::make_unique<SimpleMixedCompressor>(opts, preferred);
+  }
+};
+
+struct RoundRobinCompressor : public MultiCompressorWrapper {
+  using MultiCompressorWrapper::MultiCompressorWrapper;
+  Status CompressBlock(Slice uncompressed_data, std::string* compressed_output,
+                       CompressionType* out_compression_type,
+                       ManagedWorkingArea* wa) override {
+    const auto& compressions = GetSupportedCompressions();
+    auto counter = block_counter.FetchAddRelaxed(1);
+    auto sel_idx = counter % (compressions.size() - 1);
+    auto& compressor = compressors_[sel_idx];
+    // auto type = compressions[sel_idx];
+    // fprintf(stdout,
+    //         "[CompressorWrapper] selected compression algo: %s typeint:%d\n",
+    //         std::to_string(type).c_str(), type);
+    return compressor->CompressBlock(uncompressed_data, compressed_output,
+                                     out_compression_type, wa);
+  }
+  static RelaxedAtomic<uint64_t> block_counter;
+};
+RelaxedAtomic<uint64_t> RoundRobinCompressor::block_counter{0};
+
+class RoundRobinManager : public CompressionManagerWrapper {
+  using CompressionManagerWrapper::CompressionManagerWrapper;
+  const char* Name() const override { return wrapped_->Name(); }
+  std::unique_ptr<Compressor> GetCompressorForSST(
+      const FilterBuildingContext& context, const CompressionOptions& opts,
+      CompressionType preferred) override {
+    assert(preferred == kZSTD);
+    (void)context;
+    // fprintf(stdout,
+    //         "[CompressorWrapper] selected compression algo: %s typeint:%d\n",
+    // void)context;
+    return std::make_unique<RoundRobinCompressor>(opts, preferred);
+  }
+};
+
+}  // namespace ROCKSDB_NAMESPACE