From ef0ab584b0365f4cc46fde1320febdbea6a21cc0 Mon Sep 17 00:00:00 2001
From: Shanicky Chen <shanicky@ShanickydeMacBook-Pro.local>
Date: Mon, 15 Jan 2024 17:59:17 +0800
Subject: [PATCH] Add check_frag_opt to unreg_worker

---
 src/ctl/src/cmd_impl/meta/reschedule.rs |  3 ++-
 src/ctl/src/lib.rs                      | 16 +++++++++++++++-
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/src/ctl/src/cmd_impl/meta/reschedule.rs b/src/ctl/src/cmd_impl/meta/reschedule.rs
index e2452d408cacf..f14b5c2c30da1 100644
--- a/src/ctl/src/cmd_impl/meta/reschedule.rs
+++ b/src/ctl/src/cmd_impl/meta/reschedule.rs
@@ -241,6 +241,7 @@ pub async fn unregister_workers(
     workers: Vec<String>,
     yes: bool,
     ignore_not_found: bool,
+    check_fragment_occupied: bool,
 ) -> Result<()> {
     let meta_client = context.meta_client().await?;
 
@@ -320,7 +321,7 @@ pub async fn unregister_workers(
                 .intersection(&target_worker_ids)
                 .collect();
 
-            if !intersection_worker_ids.is_empty() {
+            if check_fragment_occupied && !intersection_worker_ids.is_empty() {
                 println!(
                     "worker ids {:?} are still occupied by fragment #{}",
                     intersection_worker_ids, fragment_id
diff --git a/src/ctl/src/lib.rs b/src/ctl/src/lib.rs
index 95c9aa25a6b7a..936a24a8d4bb6 100644
--- a/src/ctl/src/lib.rs
+++ b/src/ctl/src/lib.rs
@@ -486,6 +486,10 @@ enum MetaCommands {
         /// The worker not found will be ignored
         #[clap(long, default_value_t = false)]
         ignore_not_found: bool,
+
+        /// Checking whether the fragment is occupied by workers
+        #[clap(long, default_value_t = false)]
+        check_fragment_occupied: bool,
     },
 
     /// Validate source interface for the cloud team
@@ -706,7 +710,17 @@ pub async fn start_impl(opts: CliOpts, context: &CtlContext) -> Result<()> {
             workers,
             yes,
             ignore_not_found,
-        }) => cmd_impl::meta::unregister_workers(context, workers, yes, ignore_not_found).await?,
+            check_fragment_occupied,
+        }) => {
+            cmd_impl::meta::unregister_workers(
+                context,
+                workers,
+                yes,
+                ignore_not_found,
+                check_fragment_occupied,
+            )
+            .await?
+        }
         Commands::Meta(MetaCommands::ValidateSource { props }) => {
             cmd_impl::meta::validate_source(context, props).await?
         }