From ce7b8aefd274cef40002a52895bd434225a36ddc Mon Sep 17 00:00:00 2001
From: Alexey Milovidov <milovidov@yandex-team.ru>
Date: Wed, 7 Oct 2020 11:41:27 +0300
Subject: [PATCH] Add reconnects to zookeeper-dump-tree tool

---
 utils/zookeeper-dump-tree/main.cpp | 29 +++++++++++++++++++++++++----
 1 file changed, 25 insertions(+), 4 deletions(-)

diff --git a/utils/zookeeper-dump-tree/main.cpp b/utils/zookeeper-dump-tree/main.cpp
index 877a4a996b7..4614ee577e7 100644
--- a/utils/zookeeper-dump-tree/main.cpp
+++ b/utils/zookeeper-dump-tree/main.cpp
@@ -29,16 +29,20 @@ try
         return 1;
     }
 
-    zkutil::ZooKeeper zookeeper(options.at("address").as<std::string>());
+    zkutil::ZooKeeperPtr zookeeper = std::make_shared<zkutil::ZooKeeper>(options.at("address").as<std::string>());
 
     std::string initial_path = options.at("path").as<std::string>();
 
     std::list<std::pair<std::string, std::future<Coordination::ListResponse>>> list_futures;
-    list_futures.emplace_back(initial_path, zookeeper.asyncGetChildren(initial_path));
+    list_futures.emplace_back(initial_path, zookeeper->asyncGetChildren(initial_path));
+
+    size_t num_reconnects = 0;
+    constexpr size_t max_reconnects = 100;
 
     for (auto it = list_futures.begin(); it != list_futures.end(); ++it)
     {
         Coordination::ListResponse response;
+
         try
         {
             response = it->second.get();
@@ -46,8 +50,25 @@ try
         catch (const Coordination::Exception & e)
         {
             if (e.code == Coordination::Error::ZNONODE)
+            {
                 continue;
-            throw;
+            }
+            else if (Coordination::isHardwareError(e.code))
+            {
+                /// Reinitialize the session and move the node to the end of the queue for later retry.
+                if (zookeeper->expired())
+                {
+                    if (num_reconnects == max_reconnects)
+                        throw;
+                    ++num_reconnects;
+                    zookeeper = zookeeper->startNewSession();
+                }
+
+                list_futures.emplace_back(it->first, zookeeper->asyncGetChildren(it->first));
+                continue;
+            }
+            else
+                throw;
         }
 
         std::cout << it->first << '\t' << response.stat.numChildren << '\t' << response.stat.dataLength << '\n';
@@ -55,7 +76,7 @@ try
         for (const auto & name : response.names)
         {
             std::string child_path = it->first == "/" ? it->first + name : it->first + '/' + name;
-            list_futures.emplace_back(child_path, zookeeper.asyncGetChildren(child_path));
+            list_futures.emplace_back(child_path, zookeeper->asyncGetChildren(child_path));
         }
     }
 }