Make export part and partition respect background tasks memory limit

src/Common/ProfileEvents.cpp

@@ -332,6 +332,7 @@
     M(ExportPartitionZooKeeperRemoveRecursive, "Number of 'removeRecursive' requests to ZooKeeper made by the export partition feature.", ValueType::Number) \
     M(ExportPartitionZooKeeperMulti, "Number of 'multi' requests to ZooKeeper made by the export partition feature.", ValueType::Number) \
     M(ExportPartitionZooKeeperExists, "Number of 'exists' requests to ZooKeeper made by the export partition feature.", ValueType::Number) \
+    M(ExportPartsRejectedByMemoryLimit, "Number of background export part tasks rejected due to background memory limit.", ValueType::Number) \
     \
     M(DistributedConnectionTries, "Total count of distributed connection attempts.", ValueType::Number) \
     M(DistributedConnectionUsable, "Total count of successful distributed connections to a usable server (with required table, but maybe stale).", ValueType::Number) \

src/Storages/MergeTree/ExportPartitionTaskScheduler.cpp

@@ -3,8 +3,10 @@
 #include <Interpreters/Context.h>
 #include <Interpreters/DatabaseCatalog.h>
 #include <Common/Exception.h>
+#include <Common/MemoryTracker.h>
 #include <Common/ZooKeeper/Types.h>
 #include <Common/ProfileEvents.h>
+#include <Common/formatReadable.h>
 #include "Storages/MergeTree/ExportPartitionUtils.h"
 #include "Storages/MergeTree/MergeTreePartExportManifest.h"
 #include "Storages/MergeTree/ExportPartFromPartitionExportTask.h"
@@ -20,6 +22,7 @@ namespace ProfileEvents
     extern const Event ExportPartitionZooKeeperSet;
     extern const Event ExportPartitionZooKeeperRemove;
     extern const Event ExportPartitionZooKeeperMulti;
+    extern const Event ExportPartsRejectedByMemoryLimit;
 }
 
 
@@ -53,6 +56,20 @@ void ExportPartitionTaskScheduler::run()
         return;
     }
 
+    /// Respect the background memory soft-limit: refuse to schedule new export-part tasks when
+    /// background tasks are already pressing the limit. The task is rescheduled by the parent
+    /// background pool a few seconds later, so this just defers work without losing it.
+    if (!canEnqueueBackgroundTask())
+    {
+        ProfileEvents::increment(ProfileEvents::ExportPartsRejectedByMemoryLimit);
+        LOG_TRACE(storage.log,
+            "ExportPartition scheduler task: Reached memory limit for the background tasks ({}), "
+            "so won't select new parts to export. Current background tasks memory usage: {}.",
+            formatReadableSizeWithBinarySuffix(background_memory_tracker.getSoftLimit()),
+            formatReadableSizeWithBinarySuffix(background_memory_tracker.get()));
+        return;
+    }
+
     LOG_INFO(storage.log, "ExportPartition scheduler task: Available move executors: {}", available_move_executors);
 
     std::size_t scheduled_exports_count = 0;

src/Storages/MergeTree/MergeTreeData.cpp

@@ -16,6 +16,7 @@
 #include <Common/Config/ConfigHelper.h>
 #include <Common/CurrentMetrics.h>
 #include <Common/Increment.h>
+#include <Common/MemoryTracker.h>
 #include <Common/ProfileEventsScope.h>
 #include <Common/Stopwatch.h>
 #include <Common/StringUtils.h>
@@ -168,6 +169,7 @@ namespace ProfileEvents
     extern const Event PartsExportTotalMilliseconds;
     extern const Event PartsExportFailures;
     extern const Event PartsExportDuplicated;
+    extern const Event ExportPartsRejectedByMemoryLimit;
 }
 
 namespace CurrentMetrics
@@ -6638,6 +6640,17 @@ void MergeTreeData::exportPartToTable(
     }
 
     {
+        if (!canEnqueueBackgroundTask())
+        {
+            ProfileEvents::increment(ProfileEvents::ExportPartsRejectedByMemoryLimit);
+            throw Exception(ErrorCodes::ABORTED,
+                "Failed to schedule export part task for data part '{}'. "
+                "Reached memory limit for the background tasks ({}). Current background tasks memory usage: {}.",
+                part_name,
+                formatReadableSizeWithBinarySuffix(background_memory_tracker.getSoftLimit()),
+                formatReadableSizeWithBinarySuffix(background_memory_tracker.get()));
+        }
+
         MergeTreePartExportManifest manifest(
             dest_storage,
             part,