Merge pull request #66548 from ClickHouse/correctly-handle-realloc

Correctly track memory for `Allocator::realloc`

src/Common/Allocator.cpp

@@ -187,12 +187,9 @@ void * Allocator<clear_memory_, populate>::realloc(void * buf, size_t old_size,
 #if USE_GWP_ASAN
     if (unlikely(GWPAsan::GuardedAlloc.shouldSample()))
     {
+        auto trace_alloc = CurrentMemoryTracker::alloc(new_size);
         if (void * ptr = GWPAsan::GuardedAlloc.allocate(new_size, alignment))
         {
-            auto trace_free = CurrentMemoryTracker::free(old_size);
-            auto trace_alloc = CurrentMemoryTracker::alloc(new_size);
-            trace_free.onFree(buf, old_size);
-
             memcpy(ptr, buf, std::min(old_size, new_size));
             free(buf, old_size);
             trace_alloc.onAlloc(buf, new_size);
@@ -209,6 +206,7 @@ void * Allocator<clear_memory_, populate>::realloc(void * buf, size_t old_size,
         }
         else
         {
+            [[maybe_unused]] auto trace_free = CurrentMemoryTracker::free(old_size);
             ProfileEvents::increment(ProfileEvents::GWPAsanAllocateFailed);
         }
     }
@@ -231,13 +229,17 @@ void * Allocator<clear_memory_, populate>::realloc(void * buf, size_t old_size,
     if (alignment <= MALLOC_MIN_ALIGNMENT)
     {
         /// Resize malloc'd memory region with no special alignment requirement.
-        auto trace_free = CurrentMemoryTracker::free(old_size);
+        /// Realloc can do 2 possible things:
+        /// - expand existing memory region
+        /// - allocate new memory block and free the old one
+        /// Because we don't know which option will be picked we need to make sure there is enough
+        /// memory for all options
         auto trace_alloc = CurrentMemoryTracker::alloc(new_size);
-        trace_free.onFree(buf, old_size);
 
         void * new_buf = ::realloc(buf, new_size);
         if (nullptr == new_buf)
         {
+            [[maybe_unused]] auto trace_free = CurrentMemoryTracker::free(old_size);
             throw DB::ErrnoException(
                 DB::ErrorCodes::CANNOT_ALLOCATE_MEMORY,
                 "Allocator: Cannot realloc from {} to {}",
@@ -246,6 +248,8 @@ void * Allocator<clear_memory_, populate>::realloc(void * buf, size_t old_size,
         }
 
         buf = new_buf;
+        auto trace_free = CurrentMemoryTracker::free(old_size);
+        trace_free.onFree(buf, old_size);
         trace_alloc.onAlloc(buf, new_size);
 
         if constexpr (clear_memory)

tests/integration/test_memory_limit/configs/async_metrics_no.xml

@@ -0,0 +1,21 @@
+<clickhouse>
+    <!-- this update period also syncs MemoryTracking with RSS, disable this, by using period = 1 day -->
+    <asynchronous_metrics_update_period_s>86400</asynchronous_metrics_update_period_s>
+    <query_masking_rules remove="remove"/>
+
+    <query_thread_log remove="remove"/>
+    <query_log remove="remove" />
+    <query_views_log remove="remove" />
+    <metric_log remove="remove"/>
+    <error_log remove="remove"/>
+    <text_log remove="remove"/>
+    <trace_log remove="remove"/>
+    <asynchronous_metric_log remove="remove" />
+    <session_log remove="remove" />
+    <part_log remove="remove" />
+    <crash_log remove="remove" />
+    <opentelemetry_span_log remove="remove" />
+    <!-- just in case it will be enabled by default -->
+    <zookeeper_log remove="remove" />
+    <transactions_info_log remove="remove" />
+</clickhouse>

tests/integration/test_memory_limit/test.py

@@ -0,0 +1,54 @@
+#!/usr/bin/env python3
+import logging
+import time
+import pytest
+from helpers.cluster import ClickHouseCluster
+from multiprocessing.dummy import Pool
+
+cluster = ClickHouseCluster(__file__)
+
+node = cluster.add_instance(
+    "node",
+    main_configs=[
+        "configs/async_metrics_no.xml",
+    ],
+    mem_limit="4g",
+    env_variables={"MALLOC_CONF": "dirty_decay_ms:0"},
+)
+
+
+@pytest.fixture(scope="module", autouse=True)
+def start_cluster():
+    try:
+        cluster.start()
+        yield cluster
+    finally:
+        cluster.shutdown()
+
+
+def test_multiple_queries():
+    if node.is_built_with_sanitizer():
+        return
+
+    p = Pool(15)
+
+    def run_query(node):
+        try:
+            node.query("SELECT * FROM system.numbers GROUP BY number")
+        except Exception as ex:
+            print("Exception", ex)
+            raise ex
+
+    tasks = []
+    for i in range(30):
+        tasks.append(p.apply_async(run_query, (node,)))
+        time.sleep(i * 0.1)
+
+    for task in tasks:
+        try:
+            task.get()
+        except Exception as ex:
+            print("Exception", ex)
+
+    # test that we didn't kill the server
+    node.query("SELECT 1")

tests/queries/0_stateless/02231_buffer_aggregate_states_leak.sql

@@ -29,7 +29,7 @@ from in_02231
 group by key;
 
 set optimize_trivial_insert_select = 1;
-insert into in_02231 select * from numbers(10e6) settings max_memory_usage='310Mi', max_threads=1;
+insert into in_02231 select * from numbers(10e6) settings max_memory_usage='400Mi', max_threads=1;
 
 drop table buffer_02231;
 drop table out_02231;