Improvements to partition export

src/Common/ProfileEvents.cpp

@@ -265,6 +265,18 @@
     M(ZooKeeperBytesSent, "Number of bytes send over network while communicating with ZooKeeper.", ValueType::Bytes) \
     M(ZooKeeperBytesReceived, "Number of bytes received over network while communicating with ZooKeeper.", ValueType::Bytes) \
     \
+    M(ExportPartitionZooKeeperRequests, "Total number of ZooKeeper requests made by the export partition feature.", ValueType::Number) \
+    M(ExportPartitionZooKeeperGet, "Number of 'get' requests to ZooKeeper made by the export partition feature.", ValueType::Number) \
+    M(ExportPartitionZooKeeperGetChildren, "Number of 'getChildren' requests to ZooKeeper made by the export partition feature.", ValueType::Number) \
+    M(ExportPartitionZooKeeperGetChildrenWatch, "Number of 'getChildrenWatch' requests to ZooKeeper made by the export partition feature.", ValueType::Number) \
+    M(ExportPartitionZooKeeperGetWatch, "Number of 'getWatch' requests to ZooKeeper made by the export partition feature.", ValueType::Number) \
+    M(ExportPartitionZooKeeperCreate, "Number of 'create' requests to ZooKeeper made by the export partition feature.", ValueType::Number) \
+    M(ExportPartitionZooKeeperSet, "Number of 'set' requests to ZooKeeper made by the export partition feature.", ValueType::Number) \
+    M(ExportPartitionZooKeeperRemove, "Number of 'remove' requests to ZooKeeper made by the export partition feature.", ValueType::Number) \
+    M(ExportPartitionZooKeeperRemoveRecursive, "Number of 'removeRecursive' requests to ZooKeeper made by the export partition feature.", ValueType::Number) \
+    M(ExportPartitionZooKeeperMulti, "Number of 'multi' requests to ZooKeeper made by the export partition feature.", ValueType::Number) \
+    M(ExportPartitionZooKeeperExists, "Number of 'exists' requests to ZooKeeper made by the export partition feature.", ValueType::Number) \
+    \
     M(DistributedConnectionTries, "Total count of distributed connection attempts.", ValueType::Number) \
     M(DistributedConnectionUsable, "Total count of successful distributed connections to a usable server (with required table, but maybe stale).", ValueType::Number) \
     M(DistributedConnectionFailTry, "Total count when distributed connection fails with retry.", ValueType::Number) \

src/Core/Settings.cpp

@@ -6899,6 +6899,10 @@ Possible values:
 - `` (empty value) - use session timezone
 
 Default value is `UTC`.
+)", 0) \
+    DECLARE(Bool, export_merge_tree_partition_lock_inside_the_task, false, R"(
+Only lock a part when the task is already running. This might help with busy waiting where the scheduler locks a part, but the task ends in the pending list.
+On the other hand, there is a chance once the task executes that part has already been locked by another replica and the task will simply early exit.
 )", 0) \
     \
     /* ####################################################### */ \

src/Core/SettingsChangesHistory.cpp

@@ -55,7 +55,8 @@ const VersionToSettingsChangesMap & getSettingsChangesHistory()
             {"export_merge_tree_part_file_already_exists_policy", "skip", "skip", "New setting."},
             {"iceberg_timezone_for_timestamptz", "UTC", "UTC", "New setting."},
             {"hybrid_table_auto_cast_columns", true, true, "New setting to automatically cast Hybrid table columns when segments disagree on types. Default enabled."},
-            {"allow_experimental_hybrid_table", false, false, "Added new setting to allow the Hybrid table engine."}
+            {"allow_experimental_hybrid_table", false, false, "Added new setting to allow the Hybrid table engine."},
+            {"export_merge_tree_partition_lock_inside_the_task", false, false, "New setting."}
         });
         addSettingsChanges(settings_changes_history, "25.8",
         {

src/Storages/ExportReplicatedMergeTreePartitionManifest.h

@@ -109,6 +109,7 @@ struct ExportReplicatedMergeTreePartitionManifest
     bool parallel_formatting;
     bool parquet_parallel_encoding;
     MergeTreePartExportManifest::FileAlreadyExistsPolicy file_already_exists_policy;
+    bool lock_inside_the_task; /// todo temporary
 
     std::string toJsonString() const
     {
@@ -131,6 +132,7 @@ struct ExportReplicatedMergeTreePartitionManifest
         json.set("create_time", create_time);
         json.set("max_retries", max_retries);
         json.set("ttl_seconds", ttl_seconds);
+        json.set("lock_inside_the_task", lock_inside_the_task);
         std::ostringstream oss;     // STYLE_CHECK_ALLOW_STD_STRING_STREAM
         oss.exceptions(std::ios::failbit);
         Poco::JSON::Stringifier::stringify(json, oss);
@@ -172,6 +174,11 @@ struct ExportReplicatedMergeTreePartitionManifest
             /// what to do if it's not a valid value?
         }
 
+        if (json->has("lock_inside_the_task"))
+        {
+            manifest.lock_inside_the_task = json->getValue<bool>("lock_inside_the_task");
+        }
+
         return manifest;
     }
 };

