Sanitize places that can have big allocations w/ ignored MEMORY_LIMIT_EXCEEDED

azat · azat · commit 4117df841f4b · 2025-09-13T22:17:56.000+02:00
- Add MemoryTrackerDebugBlockerInThread for some internal places that should not throw - Add TrackedString implementation, a std::string that can throw, and use it in places where we can have a significant memory usage - Use AllocatorWithMemoryTracking in RadixSort as well (can create big allocations) I've found this from the log from CI [1]. [1]: https://s3.amazonaws.com/clickhouse-test-reports/json.html?PR=87035&sha=cf9779523fb72328052b0f541fef6cbd1cf1cf45&name_0=PR&name_1=Stateless%20tests%20%28amd_debug%2C%20AsyncInsert%2C%20s3%20storage%2C%20parallel%29
diff --git a/src/AggregateFunctions/QuantileTDigest.h b/src/AggregateFunctions/QuantileTDigest.h
@@ -119,7 +119,6 @@ class QuantileTDigest
         static constexpr size_t PART_SIZE_BITS = 8;
 
         using Transform = RadixSortFloatTransform<KeyBits>;
-        using Allocator = RadixSortAllocator;
 
         /// The function to get the key from an array element.
         static Key & extractKey(Element & elem) { return elem.mean; }
diff --git a/src/Common/RadixSort.h b/src/Common/RadixSort.h
@@ -1,6 +1,5 @@
 #pragma once
 
-#include <string.h>
 #if !defined(OS_DARWIN) && !defined(OS_FREEBSD)
 #include <malloc.h>
 #endif
@@ -16,6 +15,7 @@
 #include <base/extended_types.h>
 #include <base/sort.h>
 
+#include <Common/AllocatorWithMemoryTracking.h>
 #include <Common/TargetSpecific.h>
 
 /** Radix sort, has the following functionality:
@@ -33,22 +33,6 @@
   */
 
 
