Check malloc result for null; add TryNew* for overloads involving custom allocator

86ad7067e751c73259ce60432651ce3ecd0d057e

2 files changed, 82 insertions, 31 deletions

diff --git a/library/cpp/yt/memory/new-inl.h b/library/cpp/yt/memory/new-inl.h
index c21bd37a0c..f0646582ec 100644
--- a/library/cpp/yt/memory/new-inl.h
+++ b/library/cpp/yt/memory/new-inl.h
@@ -10,6 +10,8 @@
 
 #include <library/cpp/yt/malloc//malloc.h>
 
+#include <library/cpp/yt/system/exit.h>
+
 namespace NYT {
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -177,17 +179,41 @@ Y_FORCE_INLINE TIntrusivePtr<T> SafeConstruct(void* ptr, As&&... args)
 }
 
 template <size_t Size, size_t Alignment>
-void* AllocateConstSizeAligned()
+Y_FORCE_INLINE void* AllocateConstSizeAlignedOrCrash()
+{
+    void* ptr;
+#ifdef _win_
+    ptr = ::aligned_malloc(Size, Alignment);
+#else
+    if (Alignment <= alignof(std::max_align_t)) {
+        ptr = ::malloc(Size);
+    } else {
+        ptr = ::aligned_malloc(Size, Alignment);
+    }
+#endif
+    if (Y_UNLIKELY(!ptr)) {
+        AbortProcess(ToUnderlying(EProcessExitCode::OutOfMemory));
+    }
+    return ptr;
+}
+
+template <size_t Alignment>
+Y_FORCE_INLINE void* AllocateOrCrash(size_t size)
 {
+    void* ptr;
 #ifdef _win_
-    return ::aligned_malloc(Size, Alignment);
+    ptr = ::aligned_malloc(size, Alignment);
 #else
     if (Alignment <= alignof(std::max_align_t)) {
-        return ::malloc(Size);
+        ptr = ::malloc(size);
     } else {
-        return ::aligned_malloc(Size, Alignment);
+        ptr = ::aligned_malloc(size, Alignment);
     }
 #endif
+    if (Y_UNLIKELY(!ptr)) {
+        AbortProcess(ToUnderlying(EProcessExitCode::OutOfMemory));
+    }
+    return ptr;
 }
 
 } // namespace NDetail
@@ -198,25 +224,36 @@ template <class T, class... As, class>
 Y_FORCE_INLINE TIntrusivePtr<T> New(
     As&&... args)
 {
-    void* ptr = NYT::NDetail::AllocateConstSizeAligned<
+    void* ptr = NYT::NDetail::AllocateConstSizeAlignedOrCrash<
         NYT::NDetail::TConstructHelper<T>::Size,
         NYT::NDetail::TConstructHelper<T>::Alignment>();
-
     return NYT::NDetail::SafeConstruct<T>(ptr, std::forward<As>(args)...);
 }
 
 template <class T, class... As, class>
