Skip to content

Commit 818158c

Browse files
author
babenko
committed
Check malloc result for null; add TryNew* for overloads involving custom allocator
86ad7067e751c73259ce60432651ce3ecd0d057e
1 parent 29e274a commit 818158c

File tree

2 files changed

+82
-31
lines changed

2 files changed

+82
-31
lines changed

‎library/cpp/yt/memory/new-inl.h

Lines changed: 63 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@
1010

1111
#include <library/cpp/yt/malloc//malloc.h>
1212

13+
#include <library/cpp/yt/system/exit.h>
14+
1315
namespace NYT {
1416

1517
////////////////////////////////////////////////////////////////////////////////
@@ -177,17 +179,41 @@ Y_FORCE_INLINE TIntrusivePtr<T> SafeConstruct(void* ptr, As&&... args)
177179
}
178180

179181
template <size_t Size, size_t Alignment>
180-
void* AllocateConstSizeAligned()
182+
Y_FORCE_INLINE void* AllocateConstSizeAlignedOrCrash()
183+
{
184+
void* ptr;
185+
#ifdef _win_
186+
ptr = ::aligned_malloc(Size, Alignment);
187+
#else
188+
if (Alignment <= alignof(std::max_align_t)) {
189+
ptr = ::malloc(Size);
190+
} else {
191+
ptr = ::aligned_malloc(Size, Alignment);
192+
}
193+
#endif
194+
if (Y_UNLIKELY(!ptr)) {
195+
AbortProcess(ToUnderlying(EProcessExitCode::OutOfMemory));
196+
}
197+
return ptr;
198+
}
199+
200+
template <size_t Alignment>
201+
Y_FORCE_INLINE void* AllocateOrCrash(size_t size)
181202
{
203+
void* ptr;
182204
#ifdef _win_
183-
return ::aligned_malloc(Size, Alignment);
205+
ptr = ::aligned_malloc(size, Alignment);
184206
#else
185207
if (Alignment <= alignof(std::max_align_t)) {
186-
return ::malloc(Size);
208+
ptr = ::malloc(size);
187209
} else {
188-
return ::aligned_malloc(Size, Alignment);
210+
ptr = ::aligned_malloc(size, Alignment);
189211
}
190212
#endif
213+
if (Y_UNLIKELY(!ptr)) {
214+
AbortProcess(ToUnderlying(EProcessExitCode::OutOfMemory));
215+
}
216+
return ptr;
191217
}
192218

193219
} // namespace NDetail
@@ -198,25 +224,36 @@ template <class T, class... As, class>
198224
Y_FORCE_INLINE TIntrusivePtr<T> New(
199225
As&&... args)
200226
{
201-
void* ptr = NYT::NDetail::AllocateConstSizeAligned<
227+
void* ptr = NYT::NDetail::AllocateConstSizeAlignedOrCrash<
202228
NYT::NDetail::TConstructHelper<T>::Size,
203229
NYT::NDetail::TConstructHelper<T>::Alignment>();
204-
205230
return NYT::NDetail::SafeConstruct<T>(ptr, std::forward<As>(args)...);
206231
}
207232