src/Storages/ExportReplicatedMergeTreePartitionTaskEntry.h

@@ -30,7 +30,7 @@ struct ExportReplicatedMergeTreePartitionTaskEntry
     /// This is used to prevent the parts from being deleted before finishing the export operation
     /// It does not mean this replica will export all the parts
     /// There is also a chance this replica does not contain a given part and it is totally ok.
-    std::vector<DataPartPtr> part_references;
+    mutable std::vector<DataPartPtr> part_references;
 
     std::string getCompositeKey() const
     {

src/Storages/MergeTree/BackgroundJobsAssignee.cpp

@@ -95,6 +95,10 @@ bool BackgroundJobsAssignee::scheduleCommonTask(ExecutableTaskPtr common_task, b
     return schedule_res;
 }
 
+std::size_t BackgroundJobsAssignee::getAvailableMoveExecutors() const
+{
+    return getContext()->getMovesExecutor()->getAvailableSlots();
+}
 
 String BackgroundJobsAssignee::toString(Type type)
 {

src/Storages/MergeTree/BackgroundJobsAssignee.h

@@ -55,6 +55,8 @@ class BackgroundJobsAssignee : public WithContext
     bool scheduleMoveTask(ExecutableTaskPtr move_task);
     bool scheduleCommonTask(ExecutableTaskPtr common_task, bool need_trigger);
 
+    std::size_t getAvailableMoveExecutors() const;
+
     /// Just call finish
     ~BackgroundJobsAssignee();
 

src/Storages/MergeTree/ExportPartFromPartitionExportTask.cpp

@@ -0,0 +1,68 @@
+#include <Storages/MergeTree/ExportPartFromPartitionExportTask.h>
+#include <Common/ProfileEvents.h>
+
+namespace ProfileEvents
+{
+    extern const Event ExportPartitionZooKeeperRequests;
+    extern const Event ExportPartitionZooKeeperGetChildren;
+    extern const Event ExportPartitionZooKeeperCreate;
+}
+namespace DB
+{
+
+ExportPartFromPartitionExportTask::ExportPartFromPartitionExportTask(
+    StorageReplicatedMergeTree & storage_,
+    const std::string & key_,
+    const MergeTreePartExportManifest & manifest_)
+    : storage(storage_),
+    key(key_),
+    manifest(manifest_)
+{
+    export_part_task = std::make_shared<ExportPartTask>(storage, manifest);
+}
+
+bool ExportPartFromPartitionExportTask::executeStep()
+{
+    const auto zk = storage.getZooKeeper();
+    const auto part_name = manifest.data_part->name;
+
+    LOG_INFO(storage.log, "ExportPartFromPartitionExportTask: Attempting to lock part: {}", part_name);
+
+    ProfileEvents::increment(ProfileEvents::ExportPartitionZooKeeperRequests);
+    ProfileEvents::increment(ProfileEvents::ExportPartitionZooKeeperCreate);
+    if (Coordination::Error::ZOK == zk->tryCreate(fs::path(storage.zookeeper_path) / "exports" / key / "locks" / part_name, storage.replica_name, zkutil::CreateMode::Ephemeral))
+    {
+        LOG_INFO(storage.log, "ExportPartFromPartitionExportTask: Locked part: {}", part_name);
+        export_part_task->executeStep();
+        return false;
+    }
+
+    LOG_INFO(storage.log, "ExportPartFromPartitionExportTask: Failed to lock part {}, skipping", part_name);
+    return false;
+}
+
+void ExportPartFromPartitionExportTask::cancel() noexcept
+{
+    export_part_task->cancel();
+}
+
+void ExportPartFromPartitionExportTask::onCompleted()
+{
+    export_part_task->onCompleted();
+}
+
+StorageID ExportPartFromPartitionExportTask::getStorageID() const
+{
+    return export_part_task->getStorageID();
+}
+
+Priority ExportPartFromPartitionExportTask::getPriority() const
+{
+    return export_part_task->getPriority();
+}
+
+String ExportPartFromPartitionExportTask::getQueryId() const
+{
+    return export_part_task->getQueryId();
+}
+}

src/Storages/MergeTree/ExportPartFromPartitionExportTask.h

@@ -0,0 +1,33 @@
+#pragma once
+
+#include <Storages/MergeTree/IExecutableTask.h>
+#include <Storages/MergeTree/MergeTreePartExportManifest.h>
+#include <Storages/StorageReplicatedMergeTree.h>
+#include <Storages/MergeTree/ExportPartTask.h>
+
+namespace DB
+{
+
+class ExportPartFromPartitionExportTask : public IExecutableTask
+{
+public:
+    explicit ExportPartFromPartitionExportTask(
+        StorageReplicatedMergeTree & storage_,
+        const std::string & key_,
+        const MergeTreePartExportManifest & manifest_);
+    bool executeStep() override;
+    void onCompleted() override;
+    StorageID getStorageID() const override;
+    Priority getPriority() const override;
+    String getQueryId() const override;
+
+    void cancel() noexcept override;
+
+private:
+    StorageReplicatedMergeTree & storage;
+    std::string key;
+    MergeTreePartExportManifest manifest;
+    std::shared_ptr<ExportPartTask> export_part_task;
+};
+
+}

src/Storages/MergeTree/ExportPartTask.cpp

@@ -13,6 +13,7 @@
 #include <Common/Exception.h>
 #include <Common/ProfileEventsScope.h>
 #include <Storages/MergeTree/ExportList.h>
+#include <Formats/FormatFactory.h>
 
 namespace ProfileEvents
 {
@@ -38,15 +39,25 @@ namespace Setting
     extern const SettingsUInt64 min_bytes_to_use_direct_io;
 }
 
-ExportPartTask::ExportPartTask(MergeTreeData & storage_, const MergeTreePartExportManifest & manifest_, ContextPtr context_)
+ExportPartTask::ExportPartTask(MergeTreeData & storage_, const MergeTreePartExportManifest & manifest_)
     : storage(storage_),
-    manifest(manifest_),
-    local_context(context_)
+    manifest(manifest_)
 {
 }
 
+const MergeTreePartExportManifest & ExportPartTask::getManifest() const
+{
+    return manifest;
+}
+
 bool ExportPartTask::executeStep()
 {
+    auto local_context = Context::createCopy(storage.getContext());
+    local_context->makeQueryContextForExportPart();
+    local_context->setCurrentQueryId(manifest.transaction_id);
+    local_context->setBackgroundOperationTypeForContext(ClientInfo::BackgroundOperationType::EXPORT_PART);
+    local_context->setSettings(manifest.settings);
+
     const auto & metadata_snapshot = manifest.metadata_snapshot;
 
     Names columns_to_read = metadata_snapshot->getColumns().getNamesOfPhysical();
@@ -91,7 +102,7 @@ bool ExportPartTask::executeStep()
             block_with_partition_values,
             (*exports_list_entry)->destination_file_path,
             manifest.file_already_exists_policy == MergeTreePartExportManifest::FileAlreadyExistsPolicy::overwrite,
-            manifest.format_settings,
+            getFormatSettings(local_context),
             local_context);
     }
     catch (const Exception & e)
@@ -126,10 +137,21 @@ bool ExportPartTask::executeStep()
             }
         }
 
