Support recompress-with-CompressionManager in sst_dump (facebook#13783)

Summary:
So that we can use --command=recompress with a custom CompressionManager. (It's not required for reading files using a custom CompressionManager because those can already use ObjectLibrary for dependency injection.)

Suggested follow-up:
* These tests should not be using C arrays, snprintf, manual delete, etc. except for thin compatibility with argc/argv.

Pull Request resolved: facebook#13783

Test Plan: unit test added, some manual testing

Reviewed By: archang19

Differential Revision: D78574434

Pulled By: pdillinger

fbshipit-source-id: 609e6c6439090e6b7e9b63fbd4c2d3f04b104fcf

Author: pdillinger

.github/workflows/pr-jobs.yml

@@ -182,7 +182,7 @@ jobs:
     - uses: actions/checkout@v4.1.0
     - uses: "./.github/actions/pre-steps"
     - run: ENCRYPTED_ENV=1 ROCKSDB_DISABLE_SNAPPY=1 ROCKSDB_DISABLE_ZLIB=1 ROCKSDB_DISABLE_BZIP=1 ROCKSDB_DISABLE_LZ4=1 ROCKSDB_DISABLE_ZSTD=1 make V=1 J=32 -j32 check
-    - run: "./sst_dump --help | grep -E -q 'Supported compression types: kNoCompression$' # Verify no compiled in compression\n"
+    - run: "./sst_dump --help | grep -E -q 'Supported built-in compression types: kNoCompression$' # Verify no compiled in compression\n"
     - uses: "./.github/actions/post-steps"
   # ======================== Linux No Test Runs ======================= #
   build-linux-release:

options/options_helper.h

@@ -72,6 +72,9 @@ std::unique_ptr<Configurable> CFOptionsAsConfigurable(
 Status StringToMap(const std::string& opts_str,
                    std::unordered_map<std::string, std::string>* opts_map);
 
+Status GetStringFromCompressionType(std::string* compression_str,
+                                    CompressionType compression_type);
+
 struct OptionsHelper {
   static const std::string kCFOptionsName /*= "ColumnFamilyOptions"*/;
   static const std::string kDBOptionsName /*= "DBOptions" */;

table/sst_file_dumper.cc

@@ -260,16 +260,21 @@ Status SstFileDumper::CalculateCompressedTableSize(
 }
 
 Status SstFileDumper::ShowAllCompressionSizes(
-    size_t block_size,
-    const std::vector<std::pair<CompressionType, const char*>>&
-        compression_types,
+    size_t block_size, const std::vector<CompressionType>& compression_types,
     int32_t compress_level_from, int32_t compress_level_to,
     uint32_t max_dict_bytes, uint32_t zstd_max_train_bytes,
     uint64_t max_dict_buffer_bytes, bool use_zstd_dict_trainer) {
   fprintf(stdout, "Block Size: %" ROCKSDB_PRIszt "\n", block_size);
-  for (auto& i : compression_types) {
-    if (CompressionTypeSupported(i.first)) {
-      fprintf(stdout, "Compression: %-24s\n", i.second);
+  for (CompressionType ctype : compression_types) {
+    std::string cname;
+    if (!GetStringFromCompressionType(&cname, ctype).ok()) {
+      // Can produce names like "Reserved4F" for unrecognized values
+      cname = CompressionTypeToString(ctype);
+    }
+    if (options_.compression_manager
+            ? options_.compression_manager->SupportsCompressionType(ctype)
+            : CompressionTypeSupported(ctype)) {
+      fprintf(stdout, "Compression: %-24s\n", cname.c_str());
       CompressionOptions compress_opt;
       compress_opt.max_dict_bytes = max_dict_bytes;
       compress_opt.zstd_max_train_bytes = zstd_max_train_bytes;
@@ -278,13 +283,13 @@ Status SstFileDumper::ShowAllCompressionSizes(
       for (int32_t j = compress_level_from; j <= compress_level_to; j++) {
         fprintf(stdout, "Compression level: %d", j);
         compress_opt.level = j;
-        Status s = ShowCompressionSize(block_size, i.first, compress_opt);
+        Status s = ShowCompressionSize(block_size, ctype, compress_opt);
         if (!s.ok()) {
           return s;
         }
       }
     } else {
-      fprintf(stdout, "Unsupported compression type: %s.\n", i.second);
+      fprintf(stdout, "Unsupported compression type: %s.\n", cname.c_str());
     }
   }
   return Status::OK();

table/sst_file_dumper.h

@@ -43,9 +43,7 @@ class SstFileDumper {
   Status getStatus() { return init_result_; }
 
   Status ShowAllCompressionSizes(
-      size_t block_size,
-      const std::vector<std::pair<CompressionType, const char*>>&
-          compression_types,
+      size_t block_size, const std::vector<CompressionType>& compression_types,
       int32_t compress_level_from, int32_t compress_level_to,
       uint32_t max_dict_bytes, uint32_t zstd_max_train_bytes,
       uint64_t max_dict_buffer_bytes, bool use_zstd_dict_trainer);

tools/sst_dump_test.cc

@@ -15,6 +15,7 @@
 #include "rocksdb/convenience.h"
 #include "rocksdb/filter_policy.h"
 #include "rocksdb/sst_dump_tool.h"
+#include "rocksdb/utilities/object_registry.h"
 #include "table/block_based/block_based_table_factory.h"
 #include "table/sst_file_dumper.h"
 #include "table/table_builder.h"
@@ -346,6 +347,76 @@ TEST_F(SSTDumpToolTest, CompressedSizes) {
   }
 }
 
+namespace {
+using Compressor8A = test::CompressorCustomAlg<kCustomCompression8A>;
+class MyManager : public CompressionManager {
+ public:
+  static constexpr const char* kCompatibilityName = "SSTDumpToolTest:MyManager";
+  const char* Name() const override { return kCompatibilityName; }
+  const char* CompatibilityName() const override { return kCompatibilityName; }
+
+  bool SupportsCompressionType(CompressionType type) const override {
+    return type == kCustomCompression8A;
+  }
+
+  std::unique_ptr<Compressor> GetCompressor(const CompressionOptions& /*opts*/,
+                                            CompressionType type) override {
+    switch (static_cast<unsigned char>(type)) {
+      case kCustomCompression8A:
+        return std::make_unique<Compressor8A>();
+      default:
+        return nullptr;
+    }
+  }
+
+  std::shared_ptr<Decompressor> GetDecompressor() override {
+    return std::make_shared<test::DecompressorCustomAlg>();
+  }
+};
+}  // namespace
+
+TEST_F(SSTDumpToolTest, CompressionManager) {
+  if (!Compressor8A::Supported()) {
+    fprintf(stderr,
+            "Prerequisite compression library not supported. Skipping\n");
+    return;
+  }
+
+  // Registery in ObjectLibrary to check that sst_dump can use named
+  // CompressionManagers with dependency injection
+  auto& library = *ObjectLibrary::Default();
+  library.AddFactory<CompressionManager>(
+      MyManager::kCompatibilityName,
+      [](const std::string& /*uri*/, std::unique_ptr<CompressionManager>* guard,
+         std::string* /*errmsg*/) {
+        *guard = std::make_unique<MyManager>();
+        return guard->get();
+      });
+
+  Options opts;
+  opts.env = env();
+  BlockBasedTableOptions table_opts;
+  table_opts.filter_policy.reset(
+      ROCKSDB_NAMESPACE::NewBloomFilterPolicy(10, false));
+  opts.table_factory.reset(new BlockBasedTableFactory(table_opts));
+  std::string file_path = MakeFilePath("rocksdb_sst_test.sst");
+  createSST(opts, file_path, 10);
+
+  char* usage[5];
+  PopulateCommandArgs(file_path, "--command=recompress", usage);
+  snprintf(usage[3], kOptLength, "--compression_manager=%s",
+           MyManager::kCompatibilityName);
+  snprintf(usage[4], kOptLength, "--compression_types=kCustomCompression8A");
+
+  ROCKSDB_NAMESPACE::SSTDumpTool tool;
+  ASSERT_TRUE(!tool.Run(5, usage, opts));
+
+  cleanup(opts, file_path);
+  for (int i = 0; i < 5; i++) {
+    delete[] usage[i];
+  }
+}
+
 TEST_F(SSTDumpToolTest, MemEnv) {
   std::unique_ptr<Env> mem_env(NewMemEnv(env()));
   Options opts;

tools/sst_dump_tool.cc

@@ -17,17 +17,6 @@
 
 namespace ROCKSDB_NAMESPACE {
 
-static const std::vector<std::pair<CompressionType, const char*>>
-    kCompressions = {
-        {CompressionType::kNoCompression, "kNoCompression"},
-        {CompressionType::kSnappyCompression, "kSnappyCompression"},
-        {CompressionType::kZlibCompression, "kZlibCompression"},
-        {CompressionType::kBZip2Compression, "kBZip2Compression"},
-        {CompressionType::kLZ4Compression, "kLZ4Compression"},
-        {CompressionType::kLZ4HCCompression, "kLZ4HCCompression"},
-        {CompressionType::kXpressCompression, "kXpressCompression"},
-        {CompressionType::kZSTD, "kZSTD"}};
-
 namespace {
 
 void print_help(bool to_stderr) {
@@ -98,10 +87,15 @@ void print_help(bool to_stderr) {
       be used when trying different compression algorithms
 
     --compression_types=<comma-separated list of CompressionType members, e.g.,
-      kSnappyCompression>
+      kSnappyCompression or kCustomCompressionC4>
       Can be combined with --command=recompress to run recompression for this
       list of compression types
-      Supported compression types: %s
+      Supported built-in compression types: %s
+
+    --compression_manager=<compression manager string>
+      Used with --command=recompress to specify a compression manager to use
+      instead of the built-in compression manager, which may support a
+      different set of compression types.
 
     --parse_internal_key=<0xKEY>
       Convenience option to parse an internal key on the command line. Dumps the
@@ -178,7 +172,8 @@ int SSTDumpTool::Run(int argc, char const* const* argv, Options options) {
   std::string compression_level_to_str;
   size_t block_size = 0;
   size_t readahead_size = 2 * 1024 * 1024;
-  std::vector<std::pair<CompressionType, const char*>> compression_types;
+  std::vector<CompressionType> compression_types;
+  std::shared_ptr<CompressionManager> compression_manager;
   uint64_t total_num_files = 0;
   uint64_t total_num_data_blocks = 0;
   uint64_t total_data_block_size = 0;
@@ -244,19 +239,36 @@ int SSTDumpTool::Run(int argc, char const* const* argv, Options options) {
       std::istringstream iss(compression_types_csv);
       std::string compression_type;
       has_specified_compression_types = true;
+
       while (std::getline(iss, compression_type, ',')) {
-        auto iter = std::find_if(
-            kCompressions.begin(), kCompressions.end(),
-            [&compression_type](std::pair<CompressionType, const char*> curr) {
-              return curr.second == compression_type;
-            });
-        if (iter == kCompressions.end()) {
+        auto iter =
+            OptionsHelper::compression_type_string_map.find(compression_type);
+        if (iter == OptionsHelper::compression_type_string_map.end()) {
           fprintf(stderr, "%s is not a valid CompressionType\n",
                   compression_type.c_str());
           exit(1);
         }
-        compression_types.emplace_back(*iter);
+        compression_types.emplace_back(iter->second);
+      }
+    } else if (strncmp(argv[i], "--compression_manager=", 22) == 0) {
+      std::string compression_manager_str = argv[i] + 22;
+      ConfigOptions config_options;
+      config_options.ignore_unsupported_options = false;
+      Status s = CompressionManager::CreateFromString(
+          config_options, compression_manager_str, &compression_manager);
+      if (!s.ok()) {
+        fprintf(stderr, "Failed to create compression manager: %s\n",
+                s.ToString().c_str());
+        exit(1);
+      }
+      if (compression_manager == nullptr) {
+        fprintf(stderr, "No compression manager created: %s\n",
+                compression_manager_str.c_str());
+        exit(1);
       }
+      options.compression_manager = compression_manager;
+      printf("Using compression manager: %s\n",
+             compression_manager->GetId().c_str());
     } else if (strncmp(argv[i], "--parse_internal_key=", 21) == 0) {
       std::string in_key(argv[i] + 21);
       try {
@@ -450,9 +462,12 @@ int SSTDumpTool::Run(int argc, char const* const* argv, Options options) {
     }
 
     if (command == "recompress") {
+      // TODO: consider getting supported compressions from the compression
+      // manager
       st = dumper.ShowAllCompressionSizes(
           set_block_size ? block_size : 16384,
-          compression_types.empty() ? kCompressions : compression_types,
+          compression_types.empty() ? GetSupportedCompressions()
+                                    : compression_types,
           compress_level_from, compress_level_to, compression_max_dict_bytes,
           compression_zstd_max_train_bytes, compression_max_dict_buffer_bytes,
           !compression_use_zstd_finalize_dict);