feat: max_dict_bytes, zstd_max_train_bytes

Author: nagisa-kunhah

kvrocks.conf

@@ -937,6 +937,19 @@ rocksdb.compression snappy
 #   compression library as mentioned above)
 rocksdb.compression_level 32767
 
+# Maximum per-SST dictionary size for block compression. Set to 0 to disable
+# dictionary compression.
+#
+# Default: 0
+rocksdb.max_dict_bytes 0
+
+# Maximum amount of training data used by the ZSTD dictionary trainer. This is
+# most useful when rocksdb.compression is set to zstd and rocksdb.max_dict_bytes
+# is non-zero.
+#
+# Default: 0
+rocksdb.zstd_max_train_bytes 0
+
 # If non-zero, we perform bigger reads when doing compaction. If you're
 # running RocksDB on spinning disks, you should set this to at least 2MB.
 # That way RocksDB's compaction is doing sequential instead of random reads.

src/config/config.cc

@@ -251,8 +251,10 @@ Config::Config() {
       /* rocksdb options */
       {"rocksdb.compression", false,
        new EnumField<rocksdb::CompressionType>(&rocks_db.compression, compression_types,
-                                               rocksdb::CompressionType::kNoCompression)},
+                                               rocksdb::CompressionType::kSnappyCompression)},
       {"rocksdb.compression_level", true, new IntField(&rocks_db.compression_level, 32767, INT_MIN, INT_MAX)},
+      {"rocksdb.max_dict_bytes", true, new UInt32Field(&rocks_db.max_dict_bytes, 0, 0, UINT32_MAX)},
+      {"rocksdb.zstd_max_train_bytes", true, new UInt32Field(&rocks_db.zstd_max_train_bytes, 0, 0, UINT32_MAX)},
       {"rocksdb.compression_start_level", false,
        new IntField(&rocks_db.compression_start_level, 2, 0, KVROCKS_MAX_LSM_LEVEL - 1)},
       {"rocksdb.block_size", true, new IntField(&rocks_db.block_size, 16384, 0, INT_MAX)},

src/config/config.h

@@ -236,6 +236,8 @@ struct Config {
     rocksdb::CompressionType compression;
     int compression_start_level;
     int compression_level;
+    uint32_t max_dict_bytes;
+    uint32_t zstd_max_train_bytes;
     bool disable_auto_compactions;
     bool enable_blob_files;
     int min_blob_size;

src/storage/storage.cc

@@ -178,6 +178,8 @@ rocksdb::Options Storage::InitRocksDBOptions() {
   options.write_buffer_size = config_->rocks_db.write_buffer_size * MiB;
   options.num_levels = KVROCKS_MAX_LSM_LEVEL;
   options.compression_opts.level = config_->rocks_db.compression_level;
+  options.compression_opts.max_dict_bytes = config_->rocks_db.max_dict_bytes;
+  options.compression_opts.zstd_max_train_bytes = config_->rocks_db.zstd_max_train_bytes;
   options.compression_per_level.resize(options.num_levels);
   options.wal_compression = config_->rocks_db.wal_compression;
   for (int i = 0; i < options.num_levels; ++i) {

tests/cppunit/config_test.cc

@@ -134,6 +134,8 @@ TEST(Config, GetAndSet) {
       {"rocksdb.row_cache_size", "100"},
       {"rocksdb.rate_limiter_auto_tuned", "yes"},
       {"rocksdb.compression_level", "32767"},
+      {"rocksdb.max_dict_bytes", "16384"},
+      {"rocksdb.zstd_max_train_bytes", "262144"},
       {"rocksdb.wal_compression", "no"},
       {"histogram-bucket-boundaries", "10,100,1000,10000"},
 
@@ -149,9 +151,12 @@ TEST(Config, GetRenameCommand) {
   unlink(path);
 
   std::ofstream output_file(path, std::ios::out);
-  output_file << "rename-command KEYS KEYS_NEW" << "\n";
-  output_file << "rename-command GET GET_NEW" << "\n";
-  output_file << "rename-command SET SET_NEW" << "\n";
+  output_file << "rename-command KEYS KEYS_NEW"
+              << "\n";
+  output_file << "rename-command GET GET_NEW"
+              << "\n";
+  output_file << "rename-command SET SET_NEW"
+              << "\n";
   output_file.close();
   redis::CommandTable::Reset();
   Config config;
@@ -171,9 +176,12 @@ TEST(Config, Rewrite) {
   unlink(path);
 
   std::ofstream output_file(path, std::ios::out);
-  output_file << "rename-command KEYS KEYS_NEW" << "\n";
-  output_file << "rename-command GET GET_NEW" << "\n";
-  output_file << "rename-command SET SET_NEW" << "\n";
+  output_file << "rename-command KEYS KEYS_NEW"
+              << "\n";
+  output_file << "rename-command GET GET_NEW"
+              << "\n";
+  output_file << "rename-command SET SET_NEW"
+              << "\n";
   output_file.close();
 
   redis::CommandTable::Reset();
@@ -189,6 +197,32 @@ TEST(Config, Rewrite) {
   unlink(path);
 }
 
+TEST(Config, LoadRocksDBDictionaryCompressionOptions) {
+  const char *path = "test.conf";
+  unlink(path);
+
+  std::ofstream output_file(path, std::ios::out);
+  output_file << "rocksdb.max_dict_bytes 16384"
+              << "\n";
+  output_file << "rocksdb.zstd_max_train_bytes 262144"
+              << "\n";
+  output_file.close();
+
+  Config config;
+  ASSERT_TRUE(config.Load(CLIOptions(path)).IsOK());
+  ASSERT_EQ(config.rocks_db.max_dict_bytes, 16384U);
+  ASSERT_EQ(config.rocks_db.zstd_max_train_bytes, 262144U);
+
+  ASSERT_TRUE(config.Rewrite({}).IsOK());
+
+  Config rewritten_config;
+  ASSERT_TRUE(rewritten_config.Load(CLIOptions(path)).IsOK());
+  ASSERT_EQ(rewritten_config.rocks_db.max_dict_bytes, 16384U);
+  ASSERT_EQ(rewritten_config.rocks_db.zstd_max_train_bytes, 262144U);
+
+  unlink(path);
+}
+
 TEST(Config, ParseConfigLine) {
   ASSERT_EQ(*ParseConfigLine(""), ConfigKV{});
   ASSERT_EQ(*ParseConfigLine("# hello"), ConfigKV{});

tests/gocase/unit/config/config_test.go

@@ -123,10 +123,14 @@ func TestConfigSetCompression(t *testing.T) {
 	require.NoError(t, rdb.Do(ctx, "SET", "foo", "bar").Err())
 
 	configKey := "rocksdb.compression"
+	vals, err := rdb.ConfigGet(ctx, configKey).Result()
+	require.NoError(t, err)
+	require.EqualValues(t, "snappy", vals[configKey])
+
 	supportedCompressions := []string{"no", "snappy", "zlib", "lz4", "zstd"}
 	for _, compression := range supportedCompressions {
 		require.NoError(t, rdb.ConfigSet(ctx, configKey, compression).Err())
-		vals, err := rdb.ConfigGet(ctx, configKey).Result()
+		vals, err = rdb.ConfigGet(ctx, configKey).Result()
 		require.NoError(t, err)
 		require.EqualValues(t, compression, vals[configKey])
 	}