Fix flaky tests caused by OPTIMIZE FINAL failing memory budget check

ClickHouse · May 16, 2023 · 66cb7f8 · 66cb7f8
1 parent 3d26232
commit 66cb7f8
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 26 deletions.
diff --git a/src/Storages/StorageMergeTree.cpp b/src/Storages/StorageMergeTree.cpp
@@ -925,44 +925,70 @@ MergeMutateSelectedEntryPtr StorageMergeTree::selectPartsToMerge(
 
     SelectPartsDecision select_decision = SelectPartsDecision::CANNOT_SELECT;
 
-    if (!canEnqueueBackgroundTask())
+    auto is_background_memory_usage_ok = [](String * disable_reason) -> bool
     {
-        if (out_disable_reason)
-            *out_disable_reason = fmt::format("Current background tasks memory usage ({}) is more than the limit ({})",
+        if (canEnqueueBackgroundTask())
+            return true;
+        if (disable_reason)
+            *disable_reason = fmt::format("Current background tasks memory usage ({}) is more than the limit ({})",
                 formatReadableSizeWithBinarySuffix(background_memory_tracker.get()),
                 formatReadableSizeWithBinarySuffix(background_memory_tracker.getSoftLimit()));
-    }
-    else if (partition_id.empty())
-    {
-        UInt64 max_source_parts_size = merger_mutator.getMaxSourcePartsSizeForMerge();
-        bool merge_with_ttl_allowed = getTotalMergesWithTTLInMergeList() < data_settings->max_number_of_merges_with_ttl_in_pool;
+        return false;
+    };
 
-        /// TTL requirements is much more strict than for regular merge, so
-        /// if regular not possible, than merge with ttl is not also not
-        /// possible.
-        if (max_source_parts_size > 0)
+    if (partition_id.empty())
+    {
+        if (is_background_memory_usage_ok(out_disable_reason))
         {
-            select_decision = merger_mutator.selectPartsToMerge(
-                future_part,
-                aggressive,
-                max_source_parts_size,
-                can_merge,
-                merge_with_ttl_allowed,
-                txn,
-                out_disable_reason);
+            UInt64 max_source_parts_size = merger_mutator.getMaxSourcePartsSizeForMerge();
+            bool merge_with_ttl_allowed = getTotalMergesWithTTLInMergeList() < data_settings->max_number_of_merges_with_ttl_in_pool;
+
+            /// TTL requirements is much more strict than for regular merge, so
+            /// if regular not possible, than merge with ttl is not also not
+            /// possible.
+            if (max_source_parts_size > 0)
+            {
+                select_decision = merger_mutator.selectPartsToMerge(
+                    future_part,
+                    aggressive,
+                    max_source_parts_size,
+                    can_merge,
+                    merge_with_ttl_allowed,
+                    txn,
+                    out_disable_reason);
+            }
+            else if (out_disable_reason)
+                *out_disable_reason = "Current value of max_source_parts_size is zero";
         }
-        else if (out_disable_reason)
-            *out_disable_reason = "Current value of max_source_parts_size is zero";
     }
     else
     {
         while (true)
         {
-            select_decision = merger_mutator.selectAllPartsToMergeWithinPartition(
-                future_part, can_merge, partition_id, final, metadata_snapshot, txn, out_disable_reason, optimize_skip_merged_partitions);
             auto timeout_ms = getSettings()->lock_acquire_timeout_for_background_operations.totalMilliseconds();
             auto timeout = std::chrono::milliseconds(timeout_ms);
 
+            if (!is_background_memory_usage_ok(out_disable_reason))
+            {
+                constexpr auto poll_interval = std::chrono::seconds(1);
+                Int64 attempts = timeout / poll_interval;
+                bool ok = false;
+                for (Int64 i = 0; i < attempts; ++i)
+                {
+                    std::this_thread::sleep_for(poll_interval);
+                    if (is_background_memory_usage_ok(out_disable_reason))
+                    {
+                        ok = true;
+                        break;
+                    }
+                }
+                if (!ok)
+                    break;
+            }
+
+            select_decision = merger_mutator.selectAllPartsToMergeWithinPartition(
+                future_part, can_merge, partition_id, final, metadata_snapshot, txn, out_disable_reason, optimize_skip_merged_partitions);
+
             /// If final - we will wait for currently processing merges to finish and continue.
             if (final
                 && select_decision != SelectPartsDecision::SELECTED

diff --git a/tests/queries/0_stateless/02581_share_big_sets_between_multiple_mutations_tasks_long.sql b/tests/queries/0_stateless/02581_share_big_sets_between_multiple_mutations_tasks_long.sql
@@ -1,4 +1,6 @@
--- Tags: long, no-debug, no-tsan, no-asan, no-ubsan, no-msan
+-- Tags: long, no-debug, no-tsan, no-asan, no-ubsan, no-msan, no-parallel
+
+-- no-parallel because the sets use a lot of memory, which may interfere with other tests
 
 DROP TABLE IF EXISTS 02581_trips;
 

diff --git a/tests/queries/0_stateless/02581_share_big_sets_between_mutation_tasks_long.sql b/tests/queries/0_stateless/02581_share_big_sets_between_mutation_tasks_long.sql
@@ -1,4 +1,6 @@
--- Tags: long, no-debug, no-tsan, no-asan, no-ubsan, no-msan
+-- Tags: long, no-debug, no-tsan, no-asan, no-ubsan, no-msan, no-parallel
+
+-- no-parallel because the sets use a lot of memory, which may interfere with other tests
 
 DROP TABLE IF EXISTS 02581_trips;