diff --git a/src/Common/MemoryTracker.cpp b/src/Common/MemoryTracker.cpp index 2bbb5e89346..0204fec46fc 100644 --- a/src/Common/MemoryTracker.cpp +++ b/src/Common/MemoryTracker.cpp @@ -64,8 +64,8 @@ namespace ProfileEvents static constexpr size_t log_peak_memory_usage_every = 1ULL << 30; -thread_local bool MemoryTracker::BlockerInThread::is_blocked = false; -thread_local bool MemoryTracker::LockExceptionInThread::is_blocked = false; +thread_local uint64_t MemoryTracker::BlockerInThread::counter = 0; +thread_local uint64_t MemoryTracker::LockExceptionInThread::counter = 0; MemoryTracker total_memory_tracker(nullptr, VariableContext::Global); diff --git a/src/Common/MemoryTracker.h b/src/Common/MemoryTracker.h index 8bff1dc516f..4c20d109dfa 100644 --- a/src/Common/MemoryTracker.h +++ b/src/Common/MemoryTracker.h @@ -136,11 +136,11 @@ public: private: BlockerInThread(const BlockerInThread &) = delete; BlockerInThread & operator=(const BlockerInThread &) = delete; - static thread_local bool is_blocked; + static thread_local uint64_t counter; public: - BlockerInThread() { is_blocked = true; } - ~BlockerInThread() { is_blocked = false; } - static bool isBlocked() { return is_blocked; } + BlockerInThread() { ++counter; } + ~BlockerInThread() { --counter; } + static bool isBlocked() { return counter > 0; } }; /// To be able to avoid MEMORY_LIMIT_EXCEEDED Exception in destructors: @@ -160,11 +160,11 @@ public: private: LockExceptionInThread(const LockExceptionInThread &) = delete; LockExceptionInThread & operator=(const LockExceptionInThread &) = delete; - static thread_local bool is_blocked; + static thread_local uint64_t counter; public: - LockExceptionInThread() { is_blocked = true; } - ~LockExceptionInThread() { is_blocked = false; } - static bool isBlocked() { return is_blocked; } + LockExceptionInThread() { ++counter; } + ~LockExceptionInThread() { --counter; } + static bool isBlocked() { return counter > 0; } }; };