Merge pull request #729 from Altinity/backports/24.3/77766_stop_merges_without_written_blocks

24.3 Backport of ClickHouse#77766 and ClickHouse#59138 - Stop merges without written blocks / Fix REPLACE PARTITION waiting for mutations/merges on unrelated partitions

Author: Enmk

src/Common/ActionBlocker.h

@@ -17,7 +17,7 @@ class ActionBlocker
 
     bool isCancelled() const { return *counter > 0; }
 
-    /// Temporarily blocks corresponding actions (while the returned object is alive)
+    /// Temporary blocks corresponding actions (while the returned object is alive)
     friend class ActionLock;
     ActionLock cancel() { return ActionLock(*this); }
 

src/Common/tests/gtest_action_blocker.cpp

@@ -0,0 +1,128 @@
+#include <gtest/gtest.h>
+
+#include <chrono>
+#include <mutex>
+#include <stdexcept>
+#include <string_view>
+#include <vector>
+#include <thread>
+
+#include <Common/ActionBlocker.h>
+#include <Common/ActionLock.h>
+
+using namespace DB;
+
+TEST(ActionBlocker, TestDefaultConstructor)
+{
+    ActionBlocker blocker;
+
+    EXPECT_FALSE(blocker.isCancelled());
+    EXPECT_EQ(0, blocker.getCounter().load());
+}
+
+TEST(ActionBlocker, TestCancelForever)
+{
+    ActionBlocker blocker;
+
+    blocker.cancelForever();
+    EXPECT_TRUE(blocker.isCancelled());
+    EXPECT_EQ(1, blocker.getCounter().load());
+}
+
+TEST(ActionBlocker, TestCancel)
+{
+    ActionBlocker blocker;
+
+    {
+        auto lock = blocker.cancel();
+        EXPECT_TRUE(blocker.isCancelled());
+        EXPECT_EQ(1, blocker.getCounter().load());
+    }
+    // automatically un-cancelled on `lock` destruction
+    EXPECT_FALSE(blocker.isCancelled());
+}
+
+
+
+TEST(ActionLock, TestDefaultConstructor)
+{
+    ActionLock locker;
+    EXPECT_TRUE(locker.expired());
+}
+
+TEST(ActionLock, TestConstructorWithActionBlocker)
+{
+    ActionBlocker blocker;
+    ActionLock lock(blocker);
+
+    EXPECT_FALSE(lock.expired());
+    EXPECT_TRUE(blocker.isCancelled());
+    EXPECT_EQ(1, blocker.getCounter().load());
+}
+
+TEST(ActionLock, TestMoveAssignmentToEmpty)
+{
+    ActionBlocker blocker;
+
+    {
+        ActionLock lock;
+        lock = blocker.cancel();
+        EXPECT_TRUE(blocker.isCancelled());
+    }
+    // automatically un-cancelled on `lock` destruction
+    EXPECT_FALSE(blocker.isCancelled());
+    EXPECT_EQ(0, blocker.getCounter().load());
+}
+
+TEST(ActionLock, TestMoveAssignmentToNonEmpty)
+{
+    ActionBlocker blocker;
+    {
+        auto lock = blocker.cancel();
+        EXPECT_TRUE(blocker.isCancelled());
+
+        // cause a move
+        lock = blocker.cancel();
+
+        // blocker should remain locked
+        EXPECT_TRUE(blocker.isCancelled());
+    }
+    // automatically un-cancelled on `lock` destruction
+    EXPECT_FALSE(blocker.isCancelled());
+    EXPECT_EQ(0, blocker.getCounter().load());
+}
+
+TEST(ActionLock, TestMoveAssignmentToNonEmpty2)
+{
+    ActionBlocker blocker1;
+    ActionBlocker blocker2;
+    {
+        auto lock = blocker1.cancel();
+        EXPECT_TRUE(blocker1.isCancelled());
+
+        // cause a move
+        lock = blocker2.cancel();
+
+        // blocker2 be remain locked, blocker1 - unlocked
+        EXPECT_TRUE(blocker2.isCancelled());
+
+        EXPECT_FALSE(blocker1.isCancelled());
+    }
+    // automatically un-cancelled on `lock` destruction
+    EXPECT_FALSE(blocker1.isCancelled());
+    EXPECT_FALSE(blocker2.isCancelled());
+    EXPECT_EQ(0, blocker1.getCounter().load());
+    EXPECT_EQ(0, blocker2.getCounter().load());
+}
+
+TEST(ActionLock, TestExpiration)
+{
+    ActionLock lock;
+    {
+        ActionBlocker blocker;
+        lock = blocker.cancel();
+        EXPECT_FALSE(lock.expired());
+    }
+
+    EXPECT_TRUE(lock.expired());
+}

