add setting to control locking behavior (inside/outside task), deny export request in case pool is full

Author: arthurpassos

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/Storages/MergeTree/ExportPartFromPartitionExportTask.cpp

@@ -13,33 +13,31 @@ namespace DB
 ExportPartFromPartitionExportTask::ExportPartFromPartitionExportTask(
     StorageReplicatedMergeTree & storage_,
     const std::string & key_,
-    const MergeTreePartExportManifest & manifest_,
-    ContextPtr context_)
+    const MergeTreePartExportManifest & manifest_)
     : storage(storage_),
     key(key_),
-    manifest(manifest_),
-    local_context(context_)
+    manifest(manifest_)
 {
-    export_part_task = std::make_shared<ExportPartTask>(storage, manifest, local_context);
+    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, "ExportPartition scheduler task: Attempting to lock part: {}", 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, "ExportPartition scheduler task: Locked part: {}", part_name);
+        LOG_INFO(storage.log, "ExportPartFromPartitionExportTask: Locked part: {}", part_name);
         export_part_task->executeStep();
         return false;
     }
 
-    LOG_INFO(storage.log, "ExportPartition scheduler task: Failed to lock part {}, skipping", part_name);
+    LOG_INFO(storage.log, "ExportPartFromPartitionExportTask: Failed to lock part {}, skipping", part_name);
     return false;
 }
 

src/Storages/MergeTree/ExportPartFromPartitionExportTask.h

@@ -14,8 +14,7 @@ class ExportPartFromPartitionExportTask : public IExecutableTask
     explicit ExportPartFromPartitionExportTask(
         StorageReplicatedMergeTree & storage_,
         const std::string & key_,
-        const MergeTreePartExportManifest & manifest_,
-        ContextPtr context_);
+        const MergeTreePartExportManifest & manifest_);
     bool executeStep() override;
     void onCompleted() override;
     StorageID getStorageID() const override;
@@ -28,7 +27,6 @@ class ExportPartFromPartitionExportTask : public IExecutableTask
     StorageReplicatedMergeTree & storage;
     std::string key;
     MergeTreePartExportManifest manifest;
-    ContextPtr local_context;
     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,10 +39,9 @@ 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_)
 {
 }
 
@@ -52,6 +52,12 @@ const MergeTreePartExportManifest & ExportPartTask::getManifest() const
 
 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();
@@ -96,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)
@@ -131,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);
 
@@ -264,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,8 +12,7 @@ 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;
@@ -26,7 +25,6 @@ class ExportPartTask : public IExecutableTask
 private:
     MergeTreeData & storage;
     MergeTreePartExportManifest manifest;
-    ContextPtr local_context;
     QueryPipeline pipeline;
     std::atomic<bool> cancel_requested = false;
 

src/Storages/MergeTree/ExportPartitionManifestUpdatingTask.cpp

@@ -100,6 +100,12 @@ std::vector<ReplicatedPartitionExportInfo> ExportPartitionManifestUpdatingTask::
 {
     std::vector<ReplicatedPartitionExportInfo> infos;
     const auto zk = storage.getZooKeeper();
+
+    if (!zk->isFeatureEnabled(DB::KeeperFeatureFlag::MULTI_READ))
+    {
+        throw Exception(ErrorCodes::NOT_IMPLEMENTED, "MULTI_READ feature flag is not enabled");
+    }
+
     const auto exports_path = fs::path(storage.zookeeper_path) / "exports";
 
     ProfileEvents::increment(ProfileEvents::ExportPartitionZooKeeperRequests);
@@ -657,6 +663,7 @@ void ExportPartitionManifestUpdatingTask::handleStatusChanges()
         {
             try
             {
+                LOG_INFO(storage.log, "ExportPartition Manifest Updating task: killing export partition for task {}", key);
                 storage.killExportPart(it->manifest.transaction_id);
             }
             catch (...)