From c45bfaf6817d21300174b5db1eaeb30fdc3af02f Mon Sep 17 00:00:00 2001
From: Li0k <yuli@singularity-data.com>
Date: Thu, 30 May 2024 16:47:08 +0800
Subject: [PATCH] feat(ctl): support mutable compression_algorithm

---
 proto/hummock.proto                              |  6 ++++++
 src/ctl/src/cmd_impl/hummock/compaction_group.rs |  5 +++++
 src/ctl/src/lib.rs                               | 16 ++++++++++++++++
 .../hummock/manager/compaction_group_manager.rs  |  5 +++++
 4 files changed, 32 insertions(+)

diff --git a/proto/hummock.proto b/proto/hummock.proto
index 7caf27e155de..c3c475591420 100644
--- a/proto/hummock.proto
+++ b/proto/hummock.proto
@@ -671,6 +671,11 @@ message RiseCtlListCompactionGroupResponse {
 }
 
 message RiseCtlUpdateCompactionConfigRequest {
+  message CompressionAlgorithm {
+    uint32 level = 1;
+    string compression_algorithm = 2;
+  }
+
   message MutableConfig {
     oneof mutable_config {
       uint64 max_bytes_for_level_base = 1;
@@ -688,6 +693,7 @@ message RiseCtlUpdateCompactionConfigRequest {
       uint64 level0_max_compact_file_number = 14;
       bool enable_emergency_picker = 15;
       uint32 tombstone_reclaim_ratio = 16;
+      CompressionAlgorithm compression_algorithm = 17;
     }
   }
   repeated uint64 compaction_group_ids = 1;
diff --git a/src/ctl/src/cmd_impl/hummock/compaction_group.rs b/src/ctl/src/cmd_impl/hummock/compaction_group.rs
index 068472dd1ce4..3fb83e9d16cb 100644
--- a/src/ctl/src/cmd_impl/hummock/compaction_group.rs
+++ b/src/ctl/src/cmd_impl/hummock/compaction_group.rs
@@ -20,6 +20,7 @@ use risingwave_hummock_sdk::compaction_group::StateTableId;
 use risingwave_hummock_sdk::{CompactionGroupId, HummockContextId};
 use risingwave_pb::hummock::compact_task::TaskStatus;
 use risingwave_pb::hummock::rise_ctl_update_compaction_config_request::mutable_config::MutableConfig;
+use risingwave_pb::hummock::rise_ctl_update_compaction_config_request::CompressionAlgorithm;
 
 use crate::CtlContext;
 
@@ -63,6 +64,7 @@ pub fn build_compaction_config_vec(
     level0_overlapping_sub_level_compact_level_count: Option<u32>,
     enable_emergency_picker: Option<bool>,
     tombstone_reclaim_ratio: Option<u32>,
+    compress_algorithm: Option<CompressionAlgorithm>,
 ) -> Vec<MutableConfig> {
     let mut configs = vec![];
     if let Some(c) = max_bytes_for_level_base {
@@ -110,6 +112,9 @@ pub fn build_compaction_config_vec(
     if let Some(c) = tombstone_reclaim_ratio {
         configs.push(MutableConfig::TombstoneReclaimRatio(c))
     }
+    if let Some(c) = compress_algorithm {
+        configs.push(MutableConfig::CompressionAlgorithm(c))
+    }
 
     configs
 }
diff --git a/src/ctl/src/lib.rs b/src/ctl/src/lib.rs
index a1aaa8f48c5f..1276d4bfce43 100644
--- a/src/ctl/src/lib.rs
+++ b/src/ctl/src/lib.rs
@@ -22,6 +22,7 @@ use cmd_impl::hummock::SstDumpArgs;
 use itertools::Itertools;
 use risingwave_hummock_sdk::HummockEpoch;
 use risingwave_meta::backup_restore::RestoreOpts;
+use risingwave_pb::hummock::rise_ctl_update_compaction_config_request::CompressionAlgorithm;
 use risingwave_pb::meta::update_worker_node_schedulability_request::Schedulability;
 use thiserror_ext::AsReport;
 
@@ -254,6 +255,10 @@ enum HummockCommands {
         enable_emergency_picker: Option<bool>,
         #[clap(long)]
         tombstone_reclaim_ratio: Option<u32>,
+        #[clap(long)]
+        compression_level: Option<u32>,
+        #[clap(long)]
+        compression_algorithm: Option<String>,
     },
     /// Split given compaction group into two. Moves the given tables to the new group.
     SplitCompactionGroup {
@@ -683,6 +688,8 @@ async fn start_impl(opts: CliOpts, context: &CtlContext) -> Result<()> {
             level0_overlapping_sub_level_compact_level_count,
             enable_emergency_picker,
             tombstone_reclaim_ratio,
+            compression_level,
+            compression_algorithm,
         }) => {
             cmd_impl::hummock::update_compaction_config(
                 context,
@@ -703,6 +710,15 @@ async fn start_impl(opts: CliOpts, context: &CtlContext) -> Result<()> {
                     level0_overlapping_sub_level_compact_level_count,
                     enable_emergency_picker,
                     tombstone_reclaim_ratio,
+                    if let Some(level) = compression_level {
+                        assert!(compression_algorithm.is_some());
+                        Some(CompressionAlgorithm {
+                            level,
+                            compression_algorithm: compression_algorithm.unwrap(),
+                        })
+                    } else {
+                        None
+                    },
                 ),
             )
             .await?
diff --git a/src/meta/src/hummock/manager/compaction_group_manager.rs b/src/meta/src/hummock/manager/compaction_group_manager.rs
index d5e96c4a26d5..320ab37daedb 100644
--- a/src/meta/src/hummock/manager/compaction_group_manager.rs
+++ b/src/meta/src/hummock/manager/compaction_group_manager.rs
@@ -878,6 +878,11 @@ fn update_compaction_config(target: &mut CompactionConfig, items: &[MutableConfi
             MutableConfig::TombstoneReclaimRatio(c) => {
                 target.tombstone_reclaim_ratio = *c;
             }
+
+            MutableConfig::CompressionAlgorithm(c) => {
+                target.compression_algorithm[c.get_level() as usize]
+                    .clone_from(&c.compression_algorithm);
+            }
         }
     }
 }