use the background scheduler instead of scheduling upon every request. very hackish for now, need to improve part selection and blocking

Author: arthurpassos

src/Storages/MergeTree/MergeMutateSelectedEntry.h

@@ -55,8 +55,8 @@ struct CurrentlyExportingPartsTagger
 {
     DataPartPtr part_to_export;
     StorageMergeTree & storage;
-    CurrentlyExportingPartsTagger(DataPartPtr part_to_export_, StorageMergeTree & storage_);
     ~CurrentlyExportingPartsTagger();
+    CurrentlyExportingPartsTagger(DataPartPtr part_to_export_, StorageMergeTree & storage_);
 };
 
 using CurrentlyExportingPartsTaggerPtr = std::shared_ptr<CurrentlyExportingPartsTagger>;

src/Storages/MergeTree/MergeTreeData.h

@@ -1153,7 +1153,7 @@ class MergeTreeData : public IStorage, public WithMutableContext
     /// Schedules background job to like merge/mutate/fetch an executor
     virtual bool scheduleDataProcessingJob(BackgroundJobsAssignee & assignee) = 0;
     /// Schedules job to move parts between disks/volumes and so on.
-    bool scheduleDataMovingJob(BackgroundJobsAssignee & assignee);
+    virtual bool scheduleDataMovingJob(BackgroundJobsAssignee & assignee);
     bool areBackgroundMovesNeeded() const;
 
 

src/Storages/MergeTree/MergeTreeExportManifest.h

@@ -51,6 +51,7 @@ struct MergeTreeExportManifest
     {
         String part_name;
         String remote_path; // empty until uploaded
+        bool in_progress = false; /// this is just a hackish workaround for now
     };
 
 
@@ -174,6 +175,15 @@ struct MergeTreeExportManifest
         write();
     }
 
+    void setInProgress(const String & part_name)
+    {
+        for (auto & i : items)
+        {
+            if (i.part_name == part_name)
+                i.in_progress = true;
+        }
+    }
+
     std::vector<String> pendingParts() const
     {
         std::vector<String> res;

src/Storages/ObjectStorage/MergeTree/ExportPartPlainMergeTreeTask.cpp

@@ -11,14 +11,14 @@ namespace DB
 
 ExportPartPlainMergeTreeTask::ExportPartPlainMergeTreeTask(
     StorageMergeTree & storage_,
-    const std::shared_ptr<CurrentlyExportingPartsTagger> & exports_tagger_,
+    const DataPartPtr & part_to_export_,
     const StoragePtr & destination_storage_,
     ContextPtr context_,
     std::shared_ptr<MergeTreeExportManifest> manifest_,
     IExecutableTask::TaskResultCallback & task_result_callback_,
     size_t max_retries_)
     : storage(storage_)
-    , exports_tagger(exports_tagger_)
+    , part_to_export(part_to_export_)
     , destination_storage(destination_storage_)
     , context(std::move(context_))
     , manifest(std::move(manifest_))
@@ -63,7 +63,7 @@ bool ExportPartPlainMergeTreeTask::executeStep()
                 retry_count++;
                 LOG_INFO(getLogger("ExportPartPlainMergeTreeTask"),
                 "Retrying export attempt {} for part {}",
-                retry_count, exports_tagger->part_to_export->name);
+                retry_count, part_to_export->name);
                 state = State::NEED_EXECUTE;
             }
             else
@@ -84,7 +84,7 @@ bool ExportPartPlainMergeTreeTask::executeStep()
                 retry_count++;
                 LOG_INFO(getLogger("ExportPartPlainMergeTreeTask"),
                 "Retrying export attempt {} for part {}",
-                retry_count, exports_tagger->part_to_export->name);
+                retry_count, part_to_export->name);
                 state = State::NEED_COMMIT;
             }
             else
@@ -101,16 +101,23 @@ bool ExportPartPlainMergeTreeTask::executeStep()
             manifest->status = MergeTreeExportManifest::Status::failed;
             manifest->write();
 
-            storage.already_exported_partition_ids.erase(manifest->partition_id);
+            /// this is a mess, what if several fail? I need to re-think the architecture
+            /// I'll leave this commented out for now
+            // storage.already_exported_partition_ids.erase(manifest->partition_id);
+
+            storage.currently_merging_mutating_parts.erase(part_to_export);
 
             return false;
         }
         case State::SUCCESS:
         {
+            storage.currently_merging_mutating_parts.erase(part_to_export);
+
             return false;
         }
     }
 
+    storage.currently_merging_mutating_parts.erase(part_to_export);
     return false;
 }
 
@@ -157,15 +164,15 @@ bool ExportPartPlainMergeTreeTask::executeExport()
     {
         destination_storage->importMergeTreePart(
             storage,
-            exports_tagger->part_to_export,
+            part_to_export,
             context,
             part_log_wrapper);
 
         return true;
     }
     catch (...)
     {
-        LOG_ERROR(getLogger("ExportPartPlainMergeTreeTask"), "Failed to export part {}", exports_tagger->part_to_export->name);
+        LOG_ERROR(getLogger("ExportPartPlainMergeTreeTask"), "Failed to export part {}", part_to_export->name);
 
         return false;
     }