-/** Used as a template parameter. See below.
-  */
-struct RadixSortAllocator
-{
-    static void * allocate(size_t size)
-    {
-        return ::operator new(size);
-    }
-
-    static void deallocate(void * ptr, size_t size)
-    {
-        ::operator delete(ptr, size);
-    }
-};
-
-
 /** A transformation that transforms the bit representation of a key into an unsigned integer number,
   *  that the order relation over the keys will match the order relation over the obtained unsigned numbers.
   * For floats this conversion does the following:
@@ -97,11 +81,6 @@ struct RadixSortFloatTraits
     /// Converting a key into KeyBits is such that the order relation over the key corresponds to the order relation over KeyBits.
     using Transform = RadixSortFloatTransform<KeyBits>;
 
-    /// An object with the functions allocate and deallocate.
-    /// Can be used, for example, to allocate memory for a temporary array on the stack.
-    /// To do this, the allocator itself is created on the stack.
-    using Allocator = RadixSortAllocator;
-
     /// The function to get the key from an array element.
     static Key & extractKey(Element & elem) { return elem; }
 
@@ -144,7 +123,6 @@ struct RadixSortUIntTraits
     static constexpr size_t PART_SIZE_BITS = 8;
 
     using Transform = RadixSortIdentityTransform<KeyBits>;
-    using Allocator = RadixSortAllocator;
 
     static Key & extractKey(Element & elem) { return elem; }
     static Result & extractResult(Element & elem) { return elem; }
@@ -183,7 +161,6 @@ struct RadixSortIntTraits
     static constexpr size_t PART_SIZE_BITS = 8;
 
     using Transform = RadixSortSignedTransform<KeyBits>;
-    using Allocator = RadixSortAllocator;
 
     static Key & extractKey(Element & elem) { return elem; }
     static Result & extractResult(Element & elem) { return elem; }
@@ -289,10 +266,10 @@ struct RadixSort
         /// For each of the NUM_PASSES bit ranges of the key, consider how many times each value of this bit range met.
         std::unique_ptr<CountType[]> histograms{new CountType[HISTOGRAM_SIZE * NUM_PASSES]{}};
 
-        typename Traits::Allocator allocator;
+        AllocatorWithMemoryTracking<typename Traits::Element> allocator;
 
         /// We will do several passes through the array. On each pass, the data is transferred to another array. Let's allocate this temporary array.
-        Element * swap_buffer = reinterpret_cast<Element *>(allocator.allocate(size * sizeof(Element)));
+        Element * swap_buffer = allocator.allocate(size);
 
         /// Transform the array and calculate the histogram.
         /// NOTE This is slightly suboptimal. Look at https://github.com/powturbo/TurboHist
@@ -436,7 +413,7 @@ struct RadixSort
                 std::reverse(arr, arr + size);
         }
 
-        allocator.deallocate(swap_buffer, size * sizeof(Element));
+        allocator.deallocate(swap_buffer, size);
     }
 
 
diff --git a/src/Common/SystemLogBase.cpp b/src/Common/SystemLogBase.cpp
@@ -15,6 +15,7 @@
 #include <Interpreters/FilesystemCacheLog.h>
 #include <Interpreters/ObjectStorageQueueLog.h>
 #include <Interpreters/IcebergMetadataLog.h>
+#include <Common/MemoryTrackerDebugBlockerInThread.h>
 #if CLICKHOUSE_CLOUD
 #include <Interpreters/DistributedCacheLog.h>
 #include <Interpreters/DistributedCacheServerLog.h>
@@ -200,6 +201,8 @@ void SystemLogQueue<LogElement>::confirm(SystemLogQueue<LogElement>::Index last_
 template <typename LogElement>
 typename SystemLogQueue<LogElement>::PopResult SystemLogQueue<LogElement>::pop()
 {
+    [[maybe_unused]] MemoryTrackerDebugBlockerInThread blocker;
+
     PopResult result;
     size_t prev_ignored_logs = 0;
 
@@ -308,6 +311,7 @@ void SystemLogBase<LogElement>::stopFlushThread()
 template <typename LogElement>
 void SystemLogBase<LogElement>::add(LogElement element)
 {
+    [[maybe_unused]] MemoryTrackerDebugBlockerInThread blocker;
     queue->push(std::move(element));
 }
 
diff --git a/src/Common/TrackedString.h b/src/Common/TrackedString.h
@@ -0,0 +1,11 @@
+#pragma once
+
+#include <string>
+#include <Common/AllocatorWithMemoryTracking.h>
+
+/// String that can throw MEMORY_LIMIT_EXCEEDED
+/// Use it when the string may require significant memory.
+class TrackedString : public std::basic_string<char, std::char_traits<char>, AllocatorWithMemoryTracking<char>>
+{
+    using std::basic_string<char, std::char_traits<char>, AllocatorWithMemoryTracking<char>>::basic_string;
+};
diff --git a/src/Common/ZooKeeper/ZooKeeperImpl.cpp b/src/Common/ZooKeeper/ZooKeeperImpl.cpp
@@ -26,6 +26,7 @@
 #include <Common/logger_useful.h>
 #include <Common/setThreadName.h>
 #include <Common/thread_local_rng.h>
+#include <Common/MemoryTrackerDebugBlockerInThread.h>
 #include <Core/Settings.h>
 #include <Core/ServerSettings.h>
 
@@ -771,6 +772,8 @@ void ZooKeeper::sendAuth(const String & scheme, const String & data)
 
 void ZooKeeper::sendThread()
 {
+    [[maybe_unused]] MemoryTrackerDebugBlockerInThread blocker;
+
     setThreadName("ZooKeeperSend");
 
     scope_guard os_thread_nice_value_guard;
@@ -860,6 +863,8 @@ void ZooKeeper::sendThread()
 
 void ZooKeeper::receiveThread()
 {
+    [[maybe_unused]] MemoryTrackerDebugBlockerInThread blocker;
+
     setThreadName("ZooKeeperRecv");
 
     scope_guard os_thread_nice_value_guard;
@@ -1817,6 +1822,8 @@ void ZooKeeper::setZooKeeperLog(std::shared_ptr<DB::ZooKeeperLog> zk_log_)
 #ifdef ZOOKEEPER_LOG
 void ZooKeeper::logOperationIfNeeded(const ZooKeeperRequestPtr & request, const ZooKeeperResponsePtr & response, bool finalize, UInt64 elapsed_microseconds)
 {
+    [[maybe_unused]] MemoryTrackerDebugBlockerInThread blocker;
+
     auto maybe_zk_log = std::atomic_load(&zk_log);
     if (!maybe_zk_log)
         return;
diff --git a/src/IO/WriteBufferFromTrackedString.cpp b/src/IO/WriteBufferFromTrackedString.cpp
@@ -0,0 +1,21 @@
+#include <Common/Exception.h>
+#include <IO/WriteBufferFromTrackedString.h>
+
+namespace DB
+{
+
+/// It is safe to make them autofinalizable.
+WriteBufferFromTrackedString::~WriteBufferFromTrackedString()
+{
+    try
+    {
+        if (!this->finalized && !this->canceled)
+            this->finalize();
+    }
+    catch (...)
+    {
+        tryLogCurrentException(__PRETTY_FUNCTION__);
+    }
+}
+
+}
diff --git a/src/IO/WriteBufferFromTrackedString.h b/src/IO/WriteBufferFromTrackedString.h
@@ -0,0 +1,25 @@
+#pragma once
+
+#include <IO/WriteBufferFromVector.h>
+#include <Common/TrackedString.h>
+
+namespace DB
+{
+
+class WriteBufferFromTrackedString final : public WriteBufferFromVectorImpl<TrackedString>
+{
+    using Base = WriteBufferFromVectorImpl;
+public:
+    explicit WriteBufferFromTrackedString(TrackedString & vector_)
+        : Base(vector_)
+    {
+    }
+
+    WriteBufferFromTrackedString(TrackedString & vector_, AppendModeTag tag_)
+        : Base(vector_, tag_)
+    {
+    }
+    ~WriteBufferFromTrackedString() override;
+};
+
+}
diff --git a/src/Interpreters/AsynchronousInsertQueue.cpp b/src/Interpreters/AsynchronousInsertQueue.cpp
@@ -10,6 +10,7 @@
 #include <IO/ConcatReadBuffer.h>
 #include <IO/LimitReadBuffer.h>
 #include <IO/ReadBufferFromString.h>
+#include <IO/WriteBufferFromTrackedString.h>
 #include <IO/copyData.h>
 #include <Interpreters/ExpressionActions.h>
 #include <Interpreters/AsynchronousInsertLog.h>
@@ -399,7 +400,7 @@ AsynchronousInsertQueue::pushQueryWithInlinedData(ASTPtr query, ContextPtr query
     }
     preprocessInsertQuery(query, query_context);
 
-    String bytes;
+    TrackedString bytes;
     {
         /// Read at most 'async_insert_max_data_size' bytes of data.
         /// If limit is exceeded we will fallback to synchronous insert
@@ -424,7 +425,7 @@ AsynchronousInsertQueue::pushQueryWithInlinedData(ASTPtr query, ContextPtr query
         }
 
         {
-            WriteBufferFromString write_buf(bytes);
+            WriteBufferFromTrackedString write_buf(bytes);
             copyData(limit_buf, write_buf);
         }
 
diff --git a/src/Interpreters/AsynchronousInsertQueue.h b/src/Interpreters/AsynchronousInsertQueue.h
@@ -6,6 +6,7 @@
 #include <Common/MemoryTrackerSwitcher.h>
 #include <Common/SettingsChanges.h>
 #include <Common/ThreadPool.h>
+#include <Common/TrackedString.h>
 #include <Interpreters/AsynchronousInsertQueueDataKind.h>
 
 #include <future>
@@ -92,9 +93,9 @@ class AsynchronousInsertQueue : public WithContext
     };
 
 private:
-    struct DataChunk : public std::variant<String, Block>
+    struct DataChunk : public std::variant<TrackedString, Block>
     {
-        using std::variant<String, Block>::variant;
+        using std::variant<TrackedString, Block>::variant;
 
         size_t byteSize() const
         {
@@ -125,7 +126,7 @@ class AsynchronousInsertQueue : public WithContext
             }, *this);
         }
 
-        const String * asString() const { return std::get_if<String>(this); }
+        const TrackedString * asString() const { return std::get_if<TrackedString>(this); }
         const Block * asBlock() const { return std::get_if<Block>(this); }
     };
 
diff --git a/src/Interpreters/TraceCollector.cpp b/src/Interpreters/TraceCollector.cpp
@@ -7,6 +7,7 @@
 #include <Interpreters/TraceLog.h>
 #include <Common/MemoryTrackerBlockerInThread.h>
 #include <Common/Exception.h>
+#include <Common/MemoryTrackerDebugBlockerInThread.h>
 #include <Common/TraceSender.h>
 #include <Common/ProfileEvents.h>
 #include <Common/setThreadName.h>
@@ -99,6 +100,8 @@ TraceCollector::~TraceCollector()
 
 void TraceCollector::run()
 {
+    [[maybe_unused]] MemoryTrackerDebugBlockerInThread blocker;
+
     setThreadName("TraceCollector");
 
     MemoryTrackerBlockerInThread untrack_lock(VariableContext::Global);
diff --git a/src/Loggers/OwnSplitChannel.cpp b/src/Loggers/OwnSplitChannel.cpp
@@ -7,6 +7,7 @@
 #include <Common/DNSResolver.h>
 #include <Common/IO.h>
 #include <Common/LockMemoryExceptionInThread.h>
+#include <Common/MemoryTrackerDebugBlockerInThread.h>
 #include <Common/ProfileEvents.h>
 #include <Common/SensitiveDataMasker.h>
 #include <Common/setThreadName.h>
@@ -131,6 +132,7 @@ void logToSystemTextLogQueue(
 
 #undef SET_VALUE_IF_EXISTS
 
+    [[maybe_unused]] MemoryTrackerDebugBlockerInThread blocker;
     text_log_locked->push(std::move(elem));
 }
 }
diff --git a/src/Storages/FileLog/FileLogConsumer.h b/src/Storages/FileLog/FileLogConsumer.h
@@ -3,9 +3,8 @@
 #include <Core/BackgroundSchedulePoolTaskHolder.h>
 #include <IO/ReadBuffer.h>
 #include <Storages/FileLog/StorageFileLog.h>
-
-#include <fstream>
-#include <mutex>
+#include <Common/AllocatorWithMemoryTracking.h>
+#include <Common/TrackedString.h>
 
 namespace DB
 {
@@ -30,7 +29,7 @@ class FileLogConsumer
 
     auto getFileName() const { return current[-1].file_name; }
     auto getOffset() const { return current[-1].offset; }
-    const String & getCurrentRecord() const { return current[-1].data; }
+    const auto & getCurrentRecord() const { return current[-1].data; }
 
 private:
     enum class BufferStatus : uint8_t
@@ -56,15 +55,14 @@ class FileLogConsumer
     size_t stream_number;
     size_t max_streams_number;
 
-    using RecordData = std::string;
     struct Record
     {
-        RecordData data;
+        TrackedString data;
         std::string file_name;
         /// Offset is the start of a row, which is needed for virtual columns.
         UInt64 offset;
     };
-    using Records = std::vector<Record>;
+    using Records = std::vector<Record, AllocatorWithMemoryTracking<Record>>;
 
     Records records;
     Records::const_iterator current;

Original file line number	Diff line number	Diff line change
`@@ -10,6 +10,7 @@`
`10`	`10`	`#include <IO/ConcatReadBuffer.h>`
`11`	`11`	`#include <IO/LimitReadBuffer.h>`
`12`	`12`	`#include <IO/ReadBufferFromString.h>`
	`13`	`+#include <IO/WriteBufferFromTrackedString.h>`
`13`	`14`	`#include <IO/copyData.h>`
`14`	`15`	`#include <Interpreters/ExpressionActions.h>`
`15`	`16`	`#include <Interpreters/AsynchronousInsertLog.h>`
`@@ -399,7 +400,7 @@ AsynchronousInsertQueue::pushQueryWithInlinedData(ASTPtr query, ContextPtr query`
`399`	`400`	`}`
`400`	`401`	`preprocessInsertQuery(query, query_context);`
`401`	`402`
`402`		`- String bytes;`
	`403`	`+ TrackedString bytes;`
`403`	`404`	`{`
`404`	`405`	`/// Read at most 'async_insert_max_data_size' bytes of data.`
`405`	`406`	`/// If limit is exceeded we will fallback to synchronous insert`
`@@ -424,7 +425,7 @@ AsynchronousInsertQueue::pushQueryWithInlinedData(ASTPtr query, ContextPtr query`
`424`	`425`	`}`
`425`	`426`
`426`	`427`	`{`
`427`		`- WriteBufferFromString write_buf(bytes);`
	`428`	`+ WriteBufferFromTrackedString write_buf(bytes);`
`428`	`429`	`copyData(limit_buf, write_buf);`
`429`	`430`	`}`
`430`	`431`