Merge pull request #44363 from ClickHouse/keeper-zookeeper-test-restart-zk-cluster

Try restarting ZK cluster on failed connection in `test_keeper_zookeeper_converted`

tests/integration/test_keeper_zookeeper_converter/test.py

@@ -5,6 +5,7 @@ import helpers.keeper_utils as keeper_utils
 from kazoo.client import KazooClient
 from kazoo.retry import KazooRetry
 from kazoo.security import make_acl
+from kazoo.handlers.threading import KazooTimeoutError
 import os
 import time
 
@@ -35,6 +36,11 @@ def restart_and_clear_zookeeper():
     start_zookeeper()
 
 
+def restart_zookeeper():
+    stop_zookeeper()
+    start_zookeeper()
+
+
 def clear_clickhouse_data():
     node.exec_in_container(
         [
@@ -93,13 +99,25 @@ def get_fake_zk(timeout=60.0):
 
 
 def get_genuine_zk(timeout=60.0):
-    _genuine_zk_instance = KazooClient(
-        hosts=cluster.get_instance_ip("node") + ":2181",
-        timeout=timeout,
-        connection_retry=KazooRetry(max_tries=20),
-    )
-    _genuine_zk_instance.start()
-    return _genuine_zk_instance
+    CONNECTION_RETRIES = 100
+    for i in range(CONNECTION_RETRIES):
+        try:
+            _genuine_zk_instance = KazooClient(
+                hosts=cluster.get_instance_ip("node") + ":2181",
+                timeout=timeout,
+                connection_retry=KazooRetry(max_tries=20),
+            )
+            _genuine_zk_instance.start()
+            return _genuine_zk_instance
+        except KazooTimeoutError:
+            if i == CONNECTION_RETRIES - 1:
+                raise
+
+            print(
+                "Failed to connect to ZK cluster because of timeout. Restarting cluster and trying again."
+            )
+            time.sleep(0.2)
+            restart_zookeeper()
 
 
 def compare_stats(stat1, stat2, path, ignore_pzxid=False):