Cherry-pick: Backport #83607 to 24.8: Fix incorrect memory around max_untracked_memory

src/Common/CurrentMemoryTracker.cpp

@@ -51,19 +51,22 @@ AllocationTrace CurrentMemoryTracker::allocImpl(Int64 size, bool throw_if_memory
     {
         if (current_thread)
         {
-            Int64 will_be = current_thread->untracked_memory + size;
-
-            if (will_be > current_thread->untracked_memory_limit)
+            Int64 previous_untracked_memory = current_thread->untracked_memory;
+            current_thread->untracked_memory += size;
+            if (current_thread->untracked_memory > current_thread->untracked_memory_limit)
             {
-                auto res = memory_tracker->allocImpl(will_be, throw_if_memory_exceeded);
+                Int64 current_untracked_memory = current_thread->untracked_memory;
                 current_thread->untracked_memory = 0;
-                return res;
-            }
-            else
-            {
-                /// Update after successful allocations,
-                /// since failed allocations should not be take into account.
-                current_thread->untracked_memory = will_be;
+
+                try
+                {
+                    return memory_tracker->allocImpl(current_untracked_memory, throw_if_memory_exceeded);
+                }
+                catch (...)
+                {
+                    current_thread->untracked_memory += previous_untracked_memory;
+                    throw;
+                }
             }
         }
         /// total_memory_tracker only, ignore untracked_memory

src/Common/ThreadStatus.cpp

@@ -190,8 +190,9 @@ void ThreadStatus::flushUntrackedMemory()
     if (untracked_memory == 0)
         return;
 
-    memory_tracker.adjustWithUntrackedMemory(untracked_memory);
+    Int64 current_untracked_memory = current_thread->untracked_memory;
     untracked_memory = 0;
+    memory_tracker.adjustWithUntrackedMemory(current_untracked_memory);
 }
 
 bool ThreadStatus::isQueryCanceled() const