Avoid excessive allocation in Arena

src/Common/Arena.h

@@ -4,10 +4,8 @@
 #include <memory>
 #include <vector>
 #include <boost/noncopyable.hpp>
+#include <base/defines.h>
 #include <Core/Defines.h>
-#if __has_include(<sanitizer/asan_interface.h>) && defined(ADDRESS_SANITIZER)
-#   include <sanitizer/asan_interface.h>
-#endif
 #include <Common/memcpySmall.h>
 #include <Common/ProfileEvents.h>
 #include <Common/Allocator.h>
@@ -39,13 +37,36 @@ private:
     /// Contiguous MemoryChunk of memory and pointer to free space inside it. Member of single-linked list.
     struct alignas(16) MemoryChunk : private Allocator<false>    /// empty base optimization
     {
-        char * begin;
-        char * pos;
-        char * end; /// does not include padding.
+        char * begin = nullptr;
+        char * pos = nullptr;
+        char * end = nullptr; /// does not include padding.
 
-        MemoryChunk * prev;
+        std::unique_ptr<MemoryChunk> prev;
 
-        MemoryChunk(size_t size_, MemoryChunk * prev_)
+        MemoryChunk()
+        {
+        }
+
+        void swap(MemoryChunk & other)
+        {
+            std::swap(begin, other.begin);
+            std::swap(pos, other.pos);
+            std::swap(end, other.end);
+            prev.swap(other.prev);
+        }
+
+        MemoryChunk(MemoryChunk && other)
+        {
+            *this = std::move(other);
+        }
+
+        MemoryChunk & operator=(MemoryChunk && other)
+        {
+            swap(other);
+            return *this;
+        }
+
+        MemoryChunk(size_t size_)
         {
             ProfileEvents::increment(ProfileEvents::ArenaAllocChunks);
             ProfileEvents::increment(ProfileEvents::ArenaAllocBytes, size_);
@@ -53,7 +74,6 @@ private:
             begin = reinterpret_cast<char *>(Allocator<false>::alloc(size_));
             pos = begin;
             end = begin + size_ - pad_right;
-            prev = prev_;
 
             ASAN_POISON_MEMORY_REGION(begin, size_);
         }
@@ -67,19 +87,18 @@ private:
             ASAN_UNPOISON_MEMORY_REGION(begin, size());
 
             Allocator<false>::free(begin, size());
-
-            delete prev;
         }
 
         size_t size() const { return end + pad_right - begin; }
         size_t remaining() const { return end - pos; }
     };
 
+    size_t initial_size;
     size_t growth_factor;
     size_t linear_growth_threshold;
 
     /// Last contiguous MemoryChunk of memory.
-    MemoryChunk * head;
+    MemoryChunk head;
     size_t allocated_bytes;
     size_t used_bytes;
     size_t page_size;
