From 88a3dbe716796e5575be5b7a3c579c8e8d32f479 Mon Sep 17 00:00:00 2001
From: Azat Khuzhin <a3at.mail@gmail.com>
Date: Thu, 28 Nov 2024 18:15:14 +0100
Subject: [PATCH] Fix inaccurate MemoryTracking metric in case of allocation
 failures
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

MemoryTracking metric takes into account memory allocation even in case
of this allocation will fail with MEMORY_LIMIT_EXCEEDED, which is not
good, which eventually will lead to `amount` and `MemoryTracking`
mismatch, I found one server with **43x difference**.

<details>

```sql
SELECT
    event_time_microseconds,
    message
FROM system.text_log
WHERE (event_date = (today() - 1)) AND (logger_name = 'MemoryTracker') AND (message LIKE '%total%')
ORDER BY 1 DESC
LIMIT 1

Query id: 64d60852-fa14-4ed1-adb1-d4bbd6159475

   ┌────event_time_microseconds─┬─message───────────────────────────────────┐
1. │ 2024-11-27 05:09:48.157608 │ Current memory usage (total): 471.00 GiB. │
   └────────────────────────────┴───────────────────────────────────────────┘
```

```sql
SELECT
    metric,
    formatReadableSize(value)
FROM system.metrics
WHERE (metric ILIKE '%mem%') OR (metric ILIKE '%jemalloc%')
ORDER BY value ASC

Query id: af7908a8-956a-4684-b7c5-b2e0c6fa06f4

   ┌─metric────────────────────────┬─formatReadableSize(value)─┐
1. │ MergesMutationsMemoryTracking │ 0.00 B                    │
2. │ MemoryTracking                │ 20.37 TiB                 │
   └───────────────────────────────┴───────────────────────────┘
```

</details>

Signed-off-by: Azat Khuzhin <a.khuzhin@semrush.com>
---
 src/Common/MemoryTracker.cpp | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/Common/MemoryTracker.cpp b/src/Common/MemoryTracker.cpp
index f4af019605e..829d5dfcd79 100644
--- a/src/Common/MemoryTracker.cpp
+++ b/src/Common/MemoryTracker.cpp
@@ -256,10 +256,6 @@ AllocationTrace MemoryTracker::allocImpl(Int64 size, bool throw_if_memory_exceed
     Int64 will_be = size ? size + amount.fetch_add(size, std::memory_order_relaxed) : amount.load(std::memory_order_relaxed);
     Int64 will_be_rss = size ? size + rss.fetch_add(size, std::memory_order_relaxed) : rss.load(std::memory_order_relaxed);
 
-    auto metric_loaded = metric.load(std::memory_order_relaxed);
-    if (metric_loaded != CurrentMetrics::end() && size)
-        CurrentMetrics::add(metric_loaded, size);
-
     Int64 current_hard_limit = hard_limit.load(std::memory_order_relaxed);
     Int64 current_profiler_limit = profiler_limit.load(std::memory_order_relaxed);
 
@@ -371,6 +367,10 @@ AllocationTrace MemoryTracker::allocImpl(Int64 size, bool throw_if_memory_exceed
         }
     }
 
+    auto metric_loaded = metric.load(std::memory_order_relaxed);
+    if (metric_loaded != CurrentMetrics::end() && size)
+        CurrentMetrics::add(metric_loaded, size);
+
     if (peak_updated && allocation_traced)
     {
         MemoryTrackerBlockerInThread untrack_lock(VariableContext::Global);