src/Storages/ObjectStorage/MergeTree/ExportPartPlainMergeTreeTask.h

@@ -18,7 +18,7 @@ class ExportPartPlainMergeTreeTask : public IExecutableTask
 public:
     ExportPartPlainMergeTreeTask(
         StorageMergeTree & storage_,
-        const std::shared_ptr<CurrentlyExportingPartsTagger> & exports_tagger_,
+        const DataPartPtr & part_to_export_,
         const StoragePtr & destination_storage_,
         ContextPtr context_,
         std::shared_ptr<MergeTreeExportManifest> manifest_,
@@ -50,7 +50,7 @@ class ExportPartPlainMergeTreeTask : public IExecutableTask
     State state{State::NEED_PREPARE};
 
     StorageMergeTree & storage;
-    std::shared_ptr<CurrentlyExportingPartsTagger> exports_tagger;
+    DataPartPtr part_to_export;
     StoragePtr destination_storage;
     ContextPtr context;
     std::shared_ptr<MergeTreeExportManifest> manifest;

src/Storages/StorageMergeTree.cpp

@@ -548,8 +548,6 @@ void StorageMergeTree::exportPartitionToTable(const PartitionCommand & command,
         return;
     }
 
-    std::vector<std::shared_ptr<CurrentlyExportingPartsTagger>> taggers;
-
     {
         /// Do not put this in a scope because `CurrentlyExportingPartsTagger` instantiated above relies on this already being locked
         /// shitty design I came up with huh
@@ -559,12 +557,11 @@ void StorageMergeTree::exportPartitionToTable(const PartitionCommand & command,
         {
             throw Exception(ErrorCodes::PART_IS_LOCKED, "Partition {} has already been exported", partition_id);
         }
-        
-        taggers.reserve(all_parts.size());
 
         for (const auto & part : all_parts)
         {
-            taggers.push_back(std::make_shared<CurrentlyExportingPartsTagger>(part, *this));
+            if (!currently_merging_mutating_parts.emplace(part).second)
+                throw Exception(ErrorCodes::LOGICAL_ERROR, "Tagging already tagged part {}. This is a bug.", part->name);
         }
     }
 
@@ -584,14 +581,7 @@ void StorageMergeTree::exportPartitionToTable(const PartitionCommand & command,
         export_partition_transaction_id_to_manifest[transaction_id] = manifest;
     }
 
-    for (const auto & tagger : taggers)
-    {
-        auto task = std::make_shared<ExportPartPlainMergeTreeTask>(*this, tagger, dest_storage, getContext(), manifest, moves_assignee_trigger);
-        if (!background_moves_assignee.scheduleMoveTask(task))
-        {
-            LOG_ERROR(log, "Failed to schedule export task for part {}", tagger->part_to_export->name);
-        }
-    }
+    background_moves_assignee.trigger();
 }
 
 /// While exists, marks parts as 'currently_merging_mutating_parts' and reserves free space on filesystem.
@@ -1198,58 +1188,53 @@ void StorageMergeTree::readExportPartitionManifests()
             }
         }
     }
