Merge pull request #18976 from ClickHouse/fix_terminate_when_not_enough_memory

Finally get rid of allocations in ThreadPool::worker(...)

base/daemon/BaseDaemon.cpp

@@ -112,11 +112,13 @@ static void writeSignalIDtoSignalPipe(int sig)
 /** Signal handler for HUP / USR1 */
 static void closeLogsSignalHandler(int sig, siginfo_t *, void *)
 {
+    DENY_ALLOCATIONS_IN_SCOPE;
     writeSignalIDtoSignalPipe(sig);
 }
 
 static void terminateRequestedSignalHandler(int sig, siginfo_t *, void *)
 {
+    DENY_ALLOCATIONS_IN_SCOPE;
     writeSignalIDtoSignalPipe(sig);
 }
 
@@ -125,6 +127,7 @@ static void terminateRequestedSignalHandler(int sig, siginfo_t *, void *)
   */
 static void signalHandler(int sig, siginfo_t * info, void * context)
 {
+    DENY_ALLOCATIONS_IN_SCOPE;
     auto saved_errno = errno;   /// We must restore previous value of errno in signal handler.
 
     char buf[signal_pipe_buf_size];

src/Common/MemoryTracker.cpp

@@ -12,6 +12,10 @@
 #include <random>
 #include <cstdlib>
 
+#ifdef MEMORY_TRACKER_DEBUG_CHECKS
+thread_local bool _memory_tracker_always_throw_logical_error_on_allocation = false;
+#endif
+
 namespace
 {
 
@@ -165,6 +169,14 @@ void MemoryTracker::alloc(Int64 size)
         }
     }
 
+#ifdef MEMORY_TRACKER_DEBUG_CHECKS
+    if (unlikely(_memory_tracker_always_throw_logical_error_on_allocation))
+    {
+        _memory_tracker_always_throw_logical_error_on_allocation = false;
+        throw DB::Exception(DB::ErrorCodes::LOGICAL_ERROR, "Memory tracker: allocations not allowed.");
+    }
+#endif
+
     std::bernoulli_distribution fault(fault_probability);
     if (unlikely(fault_probability && fault(thread_local_rng)) && memoryTrackerCanThrow(level, true))
     {

src/Common/MemoryTracker.h

@@ -5,6 +5,28 @@
 #include <Common/CurrentMetrics.h>
 #include <Common/VariableContext.h>
 
+#if !defined(NDEBUG)
+#define MEMORY_TRACKER_DEBUG_CHECKS
+#endif
+
+/// DENY_ALLOCATIONS_IN_SCOPE macro makes MemoryTracker throw LOGICAL_ERROR on any allocation attempt
+/// until the end of the scope. It's useful to ensure that no allocations happen in signal handlers and
+/// outside of try/catch block of thread functions. ALLOW_ALLOCATIONS_IN_SCOPE cancels effect of
+/// DENY_ALLOCATIONS_IN_SCOPE in the inner scope. In Release builds these macros do nothing.
+#ifdef MEMORY_TRACKER_DEBUG_CHECKS
+#include <ext/scope_guard.h>
+extern thread_local bool _memory_tracker_always_throw_logical_error_on_allocation;
+#define ALLOCATIONS_IN_SCOPE_IMPL_CONCAT(n, val) \
+        bool _allocations_flag_prev_val##n = _memory_tracker_always_throw_logical_error_on_allocation; \
+        _memory_tracker_always_throw_logical_error_on_allocation = val; \
+        SCOPE_EXIT({ _memory_tracker_always_throw_logical_error_on_allocation = _allocations_flag_prev_val##n; })
+#define ALLOCATIONS_IN_SCOPE_IMPL(n, val) ALLOCATIONS_IN_SCOPE_IMPL_CONCAT(n, val)
+#define DENY_ALLOCATIONS_IN_SCOPE ALLOCATIONS_IN_SCOPE_IMPL(__LINE__, true)
+#define ALLOW_ALLOCATIONS_IN_SCOPE ALLOCATIONS_IN_SCOPE_IMPL(__LINE__, false)
+#else
+#define DENY_ALLOCATIONS_IN_SCOPE static_assert(true)
+#define ALLOW_ALLOCATIONS_IN_SCOPE static_assert(true)
+#endif
 
 /** Tracks memory consumption.
   * It throws an exception if amount of consumed memory become greater than certain limit.

src/Common/QueryProfiler.cpp

@@ -181,6 +181,7 @@ QueryProfilerReal::QueryProfilerReal(const UInt64 thread_id, const UInt32 period
 
 void QueryProfilerReal::signalHandler(int sig, siginfo_t * info, void * context)
 {
+    DENY_ALLOCATIONS_IN_SCOPE;
     writeTraceInfo(TraceType::Real, sig, info, context);
 }
 
@@ -190,6 +191,7 @@ QueryProfilerCpu::QueryProfilerCpu(const UInt64 thread_id, const UInt32 period)
 
 void QueryProfilerCpu::signalHandler(int sig, siginfo_t * info, void * context)
 {
+    DENY_ALLOCATIONS_IN_SCOPE;
     writeTraceInfo(TraceType::CPU, sig, info, context);
 }
 

src/Common/ThreadFuzzer.cpp

@@ -197,6 +197,7 @@ static void injection(
 
 void ThreadFuzzer::signalHandler(int)
 {
+    DENY_ALLOCATIONS_IN_SCOPE;
     auto saved_errno = errno;
 
     auto & fuzzer = ThreadFuzzer::instance();

src/Common/ThreadPool.cpp

@@ -208,6 +208,7 @@ size_t ThreadPoolImpl<Thread>::active() const
 template <typename Thread>
 void ThreadPoolImpl<Thread>::worker(typename std::list<Thread>::iterator thread_it)
 {
+    DENY_ALLOCATIONS_IN_SCOPE;
     CurrentMetrics::Increment metric_all_threads(
         std::is_same_v<Thread, std::thread> ? CurrentMetrics::GlobalThread : CurrentMetrics::LocalThread);
 
@@ -223,7 +224,9 @@ void ThreadPoolImpl<Thread>::worker(typename std::list<Thread>::iterator thread_
 
             if (!jobs.empty())
             {
-                job = std::move(jobs.top().job);
+                /// std::priority_queue does not provide interface for getting non-const reference to an element
+                /// to prevent us from modifying its priority. We have to use const_cast to force move semantics on JobWithPriority::job.
+                job = std::move(const_cast<Job &>(jobs.top().job));
                 jobs.pop();
             }
             else
@@ -237,6 +240,7 @@ void ThreadPoolImpl<Thread>::worker(typename std::list<Thread>::iterator thread_
         {
             try
             {
+                ALLOW_ALLOCATIONS_IN_SCOPE;
                 CurrentMetrics::Increment metric_active_threads(
                     std::is_same_v<Thread, std::thread> ? CurrentMetrics::GlobalThreadActive : CurrentMetrics::LocalThreadActive);
 

src/Storages/System/StorageSystemStackTrace.cpp

@@ -60,6 +60,7 @@ namespace
 
     void signalHandler(int, siginfo_t * info, void * context)
     {
+        DENY_ALLOCATIONS_IN_SCOPE;
         auto saved_errno = errno;   /// We must restore previous value of errno in signal handler.
 
         /// In case malicious user is sending signals manually (for unknown reason).