Add ability to block MEMORY_LIMIT_EXCEEDED exception explicitly

src/Common/MemoryTracker.cpp

@@ -33,6 +33,8 @@ MemoryTracker * getMemoryTracker()
 /// MemoryTracker cannot throw MEMORY_LIMIT_EXCEEDED (either configured memory
 /// limit reached or fault injected), in the following cases:
 ///
+/// - when it is explicitly blocked with LockExceptionInThread
+///
 /// - to avoid std::terminate(), when stack unwinding is current in progress in
 ///   this thread.
 ///
@@ -41,7 +43,7 @@ MemoryTracker * getMemoryTracker()
 ///   noexcept(false)) will cause std::terminate()
 bool inline memoryTrackerCanThrow()
 {
-    return !std::uncaught_exceptions();
+    return !MemoryTracker::LockExceptionInThread::isBlocked() && !std::uncaught_exceptions();
 }
 
 }
@@ -63,6 +65,7 @@ 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;
 
 MemoryTracker total_memory_tracker(nullptr, VariableContext::Global);
 

src/Common/MemoryTracker.h

@@ -142,6 +142,30 @@ public:
         ~BlockerInThread() { is_blocked = false; }
         static bool isBlocked() { return is_blocked; }
     };
+
+    /// To be able to avoid MEMORY_LIMIT_EXCEEDED Exception in destructors:
+    /// - either configured memory limit reached
+    /// - or fault injected
+    ///
+    /// So this will simply ignore the configured memory limit (and avoid fault injection).
+    ///
+    /// NOTE: exception will be silently ignored, no message in log
+    /// (since logging from MemoryTracker::alloc() is tricky)
+    ///
+    /// NOTE: MEMORY_LIMIT_EXCEEDED Exception implicitly blocked if
+    /// stack unwinding is currently in progress in this thread (to avoid
+    /// std::terminate()), so you don't need to use it in this case explicitly.
+    struct LockExceptionInThread
+    {
+    private:
+        LockExceptionInThread(const LockExceptionInThread &) = delete;
+        LockExceptionInThread & operator=(const LockExceptionInThread &) = delete;
+        static thread_local bool is_blocked;
+    public:
+        LockExceptionInThread() { is_blocked = true; }
+        ~LockExceptionInThread() { is_blocked = false; }
+        static bool isBlocked() { return is_blocked; }
+    };
 };
 
 extern MemoryTracker total_memory_tracker;