From b6cbd0281e699ae6237095f93fa7c7ea4d7bc1b3 Mon Sep 17 00:00:00 2001 From: yongman Date: Sun, 4 Jan 2026 19:15:07 +0800 Subject: [PATCH 1/8] metrics: refresh jemalloc stats in async metrics Signed-off-by: yongman (cherry picked from commit d54c26a4afa31b362c538394190f29299b94970c) --- dbms/src/Interpreters/AsynchronousMetrics.cpp | 3 +++ 1 file changed, 3 insertions(+) diff --git a/dbms/src/Interpreters/AsynchronousMetrics.cpp b/dbms/src/Interpreters/AsynchronousMetrics.cpp index 9d09febb740..bb683fb1fba 100644 --- a/dbms/src/Interpreters/AsynchronousMetrics.cpp +++ b/dbms/src/Interpreters/AsynchronousMetrics.cpp @@ -378,6 +378,9 @@ void AsynchronousMetrics::update() set("jemalloc." NAME, value); \ } while (0); + uint64_t epoch = 1; + size_t sz = sizeof(epoch); + je_mallctl("epoch", &epoch, &sz, &epoch, sz); FOR_EACH_METRIC(GET_JEMALLOC_METRIC); #undef GET_JEMALLOC_METRIC From 5034306a9ea1cac539e9ea34aa94cf129e78cd77 Mon Sep 17 00:00:00 2001 From: yongman Date: Tue, 6 Jan 2026 21:45:10 +0800 Subject: [PATCH 2/8] exclude mmap memory in MemoryTracker rss Signed-off-by: yongman (cherry picked from commit ec4bf8011fa46f8c9750594a01f7303f7bd96689) --- dbms/src/Common/BackgroundTask.cpp | 37 +++++++++++++++++++++++++----- dbms/src/Common/MemoryTracker.cpp | 26 +++++++++++++++------ dbms/src/Common/MemoryTracker.h | 2 +- dbms/src/Interpreters/Settings.h | 2 +- 4 files changed, 52 insertions(+), 15 deletions(-) diff --git a/dbms/src/Common/BackgroundTask.cpp b/dbms/src/Common/BackgroundTask.cpp index a3fc69bbeee..f5be5f94532 100644 --- a/dbms/src/Common/BackgroundTask.cpp +++ b/dbms/src/Common/BackgroundTask.cpp @@ -13,15 +13,38 @@ // limitations under the License. #include -#include #include +#include + #include namespace DB { namespace { +bool process_num_threads(Int64 & cur_proc_num_threads) +{ + // '/proc/self/stat' provides process thread count on Linux. + std::ifstream stat_stream("/proc/self/stat", std::ios_base::in); + if (!stat_stream.is_open()) + return false; + + std::string pid, comm, state, ppid, pgrp, session, tty_nr; + std::string tpgid, flags, minflt, cminflt, majflt, cmajflt; + std::string utime, stime, cutime, cstime, priority, nice; + std::string itrealvalue, starttime; + UInt64 cur_virt_size = 0; + Int64 rss = 0; + + stat_stream >> pid >> comm >> state >> ppid >> pgrp >> session >> tty_nr >> tpgid >> flags >> minflt >> cminflt + >> majflt >> cmajflt >> utime >> stime >> cutime >> cstime >> priority >> nice >> cur_proc_num_threads + >> itrealvalue >> starttime >> cur_virt_size >> rss; + + stat_stream.close(); + return true; +} + bool isProcStatSupported() { std::ifstream stat_stream("/proc/self/stat", std::ios_base::in); @@ -62,13 +85,15 @@ void CollectProcInfoBackgroundTask::memCheckJob() { try { + Int64 cur_proc_num_threads = 1; while (!end_syn) { - // Update the memory usage of the current process. Defined in Common/MemoryTracker.cpp - auto res = get_process_mem_usage(); - real_rss = res.resident_bytes; - proc_num_threads = res.cur_proc_num_threads; - proc_virt_size = res.cur_virt_bytes; + auto metrics = get_process_metrics(); + process_num_threads(cur_proc_num_threads); + real_rss = static_cast(metrics.rss); + real_rss_file = static_cast(metrics.rss_file); + proc_num_threads = cur_proc_num_threads; + proc_virt_size = metrics.vsize; baseline_of_query_mem_tracker = root_of_query_mem_trackers->get(); usleep(100000); // sleep 100ms } diff --git a/dbms/src/Common/MemoryTracker.cpp b/dbms/src/Common/MemoryTracker.cpp index 7e875babc57..17931e6e2e5 100644 --- a/dbms/src/Common/MemoryTracker.cpp +++ b/dbms/src/Common/MemoryTracker.cpp @@ -31,7 +31,7 @@ extern const Metric MemoryTrackingSharedColumnData; extern const Metric MemoryTrackingKVStore; } // namespace CurrentMetrics -std::atomic real_rss{0}, proc_num_threads{1}, baseline_of_query_mem_tracker{0}; +std::atomic real_rss{0}, real_rss_file{0}, proc_num_threads{1}, baseline_of_query_mem_tracker{0}; std::atomic proc_virt_size{0}; MemoryTracker::~MemoryTracker() { @@ -133,9 +133,15 @@ void MemoryTracker::alloc(Int64 size, bool check_memory_limit) { Int64 current_limit = limit.load(std::memory_order_relaxed); Int64 current_accuracy_diff_for_test = accuracy_diff_for_test.load(std::memory_order_relaxed); + Int64 current_real_rss = real_rss.load(std::memory_order_relaxed); + Int64 current_real_rss_file = real_rss_file.load(std::memory_order_relaxed); + // Exclude file-backed RSS (mmap page cache) because it can be reclaimed by OS automatically. + Int64 effective_rss = current_real_rss - current_real_rss_file; + if (unlikely(effective_rss < 0)) + effective_rss = 0; if (unlikely( !next.load(std::memory_order_relaxed) && current_accuracy_diff_for_test && current_limit - && real_rss > current_accuracy_diff_for_test + current_limit)) + && effective_rss > current_accuracy_diff_for_test + current_limit)) { DB::FmtBuffer fmt_buf; fmt_buf.append("Memory tracker accuracy "); @@ -144,9 +150,12 @@ void MemoryTracker::alloc(Int64 size, bool check_memory_limit) fmt_buf.fmtAppend(" {}", tmp_decr); fmt_buf.fmtAppend( - ": fault injected. real_rss ({}) is much larger than limit ({}). Debug info, threads of process: {}, " + ": fault injected. effective_rss ({}, excluding rss_file {}, rss_total {}) is much larger than limit " + "({}). Debug info, threads of process: {}, " "memory usage tracked by ProcessList: peak {}, current {}. Virtual memory size: {}.", - formatReadableSizeWithBinarySuffix(real_rss), + formatReadableSizeWithBinarySuffix(effective_rss), + formatReadableSizeWithBinarySuffix(current_real_rss_file), + formatReadableSizeWithBinarySuffix(current_real_rss), formatReadableSizeWithBinarySuffix(current_limit), proc_num_threads.load(), (root_of_query_mem_trackers ? formatReadableSizeWithBinarySuffix(root_of_query_mem_trackers->peak) @@ -181,7 +190,7 @@ void MemoryTracker::alloc(Int64 size, bool check_memory_limit) Int64 current_bytes_rss_larger_than_limit = bytes_rss_larger_than_limit.load(std::memory_order_relaxed); bool is_rss_too_large = (!next.load(std::memory_order_relaxed) && current_limit - && real_rss > current_limit + current_bytes_rss_larger_than_limit + && effective_rss > current_limit + current_bytes_rss_larger_than_limit && will_be > baseline_of_query_mem_tracker); if (is_rss_too_large || unlikely(current_limit && will_be > current_limit)) { @@ -208,8 +217,11 @@ void MemoryTracker::alloc(Int64 size, bool check_memory_limit) { // RSS too large fmt_buf.fmtAppend( " exceeded caused by 'RSS(Resident Set Size) much larger than limit' : process memory size would " - "be {} for (attempt to allocate chunk of {} bytes), limit of memory for data computing : {}.", - formatReadableSizeWithBinarySuffix(real_rss), + "be {} (excluding {} page cache, rss_total {}) for (attempt to allocate chunk of {} bytes), " + "limit of memory for data computing : {}.", + formatReadableSizeWithBinarySuffix(effective_rss), + formatReadableSizeWithBinarySuffix(current_real_rss_file), + formatReadableSizeWithBinarySuffix(current_real_rss), size, formatReadableSizeWithBinarySuffix(current_limit)); } diff --git a/dbms/src/Common/MemoryTracker.h b/dbms/src/Common/MemoryTracker.h index 98775a7a13b..28794b102e6 100644 --- a/dbms/src/Common/MemoryTracker.h +++ b/dbms/src/Common/MemoryTracker.h @@ -22,7 +22,7 @@ #include #include -extern std::atomic real_rss, proc_num_threads, baseline_of_query_mem_tracker; +extern std::atomic real_rss, real_rss_file, proc_num_threads, baseline_of_query_mem_tracker; extern std::atomic proc_virt_size; namespace CurrentMetrics { diff --git a/dbms/src/Interpreters/Settings.h b/dbms/src/Interpreters/Settings.h index 443b0691000..4388a78d5d5 100644 --- a/dbms/src/Interpreters/Settings.h +++ b/dbms/src/Interpreters/Settings.h @@ -112,7 +112,7 @@ struct Settings \ M(SettingFloat, memory_tracker_fault_probability, 0., "For testing of `exception safety` - throw an exception every time you allocate memory with the specified probability.") \ \ - M(SettingInt64, memory_tracker_accuracy_diff_for_test, 0, "For testing of the accuracy of the memory tracker - throw an exception when real_rss is much larger than tracked amount.") \ + M(SettingInt64, memory_tracker_accuracy_diff_for_test, 0, "For testing of the accuracy of the memory tracker - throw an exception when effective RSS (rss - rss_file) is much larger than tracked amount.") \ \ M(SettingString, count_distinct_implementation, "uniqExact", "What aggregate function to use for implementation of count(DISTINCT ...)") \ \ From c9cd4cef5aa119bc132105f11b8efb508971c153 Mon Sep 17 00:00:00 2001 From: yongman Date: Wed, 7 Jan 2026 00:41:30 +0800 Subject: [PATCH 3/8] format Signed-off-by: yongman (cherry picked from commit aeb3bf1e0300d92c6fa12ef21505b052a5fd6b41) --- dbms/src/Common/BackgroundTask.cpp | 1 - 1 file changed, 1 deletion(-) diff --git a/dbms/src/Common/BackgroundTask.cpp b/dbms/src/Common/BackgroundTask.cpp index f5be5f94532..eb39c89ed23 100644 --- a/dbms/src/Common/BackgroundTask.cpp +++ b/dbms/src/Common/BackgroundTask.cpp @@ -14,7 +14,6 @@ #include #include - #include #include From ac1d83b2227b301f6e9d8ac53a92dd852d473ab9 Mon Sep 17 00:00:00 2001 From: yongman Date: Wed, 29 Apr 2026 15:19:14 +0800 Subject: [PATCH 4/8] columnar: add mem tracker check after read block from proxy Signed-off-by: yongman (cherry picked from commit d946399942e70bf0d8aa1b7567ad5e498108f2e2) --- dbms/src/Storages/StorageDisaggregatedColumnar.cpp | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/dbms/src/Storages/StorageDisaggregatedColumnar.cpp b/dbms/src/Storages/StorageDisaggregatedColumnar.cpp index 21dbdccb486..e38d40c22e3 100644 --- a/dbms/src/Storages/StorageDisaggregatedColumnar.cpp +++ b/dbms/src/Storages/StorageDisaggregatedColumnar.cpp @@ -15,6 +15,7 @@ #include // for ENABLE_NEXT_GEN_COLUMNAR #if ENABLE_NEXT_GEN_COLUMNAR #include +#include #include #include #include @@ -845,6 +846,12 @@ Block RNProxyInputStream::readImpl([[maybe_unused]] FilterPtr & res_filter, [[ma if (rows == 0) return {}; + // Add memory tracker check hook after reading block from proxy. If the memory is over the limit, an exception + // will be thrown to stop reading more data from proxy and avoid OOM. + CurrentMemoryTracker::submitLocalDeltaMemory(); + if (current_memory_tracker != nullptr) + current_memory_tracker->alloc(0); + TableID physical_table_id = -1; Block header = getHeader(); const ColumnsWithTypeAndName col_type_and_name = header.getColumnsWithTypeAndName(); From 057b91abbeeb6066350c3ea77cc0316c640705a2 Mon Sep 17 00:00:00 2001 From: yongman Date: Thu, 30 Apr 2026 18:02:43 +0800 Subject: [PATCH 5/8] add dedicated function checkRssLimit Signed-off-by: yongman (cherry picked from commit 66c3e5a377acb05cbb9cd819e06b0fc94d8ada9a) --- dbms/src/Common/MemoryTracker.cpp | 108 +++++++++++++----- dbms/src/Common/MemoryTracker.h | 6 + .../Storages/StorageDisaggregatedColumnar.cpp | 7 +- 3 files changed, 89 insertions(+), 32 deletions(-) diff --git a/dbms/src/Common/MemoryTracker.cpp b/dbms/src/Common/MemoryTracker.cpp index 17931e6e2e5..20c2abfd3f7 100644 --- a/dbms/src/Common/MemoryTracker.cpp +++ b/dbms/src/Common/MemoryTracker.cpp @@ -187,12 +187,8 @@ void MemoryTracker::alloc(Int64 size, bool check_memory_limit) fmt_buf.fmtAppend(" Memory Usage of Storage: {}", storageMemoryUsageDetail()); throw DB::TiFlashException(fmt_buf.toString(), DB::Errors::Coprocessor::MemoryLimitExceeded); } - Int64 current_bytes_rss_larger_than_limit = bytes_rss_larger_than_limit.load(std::memory_order_relaxed); - bool is_rss_too_large - = (!next.load(std::memory_order_relaxed) && current_limit - && effective_rss > current_limit + current_bytes_rss_larger_than_limit - && will_be > baseline_of_query_mem_tracker); - if (is_rss_too_large || unlikely(current_limit && will_be > current_limit)) + checkRssLimitImpl(/* require_tracked_growth */ true, will_be, size); + if (unlikely(current_limit && will_be > current_limit)) { DB::GET_METRIC(tiflash_memory_exceed_quota_count).Increment(); amount.fetch_sub(size, std::memory_order_relaxed); @@ -204,27 +200,12 @@ void MemoryTracker::alloc(Int64 size, bool check_memory_limit) if (tmp_decr) fmt_buf.fmtAppend(" {}", tmp_decr); - if (!is_rss_too_large) - { // out of memory quota - fmt_buf.fmtAppend( - " exceeded caused by 'out of memory quota for data computing' : would use {} for data computing " - "(attempt to allocate chunk of {} bytes), limit of memory for data computing: {}.", - formatReadableSizeWithBinarySuffix(will_be), - size, - formatReadableSizeWithBinarySuffix(current_limit)); - } - else - { // RSS too large - fmt_buf.fmtAppend( - " exceeded caused by 'RSS(Resident Set Size) much larger than limit' : process memory size would " - "be {} (excluding {} page cache, rss_total {}) for (attempt to allocate chunk of {} bytes), " - "limit of memory for data computing : {}.", - formatReadableSizeWithBinarySuffix(effective_rss), - formatReadableSizeWithBinarySuffix(current_real_rss_file), - formatReadableSizeWithBinarySuffix(current_real_rss), - size, - formatReadableSizeWithBinarySuffix(current_limit)); - } + fmt_buf.fmtAppend( + " exceeded caused by 'out of memory quota for data computing' : would use {} for data computing " + "(attempt to allocate chunk of {} bytes), limit of memory for data computing: {}.", + formatReadableSizeWithBinarySuffix(will_be), + size, + formatReadableSizeWithBinarySuffix(current_limit)); fmt_buf.fmtAppend(" Memory Usage of Storage: {}", storageMemoryUsageDetail()); @@ -251,6 +232,73 @@ void MemoryTracker::alloc(Int64 size, bool check_memory_limit) } +void MemoryTracker::checkRssLimitImpl(bool require_tracked_growth, Int64 will_be, Int64 size) const +{ + Int64 current_limit = limit.load(std::memory_order_relaxed); + if (next.load(std::memory_order_relaxed) || !current_limit) + return; + + Int64 current_real_rss = real_rss.load(std::memory_order_relaxed); + Int64 current_real_rss_file = real_rss_file.load(std::memory_order_relaxed); + // Exclude file-backed RSS (mmap page cache) because it can be reclaimed by OS automatically. + Int64 effective_rss = current_real_rss - current_real_rss_file; + if (unlikely(effective_rss < 0)) + effective_rss = 0; + + Int64 current_bytes_rss_larger_than_limit = bytes_rss_larger_than_limit.load(std::memory_order_relaxed); + bool is_rss_too_large = effective_rss > current_limit + current_bytes_rss_larger_than_limit; + if (require_tracked_growth) + is_rss_too_large = is_rss_too_large && will_be > baseline_of_query_mem_tracker; + + if (!is_rss_too_large) + return; + + DB::GET_METRIC(tiflash_memory_exceed_quota_count).Increment(); + + DB::FmtBuffer fmt_buf; + fmt_buf.append("Memory limit"); + const char * tmp_decr = description.load(); + if (tmp_decr) + fmt_buf.fmtAppend(" {}", tmp_decr); + + if (require_tracked_growth) + { + fmt_buf.fmtAppend( + " exceeded caused by 'RSS(Resident Set Size) much larger than limit' : process memory size would " + "be {} (excluding {} page cache, rss_total {}) for (attempt to allocate chunk of {} bytes), " + "limit of memory for data computing : {}.", + formatReadableSizeWithBinarySuffix(effective_rss), + formatReadableSizeWithBinarySuffix(current_real_rss_file), + formatReadableSizeWithBinarySuffix(current_real_rss), + size, + formatReadableSizeWithBinarySuffix(current_limit)); + } + else + { + fmt_buf.fmtAppend( + " exceeded caused by 'RSS(Resident Set Size) much larger than limit' : process memory size is {} " + "(excluding {} page cache, rss_total {}), limit of memory for data computing : {}. " + "Detected by explicit RSS limit check.", + formatReadableSizeWithBinarySuffix(effective_rss), + formatReadableSizeWithBinarySuffix(current_real_rss_file), + formatReadableSizeWithBinarySuffix(current_real_rss), + formatReadableSizeWithBinarySuffix(current_limit)); + } + + fmt_buf.fmtAppend(" Memory Usage of Storage: {}", storageMemoryUsageDetail()); + throw DB::TiFlashException(fmt_buf.toString(), DB::Errors::Coprocessor::MemoryLimitExceeded); +} + + +void MemoryTracker::checkRssLimit() const +{ + if (auto * loaded_next = next.load(std::memory_order_relaxed)) + loaded_next->checkRssLimit(); + else + checkRssLimitImpl(/* require_tracked_growth */ false, /* will_be */ 0, /* size */ 0); +} + + void MemoryTracker::free(Int64 size) { Int64 new_amount = amount.fetch_sub(size, std::memory_order_relaxed) - size; @@ -408,6 +456,12 @@ Int64 getLocalDeltaMemory() return local_delta; } +void checkRssLimit() +{ + if (current_memory_tracker) + current_memory_tracker->checkRssLimit(); +} + void alloc(Int64 size) { checkSubmitAndUpdateLocalDelta(local_delta + size); diff --git a/dbms/src/Common/MemoryTracker.h b/dbms/src/Common/MemoryTracker.h index 28794b102e6..52668b87444 100644 --- a/dbms/src/Common/MemoryTracker.h +++ b/dbms/src/Common/MemoryTracker.h @@ -79,6 +79,7 @@ class MemoryTracker : public std::enable_shared_from_this {} void reportAmount(); + void checkRssLimitImpl(bool require_tracked_growth, Int64 will_be, Int64 size) const; public: /// Using `std::shared_ptr` and `new` instread of `std::make_shared` is because `std::make_shared` cannot call private constructors. @@ -107,6 +108,10 @@ class MemoryTracker : public std::enable_shared_from_this */ void free(Int64 size); + /// Explicitly checks whether process RSS is already much larger than the configured limit. + /// Unlike alloc(0), this probe does not require tracked memory growth in the current tracker. + void checkRssLimit() const; + Int64 get() const { return amount.load(std::memory_order_relaxed); } Int64 getPeak() const { return peak.load(std::memory_order_relaxed); } @@ -184,6 +189,7 @@ namespace CurrentMemoryTracker void disableThreshold(); void submitLocalDeltaMemory(); Int64 getLocalDeltaMemory(); +void checkRssLimit(); void alloc(Int64 size); void realloc(Int64 old_size, Int64 new_size); void free(Int64 size); diff --git a/dbms/src/Storages/StorageDisaggregatedColumnar.cpp b/dbms/src/Storages/StorageDisaggregatedColumnar.cpp index e38d40c22e3..dd59b6b1e7b 100644 --- a/dbms/src/Storages/StorageDisaggregatedColumnar.cpp +++ b/dbms/src/Storages/StorageDisaggregatedColumnar.cpp @@ -846,11 +846,8 @@ Block RNProxyInputStream::readImpl([[maybe_unused]] FilterPtr & res_filter, [[ma if (rows == 0) return {}; - // Add memory tracker check hook after reading block from proxy. If the memory is over the limit, an exception - // will be thrown to stop reading more data from proxy and avoid OOM. - CurrentMemoryTracker::submitLocalDeltaMemory(); - if (current_memory_tracker != nullptr) - current_memory_tracker->alloc(0); + // Add a memory tracker hook after reading a block from proxy. + CurrentMemoryTracker::checkRssLimit(); TableID physical_table_id = -1; Block header = getHeader(); From a60789f313fbb71f9d65c577a16aa44384af83dc Mon Sep 17 00:00:00 2001 From: yongman Date: Mon, 25 May 2026 15:43:51 +0800 Subject: [PATCH 6/8] polish Signed-off-by: yongman --- dbms/src/Common/BackgroundTask.cpp | 32 +++++------------------------- 1 file changed, 5 insertions(+), 27 deletions(-) diff --git a/dbms/src/Common/BackgroundTask.cpp b/dbms/src/Common/BackgroundTask.cpp index eb39c89ed23..4ea4139170f 100644 --- a/dbms/src/Common/BackgroundTask.cpp +++ b/dbms/src/Common/BackgroundTask.cpp @@ -13,6 +13,7 @@ // limitations under the License. #include +#include #include #include @@ -22,28 +23,6 @@ namespace DB { namespace { -bool process_num_threads(Int64 & cur_proc_num_threads) -{ - // '/proc/self/stat' provides process thread count on Linux. - std::ifstream stat_stream("/proc/self/stat", std::ios_base::in); - if (!stat_stream.is_open()) - return false; - - std::string pid, comm, state, ppid, pgrp, session, tty_nr; - std::string tpgid, flags, minflt, cminflt, majflt, cmajflt; - std::string utime, stime, cutime, cstime, priority, nice; - std::string itrealvalue, starttime; - UInt64 cur_virt_size = 0; - Int64 rss = 0; - - stat_stream >> pid >> comm >> state >> ppid >> pgrp >> session >> tty_nr >> tpgid >> flags >> minflt >> cminflt - >> majflt >> cmajflt >> utime >> stime >> cutime >> cstime >> priority >> nice >> cur_proc_num_threads - >> itrealvalue >> starttime >> cur_virt_size >> rss; - - stat_stream.close(); - return true; -} - bool isProcStatSupported() { std::ifstream stat_stream("/proc/self/stat", std::ios_base::in); @@ -84,15 +63,14 @@ void CollectProcInfoBackgroundTask::memCheckJob() { try { - Int64 cur_proc_num_threads = 1; while (!end_syn) { + auto mem_usage = get_process_mem_usage(); auto metrics = get_process_metrics(); - process_num_threads(cur_proc_num_threads); - real_rss = static_cast(metrics.rss); + real_rss = static_cast(mem_usage.resident_bytes); real_rss_file = static_cast(metrics.rss_file); - proc_num_threads = cur_proc_num_threads; - proc_virt_size = metrics.vsize; + proc_num_threads = mem_usage.cur_proc_num_threads; + proc_virt_size = mem_usage.cur_virt_bytes; baseline_of_query_mem_tracker = root_of_query_mem_trackers->get(); usleep(100000); // sleep 100ms } From 6b4336a0bc5f04d9f130d901a71d82612f63b147 Mon Sep 17 00:00:00 2001 From: yongman Date: Mon, 25 May 2026 16:40:49 +0800 Subject: [PATCH 7/8] rollback alloc when exception Signed-off-by: yongman --- dbms/src/Common/MemoryTracker.cpp | 24 +++++++++++++++++++----- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/dbms/src/Common/MemoryTracker.cpp b/dbms/src/Common/MemoryTracker.cpp index 20c2abfd3f7..64f6f47de86 100644 --- a/dbms/src/Common/MemoryTracker.cpp +++ b/dbms/src/Common/MemoryTracker.cpp @@ -120,6 +120,12 @@ void MemoryTracker::alloc(Int64 size, bool check_memory_limit) */ Int64 will_be = size + amount.fetch_add(size, std::memory_order_relaxed); reportAmount(); + auto rollbackCurrentAlloc = [&] { + amount.fetch_sub(size, std::memory_order_relaxed); + reportAmount(); + if (!next.load(std::memory_order_relaxed)) + CurrentMetrics::sub(metric, size); + }; if (!next.load(std::memory_order_relaxed)) { @@ -143,6 +149,8 @@ void MemoryTracker::alloc(Int64 size, bool check_memory_limit) !next.load(std::memory_order_relaxed) && current_accuracy_diff_for_test && current_limit && effective_rss > current_accuracy_diff_for_test + current_limit)) { + rollbackCurrentAlloc(); + DB::FmtBuffer fmt_buf; fmt_buf.append("Memory tracker accuracy "); const char * tmp_decr = description.load(); @@ -171,8 +179,7 @@ void MemoryTracker::alloc(Int64 size, bool check_memory_limit) /// In this case, it doesn't matter. if (unlikely(fault_probability && drand48() < fault_probability)) { - amount.fetch_sub(size, std::memory_order_relaxed); - reportAmount(); + rollbackCurrentAlloc(); DB::FmtBuffer fmt_buf; fmt_buf.append("Memory tracker"); @@ -187,12 +194,19 @@ void MemoryTracker::alloc(Int64 size, bool check_memory_limit) fmt_buf.fmtAppend(" Memory Usage of Storage: {}", storageMemoryUsageDetail()); throw DB::TiFlashException(fmt_buf.toString(), DB::Errors::Coprocessor::MemoryLimitExceeded); } - checkRssLimitImpl(/* require_tracked_growth */ true, will_be, size); + try + { + checkRssLimitImpl(/* require_tracked_growth */ true, will_be, size); + } + catch (...) + { + rollbackCurrentAlloc(); + throw; + } if (unlikely(current_limit && will_be > current_limit)) { DB::GET_METRIC(tiflash_memory_exceed_quota_count).Increment(); - amount.fetch_sub(size, std::memory_order_relaxed); - reportAmount(); + rollbackCurrentAlloc(); DB::FmtBuffer fmt_buf; fmt_buf.append("Memory limit"); From dd94f94a4e868ce12c5839bbff4a9ff906c056f4 Mon Sep 17 00:00:00 2001 From: yongman Date: Mon, 25 May 2026 17:53:42 +0800 Subject: [PATCH 8/8] tidy Signed-off-by: yongman --- dbms/src/Common/MemoryTracker.cpp | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/dbms/src/Common/MemoryTracker.cpp b/dbms/src/Common/MemoryTracker.cpp index 64f6f47de86..d84f6c157bf 100644 --- a/dbms/src/Common/MemoryTracker.cpp +++ b/dbms/src/Common/MemoryTracker.cpp @@ -120,7 +120,7 @@ void MemoryTracker::alloc(Int64 size, bool check_memory_limit) */ Int64 will_be = size + amount.fetch_add(size, std::memory_order_relaxed); reportAmount(); - auto rollbackCurrentAlloc = [&] { + auto rollback_current_alloc = [&] { amount.fetch_sub(size, std::memory_order_relaxed); reportAmount(); if (!next.load(std::memory_order_relaxed)) @@ -149,7 +149,7 @@ void MemoryTracker::alloc(Int64 size, bool check_memory_limit) !next.load(std::memory_order_relaxed) && current_accuracy_diff_for_test && current_limit && effective_rss > current_accuracy_diff_for_test + current_limit)) { - rollbackCurrentAlloc(); + rollback_current_alloc(); DB::FmtBuffer fmt_buf; fmt_buf.append("Memory tracker accuracy "); @@ -179,7 +179,7 @@ void MemoryTracker::alloc(Int64 size, bool check_memory_limit) /// In this case, it doesn't matter. if (unlikely(fault_probability && drand48() < fault_probability)) { - rollbackCurrentAlloc(); + rollback_current_alloc(); DB::FmtBuffer fmt_buf; fmt_buf.append("Memory tracker"); @@ -200,13 +200,13 @@ void MemoryTracker::alloc(Int64 size, bool check_memory_limit) } catch (...) { - rollbackCurrentAlloc(); + rollback_current_alloc(); throw; } if (unlikely(current_limit && will_be > current_limit)) { DB::GET_METRIC(tiflash_memory_exceed_quota_count).Increment(); - rollbackCurrentAlloc(); + rollback_current_alloc(); DB::FmtBuffer fmt_buf; fmt_buf.append("Memory limit");