diff --git a/src/Common/CurrentMemoryTracker.cpp b/src/Common/CurrentMemoryTracker.cpp index fed8e341b17..a564dba094f 100644 --- a/src/Common/CurrentMemoryTracker.cpp +++ b/src/Common/CurrentMemoryTracker.cpp @@ -52,13 +52,16 @@ void CurrentMemoryTracker::allocImpl(Int64 size, bool throw_if_memory_exceeded) if (current_thread) { Int64 will_be = current_thread->untracked_memory + size; + Int64 limit = current_thread->untracked_memory_limit + current_thread->untracked_memory_limit_increase; - if (will_be > 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. - current_thread->untracked_memory = 0; + 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; } else { diff --git a/src/Common/ThreadStatus.h b/src/Common/ThreadStatus.h index 0b01f43a226..2a4ffd229f2 100644 --- a/src/Common/ThreadStatus.h +++ b/src/Common/ThreadStatus.h @@ -135,6 +135,8 @@ public: 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;