Add new/delete overloads

src/Common/Concepts.h

@@ -0,0 +1,14 @@
+#pragma once
+
+#include <concepts>
+
+namespace DB
+{
+
+template <typename... T>
+concept OptionalArgument = requires(T &&...)
+{
+    requires(sizeof...(T) == 0 || sizeof...(T) == 1);
+};
+
+}

src/Common/memory.h

@@ -1,8 +1,11 @@
 #pragma once
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wold-style-cast"
 
 #include <new>
 #include <base/defines.h>
 
+#include <Common/Concepts.h>
 #include <Common/CurrentMemoryTracker.h>
 #include <Common/config.h>
 
@@ -14,13 +17,24 @@
 #    include <cstdlib>
 #endif
 
-
 namespace Memory
 {
 
-inline ALWAYS_INLINE void * newImpl(std::size_t size)
+inline ALWAYS_INLINE size_t alignToSizeT(std::align_val_t align) noexcept
 {
-    auto * ptr = malloc(size);
+    return static_cast<size_t>(align);
+}
+
+template <std::same_as<std::align_val_t>... TAlign>
+requires DB::OptionalArgument<TAlign...>
+inline ALWAYS_INLINE void * newImpl(std::size_t size, TAlign... align)
+{
+    void * ptr = nullptr;
+    if constexpr (sizeof...(TAlign) == 1)
+        ptr = aligned_alloc(alignToSizeT(align...), size);
+    else
+        ptr = malloc(size);
+
     if (likely(ptr != nullptr))
         return ptr;
 
@@ -33,6 +47,11 @@ inline ALWAYS_INLINE void * newNoExept(std::size_t size) noexcept
     return malloc(size);
 }
 
+inline ALWAYS_INLINE void * newNoExept(std::size_t size, std::align_val_t align) noexcept
+{
+    return aligned_alloc(static_cast<size_t>(align), size);
+}
+
 inline ALWAYS_INLINE void deleteImpl(void * ptr) noexcept
 {
     free(ptr);
@@ -40,17 +59,24 @@ inline ALWAYS_INLINE void deleteImpl(void * ptr) noexcept
 
 #if USE_JEMALLOC
 
-inline ALWAYS_INLINE void deleteSized(void * ptr, std::size_t size) noexcept
+template <std::same_as<std::align_val_t>... TAlign>
+requires DB::OptionalArgument<TAlign...>
+inline ALWAYS_INLINE void deleteSized(void * ptr, std::size_t size, TAlign... align) noexcept
 {
     if (unlikely(ptr == nullptr))
         return;
 
+    if constexpr (sizeof...(TAlign) == 1)
+        sdallocx(ptr, size, MALLOCX_ALIGN(alignToSizeT(align...)));
+    else
         sdallocx(ptr, size, 0);
 }
 
 #else
 
-inline ALWAYS_INLINE void deleteSized(void * ptr, std::size_t size [[maybe_unused]]) noexcept
+template <std::same_as<std::align_val_t>... TAlign>
+requires DB::OptionalArgument<TAlign...>
+inline ALWAYS_INLINE void deleteSized(void * ptr, std::size_t size [[maybe_unused]], TAlign... /* align */) noexcept
 {
     free(ptr);
 }
@@ -63,8 +89,9 @@ inline ALWAYS_INLINE void deleteSized(void * ptr, std::size_t size [[maybe_unuse
 #    include <malloc/malloc.h>
 #endif
 
-
+template <std::same_as<std::align_val_t>... TAlign>
-inline ALWAYS_INLINE size_t getActualAllocationSize(size_t size)
+requires DB::OptionalArgument<TAlign...>
+inline ALWAYS_INLINE size_t getActualAllocationSize(size_t size, TAlign... align)
 {
     size_t actual_size = size;
 
@@ -72,26 +99,41 @@ inline ALWAYS_INLINE size_t getActualAllocationSize(size_t size)
     /// The nallocx() function allocates no memory, but it performs the same size computation as the mallocx() function
     /// @note je_mallocx() != je_malloc(). It's expected they don't differ much in allocation logic.
     if (likely(size != 0))
+    {
+        if constexpr (sizeof...(TAlign) == 1)
+            actual_size = nallocx(size, MALLOCX_ALIGN(alignToSizeT(align...)));
+        else
             actual_size = nallocx(size, 0);
+    }
 #endif
 
     return actual_size;
 }
 
-inline ALWAYS_INLINE void trackMemory(std::size_t size)
+template <std::same_as<std::align_val_t>... TAlign>
+requires DB::OptionalArgument<TAlign...>
+inline ALWAYS_INLINE void trackMemory(std::size_t size, TAlign... align)
 {
-    std::size_t actual_size = getActualAllocationSize(size);
+    std::size_t actual_size = getActualAllocationSize(size, align...);
     CurrentMemoryTracker::allocNoThrow(actual_size);
 }
 
-inline ALWAYS_INLINE void untrackMemory(void * ptr [[maybe_unused]], std::size_t size [[maybe_unused]] = 0) noexcept
+template <std::same_as<std::align_val_t>... TAlign>
+requires DB::OptionalArgument<TAlign...>
+inline ALWAYS_INLINE void untrackMemory(void * ptr [[maybe_unused]], std::size_t size [[maybe_unused]] = 0, TAlign... align) noexcept
 {
     try
     {
 #if USE_JEMALLOC
+
         /// @note It's also possible to use je_malloc_usable_size() here.
         if (likely(ptr != nullptr))
+        {
+            if constexpr (sizeof...(TAlign) == 1)
+                CurrentMemoryTracker::free(sallocx(ptr, MALLOCX_ALIGN(alignToSizeT(align...))));
+            else
                 CurrentMemoryTracker::free(sallocx(ptr, 0));
+        }
 #else
         if (size)
             CurrentMemoryTracker::free(size);
@@ -103,7 +145,10 @@ inline ALWAYS_INLINE void untrackMemory(void * ptr [[maybe_unused]], std::size_t
 #endif
     }
     catch (...)
-    {}
+    {
+    }
 }
 
 }
+
+#pragma GCC diagnostic pop

src/Common/new_delete.cpp

@@ -1,6 +1,7 @@
-#include <Common/memory.h>
+#include <cassert>
-#include <Common/config.h>
 #include <new>
+#include <Common/config.h>
+#include <Common/memory.h>
 
 #if defined(OS_DARWIN) && (USE_JEMALLOC)
 /// In case of OSX jemalloc register itself as a default zone allocator.
@@ -53,12 +54,24 @@ void * operator new(std::size_t size)
     return Memory::newImpl(size);
 }
 
+void * operator new(std::size_t size, std::align_val_t align)
+{
+    Memory::trackMemory(size, align);
+    return Memory::newImpl(size, align);
+}
+
 void * operator new[](std::size_t size)
 {
     Memory::trackMemory(size);
     return Memory::newImpl(size);
 }
 
+void * operator new[](std::size_t size, std::align_val_t align)
+{
+    Memory::trackMemory(size, align);
+    return Memory::newImpl(size, align);
+}
+
 void * operator new(std::size_t size, const std::nothrow_t &) noexcept
 {
     Memory::trackMemory(size);
@@ -71,6 +84,18 @@ void * operator new[](std::size_t size, const std::nothrow_t &) noexcept
     return Memory::newNoExept(size);
 }
 
+void * operator new(std::size_t size, std::align_val_t align, const std::nothrow_t &) noexcept
+{
+    Memory::trackMemory(size, align);
+    return Memory::newNoExept(size, align);
+}
+
+void * operator new[](std::size_t size, std::align_val_t align, const std::nothrow_t &) noexcept
+{
+    Memory::trackMemory(size, align);
+    return Memory::newNoExept(size, align);
+}
+
 /// delete
 
 /// C++17 std 21.6.2.1 (11)
@@ -81,26 +106,51 @@ void * operator new[](std::size_t size, const std::nothrow_t &) noexcept
 /// It's unspecified whether size-aware or size-unaware version is called when deleting objects of
 /// incomplete type and arrays of non-class and trivially-destructible class types.
 
+
 void operator delete(void * ptr) noexcept
 {
     Memory::untrackMemory(ptr);
     Memory::deleteImpl(ptr);
 }
 
+void operator delete(void * ptr, std::align_val_t align) noexcept
+{
+    Memory::untrackMemory(ptr, 0, align);
+    Memory::deleteImpl(ptr);
+}
+
 void operator delete[](void * ptr) noexcept
 {
     Memory::untrackMemory(ptr);
     Memory::deleteImpl(ptr);
 }
 
+void operator delete[](void * ptr, std::align_val_t align) noexcept
+{
+    Memory::untrackMemory(ptr, 0, align);
+    Memory::deleteImpl(ptr);
+}
+
 void operator delete(void * ptr, std::size_t size) noexcept
 {
     Memory::untrackMemory(ptr, size);
     Memory::deleteSized(ptr, size);
 }
 
+void operator delete(void * ptr, std::size_t size, std::align_val_t align) noexcept
+{
+    Memory::untrackMemory(ptr, size, align);
+    Memory::deleteSized(ptr, size, align);
+}
+
 void operator delete[](void * ptr, std::size_t size) noexcept
 {
     Memory::untrackMemory(ptr, size);
     Memory::deleteSized(ptr, size);
 }
+
+void operator delete[](void * ptr, std::size_t size, std::align_val_t align) noexcept
+{
+    Memory::untrackMemory(ptr, size, align);
+    Memory::deleteSized(ptr, size, align);
+}