Tracking memory used in merges [#METR-23888].

dbms/include/DB/Common/MemoryTracker.h

@@ -2,6 +2,13 @@
 
 #include <atomic>
 #include <common/Common.h>
+#include <DB/Common/CurrentMetrics.h>
+
+
+namespace CurrentMetrics
+{
+	extern const Metric MemoryTracking;
+}
 
 
 /** Отслеживает потребление памяти.
@@ -20,6 +27,9 @@ class MemoryTracker
 	/// Односвязный список. Вся информация будет передаваться в следующие MemoryTracker-ы тоже. Они должны жить во время жизни данного MemoryTracker.
 	MemoryTracker * next = nullptr;
 
+	/// You could specify custom metric to track memory usage.
+	CurrentMetrics::Metric metric = CurrentMetrics::MemoryTracking;
+
 	/// Если задано (например, "for user") - в сообщениях в логе будет указываться это описание.
 	const char * description = nullptr;
 
@@ -67,6 +77,11 @@ public:
 		next = elem;
 	}
 
+	void setMetric(CurrentMetrics::Metric metric_)
+	{
+		metric = metric_;
+	}
+
 	void setDescription(const char * description_)
 	{
 		description = description_;

dbms/src/Common/CurrentMetrics.cpp

@@ -24,6 +24,7 @@
 	M(ReadonlyReplica) \
 	M(LeaderReplica) \
 	M(MemoryTracking) \
+	M(MemoryTrackingInBackgroundProcessingPool) \
 	M(LeaderElection) \
 	M(EphemeralNode) \
 	M(ZooKeeperWatch) \

dbms/src/Common/MemoryTracker.cpp

@@ -1,7 +1,6 @@
 #include <common/likely.h>
 #include <common/logger_useful.h>
 #include <DB/Common/Exception.h>
-#include <DB/Common/CurrentMetrics.h>
 #include <DB/Common/formatReadable.h>
 #include <DB/IO/WriteHelpers.h>
 #include <iomanip>
@@ -9,11 +8,6 @@
 #include <DB/Common/MemoryTracker.h>
 
 
-namespace CurrentMetrics
-{
-	extern const Metric MemoryTracking;
-}
-
 namespace DB
 {
 	namespace ErrorCodes
@@ -29,7 +23,7 @@ MemoryTracker::~MemoryTracker()
 		logPeakMemoryUsage();
 
 	if (amount && !next)
-		CurrentMetrics::sub(CurrentMetrics::MemoryTracking, amount);
+		CurrentMetrics::sub(metric, amount);
 }
 
 
@@ -46,7 +40,7 @@ void MemoryTracker::alloc(Int64 size)
 	Int64 will_be = amount += size;
 
 	if (!next)
-		CurrentMetrics::add(CurrentMetrics::MemoryTracking, size);
+		CurrentMetrics::add(metric, size);
 
 	/// Using non-thread-safe random number generator. Joint distribution in different threads would not be uniform.
 	/// In this case, it doesn't matter.
@@ -95,14 +89,14 @@ void MemoryTracker::free(Int64 size)
 	if (next)
 		next->free(size);
 	else
-		CurrentMetrics::sub(CurrentMetrics::MemoryTracking, size);
+		CurrentMetrics::sub(metric, size);
 }
 
 
 void MemoryTracker::reset()
 {
 	if (!next)
-		CurrentMetrics::sub(CurrentMetrics::MemoryTracking, amount);
+		CurrentMetrics::sub(metric, amount);
 
 	amount = 0;
 	peak = 0;

dbms/src/Storages/MergeTree/BackgroundProcessingPool.cpp

@@ -1,6 +1,7 @@
 #include <DB/Common/Exception.h>
 #include <DB/Common/setThreadName.h>
 #include <DB/Common/CurrentMetrics.h>
+#include <DB/Common/MemoryTracker.h>
 #include <DB/IO/WriteHelpers.h>
 #include <common/logger_useful.h>
 #include <DB/Storages/MergeTree/BackgroundProcessingPool.h>
@@ -11,6 +12,7 @@
 namespace CurrentMetrics
 {
 	extern const Metric BackgroundPoolTask;
+	extern const Metric MemoryTrackingInBackgroundProcessingPool;
 }
 
 namespace DB
@@ -109,6 +111,10 @@ void BackgroundProcessingPool::threadFunction()
 {
 	setThreadName("BackgrProcPool");
 
+	MemoryTracker memory_tracker;
+	memory_tracker.setMetric(CurrentMetrics::MemoryTrackingInBackgroundProcessingPool);
+	current_memory_tracker = &memory_tracker;
+
 	std::mt19937 rng(reinterpret_cast<intptr_t>(&rng));
 	std::this_thread::sleep_for(std::chrono::duration<double>(std::uniform_real_distribution<double>(0, sleep_seconds_random_part)(rng)));
 
@@ -191,6 +197,8 @@ void BackgroundProcessingPool::threadFunction()
 			task->iterator = tasks.emplace(next_time_to_execute, task);
 		}
 	}
+
+	current_memory_tracker = nullptr;
 }
 
 }