-        tryLogCurrentException(__PRETTY_FUNCTION__);
+        LOG_INFO(getLogger("ExportPartTask"), "Export part {} failed: {}", manifest.data_part->name, e.message());
 
         ProfileEvents::increment(ProfileEvents::PartsExportFailures);
 
+        storage.writePartLog(
+            PartLogElement::Type::EXPORT_PART,
+            ExecutionStatus::fromCurrentException("", true),
+            static_cast<UInt64>((*exports_list_entry)->elapsed * 1000000000),
+            manifest.data_part->name,
+            manifest.data_part,
+            {manifest.data_part},
+            nullptr,
+            nullptr,
+            exports_list_entry.get());
+
         std::lock_guard inner_lock(storage.export_manifests_mutex);
         storage.export_manifests.erase(manifest);
 
@@ -259,6 +281,7 @@ bool ExportPartTask::executeStep()
 
 void ExportPartTask::cancel() noexcept
 {
+    LOG_INFO(getLogger("ExportPartTask"), "Export part {} task cancel() method called", manifest.data_part->name);
     cancel_requested.store(true);
     pipeline.cancel();
 }

src/Storages/MergeTree/ExportPartTask.h

@@ -12,20 +12,19 @@ class ExportPartTask : public IExecutableTask
 public:
     explicit ExportPartTask(
         MergeTreeData & storage_,
-        const MergeTreePartExportManifest & manifest_,
-        ContextPtr context_);
+        const MergeTreePartExportManifest & manifest_);
     bool executeStep() override;
     void onCompleted() override;
     StorageID getStorageID() const override;
     Priority getPriority() const override;
     String getQueryId() const override;
+    const MergeTreePartExportManifest & getManifest() const;
 
     void cancel() noexcept override;
 
 private:
     MergeTreeData & storage;
     MergeTreePartExportManifest manifest;
-    ContextPtr local_context;
     QueryPipeline pipeline;
     std::atomic<bool> cancel_requested = false;