From a8f515b454ff59063561d5228163d68d777330d8 Mon Sep 17 00:00:00 2001
From: Alexandr Sorokin <a.sorokin@picodata.io>
Date: Fri, 19 Apr 2024 15:21:15 +0300
Subject: [PATCH] Set the level in hosts as zone

---
 src/task/args.rs             |  5 ++++-
 src/task/cluster.rs          |  4 ++--
 src/task/cluster/hst/test.rs |  8 ++++++--
 src/task/cluster/hst/v2.rs   | 29 +++++++++++++++--------------
 4 files changed, 27 insertions(+), 19 deletions(-)

diff --git a/src/task/args.rs b/src/task/args.rs
index 645bd0c..e3eb2f2 100644
--- a/src/task/args.rs
+++ b/src/task/args.rs
@@ -316,6 +316,9 @@ pub(super) fn read() -> ArgMatches {
 fn fd_as_zone_arg() -> Arg {
     Arg::new("fd-as-zone")
         .long("fd-as-zone")
-        .action(ArgAction::SetTrue)
+        .action(ArgAction::Set)
+        .num_args(0..=1)
+        .default_missing_value("1")
+        .value_parser(clap::value_parser!(u8))
         .help("Used to insert 'failure_domain' field's value of instances in their 'zone' field.")
 }
diff --git a/src/task/cluster.rs b/src/task/cluster.rs
index 9ac9398..b22a1e6 100644
--- a/src/task/cluster.rs
+++ b/src/task/cluster.rs
@@ -627,8 +627,8 @@ impl Cluster {
     /// Note that method is intended to be called after cluster is spread
     /// - that's it, when there may only be single domain name in instance's `failure_domains`.
     pub fn use_failure_domain_as_zone_for_instances(mut self, args: &ArgMatches) -> Self {
-        if args.get_flag("fd-as-zone") {
-            self.hosts.use_failure_domain_as_zone(None);
+        if let Ok(Some(lvl)) = args.try_get_one::<u8>("fd-as-zone") {
+            self.hosts.use_failure_domain_as_zone(*lvl);
         }
         self
     }
diff --git a/src/task/cluster/hst/test.rs b/src/task/cluster/hst/test.rs
index 2e126de..25c6673 100644
--- a/src/task/cluster/hst/test.rs
+++ b/src/task/cluster/hst/test.rs
@@ -515,14 +515,18 @@ fn hosts_use_failure_domain_as_zone() {
     assert_eq!(failure_domain_instance_zone(&host, "cache-2-1"), None);
     assert_eq!(failure_domain_instance_zone(&host, "cache-2-2"), None);
 
-    host.use_failure_domain_as_zone(None);
+    host.use_failure_domain_as_zone(1);
     assert_eq!(
         failure_domain_instance_zone(&host, "cache-2-1"),
         Some("dc-2")
     );
     assert_eq!(
         failure_domain_instance_zone(&host, "cache-2-2"),
-        Some("server-5")
+        Some("dc-2")
+    );
+    assert_eq!(
+        failure_domain_instance_zone(&host, "stateboard-1-1"),
+        Some("dc-1")
     );
 }
 
diff --git a/src/task/cluster/hst/v2.rs b/src/task/cluster/hst/v2.rs
index c9c35b7..a62d88d 100644
--- a/src/task/cluster/hst/v2.rs
+++ b/src/task/cluster/hst/v2.rs
@@ -774,25 +774,26 @@ impl HostV2 {
     }
 
     /// For every instance that has finalized failure domain, replace its zone with that domain name.
-    pub fn use_failure_domain_as_zone(&mut self, dc: Option<String>) {
-        let dc = if dc.is_some() {
-            dc
-        } else if self.name == Name::from("cluster") {
-            None
-        } else {
-            Some(self.name.to_string())
-        };
+    pub fn use_failure_domain_as_zone(&mut self, dc_lvl: u8) {
+        let lvl: u8 = 0;
+        let zone: Option<String> = None;
+        self.set_zone(lvl, dc_lvl, zone)
+    }
+
+    pub fn set_zone(&mut self, mut lvl: u8, dc_lvl: u8, mut zone: Option<String>) {
+        if dc_lvl == lvl {
+            zone = Some(self.name.to_string());
+        }
+
         for instance in self.instances.iter_mut() {
             if let FailureDomains::Finished(failure_domain) = &instance.failure_domains {
-                if let Some(dc) = &dc {
-                    instance.config.zone = Some(dc.clone());
-                } else {
-                    instance.config.zone = Some(failure_domain.clone());
-                }
+                instance.config.zone = zone.clone().or(Some(failure_domain.clone()));
             }
         }
+
+        lvl = lvl + 1;
         for sub_host in self.hosts.iter_mut() {
-            sub_host.use_failure_domain_as_zone(dc.clone())
+            sub_host.set_zone(lvl, dc_lvl, zone.clone())
         }
     }