@@ -95,9 +114,13 @@ private:
     {
         size_t size_after_grow = 0;
 
-        if (head->size() < linear_growth_threshold)
+        if (head.size() == 0)
         {
-            size_after_grow = std::max(min_next_size, head->size() * growth_factor);
+            size_after_grow = initial_size;
+        }
+        else if (head.size() < linear_growth_threshold)
+        {
+            size_after_grow = std::max(min_next_size, head.size() * growth_factor);
         }
         else
         {
@@ -119,8 +142,18 @@ private:
     /// Add next contiguous MemoryChunk of memory with size not less than specified.
     void NO_INLINE addMemoryChunk(size_t min_size)
     {
-        head = new MemoryChunk(nextSize(min_size + pad_right), head);
-        allocated_bytes += head->size();
+        size_t next_size = nextSize(min_size + pad_right);
+        if (!head.begin)
+        {
+            head = MemoryChunk(next_size);
+        }
+        else
+        {
+            auto chunk = std::make_unique<MemoryChunk>(next_size);
+            head.swap(*chunk);
+            head.prev = std::move(chunk);
+            allocated_bytes += head.size();
+        }
     }
 
     friend class ArenaAllocator;
@@ -128,29 +161,24 @@ private:
 
 public:
     explicit Arena(size_t initial_size_ = 4096, size_t growth_factor_ = 2, size_t linear_growth_threshold_ = 128 * 1024 * 1024)
-        : growth_factor(growth_factor_)
+        : initial_size(initial_size_)
+        , growth_factor(growth_factor_)
         , linear_growth_threshold(linear_growth_threshold_)
-        , head(new MemoryChunk(initial_size_, nullptr))
-        , allocated_bytes(head->size())
+        , allocated_bytes(head.size())
         , used_bytes(0)
         , page_size(static_cast<size_t>(::getPageSize()))
     {
     }
 
-    ~Arena()
-    {
-        delete head;
-    }
-
     /// Get piece of memory, without alignment.
     char * alloc(size_t size)
     {
         used_bytes += size;
-        if (unlikely(static_cast<std::ptrdiff_t>(size) > head->end - head->pos))
+        if (unlikely(static_cast<std::ptrdiff_t>(size) > head.end - head.pos))
             addMemoryChunk(size);
 
-        char * res = head->pos;
-        head->pos += size;
+        char * res = head.pos;
+        head.pos += size;
         ASAN_UNPOISON_MEMORY_REGION(res, size + pad_right);
         return res;
     }
@@ -161,14 +189,14 @@ public:
         used_bytes += size;
         do
         {
-            void * head_pos = head->pos;
-            size_t space = head->end - head->pos;
+            void * head_pos = head.pos;
+            size_t space = head.end - head.pos;
 
             auto * res = static_cast<char *>(std::align(alignment, size, head_pos, space));
             if (res)
             {
-                head->pos = static_cast<char *>(head_pos);
-                head->pos += size;
+                head.pos = static_cast<char *>(head_pos);
+                head.pos += size;
                 ASAN_UNPOISON_MEMORY_REGION(res, size + pad_right);
                 return res;
             }
@@ -191,9 +219,9 @@ public:
     void * rollback(size_t size)
     {
         used_bytes -= size;
-        head->pos -= size;
-        ASAN_POISON_MEMORY_REGION(head->pos, size + pad_right);
-        return head->pos;
+        head.pos -= size;
+        ASAN_POISON_MEMORY_REGION(head.pos, size + pad_right);
+        return head.pos;
     }
 
     /** Begin or expand a contiguous range of memory.
@@ -234,10 +262,10 @@ public:
         // This method only works for extending the last allocation. For lack of
         // original size, check a weaker condition: that 'begin' is at least in
         // the current MemoryChunk.
-        assert(range_start >= head->begin);
-        assert(range_start < head->end);
+        assert(range_start >= head.begin);
+        assert(range_start < head.end);
 
-        if (head->pos + additional_bytes <= head->end)
+        if (head.pos + additional_bytes <= head.end)
         {
             // The new size fits into the last MemoryChunk, so just alloc the
             // additional size. We can alloc without alignment here, because it
@@ -254,7 +282,7 @@ public:
         // solved not by complicating this method, but by rethinking the
         // approach to memory management for aggregate function states, so that
         // we can provide a proper realloc().
-        const size_t existing_bytes = head->pos - range_start;
+        const size_t existing_bytes = head.pos - range_start;
         const size_t new_bytes = existing_bytes + additional_bytes;
         const char * old_range = range_start;
 
@@ -317,12 +345,11 @@ public:
     /// yourself having to use this method, probably you're doing something wrong.
     size_t remainingSpaceInCurrentMemoryChunk() const
     {
-        return head->remaining();
+        return head.remaining();
     }
 };
 
 using ArenaPtr = std::shared_ptr<Arena>;
 using Arenas = std::vector<ArenaPtr>;
 
-
 }

src/Common/ArenaAllocator.h

@@ -20,7 +20,7 @@ public:
         char const * data = reinterpret_cast<char *>(buf);
 
         // Invariant should be maintained: new_size > old_size
-        if (data + old_size == arena->head->pos)
+        if (data + old_size == arena->head.pos)
         {
             // Consecutive optimization
             arena->allocContinue(new_size - old_size, data);
@@ -59,7 +59,7 @@ public:
     {
         char const * data = reinterpret_cast<char *>(buf);
 
-        if (data + old_size == arena->head->pos)
+        if (data + old_size == arena->head.pos)
         {
             arena->allocContinue(new_size - old_size, data, alignment);
             return reinterpret_cast<void *>(const_cast<char *>(data));