From 068ada57baace782d98923de2174d50aebc16bf0 Mon Sep 17 00:00:00 2001
From: Mikhail Artemenko <mikhail.artemenko@clickhouse.com>
Date: Fri, 6 Sep 2024 12:23:57 +0000
Subject: [PATCH] update api

---
 src/Common/ZooKeeper/ZooKeeper.cpp | 33 ++++++++++++++++++++----------
 src/Common/ZooKeeper/ZooKeeper.h   | 13 ++++++------
 2 files changed, 29 insertions(+), 17 deletions(-)

diff --git a/src/Common/ZooKeeper/ZooKeeper.cpp b/src/Common/ZooKeeper/ZooKeeper.cpp
index 65c3fdba8d2..ae60520affb 100644
--- a/src/Common/ZooKeeper/ZooKeeper.cpp
+++ b/src/Common/ZooKeeper/ZooKeeper.cpp
@@ -979,8 +979,26 @@ bool ZooKeeper::tryRemoveChildrenRecursive(const std::string & path, bool probab
     return removed_as_expected;
 }
 
-void ZooKeeper::removeRecursive(const std::string & path) // TODO(michicosun) rewrite
+void ZooKeeper::removeRecursive(const std::string & path, uint32_t remove_nodes_limit)
 {
+    if (!isFeatureEnabled(DB::KeeperFeatureFlag::REMOVE_RECURSIVE))
+    {
+        removeChildrenRecursive(path);
+        remove(path);
+        return;
+    }
+
+    check(tryRemoveRecursive(path, remove_nodes_limit), path);
+}
+
+Coordination::Error ZooKeeper::tryRemoveRecursive(const std::string & path, uint32_t remove_nodes_limit)
+{
+    if (!isFeatureEnabled(DB::KeeperFeatureFlag::REMOVE_RECURSIVE))
+    {
+        tryRemoveChildrenRecursive(path);
+        return tryRemove(path);
+    }
+
     auto promise = std::make_shared<std::promise<Coordination::RemoveRecursiveResponse>>();
     auto future = promise->get_future();
 
@@ -989,27 +1007,20 @@ void ZooKeeper::removeRecursive(const std::string & path) // TODO(michicosun) re
         promise->set_value(response);
     };
 
-    impl->removeRecursive(path, 100, std::move(callback));
+    impl->removeRecursive(path, remove_nodes_limit, std::move(callback));
 
     if (future.wait_for(std::chrono::milliseconds(args.operation_timeout_ms)) != std::future_status::ready)
     {
         impl->finalize(fmt::format("Operation timeout on {} {}", Coordination::OpNum::RemoveRecursive, path));
-        check(Coordination::Error::ZOPERATIONTIMEOUT, path);
+        return Coordination::Error::ZOPERATIONTIMEOUT;
     }
     else
     {
         auto response = future.get();
-        check(response.error, path);
+        return response.error;
     }
 }
 
-Coordination::Error ZooKeeper::tryRemoveRecursive(const std::string & path)
-{
-    tryRemoveChildrenRecursive(path);
-    return tryRemove(path);
-}
-
-
 namespace
 {
     struct WaitForDisappearState
diff --git a/src/Common/ZooKeeper/ZooKeeper.h b/src/Common/ZooKeeper/ZooKeeper.h
index 3b94e6004cb..29c4fbc9b74 100644
--- a/src/Common/ZooKeeper/ZooKeeper.h
+++ b/src/Common/ZooKeeper/ZooKeeper.h
@@ -479,15 +479,16 @@ public:
 
     Int64 getClientID();
 
-    /// Remove the node with the subtree. If someone concurrently adds or removes a node
-    /// in the subtree, the result is undefined.
-    void removeRecursive(const std::string & path);
+    /// Remove the node with the subtree.
+    /// If Keeper supports RemoveRecursive operation then it will be performed atomically.
+    /// Otherwise if someone concurrently adds or removes a node in the subtree, the result is undefined.
+    void removeRecursive(const std::string & path, uint32_t remove_nodes_limit = 100);
 
-    /// Remove the node with the subtree. If someone concurrently removes a node in the subtree,
-    /// this will not cause errors.
+    /// Same as removeRecursive but in case if Keeper does not supports RemoveRecursive and
+    /// if someone concurrently removes a node in the subtree, this will not cause errors.
     /// For instance, you can call this method twice concurrently for the same node and the end
     /// result would be the same as for the single call.
-    Coordination::Error tryRemoveRecursive(const std::string & path);
+    Coordination::Error tryRemoveRecursive(const std::string & path, uint32_t remove_nodes_limit = 100);
 
     /// Similar to removeRecursive(...) and tryRemoveRecursive(...), but does not remove path itself.
     /// Node defined as RemoveException will not be deleted.