+
+    background_moves_assignee.trigger();
 }
 
 void StorageMergeTree::resumeExportPartitionTasks()
 {
     /// Initially I opted for having two separate methods: read and resume because I wanted to schedule the tasks in order
     /// but it turns out the background executor schedules tasks based on their priority, so it is likely this is not needed anymore.
-    for (const auto & [transaction_id, manifest] : export_partition_transaction_id_to_manifest)
-    {
-        if (manifest->status != MergeTreeExportManifest::Status::pending)
-            continue;
+    // for (const auto & [transaction_id, manifest] : export_partition_transaction_id_to_manifest)
+    // {
+    //     if (manifest->status != MergeTreeExportManifest::Status::pending)
+    //         continue;
 
-        auto destination_storage = DatabaseCatalog::instance().tryGetTable(manifest->destination_storage_id, getContext());
-        if (!destination_storage)
-        {
-            LOG_ERROR(log, "Failed to reconstruct destination storage: {}", manifest->destination_storage_id.getNameForLogs());
-            continue;
-        }
+    //     auto destination_storage = DatabaseCatalog::instance().tryGetTable(manifest->destination_storage_id, getContext());
+    //     if (!destination_storage)
+    //     {
+    //         LOG_ERROR(log, "Failed to reconstruct destination storage: {}", manifest->destination_storage_id.getNameForLogs());
+    //         continue;
+    //     }
 
-        auto pending_part_names = manifest->pendingParts();
+    //     auto pending_part_names = manifest->pendingParts();
 
-        /// apparently, it is possible that pending parts are empty
-        /// if it is empty, I have to somehow commit and mark as completed..
+    //     /// apparently, it is possible that pending parts are empty
+    //     /// if it is empty, I have to somehow commit and mark as completed..
 
-        std::vector<DataPartPtr> parts_to_export;
+    //     std::vector<DataPartPtr> parts_to_export;
 
-        for (const auto & part_name : pending_part_names)
-        {
-            auto part = getPartIfExists(part_name, {MergeTreeDataPartState::Active});
+    //     for (const auto & part_name : pending_part_names)
+    //     {
+    //         auto part = getPartIfExists(part_name, {MergeTreeDataPartState::Active});
 
-            if (!part)
-            {
-                LOG_ERROR(log, "Part {} is present in the manifest file {}, but not found in the storage {}",
-                    part_name,
-                    manifest->transaction_id,
-                    getStorageID().getNameForLogs());
-                manifest->status = MergeTreeExportManifest::Status::failed;
-                manifest->write();
-
-                already_exported_partition_ids.erase(manifest->partition_id);
-                continue;
-            }
+    //         if (!part)
+    //         {
+    //             LOG_ERROR(log, "Part {} is present in the manifest file {}, but not found in the storage {}",
+    //                 part_name,
+    //                 manifest->transaction_id,
+    //                 getStorageID().getNameForLogs());
+    //             manifest->status = MergeTreeExportManifest::Status::failed;
+    //             manifest->write();
 
-            parts_to_export.emplace_back(part);
-        }
+    //             already_exported_partition_ids.erase(manifest->partition_id);
+    //             continue;
+    //         }
 
-        for (const auto & part : parts_to_export)
-        {
-            auto tagger = std::make_shared<CurrentlyExportingPartsTagger>(part, *this);
-            auto task = std::make_shared<ExportPartPlainMergeTreeTask>(*this, tagger, destination_storage, getContext(), manifest, moves_assignee_trigger);
-            background_moves_assignee.scheduleMoveTask(task);
-        }
-    }
+    //         parts_to_export.emplace_back(part);
+    //     }
+    // }
 }
 
 void StorageMergeTree::loadExportPartition()
@@ -1662,6 +1647,57 @@ UInt32 StorageMergeTree::getMaxLevelInBetween(const PartProperties & left, const
     return level;
 }
 
+
+bool StorageMergeTree::scheduleDataMovingJob(BackgroundJobsAssignee & assignee)
+{
+    if (MergeTreeData::scheduleDataMovingJob(assignee))
+    {
+        return true;
+    }
+
+    /// Try to schedule one export part task if any pending export exists
+    {
+        std::lock_guard lock(export_partition_transaction_id_to_manifest_mutex);
+        for (const auto & [transaction_id, manifest] : export_partition_transaction_id_to_manifest)
+        {
+            if (manifest->status != MergeTreeExportManifest::Status::pending)
+                continue;
+
+            auto destination_storage = DatabaseCatalog::instance().tryGetTable(manifest->destination_storage_id, getContext());
+            if (!destination_storage)
+            {
+                LOG_ERROR(log, "Failed to reconstruct destination storage: {}", manifest->destination_storage_id.getNameForLogs());
+                continue;
+            }
+
+            for (auto & item : manifest->items)
+            {
+                if (item.in_progress)
+                    continue;
+
+                auto part = getPartIfExists(item.part_name, {MergeTreeDataPartState::Active});
+                if (!part)
+                {
+                    LOG_ERROR(log, "Part {} is present in the manifest file {}, but not found in the storage {}",
+                        item.part_name,
+                        manifest->transaction_id,
+                        getStorageID().getNameForLogs());
+                    manifest->status = MergeTreeExportManifest::Status::failed;
+                    manifest->write();
+                    already_exported_partition_ids.erase(manifest->partition_id);
+                    continue;
+                }
+
+                auto task = std::make_shared<ExportPartPlainMergeTreeTask>(*this, part, destination_storage, getContext(), manifest, moves_assignee_trigger);
+                item.in_progress = background_moves_assignee.scheduleMoveTask(task);
+                return true;
+
+            }
+        }
+    }
+    return false;
+}
+
 bool StorageMergeTree::scheduleDataProcessingJob(BackgroundJobsAssignee & assignee)
 {
     if (shutdown_called)
@@ -2949,8 +2985,8 @@ MutationCounters StorageMergeTree::getMutationCounters() const
 
 void StorageMergeTree::startBackgroundMovesIfNeeded()
 {
-    if (areBackgroundMovesNeeded())
-        background_moves_assignee.start();
+    // if (areBackgroundMovesNeeded())
+    background_moves_assignee.start();
 }
 
 std::unique_ptr<MergeTreeSettings> StorageMergeTree::getDefaultSettings() const

src/Storages/StorageMergeTree.h

@@ -119,6 +119,8 @@ class StorageMergeTree final : public MergeTreeData
 
     bool scheduleDataProcessingJob(BackgroundJobsAssignee & assignee) override;
 
+    bool scheduleDataMovingJob(BackgroundJobsAssignee & assignee) override;
+
     std::map<std::string, MutationCommands> getUnfinishedMutationCommands() const override;
 
     MergeTreeDeduplicationLog * getDeduplicationLog() { return deduplication_log.get(); }