208233
template <class T, class... As, class>
209-
Y_FORCE_INLINE TIntrusivePtr<T> New(
234+
Y_FORCE_INLINE TIntrusivePtr<T> TryNew(
210235
typename T::TAllocator* allocator,
211236
As&&... args)
212237
{
213238
auto* ptr = allocator->Allocate(NYT::NDetail::TConstructHelper<T>::Size);
214-
if (!ptr) {
239+
if (Y_UNLIKELY(!ptr)) {
215240
return nullptr;
216241
}
217242
return NYT::NDetail::SafeConstruct<T>(ptr, std::forward<As>(args)...);
218243
}
219244

245+
template <class T, class... As, class>
246+
Y_FORCE_INLINE TIntrusivePtr<T> New(
247+
typename T::TAllocator* allocator,
248+
As&&... args)
249+
{
250+
auto obj = TryNew<T>(allocator, std::forward<As>(args)...);
251+
if (Y_UNLIKELY(!obj)) {
252+
AbortProcess(ToUnderlying(EProcessExitCode::OutOfMemory));
253+
}
254+
return obj;
255+
}
256+
220257
////////////////////////////////////////////////////////////////////////////////
221258

222259
template <class T, class... As, class>
@@ -225,48 +262,48 @@ Y_FORCE_INLINE TIntrusivePtr<T> NewWithExtraSpace(
225262
As&&... args)
226263
{
227264
auto totalSize = NYT::NDetail::TConstructHelper<T>::Size + extraSpaceSize;
228-
void* ptr = nullptr;
229-
230-
#ifdef _win_
231-
ptr = ::aligned_malloc(totalSize, NYT::NDetail::TConstructHelper<T>::Alignment);
232-
#else
233-
if (NYT::NDetail::TConstructHelper<T>::Alignment <= alignof(std::max_align_t)) {
234-
ptr = ::malloc(totalSize);
235-
} else {
236-
ptr = ::aligned_malloc(totalSize, NYT::NDetail::TConstructHelper<T>::Alignment);
237-
}
238-
#endif
239-
265+
void* ptr = NYT::NDetail::AllocateOrCrash<NYT::NDetail::TConstructHelper<T>::Alignment>(totalSize);
240266
return NYT::NDetail::SafeConstruct<T>(ptr, std::forward<As>(args)...);
241267
}
242268

243269
template <class T, class... As, class>
244-
Y_FORCE_INLINE TIntrusivePtr<T> NewWithExtraSpace(
270+
Y_FORCE_INLINE TIntrusivePtr<T> TryNewWithExtraSpace(
245271
typename T::TAllocator* allocator,
246272
size_t extraSpaceSize,
247273
As&&... args)
248274
{
249275
auto totalSize = NYT::NDetail::TConstructHelper<T>::Size + extraSpaceSize;
250276
auto* ptr = allocator->Allocate(totalSize);
251-
if (!ptr) {
277+
if (Y_UNLIKELY(!ptr)) {
252278
return nullptr;
253279
}
254280
return NYT::NDetail::SafeConstruct<T>(ptr, std::forward<As>(args)...);
255281
}
256282

283+
template <class T, class... As, class>
284+
Y_FORCE_INLINE TIntrusivePtr<T> NewWithExtraSpace(
285+
typename T::TAllocator* allocator,
286+
size_t extraSpaceSize,
287+
As&&... args)
288+
{
289+
auto obj = TryNewWithExtraSpace<T>(allocator, extraSpaceSize, std::forward<As>(args)...);
290+
if (Y_UNLIKELY(!obj)) {
291+
AbortProcess(ToUnderlying(EProcessExitCode::OutOfMemory));
292+
}
293+
return obj;
294+
}
295+
257296
////////////////////////////////////////////////////////////////////////////////
258297

259298
template <class T, class TDeleter, class... As>
260299
Y_FORCE_INLINE TIntrusivePtr<T> NewWithDeleter(TDeleter deleter, As&&... args)
261300
{
262301
using TWrapper = TRefCountedWrapperWithDeleter<T, TDeleter>;
263-
void* ptr = NYT::NDetail::AllocateConstSizeAligned<sizeof(TWrapper), alignof(TWrapper)>();
264-
302+
void* ptr = NYT::NDetail::AllocateConstSizeAlignedOrCrash<sizeof(TWrapper), alignof(TWrapper)>();
265303
auto* instance = NYT::NDetail::NewEpilogue<TWrapper>(
266304
ptr,
267305
std::move(deleter),
268306
std::forward<As>(args)...);
269-
270307
return TIntrusivePtr<T>(instance, /*addReference*/ false);
271308
}
272309

@@ -278,16 +315,13 @@ Y_FORCE_INLINE TIntrusivePtr<T> NewWithLocation(
278315
As&&... args)
279316
{
280317
using TWrapper = TRefCountedWrapperWithCookie<T>;
281-
void* ptr = NYT::NDetail::AllocateConstSizeAligned<sizeof(TWrapper), alignof(TWrapper)>();
282-
318+
void* ptr = NYT::NDetail::AllocateConstSizeAlignedOrCrash<sizeof(TWrapper), alignof(TWrapper)>();
283319
auto* instance = NYT::NDetail::NewEpilogue<TWrapper>(ptr, std::forward<As>(args)...);
284-
285320
#ifdef YT_ENABLE_REF_COUNTED_TRACKING
286321
instance->InitializeTracking(GetRefCountedTypeCookieWithLocation<T, TTag, Counter>(location));
287322
#else
288323
Y_UNUSED(location);
289324
#endif
290-
291325
return TIntrusivePtr<T>(instance, /*addReference*/ false);
292326
}
293327

‎library/cpp/yt/memory/new.h

Lines changed: 19 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -78,26 +78,43 @@ struct THasAllocator<T, std::void_t<typename T::TAllocator>>
7878

7979
////////////////////////////////////////////////////////////////////////////////
8080

81-
//! Allocates a new instance of |T|.
81+
//! Allocates a new instance of |T| using the standard allocator.
82+
//! Aborts the process on out-of-memory condition.
8283
template <class T, class... As, class = typename THasAllocator<T>::TFalse>
8384
TIntrusivePtr<T> New(As&&... args);
8485

86+
//! Allocates a new instance of |T| using a custom #allocator.
87+
//! Returns null on allocation failure.
88+
template <class T, class... As, class = typename THasAllocator<T>::TTrue>
89+
TIntrusivePtr<T> TryNew(typename T::TAllocator* allocator, As&&... args);
90+
91+
//! Same as #TryNewWit but aborts on allocation failure.
8592
template <class T, class... As, class = typename THasAllocator<T>::TTrue>
8693
TIntrusivePtr<T> New(typename T::TAllocator* allocator, As&&... args);
8794

88-
//! Allocates an instance of |T| with additional storage of #extraSpaceSize bytes.
95+
//! Allocates an instance of |T|
96+
//! Aborts the process on out-of-memory condition.
8997
template <class T, class... As, class = typename THasAllocator<T>::TFalse>
9098
TIntrusivePtr<T> NewWithExtraSpace(size_t extraSpaceSize, As&&... args);
9199

100+
//! Allocates a new instance of |T| with additional storage of #extraSpaceSize bytes
101+
//! using a custom #allocator.
102+
//! Returns null on allocation failure.
103+
template <class T, class... As, class = typename THasAllocator<T>::TTrue>
104+
TIntrusivePtr<T> TryNewWithExtraSpace(typename T::TAllocator* allocator, size_t extraSpaceSize, As&&... args);
105+
106+
//! Same as #TryNewWithExtraSpace but aborts on allocation failure.
92107
template <class T, class... As, class = typename THasAllocator<T>::TTrue>
93108
TIntrusivePtr<T> NewWithExtraSpace(typename T::TAllocator* allocator, size_t extraSpaceSize, As&&... args);
94109

95110
//! Allocates a new instance of |T| with a custom #deleter.
111+
//! Aborts the process on out-of-memory condition.
96112
template <class T, class TDeleter, class... As>
97113
TIntrusivePtr<T> NewWithDeleter(TDeleter deleter, As&&... args);
98114

99115
//! Allocates a new instance of |T|.
100116
//! The allocation is additionally marked with #location.
117+
//! Aborts the process on out-of-memory condition.
101118
template <class T, class TTag, int Counter, class... As>
102119
TIntrusivePtr<T> NewWithLocation(const TSourceLocation& location, As&&... args);
103120

0 commit comments

Comments
 (0)