lock part inside task

Author: arthurpassos

src/Storages/MergeTree/ExportPartFromPartitionExportTask.cpp

@@ -0,0 +1,67 @@
+#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_,
+    ContextPtr context_)
+    : storage(storage_),
+    key(key_),
+    manifest(manifest_),
+    local_context(context_)
+{
+    export_part_task = std::make_shared<ExportPartTask>(storage, manifest, local_context);
+}
+
+bool ExportPartFromPartitionExportTask::executeStep()
+{
+    const auto zk = storage.getZooKeeper();
+    const auto part_name = manifest.data_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))
+    {
+        export_part_task->executeStep();
+        return false;
+    }
+
+    LOG_INFO(storage.log, "ExportPartition scheduler task: 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,35 @@
+#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_,
+        ContextPtr context_);
+    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;
+    ContextPtr local_context;
+    std::shared_ptr<ExportPartTask> export_part_task;
+};
+
+}

src/Storages/MergeTree/ExportPartTask.cpp

@@ -101,6 +101,11 @@ ExportPartTask::ExportPartTask(MergeTreeData & storage_, const MergeTreePartExpo
 {
 }
 
+const MergeTreePartExportManifest & ExportPartTask::getManifest() const
+{
+    return manifest;
+}
+
 bool ExportPartTask::executeStep()
 {
     auto local_context = Context::createCopy(storage.getContext());

src/Storages/MergeTree/ExportPartTask.h

@@ -18,6 +18,7 @@ class ExportPartTask : public IExecutableTask
     StorageID getStorageID() const override;
     Priority getPriority() const override;
     String getQueryId() const override;
+    const MergeTreePartExportManifest & getManifest() const;
 
     void cancel() noexcept override;
 

src/Storages/MergeTree/ExportPartitionTaskScheduler.cpp

@@ -7,6 +7,9 @@
 #include <Common/ProfileEvents.h>
 #include "Storages/MergeTree/ExportPartitionUtils.h"
 #include "Storages/MergeTree/MergeTreePartExportManifest.h"
+#include "Storages/MergeTree/ExportPartFromPartitionExportTask.h"
+#include "Formats/FormatFactory.h"
+#include <Core/Settings.h>
 
 namespace ProfileEvents
 {
@@ -24,6 +27,11 @@ namespace ProfileEvents
 namespace DB
 {
 
+namespace Setting
+{
+    extern const SettingsMergeTreePartExportFileAlreadyExistsPolicy export_merge_tree_part_file_already_exists_policy;
+}
+
 namespace ErrorCodes
 {
     extern const int QUERY_WAS_CANCELLED;
@@ -156,41 +164,29 @@ void ExportPartitionTaskScheduler::run()
                 continue;
             }
 
-            ProfileEvents::increment(ProfileEvents::ExportPartitionZooKeeperRequests);
-            ProfileEvents::increment(ProfileEvents::ExportPartitionZooKeeperCreate);
-            if (Coordination::Error::ZOK != zk->tryCreate(fs::path(storage.zookeeper_path) / "exports" / key / "locks" / zk_part_name, storage.replica_name, zkutil::CreateMode::Ephemeral))
-            {
-                LOG_INFO(storage.log, "ExportPartition scheduler task: Failed to lock part {}, skipping", zk_part_name);
-                continue;
-            }
+            std::lock_guard part_export_lock(storage.export_manifests_mutex);
 
-            try
-            {
-                storage.exportPartToTable(
-                    part->name,
-                    destination_storage_id,
-                    manifest.transaction_id,
-                    getContextCopyWithTaskSettings(storage.getContext(), manifest),
-                    /*allow_outdated_parts*/ true,
-                    [this, key, zk_part_name, manifest, destination_storage]
-                    (MergeTreePartExportManifest::CompletionCallbackResult result)
-                    {
-                        handlePartExportCompletion(key, zk_part_name, manifest, destination_storage, result);
-                    });
-            }
-            catch (const Exception &)
-            {
-                tryLogCurrentException(__PRETTY_FUNCTION__);
-                zk->tryRemove(fs::path(storage.zookeeper_path) / "exports" / key / "locks" / zk_part_name);
-                ProfileEvents::increment(ProfileEvents::ExportPartitionZooKeeperRequests);
-                ProfileEvents::increment(ProfileEvents::ExportPartitionZooKeeperRemove);
-                /// we should not increment retry_count because the node might just be full
-            }
+            auto context = getContextCopyWithTaskSettings(storage.getContext(), manifest);
+
+            const auto format_settings = getFormatSettings(context);
+
+            MergeTreePartExportManifest part_export_manifest(
+                destination_storage->getStorageID(),
+                part,
+                manifest.transaction_id,
+                context->getSettingsRef()[Setting::export_merge_tree_part_file_already_exists_policy].value,
+                format_settings,
+                storage.getInMemoryMetadataPtr(),
+                [this, key, zk_part_name, manifest, destination_storage]
+                (MergeTreePartExportManifest::CompletionCallbackResult result)
+                {
+                    handlePartExportCompletion(key, zk_part_name, manifest, destination_storage, result);
+                });
+
+            storage.background_moves_assignee.scheduleMoveTask(
+                std::make_shared<ExportPartFromPartitionExportTask>(storage, key, part_export_manifest, getContextCopyWithTaskSettings(storage.getContext(), manifest)));
         }
     }
-
-    /// 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(

src/Storages/MergeTree/MergeTreeData.h

@@ -1404,6 +1404,7 @@ class MergeTreeData : public IStorage, public WithMutableContext
     friend class IMergedBlockOutputStream; // for access to log
     friend struct DataPartsLock; // for access to shared_parts_list/shared_ranges_in_parts
     friend class ExportPartTask;
+    friend class ExportPartFromPartitionExportTask;
 
     bool require_part_metadata;
 

src/Storages/StorageReplicatedMergeTree.h

@@ -406,6 +406,7 @@ class StorageReplicatedMergeTree final : public MergeTreeData
     friend class ReplicatedMergeMutateTaskBase;
     friend class ExportPartitionManifestUpdatingTask;
     friend class ExportPartitionTaskScheduler;
+    friend class ExportPartFromPartitionExportTask;
 
     using MergeStrategyPicker = ReplicatedMergeTreeMergeStrategyPicker;
     using LogEntry = ReplicatedMergeTreeLogEntry;