From 0bf6ed2e94481814fc8566b8c2170936af17302c Mon Sep 17 00:00:00 2001
From: Azat Khuzhin <a3at.mail@gmail.com>
Date: Tue, 1 Dec 2020 10:34:25 +0300
Subject: [PATCH] Make LockExceptionInThread/BlockerInThread nested

---
 src/Common/MemoryTracker.cpp |  4 ++--
 src/Common/MemoryTracker.h   | 16 ++++++++--------
 2 files changed, 10 insertions(+), 10 deletions(-)

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; }
     };
 };