From 7b1890f0b3e096a43006e4163b98f6316ac1e91b Mon Sep 17 00:00:00 2001
From: Alexey Milovidov <milovidov@yandex-team.ru>
Date: Wed, 31 Mar 2021 06:19:34 +0300
Subject: [PATCH] Fix some OOMs in stress tests

---
 src/Common/AllocatorWithMemoryTracking.h      | 62 +++++++++++++++++++
 src/Core/Field.h                              |  3 +-
 .../0_stateless/01782_field_oom.reference     |  0
 tests/queries/0_stateless/01782_field_oom.sql |  2 +
 4 files changed, 66 insertions(+), 1 deletion(-)
 create mode 100644 src/Common/AllocatorWithMemoryTracking.h
 create mode 100644 tests/queries/0_stateless/01782_field_oom.reference
 create mode 100644 tests/queries/0_stateless/01782_field_oom.sql

diff --git a/src/Common/AllocatorWithMemoryTracking.h b/src/Common/AllocatorWithMemoryTracking.h
new file mode 100644
index 00000000000..e9597e4bf5a
--- /dev/null
+++ b/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;
+}
+
diff --git a/src/Core/Field.h b/src/Core/Field.h
index 81d06693a7f..77549854982 100644
--- a/src/Core/Field.h
+++ b/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
diff --git a/tests/queries/0_stateless/01782_field_oom.reference b/tests/queries/0_stateless/01782_field_oom.reference
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/queries/0_stateless/01782_field_oom.sql b/tests/queries/0_stateless/01782_field_oom.sql
new file mode 100644
index 00000000000..60326b09a7a
--- /dev/null
+++ b/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 }