refactor kill so that it sets status instead of removing entry, adds some logs

Author: arthurpassos

src/Storages/ExportReplicatedMergeTreePartitionTaskEntry.h

@@ -19,7 +19,8 @@ struct ExportReplicatedMergeTreePartitionTaskEntry
     {
         PENDING,
         COMPLETED,
-        FAILED
+        FAILED,
+        KILLED
     };
 
     /// Allows us to skip completed / failed entries during scheduling
@@ -37,6 +38,11 @@ struct ExportReplicatedMergeTreePartitionTaskEntry
         return manifest.partition_id + "_" + qualified_table_name.getFullName();
     }
 
+    std::string getTransactionId() const
+    {
+        return manifest.transaction_id;
+    }
+
     /// Get create_time for sorted iteration
     time_t getCreateTime() const
     {
@@ -46,6 +52,7 @@ struct ExportReplicatedMergeTreePartitionTaskEntry
 
 struct ExportPartitionTaskEntryTagByCompositeKey {};
 struct ExportPartitionTaskEntryTagByCreateTime {};
+struct ExportPartitionTaskEntryTagByTransactionId {};
 
 // Multi-index container for export partition task entries
 // - Index 0 (TagByCompositeKey): hashed_unique on composite key for O(1) lookup
@@ -60,6 +67,10 @@ using ExportPartitionTaskEntriesContainer = boost::multi_index_container<
         boost::multi_index::ordered_non_unique<
             boost::multi_index::tag<ExportPartitionTaskEntryTagByCreateTime>,
             boost::multi_index::const_mem_fun<ExportReplicatedMergeTreePartitionTaskEntry, time_t, &ExportReplicatedMergeTreePartitionTaskEntry::getCreateTime>
+        >,
+        boost::multi_index::hashed_unique<
+            boost::multi_index::tag<ExportPartitionTaskEntryTagByTransactionId>,
+            boost::multi_index::const_mem_fun<ExportReplicatedMergeTreePartitionTaskEntry, std::string, &ExportReplicatedMergeTreePartitionTaskEntry::getTransactionId>
         >
     >
 >;

src/Storages/MergeTree/ExportPartTask.cpp

@@ -30,6 +30,7 @@ namespace ErrorCodes
     extern const int UNKNOWN_TABLE;
     extern const int FILE_ALREADY_EXISTS;
     extern const int LOGICAL_ERROR;
+    extern const int QUERY_WAS_CANCELLED;
 }
 
 namespace Setting
@@ -201,12 +202,18 @@ bool ExportPartTask::executeStep()
 
         auto is_cancelled_callback = [this]()
         {
-            return storage.parts_mover.moves_blocker.isCancelled();
+            return isCancelled();
         };
 
         exec.setCancelCallback(is_cancelled_callback, 100);
 
         exec.execute();
+
+        if (isCancelled())
+        {
+            throw Exception(ErrorCodes::QUERY_WAS_CANCELLED, "Export part was cancelled");
+        }
+
         std::lock_guard inner_lock(storage.export_manifests_mutex);
         storage.writePartLog(
             PartLogElement::Type::EXPORT_PART,
@@ -257,9 +264,15 @@ bool ExportPartTask::executeStep()
 
 void ExportPartTask::cancel() noexcept
 {
+    cancel_requested.store(true);
     pipeline.cancel();
 }
 
+bool ExportPartTask::isCancelled() const
+{
+    return cancel_requested.load() || storage.parts_mover.moves_blocker.isCancelled();
+}
+
 void ExportPartTask::onCompleted()
 {
 }

src/Storages/MergeTree/ExportPartTask.h

@@ -27,6 +27,9 @@ class ExportPartTask : public IExecutableTask
     MergeTreePartExportManifest manifest;
     ContextPtr local_context;
     QueryPipeline pipeline;
+    std::atomic<bool> cancel_requested = false;
+
+    bool isCancelled() const;
 };
 
 }

src/Storages/MergeTree/ExportPartitionManifestUpdatingTask.cpp

@@ -2,6 +2,7 @@
 #include <Storages/StorageReplicatedMergeTree.h>
 #include <Storages/ExportReplicatedMergeTreePartitionTaskEntry.h>
 #include "Storages/MergeTree/ExportPartitionUtils.h"
+#include "Common/logger_useful.h"
 #include <Common/ZooKeeper/Types.h>
 #include <Common/ZooKeeper/ZooKeeper.h>
 #include <Interpreters/DatabaseCatalog.h>
