From 9aaa2c10c14ab4dc916b1d48e2e46df1f593cbc7 Mon Sep 17 00:00:00 2001
From: Maxim Yurchuk <maxim-yurchuk@ydb.tech>
Date: Mon, 15 Apr 2024 11:18:08 +0000
Subject: [PATCH] Use drives from config for formatting in ydbd_slice (#3729)

---
 ydb/tools/ydbd_slice/handlers.py      | 34 ++++++++++++++++++++-------
 ydb/tools/ydbd_slice/kube/handlers.py |  3 +++
 2 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/ydb/tools/ydbd_slice/handlers.py b/ydb/tools/ydbd_slice/handlers.py
index 32627706efa6..dd8a2e0884a6 100644
--- a/ydb/tools/ydbd_slice/handlers.py
+++ b/ydb/tools/ydbd_slice/handlers.py
@@ -20,12 +20,6 @@ def __str__(self):
         )
 
 
-def format_drivers(nodes):
-    cmd = r"sudo find /dev/disk/ -path '*/by-partlabel/kikimr_*' " \
-          r"-exec dd if=/dev/zero of={} bs=1M count=1 status=none \;"
-    nodes.execute_async(cmd)
-
-
 class Slice:
     def __init__(self, components, nodes, cluster_details, configurator, do_clear_logs, args, walle_provider):
         self.slice_kikimr_path = '/Berkanavt/kikimr/bin/kikimr'
@@ -56,9 +50,30 @@ def _clear_logs(self):
             "sudo service rsyslog start;"
         self.nodes.execute_async(cmd)
 
+    def _get_all_drives(self):
+        result = []
+        for host in self.cluster_details.hosts:
+            for drive in host.drives:
+                result.append((host.hostname, drive.path))
+        return result
+
+    def _format_drives(self):
+        tasks = []
+        for (host_name, drive_path) in self._get_all_drives():
+            cmd = "dd if=/dev/zero of={} bs=1M count=1 status=none conv=notrunc".format(drive_path)
+            tasks.extend(self.nodes.execute_async_ret(cmd, nodes=[host_name]))
+        self.nodes._check_async_execution(tasks)
+
+    def _check_drives_exist(self):
+        tasks = []
+        for (host_name, drive_path) in self._get_all_drives():
+            cmd = "echo 'Check existance of drive' && test -f {}".format(drive_path)
+            tasks.extend(self.nodes.execute_async_ret(cmd, nodes=[host_name]))
+        self.nodes._check_async_execution(tasks)
+
     def slice_format(self):
         self.slice_stop()
-        format_drivers(self.nodes)
+        self._format_drives()
         self.slice_start()
 
     def slice_clear(self):
@@ -69,7 +84,7 @@ def slice_clear(self):
                 self._clear_slot(slot)
 
         if 'kikimr' in self.components:
-            format_drivers(self.nodes)
+            self._format_drives()
 
     def _invoke_scripts(self, dynamic_cfg_path, scripts):
         for script_name in scripts:
@@ -117,7 +132,8 @@ def slice_install(self):
             self._clear_logs()
 
         if 'kikimr' in self.components:
-            format_drivers(self.nodes)
+            self._check_drives_exist()
+            self._format_drives()
 
             if 'bin' in self.components.get('kikimr', []):
                 self._update_kikimr(self.nodes, self.configurator.kikimr_bin, self.configurator.kikimr_compressed_bin)
diff --git a/ydb/tools/ydbd_slice/kube/handlers.py b/ydb/tools/ydbd_slice/kube/handlers.py
index 05080b8b93a6..e627e903a261 100644
--- a/ydb/tools/ydbd_slice/kube/handlers.py
+++ b/ydb/tools/ydbd_slice/kube/handlers.py
@@ -231,6 +231,9 @@ def slice_nodeclaim_format(api_client, project_path, manifests):
         return
     node_list = nodes.Nodes(node_list)
     handlers.format_drivers(node_list)
+    cmd = r"sudo find /dev/disk/ -path '*/by-partlabel/kikimr_*' " \
+          r"-exec dd if=/dev/zero of={} bs=1M count=1 status=none \;"
+    nodes.execute_async(cmd)
 
 
 def slice_nodeclaim_delete(api_client, project_path, manifests):