src/Storages/MergeTree/DataPartsExchange.h

@@ -8,6 +8,7 @@
 #include <IO/copyData.h>
 #include <IO/ConnectionTimeouts.h>
 #include <Common/Throttler.h>
+#include <Common/ActionBlocker.h>
 
 
 namespace zkutil

src/Storages/MergeTree/MergeProgress.h

@@ -1,9 +1,9 @@
 #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>
 
 
 namespace ProfileEvents
@@ -47,27 +47,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 +73,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 +88,19 @@ 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;
+    }
+
 };
 
 }

src/Storages/MergeTree/MergeTask.cpp

@@ -128,6 +128,19 @@ static void addMissedColumnsToSerializationInfos(
     }
 }
 
+bool MergeTask::GlobalRuntimeContext::isCancelled() const
+{
+    return (future_part ? merges_blocker->isCancelledForPartition(future_part->part_info.partition_id) : merges_blocker->isCancelled())
+        || merge_list_element_ptr->is_cancelled.load(std::memory_order_relaxed);
+}
+
+void MergeTask::GlobalRuntimeContext::checkOperationIsNotCanceled() const
+{
+    if (isCancelled())
+    {
+        throw Exception(ErrorCodes::ABORTED, "Cancelled merging parts");
+    }
+}
 
 bool MergeTask::ExecuteAndFinalizeHorizontalPart::prepare()
 {
@@ -140,8 +153,7 @@ bool MergeTask::ExecuteAndFinalizeHorizontalPart::prepare()
     }
     const String local_tmp_suffix = global_ctx->parent_part ? ctx->suffix : "";
 
-    if (global_ctx->merges_blocker->isCancelled() || global_ctx->merge_list_element_ptr->is_cancelled.load(std::memory_order_relaxed))
-        throw Exception(ErrorCodes::ABORTED, "Cancelled merging parts");
+    global_ctx->checkOperationIsNotCanceled();
 
     /// We don't want to perform merge assigned with TTL as normal merge, so
     /// throw exception
@@ -391,9 +403,10 @@ bool MergeTask::ExecuteAndFinalizeHorizontalPart::prepare()
     ctx->is_cancelled = [merges_blocker = global_ctx->merges_blocker,
         ttl_merges_blocker = global_ctx->ttl_merges_blocker,
         need_remove = ctx->need_remove_expired_values,
-        merge_list_element = global_ctx->merge_list_element_ptr]() -> bool
+        merge_list_element = global_ctx->merge_list_element_ptr,
+        partition_id = global_ctx->future_part->part_info.partition_id]() -> bool
     {
-        return merges_blocker->isCancelled()
+        return merges_blocker->isCancelledForPartition(partition_id)
             || (need_remove && ttl_merges_blocker->isCancelled())
             || merge_list_element->is_cancelled.load(std::memory_order_relaxed);
     };
@@ -478,8 +491,7 @@ bool MergeTask::ExecuteAndFinalizeHorizontalPart::executeImpl()
     global_ctx->merging_executor.reset();
     global_ctx->merged_pipeline.reset();
 