@@ -154,10 +155,12 @@ ExportPartitionManifestUpdatingTask::ExportPartitionManifestUpdatingTask(Storage
 {
 }
 
-void ExportPartitionManifestUpdatingTask::run()
+void ExportPartitionManifestUpdatingTask::poll()
 {
     std::lock_guard lock(storage.export_merge_tree_partition_mutex);
 
+    LOG_INFO(storage.log, "ExportPartition Manifest Updating Task: Polling for new entries for table {}. Current number of entries: {}", storage.getStorageID().getNameForLogs(), storage.export_merge_tree_partition_task_entries_by_key.size());
+
     auto zk = storage.getZooKeeper();
 
     const std::string exports_path = fs::path(storage.zookeeper_path) / "exports";
@@ -200,8 +203,13 @@ void ExportPartitionManifestUpdatingTask::run()
         if (!cleanup_lock.is_locked && has_local_entry_and_is_up_to_date)
             continue;
 
+        auto status_watch_callback = std::make_shared<Coordination::WatchCallback>([this, key](const Coordination::WatchResponse &) {
+            storage.export_merge_tree_partition_manifest_updater->addStatusChange(key);
+            storage.export_merge_tree_partition_status_handling_task->schedule();
+        });
+
         std::string status;
-        if (!zk->tryGet(fs::path(entry_path) / "status", status))
+        if (!zk->tryGetWatch(fs::path(entry_path) / "status", status, nullptr, status_watch_callback))
         {
             LOG_INFO(storage.log, "ExportPartition Manifest Updating Task: Skipping {}: missing status", key);
             continue;
@@ -245,6 +253,8 @@ void ExportPartitionManifestUpdatingTask::run()
     /// Remove entries that were deleted by someone else
     removeStaleEntries(zk_children, entries_by_key);
 
+    LOG_INFO(storage.log, "ExportPartition Manifest Updating task: finished polling for new entries. Number of entries: {}", entries_by_key.size());
+
     storage.export_merge_tree_partition_select_task->schedule();
 }
 
@@ -303,4 +313,62 @@ void ExportPartitionManifestUpdatingTask::removeStaleEntries(
     }
 }
 
+void ExportPartitionManifestUpdatingTask::addStatusChange(const std::string & key)
+{
+    std::lock_guard lock(status_changes_mutex);
+    status_changes.emplace(key);
+}
+
+void ExportPartitionManifestUpdatingTask::handleStatusChanges()
+{
+    std::lock_guard lock(status_changes_mutex);
+    std::lock_guard task_entries_lock(storage.export_merge_tree_partition_mutex);
+    auto zk = storage.getZooKeeper();
+
+    LOG_INFO(storage.log, "ExportPartition Manifest Updating task: handling status changes. Number of status changes: {}", status_changes.size());
+
+    while (!status_changes.empty())
+    {
+        LOG_INFO(storage.log, "ExportPartition Manifest Updating task: handling status change for task {}", status_changes.front());
+        const auto key = status_changes.front();
+        status_changes.pop();
+
+        auto it = storage.export_merge_tree_partition_task_entries_by_key.find(key);
+        if (it == storage.export_merge_tree_partition_task_entries_by_key.end())
+            continue;
+
+        /// get new status from zk
+        std::string new_status_string;
+        if (!zk->tryGet(fs::path(storage.zookeeper_path) / "exports" / key / "status", new_status_string))
+        {
+            LOG_INFO(storage.log, "ExportPartition Manifest Updating Task: Failed to get new status for task {}, skipping", key);
+            continue;
+        }
+
+        const auto new_status = magic_enum::enum_cast<ExportReplicatedMergeTreePartitionTaskEntry::Status>(new_status_string);
+        if (!new_status)
+        {
+            LOG_INFO(storage.log, "ExportPartition Manifest Updating Task: Invalid status {} for task {}, skipping", new_status_string, key);
+            continue;
+        }
+
+        LOG_INFO(storage.log, "ExportPartition Manifest Updating task: status changed for task {}. New status: {}", key, magic_enum::enum_name(*new_status).data());
+
+        /// If status changed to KILLED, cancel local export operations
+        if (*new_status == ExportReplicatedMergeTreePartitionTaskEntry::Status::KILLED)
+        {
+            try
+            {
+                storage.killExportPart(it->manifest.transaction_id);
+            }
+            catch (...)
+            {
+                tryLogCurrentException(storage.log, __PRETTY_FUNCTION__);
+            }
+        }
+
+        it->status = *new_status;
+    }
+}
+
 }

src/Storages/MergeTree/ExportPartitionManifestUpdatingTask.h

@@ -1,5 +1,7 @@
 #pragma once
 
+#include <mutex>
+#include <queue>
 #include <string>
 #include <unordered_set>
 namespace DB
@@ -13,7 +15,11 @@ class ExportPartitionManifestUpdatingTask
 public:
     ExportPartitionManifestUpdatingTask(StorageReplicatedMergeTree & storage);
 
-    void run();
+    void poll();
+
+    void handleStatusChanges();
+
+    void addStatusChange(const std::string & key);
 
 private:
     StorageReplicatedMergeTree & storage;
@@ -28,6 +34,9 @@ class ExportPartitionManifestUpdatingTask
         const std::unordered_set<std::string> & zk_children,
         auto & entries_by_key
     );
+
+    std::mutex status_changes_mutex;
+    std::queue<std::string> status_changes;
 };
 
 }

src/Storages/MergeTree/ReplicatedMergeTreeRestartingThread.cpp

@@ -176,6 +176,7 @@ bool ReplicatedMergeTreeRestartingThread::runImpl()
     storage.part_check_thread.start();
     storage.export_merge_tree_partition_updating_task->activateAndSchedule();
     storage.export_merge_tree_partition_select_task->activateAndSchedule();
+    storage.export_merge_tree_partition_status_handling_task->activateAndSchedule();
 
     LOG_DEBUG(log, "Table started successfully");
     return true;