Check cancellation in progress callback

Author: antaljanosbenjamin

src/Storages/MergeTree/MergeProgress.h

@@ -1,9 +1,10 @@
 #pragma once
 
-#include <base/types.h>
-#include <Common/ProfileEvents.h>
+#include <functional>
 #include <IO/Progress.h>
 #include <Storages/MergeTree/MergeList.h>
+#include <base/types.h>
+#include <Common/ProfileEvents.h>
 
 
 namespace ProfileEvents
@@ -47,27 +48,23 @@ struct MergeStageProgress
 class MergeProgressCallback
 {
 public:
+    // It should throw an exception in case the operation should be cancelled
+    using CancellationChecker = std::function<void()>;
+
     MergeProgressCallback(
-        MergeListElement * merge_list_element_ptr_, UInt64 & watch_prev_elapsed_, MergeStageProgress & stage_)
+        MergeListElement * merge_list_element_ptr_,
+        UInt64 & watch_prev_elapsed_,
+        MergeStageProgress & stage_,
+        CancellationChecker && cancellation_checker_)
         : merge_list_element_ptr(merge_list_element_ptr_)
         , watch_prev_elapsed(watch_prev_elapsed_)
         , stage(stage_)
+        , cancellation_checker(std::move(cancellation_checker_))
     {
         updateWatch();
     }
 
-    MergeListElement * merge_list_element_ptr;
-    UInt64 & watch_prev_elapsed;
-    MergeStageProgress & stage;
-
-    void updateWatch()
-    {
-        UInt64 watch_curr_elapsed = merge_list_element_ptr->watch.elapsed();
-        ProfileEvents::increment(ProfileEvents::MergesTimeMilliseconds, (watch_curr_elapsed - watch_prev_elapsed) / 1000000);
-        watch_prev_elapsed = watch_curr_elapsed;
-    }
-
-    void operator() (const Progress & value)
+    void operator()(const Progress & value)
     {
         ProfileEvents::increment(ProfileEvents::MergedUncompressedBytes, value.read_bytes);
         if (stage.is_first)
@@ -77,6 +74,8 @@ class MergeProgressCallback
         }
         updateWatch();
 
+        cancellation_checker();
+
         merge_list_element_ptr->bytes_read_uncompressed += value.read_bytes;
         if (stage.is_first)
             merge_list_element_ptr->rows_read += value.read_rows;
@@ -90,6 +89,25 @@ class MergeProgressCallback
                 std::memory_order_relaxed);
         }
     }
+
+private:
+    MergeListElement * merge_list_element_ptr;
+    UInt64 & watch_prev_elapsed;
+    MergeStageProgress & stage;
+    CancellationChecker cancellation_checker;
+
+    void updateWatch()
+    {
+        UInt64 watch_curr_elapsed = merge_list_element_ptr->watch.elapsed();
+        watch_prev_elapsed = watch_curr_elapsed;
+    }
+
+    void updateProfileEvents(const Progress & value, ProfileEvents::Event rows_event, ProfileEvents::Event bytes_event) const
+    {
+        ProfileEvents::increment(bytes_event, value.read_bytes);
+        if (stage.is_first)
+            ProfileEvents::increment(rows_event, value.read_rows);
+    }
 };
 
 }

src/Storages/MergeTree/MergeTask.cpp

@@ -607,7 +607,8 @@ void MergeTask::VerticalMergeStage::prepareVerticalMergeForOneColumn() const
     ctx->column_parts_pipeline.setProgressCallback(MergeProgressCallback(
         global_ctx->merge_list_element_ptr,
         global_ctx->watch_prev_elapsed,
-        *global_ctx->column_progress));
+        *global_ctx->column_progress,
+        [&my_ctx = *global_ctx]() { my_ctx.checkOperationIsNotCanceled(); }));
 
     /// Is calculated inside MergeProgressCallback.
     ctx->column_parts_pipeline.disableProfileEventUpdate();
@@ -1120,7 +1121,11 @@ void MergeTask::ExecuteAndFinalizeHorizontalPart::createMergedStream()
 
     global_ctx->merged_pipeline = QueryPipelineBuilder::getPipeline(std::move(*builder));
     /// Dereference unique_ptr and pass horizontal_stage_progress by reference
-    global_ctx->merged_pipeline.setProgressCallback(MergeProgressCallback(global_ctx->merge_list_element_ptr, global_ctx->watch_prev_elapsed, *global_ctx->horizontal_stage_progress));
+    global_ctx->merged_pipeline.setProgressCallback(MergeProgressCallback(
+        global_ctx->merge_list_element_ptr,
+        global_ctx->watch_prev_elapsed,
+        *global_ctx->horizontal_stage_progress,
+        [&my_ctx = *global_ctx]() { my_ctx.checkOperationIsNotCanceled(); }));
     /// Is calculated inside MergeProgressCallback.
     global_ctx->merged_pipeline.disableProfileEventUpdate();
 

src/Storages/MergeTree/MutateTask.cpp

@@ -2190,7 +2190,11 @@ bool MutateTask::prepare()
         ctx->materialized_projections = ctx->interpreter->grabMaterializedProjections();
         ctx->mutating_pipeline_builder = ctx->interpreter->execute();
         ctx->updated_header = ctx->interpreter->getUpdatedHeader();
-        ctx->progress_callback = MergeProgressCallback((*ctx->mutate_entry)->ptr(), ctx->watch_prev_elapsed, *ctx->stage_progress);
+        ctx->progress_callback = MergeProgressCallback(
+            (*ctx->mutate_entry)->ptr(),
+            ctx->watch_prev_elapsed,
+            *ctx->stage_progress,
+            [&my_ctx = *ctx]() { my_ctx.checkOperationIsNotCanceled(); });
     }
 
     auto single_disk_volume = std::make_shared<SingleDiskVolume>("volume_" + ctx->future_part->name, ctx->space_reservation->getDisk(), 0);