Fix memory accounting in case of MEMORY_LIMIT_EXCEEDED errors

src/Common/CurrentMemoryTracker.cpp

@@ -51,15 +51,23 @@ void CurrentMemoryTracker::allocImpl(Int64 size, bool throw_if_memory_exceeded)
     {
         if (current_thread)
         {
-            current_thread->untracked_memory += size;
+            Int64 will_be = current_thread->untracked_memory + size;
+            Int64 limit = current_thread->untracked_memory_limit + current_thread->untracked_memory_limit_increase;
 
-            if (current_thread->untracked_memory > current_thread->untracked_memory_limit)
+            if (will_be > limit)
             {
-                /// Zero untracked before track. If tracker throws out-of-limit we would be able to alloc up to untracked_memory_limit bytes
+                /// Increase limit before track. If tracker throws out-of-limit we would be able to alloc up to untracked_memory_limit bytes
                 /// more. It could be useful to enlarge Exception message in rethrow logic.
-                Int64 tmp = current_thread->untracked_memory;
+                current_thread->untracked_memory_limit_increase = current_thread->untracked_memory_limit;
+                memory_tracker->allocImpl(will_be, throw_if_memory_exceeded);
+                current_thread->untracked_memory_limit_increase = 0;
                 current_thread->untracked_memory = 0;
-                memory_tracker->allocImpl(tmp, throw_if_memory_exceeded);
+            }
+            else
+            {
+                /// Update after successful allocations,
+                /// since failed allocations should not be take into account.
+                current_thread->untracked_memory = will_be;
             }
         }
         /// total_memory_tracker only, ignore untracked_memory

src/Common/MemoryTracker.cpp

@@ -169,12 +169,14 @@ void MemoryTracker::allocImpl(Int64 size, bool throw_if_memory_exceeded, MemoryT
     std::bernoulli_distribution fault(fault_probability);
     if (unlikely(fault_probability && fault(thread_local_rng)) && memoryTrackerCanThrow(level, true) && throw_if_memory_exceeded)
     {
+        /// Revert
+        amount.fetch_sub(size, std::memory_order_relaxed);
+
         /// Prevent recursion. Exception::ctor -> std::string -> new[] -> MemoryTracker::alloc
         MemoryTrackerBlockerInThread untrack_lock(VariableContext::Global);
 
         ProfileEvents::increment(ProfileEvents::QueryMemoryLimitExceeded);
         const auto * description = description_ptr.load(std::memory_order_relaxed);
-        amount.fetch_sub(size, std::memory_order_relaxed);
         throw DB::Exception(
             DB::ErrorCodes::MEMORY_LIMIT_EXCEEDED,
             "Memory tracker{}{}: fault injected. Would use {} (attempt to allocate chunk of {} bytes), maximum: {}",
@@ -211,6 +213,9 @@ void MemoryTracker::allocImpl(Int64 size, bool throw_if_memory_exceeded, MemoryT
 
         if (overcommit_result != OvercommitResult::MEMORY_FREED)
         {
+            /// Revert
+            amount.fetch_sub(size, std::memory_order_relaxed);
+
             /// Prevent recursion. Exception::ctor -> std::string -> new[] -> MemoryTracker::alloc
             MemoryTrackerBlockerInThread untrack_lock(VariableContext::Global);
             ProfileEvents::increment(ProfileEvents::QueryMemoryLimitExceeded);

src/Common/ThreadStatus.h

@@ -135,6 +135,8 @@ class ThreadStatus : public boost::noncopyable
     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;
+    /// Increase limit in case of exception.
+    Int64 untracked_memory_limit_increase = 0;
 
     /// Statistics of read and write rows/bytes
     Progress progress_in;

tests/queries/0_stateless/02400_memory_accounting_on_error.sql

@@ -0,0 +1,4 @@
+-- max_block_size to avoid randomization
+-- SELECT * FROM generateRandom('i Array(Int8)', 1, 1, 1048577) LIMIT 65536 SETTINGS max_memory_usage='1Gi', max_block_size=65505, log_queries=1; -- { serverError MEMORY_LIMIT_EXCEEDED }
+-- SYSTEM FLUSH LOGS;
+-- SELECT * FROM system.query_log WHERE event_date >= yesterday() AND current_database = currentDatabase() AND memory_usage > 100e6;