-    if (global_ctx->merges_blocker->isCancelled() || global_ctx->merge_list_element_ptr->is_cancelled.load(std::memory_order_relaxed))
-        throw Exception(ErrorCodes::ABORTED, "Cancelled merging parts");
+    global_ctx->checkOperationIsNotCanceled();
 
     if (ctx->need_remove_expired_values && global_ctx->ttl_merges_blocker->isCancelled())
         throw Exception(ErrorCodes::ABORTED, "Cancelled merging parts with expired TTL");
@@ -607,7 +619,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();
@@ -634,8 +647,7 @@ void MergeTask::VerticalMergeStage::prepareVerticalMergeForOneColumn() const
 bool MergeTask::VerticalMergeStage::executeVerticalMergeForOneColumn() const
 {
     Block block;
-    if (!global_ctx->merges_blocker->isCancelled() && !global_ctx->merge_list_element_ptr->is_cancelled.load(std::memory_order_relaxed)
-        && ctx->executor->pull(block))
+    if (!global_ctx->isCancelled() && ctx->executor->pull(block))
     {
         ctx->column_elems_written += block.rows();
         ctx->column_to->write(block);
@@ -650,8 +662,7 @@ bool MergeTask::VerticalMergeStage::executeVerticalMergeForOneColumn() const
 void MergeTask::VerticalMergeStage::finalizeVerticalMergeForOneColumn() const
 {
     const String & column_name = ctx->it_name_and_type->name;
-    if (global_ctx->merges_blocker->isCancelled() || global_ctx->merge_list_element_ptr->is_cancelled.load(std::memory_order_relaxed))
-        throw Exception(ErrorCodes::ABORTED, "Cancelled merging parts");
+    global_ctx->checkOperationIsNotCanceled();
 
     ctx->executor.reset();
     auto changed_checksums = ctx->column_to->fillChecksums(global_ctx->new_data_part, global_ctx->checksums_gathered_columns);
@@ -1120,7 +1131,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/MergeTask.h

@@ -24,6 +24,8 @@
 #include <Storages/MergeTree/MergedColumnOnlyOutputStream.h>
 #include <Storages/MergeTree/MergeProgress.h>
 #include <Storages/MergeTree/MergeTreeData.h>
+#include <Storages/MergeTree/MergeTreeIndices.h>
+#include <Storages/MergeTree/PartitionActionBlocker.h>
 
 
 namespace DB
@@ -75,7 +77,7 @@ class MergeTask
         MergeTreeTransactionPtr txn,
         MergeTreeData * data_,
         MergeTreeDataMergerMutator * mutator_,
-        ActionBlocker * merges_blocker_,
+        PartitionActionBlocker * merges_blocker_,
         ActionBlocker * ttl_merges_blocker_)
         {
             global_ctx = std::make_shared<GlobalRuntimeContext>();
@@ -148,7 +150,7 @@ class MergeTask
         MergeListElement * merge_list_element_ptr{nullptr};
         MergeTreeData * data{nullptr};
         MergeTreeDataMergerMutator * mutator{nullptr};
-        ActionBlocker * merges_blocker{nullptr};
+        PartitionActionBlocker * merges_blocker{nullptr};
         ActionBlocker * ttl_merges_blocker{nullptr};
         StorageSnapshotPtr storage_snapshot{nullptr};
         StorageMetadataPtr metadata_snapshot{nullptr};
@@ -195,6 +197,10 @@ class MergeTask
         bool need_prefix;
 
         scope_guard temporary_directory_lock;
+
+        // will throw an exception if merge was cancelled in any way.
+        void checkOperationIsNotCanceled() const;
+        bool isCancelled() const;
     };
 
     using GlobalRuntimeContextPtr = std::shared_ptr<GlobalRuntimeContext>;

src/Storages/MergeTree/MergeTreeDataMergerMutator.h

@@ -209,7 +209,7 @@ public :
     /** Is used to cancel all merges and mutations. On cancel() call all currently running actions will throw exception soon.
       * All new attempts to start a merge or mutation will throw an exception until all 'LockHolder' objects will be destroyed.
       */
-    ActionBlocker merges_blocker;
+    PartitionActionBlocker merges_blocker;
     ActionBlocker ttl_merges_blocker;
 
 private: