looking good before ordering

arthurpassos · arthurpassos · commit 44874313b798 · 2025-11-05T00:28:26.000+01:00
diff --git a/src/Core/Settings.cpp b/src/Core/Settings.cpp
@@ -6876,6 +6876,9 @@ Overwrite file if it already exists when exporting a merge tree part
 )", 0) \
     DECLARE(Bool, export_merge_tree_partition_force_export, false, R"(
 Ignore existing partition export and overwrite the zookeeper entry
+)", 0) \
+    DECLARE(UInt64, export_merge_tree_partition_max_retries, 3, R"(
+Maximum number of retries for exporting a merge tree part in an export partition task
 )", 0) \
     \
     /* ####################################################### */ \
diff --git a/src/Storages/ExportReplicatedMergeTreePartitionManifest.h b/src/Storages/ExportReplicatedMergeTreePartitionManifest.h
@@ -19,6 +19,7 @@ struct ExportReplicatedMergeTreePartitionManifest
     size_t number_of_parts;
     std::vector<String> parts;
     time_t create_time;
+    size_t max_retries;
 
     std::string toJsonString() const
     {
@@ -36,6 +37,7 @@ struct ExportReplicatedMergeTreePartitionManifest
         json.set("parts", parts_array);
         
         json.set("create_time", create_time);
+        json.set("max_retries", max_retries);
         std::ostringstream oss;     // STYLE_CHECK_ALLOW_STD_STRING_STREAM
         oss.exceptions(std::ios::failbit);
         Poco::JSON::Stringifier::stringify(json, oss);
@@ -55,7 +57,7 @@ struct ExportReplicatedMergeTreePartitionManifest
         manifest.destination_table = json->getValue<String>("destination_table");
         manifest.source_replica = json->getValue<String>("source_replica");
         manifest.number_of_parts = json->getValue<size_t>("number_of_parts");
-        
+        manifest.max_retries = json->getValue<size_t>("max_retries");
         auto parts_array = json->getArray("parts");
         for (size_t i = 0; i < parts_array->size(); ++i)
             manifest.parts.push_back(parts_array->getElement<String>(static_cast<unsigned int>(i)));
diff --git a/src/Storages/MergeTree/ExportPartitionTaskScheduler.cpp b/src/Storages/MergeTree/ExportPartitionTaskScheduler.cpp
@@ -22,7 +22,6 @@ void ExportPartitionTaskScheduler::run()
 
     for (const auto & [key, entry] : storage.export_merge_tree_partition_task_entries)
     {
-
         const auto & manifest = entry.manifest;
         const auto & database = storage.getContext()->resolveDatabase(manifest.destination_database);
         const auto & table = manifest.destination_table;
@@ -112,11 +111,13 @@ void ExportPartitionTaskScheduler::run()
             {
                 tryLogCurrentException(__PRETTY_FUNCTION__);
                 zk->tryRemove(fs::path(storage.zookeeper_path) / "exports" / key / "locks" / zk_part_name);
-                /// todo arthur re-schedule this so we can try later
                 /// we should not increment retry_count because the node might just be full
             }
         }
     }
+
+    /// maybe we failed to schedule or failed to export, need to retry eventually
+    storage.export_merge_tree_partition_select_task->scheduleAfter(1000 * 5);
 }
 
 void ExportPartitionTaskScheduler::handlePartExportCompletion(
@@ -137,7 +138,7 @@ void ExportPartitionTaskScheduler::handlePartExportCompletion(
     }
     else
     {
-        handlePartExportFailure(processing_parts_path, processed_part_path, part_name, export_path, zk, result.exception);
+        handlePartExportFailure(processing_parts_path, part_name, export_path, zk, result.exception, manifest.max_retries);
     }
 }
 
@@ -182,7 +183,6 @@ void ExportPartitionTaskScheduler::handlePartExportSuccess(
         // Remove children before parent (order matters for multi operations)
         // Maybe a ls + multi rm..
         requests.emplace_back(zkutil::makeRemoveRequest(processing_parts_path / part_name / "retry_count", -1));
-        requests.emplace_back(zkutil::makeRemoveRequest(processing_parts_path / part_name / "max_retry", -1));
         requests.emplace_back(zkutil::makeRemoveRequest(processing_parts_path / part_name / "status", -1));
         requests.emplace_back(zkutil::makeRemoveRequest(processing_parts_path / part_name, -1));
     }
@@ -223,15 +223,30 @@ void ExportPartitionTaskScheduler::handlePartExportSuccess(
 
 void ExportPartitionTaskScheduler::handlePartExportFailure(
     const std::filesystem::path & processing_parts_path,
-    const std::filesystem::path & processed_part_path,
     const std::string & part_name,
     const std::filesystem::path & export_path,
     const zkutil::ZooKeeperPtr & zk,
-    const String & exception
+    const String & exception,
+    size_t max_retries
 )
 {
     tryLogCurrentException(__PRETTY_FUNCTION__);
 
+    Coordination::Stat locked_by_stat;
+    std::string locked_by;
+
+    if (!zk->tryGet(export_path / "locks" / part_name, locked_by, &locked_by_stat))
+    {
+        LOG_INFO(storage.log, "ExportPartition: Part {} is not locked by any replica, will not increment error counts", part_name);
+        return;
+    }
+
+    if (locked_by != storage.replica_name)
+    {
+        LOG_INFO(storage.log, "ExportPartition: Part {} is locked by another replica, will not increment error counts", part_name);
+        return;
+    }
+
     Coordination::Requests ops;
 
     const auto processing_part_path = processing_parts_path / part_name;
@@ -240,21 +255,17 @@ void ExportPartitionTaskScheduler::handlePartExportFailure(
     {
         std::size_t retry_count = std::stoull(retry_count_string.c_str()) + 1;
 
-
         ops.emplace_back(zkutil::makeSetRequest(processing_part_path / "retry_count", std::to_string(retry_count), -1));
-        ops.emplace_back(zkutil::makeRemoveRequest(export_path / "locks" / part_name, -1));
+        ops.emplace_back(zkutil::makeRemoveRequest(export_path / "locks" / part_name, locked_by_stat.version));
 
-        if (retry_count >= 3)
+        if (retry_count >= max_retries)
         {
-            /// remove from processing and create in processed with status FAILED
-
-            ops.emplace_back(zkutil::makeRemoveRequest(processing_parts_path / part_name, -1));
-            ops.emplace_back(zkutil::makeCreateRequest(processed_part_path, "", zkutil::CreateMode::Persistent));
-            ops.emplace_back(zkutil::makeCreateRequest(processed_part_path / "status", "FAILED", zkutil::CreateMode::Persistent));
-            ops.emplace_back(zkutil::makeCreateRequest(processed_part_path / "finished_by", storage.replica_name, zkutil::CreateMode::Persistent));
+            /// just set status in processing_part_path and finished_by
+            ops.emplace_back(zkutil::makeSetRequest(processing_part_path / "status", "FAILED", -1));
+            ops.emplace_back(zkutil::makeCreateRequest(processing_part_path / "finished_by", storage.replica_name, zkutil::CreateMode::Persistent));
             ops.emplace_back(zkutil::makeSetRequest(export_path / "status", "FAILED", -1));
 
-            LOG_INFO(storage.log, "ExportPartition: Marked part export {} as failed", part_name);
+            LOG_INFO(storage.log, "ExportPartition: Retry count limit exceeded for part {}, will try to fail the entire task", part_name);
         }
 
         std::size_t num_exceptions = 0;
diff --git a/src/Storages/MergeTree/ExportPartitionTaskScheduler.h b/src/Storages/MergeTree/ExportPartitionTaskScheduler.h
@@ -41,11 +41,11 @@ class ExportPartitionTaskScheduler
 
     void handlePartExportFailure(
         const std::filesystem::path & processing_parts_path,
-        const std::filesystem::path & processed_part_path,
         const std::string & part_name,
         const std::filesystem::path & export_path,
         const zkutil::ZooKeeperPtr & zk,
-        const String & exception);
+        const String & exception,
+        size_t max_retries);
 };
 
 }
diff --git a/src/Storages/StorageReplicatedMergeTree.cpp b/src/Storages/StorageReplicatedMergeTree.cpp
@@ -193,6 +193,7 @@ namespace Setting
     extern const SettingsBool update_sequential_consistency;
     extern const SettingsBool allow_experimental_export_merge_tree_part;
     extern const SettingsBool export_merge_tree_partition_force_export;
+    extern const SettingsUInt64 export_merge_tree_partition_max_retries;
 }
 
 namespace MergeTreeSetting
@@ -8124,6 +8125,7 @@ void StorageReplicatedMergeTree::exportPartitionToTable(const PartitionCommand &
     manifest.number_of_parts = part_names.size();
     manifest.parts = part_names;
     manifest.create_time = time(nullptr);
+    manifest.max_retries = query_context->getSettingsRef()[Setting::export_merge_tree_partition_max_retries];
 
     ops.emplace_back(zkutil::makeCreateRequest(
         fs::path(partition_exports_path) / "metadata.json", 
@@ -8147,11 +8149,6 @@ void StorageReplicatedMergeTree::exportPartitionToTable(const PartitionCommand &
             "", 
             zkutil::CreateMode::Persistent));
 
-        ops.emplace_back(zkutil::makeCreateRequest(
-            fs::path(partition_exports_path) / "processing" / part / "max_retry", 
-            "3",
-            zkutil::CreateMode::Persistent));
-
         ops.emplace_back(zkutil::makeCreateRequest(
             fs::path(partition_exports_path) / "processing" / part / "status", 
             "PENDING", 
diff --git a/tests/integration/test_export_replicated_mt_partition_to_object_storage/configs/users.d/profile.xml b/tests/integration/test_export_replicated_mt_partition_to_object_storage/configs/users.d/profile.xml
@@ -0,0 +1,8 @@
+<clickhouse>
+    <profiles>
+        <default>
+            <s3_retry_attempts>3</s3_retry_attempts>
+        </default>
+    </profiles>
+</clickhouse>
+
diff --git a/tests/integration/test_export_replicated_mt_partition_to_object_storage/test.py b/tests/integration/test_export_replicated_mt_partition_to_object_storage/test.py
