Fix some OOMs in stress tests

src/Common/AllocatorWithMemoryTracking.h

@@ -0,0 +1,62 @@
+#pragma once
+
+#include <stdexcept>
+#include <cstddef>
+#include <cstdlib>
+
+#include <Common/CurrentMemoryTracker.h>
+
+
+/// Implementation of std::allocator interface that tracks memory with MemoryTracker.
+/// NOTE We already plug MemoryTracker into new/delete operators. So, everything works even with default allocator.
+/// But it is enabled only if jemalloc is used (to obtain the size of the allocation on call to delete).
+/// And jemalloc is disabled for builds with sanitizers. In these cases memory was not always tracked.
+
+template <typename T>
+struct AllocatorWithMemoryTracking
+{
+    typedef T value_type;
+
+    AllocatorWithMemoryTracking() = default;
+
+    template <typename U>
+    constexpr AllocatorWithMemoryTracking(const AllocatorWithMemoryTracking<U> &) noexcept
+    {
+    }
+
+    [[nodiscard]] T * allocate(size_t n)
+    {
+        if (n > std::numeric_limits<size_t>::max() / sizeof(T))
+            throw std::bad_alloc();
+
+        size_t bytes = n * sizeof(T);
+        CurrentMemoryTracker::alloc(bytes);
+
+        T * p = static_cast<T *>(malloc(bytes));
+        if (!p)
+            throw std::bad_alloc();
+
+        return p;
+    }
+
+    void deallocate(T * p, size_t n) noexcept
+    {
+        free(p);
+
+        size_t bytes = n * sizeof(T);
+        CurrentMemoryTracker::free(bytes);
+    }
+};
+
+template <typename T, typename U>
+bool operator==(const AllocatorWithMemoryTracking <T> &, const AllocatorWithMemoryTracking <U> &)
+{
+    return true;
+}
+
+template <typename T, typename U>
+bool operator!=(const AllocatorWithMemoryTracking <T> &, const AllocatorWithMemoryTracking <U> &)
+{
+    return false;
+}
+

src/Core/Field.h

@@ -8,6 +8,7 @@
 
 #include <Common/Exception.h>
 #include <Common/UInt128.h>
+#include <Common/AllocatorWithMemoryTracking.h>
 #include <Core/Types.h>
 #include <Core/Defines.h>
 #include <Core/DecimalFunctions.h>
@@ -35,7 +36,7 @@ template <typename T>
 using NearestFieldType = typename NearestFieldTypeImpl<T>::Type;
 
 class Field;
-using FieldVector = std::vector<Field>;
+using FieldVector = std::vector<Field, AllocatorWithMemoryTracking<Field>>;
 
 /// Array and Tuple use the same storage type -- FieldVector, but we declare
 /// distinct types for them, so that the caller can choose whether it wants to

tests/queries/0_stateless/01782_field_oom.sql

@@ -0,0 +1,2 @@
+SET max_memory_usage = '1G';
+SELECT sumMap([number], [number]) FROM system.numbers_mt; -- { serverError 241 }