From 9fe9911f703c02f7f5e0e83bf900e9d545f21684 Mon Sep 17 00:00:00 2001
From: xiaobiaozhao <zhaoxiaobiaogm@gmail.com>
Date: Sun, 9 Jul 2023 15:04:26 +0000
Subject: [PATCH] feat: enable dynamin compression

---
 src/config/config.cc                    |  5 +++++
 src/config/config.h                     |  1 +
 src/storage/storage.cc                  | 23 ++++++++++++++++++++
 tests/gocase/unit/config/config_test.go | 28 +++++++++++++++++++++++++
 4 files changed, 57 insertions(+)

diff --git a/src/config/config.cc b/src/config/config.cc
index e278aebf862..1c830b636ae 100644
--- a/src/config/config.cc
+++ b/src/config/config.cc
@@ -601,6 +601,11 @@ void Config::initFieldCallback() {
       {"rocksdb.level0_slowdown_writes_trigger", set_cf_option_cb},
       {"rocksdb.level0_stop_writes_trigger", set_cf_option_cb},
       {"rocksdb.level0_file_num_compaction_trigger", set_cf_option_cb},
+      {"rocksdb.compression",
+       [](Server *srv, const std::string &k, const std::string &v) -> Status {
+         if (!srv) return Status::OK();
+         return srv->storage->SetOptionForAllColumnFamilies(TrimRocksDbPrefix(k), v);
+       }},
 #ifdef ENABLE_OPENSSL
       {"tls-cert-file", set_tls_option},
       {"tls-key-file", set_tls_option},
diff --git a/src/config/config.h b/src/config/config.h
index 3eabc3bed05..b9ef433a941 100644
--- a/src/config/config.h
+++ b/src/config/config.h
@@ -39,6 +39,7 @@ namespace engine {
 class Storage;
 }
 
+extern ConfigEnum compression_types[];
 constexpr const uint32_t PORT_LIMIT = 65535;
 
 enum SupervisedMode { kSupervisedNone = 0, kSupervisedAutoDetect, kSupervisedSystemd, kSupervisedUpStart };
diff --git a/src/storage/storage.cc b/src/storage/storage.cc
index b2b6757f770..f46ef209f52 100644
--- a/src/storage/storage.cc
+++ b/src/storage/storage.cc
@@ -180,6 +180,29 @@ rocksdb::Options Storage::InitRocksDBOptions() {
 }
 
 Status Storage::SetOptionForAllColumnFamilies(const std::string &key, const std::string &value) {
+  if ("compression" == key) {
+    std::unordered_map<std::string, std::string> compression_map = {{"no", "kNoCompression"},
+                                                                    {"snappy", "kSnappyCompression"},
+                                                                    {"lz4", "kLZ4Compression"},
+                                                                    {"zstd", "kZSTD"},
+                                                                    {"zlib", "kZlibCompression"}};
+    if (compression_map.find(value) == compression_map.end()) {
+      return {Status::NotOK, "Invalid compression type"};
+    }
+
+    std::string compression_levels = "kNoCompression:kNoCompression";
+    for (size_t i = 2; i < db_->GetOptions().compression_per_level.size(); i++) {
+      compression_levels += ":";
+      compression_levels += compression_map[value];
+    }
+
+    for (auto &cf_handle : cf_handles_) {
+      auto s = db_->SetOptions(cf_handle, {{"compression_per_level", compression_levels}});
+      if (!s.ok()) return {Status::NotOK, s.ToString()};
+    }
+    return Status::OK();
+  }
+
   for (auto &cf_handle : cf_handles_) {
     auto s = db_->SetOptions(cf_handle, {{key, value}});
     if (!s.ok()) return {Status::NotOK, s.ToString()};
diff --git a/tests/gocase/unit/config/config_test.go b/tests/gocase/unit/config/config_test.go
index 1642b763963..1c23509965f 100644
--- a/tests/gocase/unit/config/config_test.go
+++ b/tests/gocase/unit/config/config_test.go
@@ -108,3 +108,31 @@ func TestSetConfigBackupDir(t *testing.T) {
 	require.True(t, hasCompactionFiles(newBackupDir))
 	require.True(t, hasCompactionFiles(originBackupDir))
 }
+
+func TestSetConfigCompression(t *testing.T) {
+	configs := map[string]string{}
+	srv := util.StartServer(t, configs)
+	defer srv.Close()
+
+	ctx := context.Background()
+	rdb := srv.NewClient()
+	defer func() { require.NoError(t, rdb.Close()) }()
+	require.NoError(t, rdb.Do(ctx, "SET", "foo", "bar").Err())
+
+	require.Equal(t, "bar", rdb.Get(ctx, "foo").Val())
+	r := rdb.Do(ctx, "CONFIG", "GET", "rocksdb.compression")
+	rList := r.Val().([]interface{})
+	require.EqualValues(t, rList[1], "no")
+
+	require.NoError(t, rdb.Do(ctx, "CONFIG", "SET", "rocksdb.compression", "lz4").Err())
+	require.Equal(t, "bar", rdb.Get(ctx, "foo").Val())
+	r = rdb.Do(ctx, "CONFIG", "GET", "rocksdb.compression")
+	rList = r.Val().([]interface{})
+	require.EqualValues(t, rList[1], "lz4")
+
+	require.NoError(t, rdb.Do(ctx, "CONFIG", "SET", "rocksdb.compression", "zstd").Err())
+	require.Equal(t, "bar", rdb.Get(ctx, "foo").Val())
+	r = rdb.Do(ctx, "CONFIG", "GET", "rocksdb.compression")
+	rList = r.Val().([]interface{})
+	require.EqualValues(t, rList[1], "zstd")
+}