diff --git a/src/Common/CurrentMemoryTracker.cpp b/src/Common/CurrentMemoryTracker.cpp index 02c7dc6e2248..5ca6bd3d80d6 100644 --- a/src/Common/CurrentMemoryTracker.cpp +++ b/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 diff --git a/src/Common/ThreadStatus.cpp b/src/Common/ThreadStatus.cpp index 77ad2769322b..055f9cede279 100644 --- a/src/Common/ThreadStatus.cpp +++ b/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