-Y_FORCE_INLINE TIntrusivePtr<T> New(
+Y_FORCE_INLINE TIntrusivePtr<T> TryNew(
     typename T::TAllocator* allocator,
     As&&... args)
 {
     auto* ptr = allocator->Allocate(NYT::NDetail::TConstructHelper<T>::Size);
-    if (!ptr) {
+    if (Y_UNLIKELY(!ptr)) {
         return nullptr;
     }
     return NYT::NDetail::SafeConstruct<T>(ptr, std::forward<As>(args)...);
 }
 
+template <class T, class... As, class>
+Y_FORCE_INLINE TIntrusivePtr<T> New(
+    typename T::TAllocator* allocator,
+    As&&... args)
+{
+    auto obj = TryNew<T>(allocator, std::forward<As>(args)...);
+    if (Y_UNLIKELY(!obj)) {
+        AbortProcess(ToUnderlying(EProcessExitCode::OutOfMemory));
+    }
+    return obj;
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 template <class T, class... As, class>
@@ -225,48 +262,48 @@ Y_FORCE_INLINE TIntrusivePtr<T> NewWithExtraSpace(
     As&&... args)
 {
     auto totalSize = NYT::NDetail::TConstructHelper<T>::Size + extraSpaceSize;
-    void* ptr = nullptr;
-
-#ifdef _win_
-    ptr = ::aligned_malloc(totalSize, NYT::NDetail::TConstructHelper<T>::Alignment);
-#else
-    if (NYT::NDetail::TConstructHelper<T>::Alignment <= alignof(std::max_align_t)) {
-        ptr = ::malloc(totalSize);
-    } else {
-        ptr = ::aligned_malloc(totalSize, NYT::NDetail::TConstructHelper<T>::Alignment);
-    }
-#endif
-
+    void* ptr = NYT::NDetail::AllocateOrCrash<NYT::NDetail::TConstructHelper<T>::Alignment>(totalSize);
     return NYT::NDetail::SafeConstruct<T>(ptr, std::forward<As>(args)...);
 }
 
 template <class T, class... As, class>
-Y_FORCE_INLINE TIntrusivePtr<T> NewWithExtraSpace(
+Y_FORCE_INLINE TIntrusivePtr<T> TryNewWithExtraSpace(
     typename T::TAllocator* allocator,
     size_t extraSpaceSize,
     As&&... args)
 {
     auto totalSize = NYT::NDetail::TConstructHelper<T>::Size + extraSpaceSize;
     auto* ptr = allocator->Allocate(totalSize);
-    if (!ptr) {
+    if (Y_UNLIKELY(!ptr)) {
         return nullptr;
     }
     return NYT::NDetail::SafeConstruct<T>(ptr, std::forward<As>(args)...);
 }
 
+template <class T, class... As, class>
+Y_FORCE_INLINE TIntrusivePtr<T> NewWithExtraSpace(
+    typename T::TAllocator* allocator,
+    size_t extraSpaceSize,
+    As&&... args)
+{
+    auto obj = TryNewWithExtraSpace<T>(allocator, extraSpaceSize, std::forward<As>(args)...);
+    if (Y_UNLIKELY(!obj)) {
+        AbortProcess(ToUnderlying(EProcessExitCode::OutOfMemory));
+    }
+    return obj;
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 template <class T, class TDeleter, class... As>
 Y_FORCE_INLINE TIntrusivePtr<T> NewWithDeleter(TDeleter deleter, As&&... args)
 {
     using TWrapper = TRefCountedWrapperWithDeleter<T, TDeleter>;
-    void* ptr = NYT::NDetail::AllocateConstSizeAligned<sizeof(TWrapper), alignof(TWrapper)>();
-
+    void* ptr = NYT::NDetail::AllocateConstSizeAlignedOrCrash<sizeof(TWrapper), alignof(TWrapper)>();
     auto* instance = NYT::NDetail::NewEpilogue<TWrapper>(
         ptr,
         std::move(deleter),
         std::forward<As>(args)...);
-
     return TIntrusivePtr<T>(instance, /*addReference*/ false);
 }
 
@@ -278,16 +315,13 @@ Y_FORCE_INLINE TIntrusivePtr<T> NewWithLocation(
     As&&... args)
 {
     using TWrapper = TRefCountedWrapperWithCookie<T>;
-    void* ptr = NYT::NDetail::AllocateConstSizeAligned<sizeof(TWrapper), alignof(TWrapper)>();
-
+    void* ptr = NYT::NDetail::AllocateConstSizeAlignedOrCrash<sizeof(TWrapper), alignof(TWrapper)>();
     auto* instance = NYT::NDetail::NewEpilogue<TWrapper>(ptr, std::forward<As>(args)...);
-
 #ifdef YT_ENABLE_REF_COUNTED_TRACKING
     instance->InitializeTracking(GetRefCountedTypeCookieWithLocation<T, TTag, Counter>(location));
 #else
     Y_UNUSED(location);
 #endif
-
     return TIntrusivePtr<T>(instance, /*addReference*/ false);
 }
 
diff --git a/library/cpp/yt/memory/new.h b/library/cpp/yt/memory/new.h
index a9591724ed..d022f0630f 100644
--- a/library/cpp/yt/memory/new.h
+++ b/library/cpp/yt/memory/new.h
@@ -78,26 +78,43 @@ struct THasAllocator<T, std::void_t<typename T::TAllocator>>
 
 ////////////////////////////////////////////////////////////////////////////////
 
-//! Allocates a new instance of |T|.
+//! Allocates a new instance of |T| using the standard allocator.
+//! Aborts the process on out-of-memory condition.
 template <class T, class... As, class = typename THasAllocator<T>::TFalse>
 TIntrusivePtr<T> New(As&&... args);
 
+//! Allocates a new instance of |T| using a custom #allocator.
+//! Returns null on allocation failure.
+template <class T, class... As, class = typename THasAllocator<T>::TTrue>
+TIntrusivePtr<T> TryNew(typename T::TAllocator* allocator, As&&... args);
+
+//! Same as #TryNewWit but aborts on allocation failure.
 template <class T, class... As, class = typename THasAllocator<T>::TTrue>
 TIntrusivePtr<T> New(typename T::TAllocator* allocator, As&&... args);
 
-//! Allocates an instance of |T| with additional storage of #extraSpaceSize bytes.
+//! Allocates an instance of |T|
+//! Aborts the process on out-of-memory condition.
 template <class T, class... As, class = typename THasAllocator<T>::TFalse>
 TIntrusivePtr<T> NewWithExtraSpace(size_t extraSpaceSize, As&&... args);
 
+//! Allocates a new instance of |T| with additional storage of #extraSpaceSize bytes
+//! using a custom #allocator.
+//! Returns null on allocation failure.
+template <class T, class... As, class = typename THasAllocator<T>::TTrue>
+TIntrusivePtr<T> TryNewWithExtraSpace(typename T::TAllocator* allocator, size_t extraSpaceSize, As&&... args);
+
+//! Same as #TryNewWithExtraSpace but aborts on allocation failure.
 template <class T, class... As, class = typename THasAllocator<T>::TTrue>
 TIntrusivePtr<T> NewWithExtraSpace(typename T::TAllocator* allocator, size_t extraSpaceSize, As&&... args);
 
 //! Allocates a new instance of |T| with a custom #deleter.
+//! Aborts the process on out-of-memory condition.
 template <class T, class TDeleter, class... As>
 TIntrusivePtr<T> NewWithDeleter(TDeleter deleter, As&&... args);
 
 //! Allocates a new instance of |T|.
 //! The allocation is additionally marked with #location.
+//! Aborts the process on out-of-memory condition.
 template <class T, class TTag, int Counter, class... As>
 TIntrusivePtr<T> NewWithLocation(const TSourceLocation& location, As&&... args);