add dynamic untracked memory limits for more precise memory tracking

src/Common/CurrentMemoryTracker.cpp

@@ -57,6 +57,7 @@ AllocationTrace CurrentMemoryTracker::allocImpl(Int64 size, bool throw_if_memory
             {
                 auto res = memory_tracker->allocImpl(will_be, throw_if_memory_exceeded);
                 current_thread->untracked_memory = 0;
+                current_thread->updateUntrackedMemoryLimit(memory_tracker->get());
                 return res;
             }
             else
@@ -84,6 +85,13 @@ void CurrentMemoryTracker::check()
         std::ignore = memory_tracker->allocImpl(0, true);
 }
 
+Int64 CurrentMemoryTracker::get()
+{
+    if (auto * memory_tracker = getMemoryTracker())
+        return memory_tracker->get();
+    return 0;
+}
+
 AllocationTrace CurrentMemoryTracker::alloc(Int64 size)
 {
     bool throw_if_memory_exceeded = true;
@@ -107,6 +115,7 @@ AllocationTrace CurrentMemoryTracker::free(Int64 size)
             {
                 Int64 untracked_memory = current_thread->untracked_memory;
                 current_thread->untracked_memory = 0;
+                current_thread->updateUntrackedMemoryLimit(memory_tracker->get() + untracked_memory);
                 return memory_tracker->free(-untracked_memory);
             }
         }

src/Common/CurrentMemoryTracker.h

@@ -12,7 +12,9 @@ struct CurrentMemoryTracker
 
     /// This function should be called after memory deallocation.
     [[nodiscard]] static AllocationTrace free(Int64 size);
+
     static void check();
+    [[nodiscard]] static Int64 get();
 
     /// Throws MEMORY_LIMIT_EXCEEDED (if it's allowed to throw exceptions)
     static void injectFault();

src/Common/ThreadStatus.h

@@ -183,6 +183,12 @@ public:
     Int64 untracked_memory = 0;
     /// Each thread could new/delete memory in range of (-untracked_memory_limit, untracked_memory_limit) without access to common counters.
     Int64 untracked_memory_limit = 4 * 1024 * 1024;
+    /// To keep total untracked memory limited to `untracked_memory_ratio * RSS` we have to account threads with small and large memory footprint differently.
+    /// For this purpose we dynamically change `untracked_memory_limit` after every tracking event using a simple formula:
+    /// untracked_memory_limit = clamp(untracked_memory_ratio * cur_memory_bytes, min_untracked_memory, max_untracked_memory)
+    /// Note that this values are updated when thread is attached to a group
+    Int64 min_untracked_memory = 4 * 1024 * 1024;
+    Int64 max_untracked_memory = 4 * 1024;
 
     /// Statistics of read and write rows/bytes
     Progress progress_in;
@@ -309,6 +315,12 @@ public:
 
     void initGlobalProfiler(UInt64 global_profiler_real_time_period, UInt64 global_profiler_cpu_time_period);
 
+    void updateUntrackedMemoryLimit(Int64 current)
+    {
+        constexpr Int64 untracked_memory_ratio_bits = 4; // untracked_memory_ratio = 1.0 / (1 << untracked_memory_ratio_bits) = 1.0 / 16 = 6.25%
+        untracked_memory_limit = std::clamp<Int64>(current >> untracked_memory_ratio_bits, min_untracked_memory, max_untracked_memory);
+    }
+
 private:
     void applyGlobalSettings();
     void applyQuerySettings();

src/Core/Settings.h

@@ -489,6 +489,7 @@ class IColumn;
     M(UInt64, max_memory_usage_for_user, 0, "Maximum memory usage for processing all concurrently running queries for the user. Zero means unlimited.", 0) \
     M(UInt64, memory_overcommit_ratio_denominator_for_user, 1_GiB, "It represents soft memory limit on the global level. This value is used to compute query overcommit ratio.", 0) \
     M(UInt64, max_untracked_memory, (4 * 1024 * 1024), "Small allocations and deallocations are grouped in thread local variable and tracked or profiled only when amount (in absolute value) becomes larger than specified value. If the value is higher than 'memory_profiler_step' it will be effectively lowered to 'memory_profiler_step'.", 0) \
+    M(UInt64, min_untracked_memory, (4 * 1024), "Lower bound for untracked memory limit which is applied to threads with low memory consumption. Untracked memory limit equals thread_memory_usage/16 and clamped between min_untracked_memory and max_untracked_memory for every thread.", 0) \
     M(UInt64, memory_profiler_step, (4 * 1024 * 1024), "Whenever query memory usage becomes larger than every next step in number of bytes the memory profiler will collect the allocating stack trace. Zero means disabled memory profiler. Values lower than a few megabytes will slow down query processing.", 0) \
     M(Float, memory_profiler_sample_probability, 0., "Collect random allocations and deallocations and write them into system.trace_log with 'MemorySample' trace_type. The probability is for every alloc/free regardless to the size of the allocation (can be changed with `memory_profiler_sample_min_allocation_size` and `memory_profiler_sample_max_allocation_size`). Note that sampling happens only when the amount of untracked memory exceeds 'max_untracked_memory'. You may want to set 'max_untracked_memory' to 0 for extra fine grained sampling.", 0) \
     M(UInt64, memory_profiler_sample_min_allocation_size, 0, "Collect random allocations of size greater or equal than specified value with probability equal to `memory_profiler_sample_probability`. 0 means disabled. You may want to set 'max_untracked_memory' to 0 to make this threshold to work as expected.", 0) \

src/Interpreters/ThreadStatusExt.cpp

@@ -11,6 +11,7 @@
 #include <Parsers/formatAST.h>
 #include <Parsers/queryNormalization.h>
 #include <Common/CurrentThread.h>
+#include <Common/CurrentMemoryTracker.h>
 #include <Common/Exception.h>
 #include <Common/ProfileEvents.h>
 #include <Common/QueryProfiler.h>
@@ -210,9 +211,12 @@ void ThreadStatus::applyQuerySettings()
     query_id_from_query_context = query_context_ptr->getCurrentQueryId();
     initQueryProfiler();
 
-    untracked_memory_limit = settings.max_untracked_memory;
-    if (settings.memory_profiler_step && settings.memory_profiler_step < static_cast<UInt64>(untracked_memory_limit))
-        untracked_memory_limit = settings.memory_profiler_step;
+    max_untracked_memory = settings.max_untracked_memory;
+    if (settings.memory_profiler_step && settings.memory_profiler_step < static_cast<UInt64>(max_untracked_memory))
+        max_untracked_memory = settings.memory_profiler_step;
+    min_untracked_memory = std::min<Int64>(settings.min_untracked_memory, max_untracked_memory);
+
+    updateUntrackedMemoryLimit(CurrentMemoryTracker::get());
 
 #if defined(OS_LINUX)
     /// Set "nice" value if required.