Pass continuations for success and failure cases (#788)

This PR provides a templated parameter to the allocation routines. This can be used to add special behaviour in
both the successful allocation behaviour, and in the failing
to allocate cases.

The intent of this is two enable two future features
* set_new_handler - so that the failure case doesn't just set ENOMEM, and return nullptr.  But can handle both the Windows and C++ versions of (_)set_new_handler.
* The success handler can be used to add checking, zeroing and in the future storing precise size information in metadata for each allocation.
This commit is contained in:
Matthew Parkinson
2025-07-03 20:05:52 +01:00
committed by GitHub
parent 012138e29f
commit d4c7e01cd7
14 changed files with 145 additions and 119 deletions

View File

@@ -166,7 +166,7 @@ void snmalloc_random_walk(
case EventKind::AllocZero:
{
auto ptr =
static_cast<char*>(scoped->alloc<snmalloc::YesZero>(e.size_or_index));
static_cast<char*>(scoped->alloc<snmalloc::Zero>(e.size_or_index));
results.emplace_back(0, ptr, e.size_or_index);
break;
}
@@ -174,7 +174,7 @@ void snmalloc_random_walk(
case EventKind::AllocNoZero:
{
auto ptr =
static_cast<char*>(scoped->alloc<snmalloc::NoZero>(e.size_or_index));
static_cast<char*>(scoped->alloc<snmalloc::Uninit>(e.size_or_index));
std::fill(ptr, ptr + e.size_or_index, e.filler);
results.emplace_back(e.filler, ptr, e.size_or_index);
break;

View File

@@ -59,7 +59,6 @@ namespace snmalloc
if (p == nullptr)
{
errno = ENOMEM;
return nullptr;
}
@@ -112,7 +111,6 @@ namespace snmalloc
if (meta == nullptr)
{
errno = ENOMEM;
return {nullptr, nullptr};
}
@@ -124,7 +122,6 @@ namespace snmalloc
if (p == nullptr)
{
local_state.get_meta_range().dealloc_range(meta_cap, meta_size);
errno = ENOMEM;
#ifdef SNMALLOC_TRACING
message<1024>("Out of memory");
#endif

View File

@@ -321,24 +321,24 @@ namespace snmalloc
return sizeclass_full_to_size(entry.get_sizeclass());
}
template<size_t size, ZeroMem zero_mem = NoZero, size_t align = 1>
template<size_t size, typename Conts = Uninit, size_t align = 1>
SNMALLOC_FAST_PATH_INLINE void* alloc()
{
return ThreadAlloc::get().alloc<zero_mem, ThreadAlloc::CheckInit>(
return ThreadAlloc::get().alloc<Conts, ThreadAlloc::CheckInit>(
aligned_size(align, size));
}
template<ZeroMem zero_mem = NoZero, size_t align = 1>
template<typename Conts = Uninit, size_t align = 1>
SNMALLOC_FAST_PATH_INLINE void* alloc(size_t size)
{
return ThreadAlloc::get().alloc<zero_mem, ThreadAlloc::CheckInit>(
return ThreadAlloc::get().alloc<Conts, ThreadAlloc::CheckInit>(
aligned_size(align, size));
}
template<ZeroMem zero_mem = NoZero>
template<typename Conts = Uninit>
SNMALLOC_FAST_PATH_INLINE void* alloc_aligned(size_t align, size_t size)
{
return ThreadAlloc::get().alloc<zero_mem, ThreadAlloc::CheckInit>(
return ThreadAlloc::get().alloc<Conts, ThreadAlloc::CheckInit>(
aligned_size(align, size));
}

View File

@@ -47,7 +47,7 @@ namespace snmalloc::libc
{
return set_error();
}
return alloc<ZeroMem::YesZero>(sz);
return alloc<Zero>(sz);
}
SNMALLOC_FAST_PATH_INLINE void* realloc(void* ptr, size_t size)

View File

@@ -20,29 +20,55 @@
namespace snmalloc
{
inline static SNMALLOC_FAST_PATH capptr::Alloc<void>
finish_alloc_no_zero(freelist::HeadPtr p, smallsizeclass_t sizeclass)
/**
* @brief This is the default continuation class used for handling successful,
* and failing allocations.
*
* @tparam zero_mem Specifies whether the allocated memory should be zeroed.
*
* The failure case sets errno to ENOMEM, as per the C standard.
*/
template<ZeroMem zero_mem = ZeroMem::NoZero>
class DefaultConts
{
SNMALLOC_ASSERT(is_start_of_object(
sizeclass_t::from_small_class(sizeclass), address_cast(p)));
UNUSED(sizeclass);
public:
static void* success(void* p, size_t size, bool secondary_allocator = false)
{
UNUSED(secondary_allocator);
SNMALLOC_ASSERT(p != nullptr);
return p.as_void();
}
SNMALLOC_ASSERT(
secondary_allocator ||
is_start_of_object(size_to_sizeclass_full(size), address_cast(p)));
template<ZeroMem zero_mem, typename Config>
inline static SNMALLOC_FAST_PATH capptr::Alloc<void>
finish_alloc(freelist::HeadPtr p, smallsizeclass_t sizeclass)
if (zero_mem == ZeroMem::YesZero)
{
// Zero the memory if requested.
return ::memset(p, 0, round_size(size));
}
return p;
}
static void* failure(size_t size)
{
UNUSED(size);
// If we are here, then the allocation failed.
// Set errno to ENOMEM, as per the C standard.
errno = ENOMEM;
// Return nullptr on failure.
return nullptr;
}
};
using Zero = DefaultConts<ZeroMem::YesZero>;
using Uninit = DefaultConts<ZeroMem::NoZero>;
template<typename Conts>
inline static SNMALLOC_FAST_PATH void*
finish_alloc(freelist::HeadPtr p, size_t size)
{
auto r = finish_alloc_no_zero(p, sizeclass);
if constexpr (zero_mem == YesZero)
Config::Pal::zero(r.unsafe_ptr(), sizeclass_to_size(sizeclass));
// TODO: Should this be zeroing the free Object state, in the non-zeroing
// case?
return r;
return Conts::success(capptr_reveal(p.as_void()), size, false);
}
struct FastFreeLists
@@ -568,11 +594,14 @@ namespace snmalloc
* @param size The size of the block to allocate.
* @return A pointer to the allocated block, or nullptr if the allocation
* failed.
* @tparam zero_mem Whether to zero the allocated memory.
* @tparam Conts : Carries two function success and failure, that are used
* to either apply additional operations in the case of success, such as
* zeroing, or to handle the failure case, such as set_new_handler, or
* throwing an exception or setting errno to ENOMEM.
* @tparam CheckInit is a class that can handle lazy initiasation if
* required. It is defaulted to do nothing.
*/
template<ZeroMem zero_mem = NoZero, typename CheckInit = CheckInitNoOp>
template<typename Conts = Uninit, typename CheckInit = CheckInitNoOp>
SNMALLOC_FAST_PATH ALLOCATOR void* alloc(size_t size)
{
// Perform the - 1 on size, so that zero wraps around and ends up on
@@ -582,17 +611,17 @@ namespace snmalloc
{
// Small allocations are more likely. Improve
// branch prediction by placing this case first.
return capptr_reveal(small_alloc<zero_mem, CheckInit>(size));
return small_alloc<Conts, CheckInit>(size);
}
return capptr_reveal(alloc_not_small<zero_mem, CheckInit>(size, this));
return alloc_not_small<Conts, CheckInit>(size, this);
}
/**
* Fast allocation for small objects.
*/
template<ZeroMem zero_mem, typename CheckInit>
SNMALLOC_FAST_PATH capptr::Alloc<void> small_alloc(size_t size)
template<typename Conts, typename CheckInit>
SNMALLOC_FAST_PATH void* small_alloc(size_t size)
{
auto domesticate =
[this](freelist::QueuePtr p) SNMALLOC_FAST_PATH_LAMBDA {
@@ -605,17 +634,21 @@ namespace snmalloc
if (SNMALLOC_LIKELY(!fl->empty()))
{
auto p = fl->take(key, domesticate);
return finish_alloc<zero_mem, Config>(p, sizeclass);
return finish_alloc<Conts>(p, size);
}
return handle_message_queue(
[](Allocator* alloc, smallsizeclass_t sizeclass, freelist::Iter<>* fl)
-> capptr::Alloc<void> {
return alloc->small_refill<zero_mem, CheckInit>(sizeclass, *fl);
[](
Allocator* alloc,
smallsizeclass_t sizeclass,
freelist::Iter<>* fl,
size_t size) -> void* {
return alloc->small_refill<Conts, CheckInit>(sizeclass, *fl, size);
},
this,
sizeclass,
fl);
fl,
size);
}
/**
@@ -627,8 +660,8 @@ namespace snmalloc
* code quality as the calling code already has size in the first argument
* register.
*/
template<ZeroMem zero_mem, typename CheckInit>
static SNMALLOC_SLOW_PATH capptr::Alloc<void>
template<typename Conts, typename CheckInit>
static SNMALLOC_SLOW_PATH void*
alloc_not_small(size_t size, Allocator* self)
{
if (size == 0)
@@ -636,19 +669,18 @@ namespace snmalloc
// Deal with alloc zero of with a small object here.
// Alternative semantics giving nullptr is also allowed by the
// standard.
return self->small_alloc<NoZero, CheckInit>(1);
return self->small_alloc<Conts, CheckInit>(1);
}
return self->handle_message_queue(
[](Allocator* self, size_t size) -> capptr::Alloc<void> {
[](Allocator* self, size_t size) -> void* {
return CheckInit::check_init(
[self, size]() {
if (size > bits::one_at_bit(bits::BITS - 1))
{
// Cannot allocate something that is more that half the size of
// the address space
errno = ENOMEM;
return capptr::Alloc<void>{nullptr};
return Conts::failure(size);
}
// Check if secondary allocator wants to offer the memory
@@ -658,9 +690,7 @@ namespace snmalloc
});
if (result != nullptr)
{
if constexpr (zero_mem == YesZero)
Config::Pal::zero(result, size);
return capptr::Alloc<void>::unsafe_from(result);
return Conts::success(result, size, true);
}
// Grab slab of correct size
@@ -677,6 +707,9 @@ namespace snmalloc
"size {} pow2size {}", size, bits::next_pow2_bits(size));
#endif
SNMALLOC_ASSERT(
(chunk.unsafe_ptr() == nullptr) == (meta == nullptr));
// set up meta data so sizeclass is correct, and hence alloc size,
// and external pointer. Initialise meta data for a successful
// large allocation.
@@ -685,19 +718,18 @@ namespace snmalloc
meta->initialise_large(
address_cast(chunk), freelist::Object::key_root);
self->laden.insert(meta);
// Make capptr system happy we are allowed to pass this to
// `success`.
auto p = capptr_reveal(
capptr_chunk_is_alloc(capptr_to_user_address_control(chunk)));
return Conts::success(p, size);
}
if (zero_mem == YesZero && chunk.unsafe_ptr() != nullptr)
{
Config::Pal::template zero<false>(
chunk.unsafe_ptr(), bits::next_pow2(size));
}
return capptr_chunk_is_alloc(
capptr_to_user_address_control(chunk));
return Conts::failure(size);
},
[](Allocator* a, size_t size) {
return alloc_not_small<zero_mem, CheckInitNoOp>(size, a);
return alloc_not_small<Conts, CheckInitNoOp>(size, a);
},
size);
},
@@ -705,20 +737,18 @@ namespace snmalloc
size);
}
template<ZeroMem zero_mem, typename CheckInit>
SNMALLOC_FAST_PATH capptr::Alloc<void>
small_refill(smallsizeclass_t sizeclass, freelist::Iter<>& fast_free_list)
template<typename Conts, typename CheckInit>
SNMALLOC_FAST_PATH void* small_refill(
smallsizeclass_t sizeclass, freelist::Iter<>& fast_free_list, size_t size)
{
void* result = Config::SecondaryAllocator::allocate(
[sizeclass]() -> stl::Pair<size_t, size_t> {
auto size = sizeclass_to_size(sizeclass);
[size]() -> stl::Pair<size_t, size_t> {
return {size, natural_alignment(size)};
});
if (result != nullptr)
{
if constexpr (zero_mem == YesZero)
Config::Pal::zero(result, sizeclass_to_size(sizeclass));
result = Conts::success(result, size, true);
// We need to check for initialisation here in the case where
// this is the first allocation in the system, so snmalloc has
@@ -726,10 +756,8 @@ namespace snmalloc
// deallocated, before snmalloc is initialised, then it will fail
// to access the pagemap.
return CheckInit::check_init(
[result]() { return capptr::Alloc<void>::unsafe_from(result); },
[](Allocator*, void* result) {
return capptr::Alloc<void>::unsafe_from(result);
},
[result]() { return result; },
[](Allocator*, void* result) { return result; },
result);
}
@@ -743,8 +771,8 @@ namespace snmalloc
if (SNMALLOC_UNLIKELY(alloc_classes[sizeclass].length == 1))
{
if (entropy.next_bit() == 0)
return small_refill_slow<zero_mem, CheckInit>(
sizeclass, fast_free_list);
return small_refill_slow<Conts, CheckInit>(
sizeclass, fast_free_list, size);
}
}
@@ -772,18 +800,19 @@ namespace snmalloc
laden.insert(meta);
}
auto r = finish_alloc<zero_mem, Config>(p, sizeclass);
auto r = finish_alloc<Conts>(p, size);
return ticker.check_tick(r);
}
return small_refill_slow<zero_mem, CheckInit>(sizeclass, fast_free_list);
return small_refill_slow<Conts, CheckInit>(
sizeclass, fast_free_list, size);
}
template<ZeroMem zero_mem, typename CheckInit>
SNMALLOC_SLOW_PATH capptr::Alloc<void> small_refill_slow(
smallsizeclass_t sizeclass, freelist::Iter<>& fast_free_list)
template<typename Conts, typename CheckInit>
SNMALLOC_SLOW_PATH void* small_refill_slow(
smallsizeclass_t sizeclass, freelist::Iter<>& fast_free_list, size_t size)
{
return CheckInit::check_init(
[this, sizeclass, &fast_free_list]() -> capptr::Alloc<void> {
[this, size, sizeclass, &fast_free_list]() -> void* {
size_t rsize = sizeclass_to_size(sizeclass);
// No existing free list get a new slab.
@@ -803,7 +832,7 @@ namespace snmalloc
if (slab == nullptr)
{
return nullptr;
return Conts::failure(sizeclass_to_size(sizeclass));
}
// Set meta slab to empty.
@@ -830,14 +859,13 @@ namespace snmalloc
laden.insert(meta);
}
auto r = finish_alloc<zero_mem, Config>(p, sizeclass);
auto r = finish_alloc<Conts>(p, size);
return ticker.check_tick(r);
},
[](Allocator* a, smallsizeclass_t sizeclass) {
return a->small_alloc<zero_mem, CheckInitNoOp>(
sizeclass_to_size(sizeclass));
[](Allocator* a, size_t size) {
return a->small_alloc<Conts, CheckInitNoOp>(size);
},
sizeclass);
size);
}
static SNMALLOC_FAST_PATH void alloc_new_list(
@@ -1191,8 +1219,10 @@ namespace snmalloc
SNMALLOC_FAST_PATH_LAMBDA {
return capptr_domesticate<Config>(local_state, p);
};
capptr::Alloc<void> p = finish_alloc_no_zero(
fl.take(freelist::Object::key_root, domesticate), sizeclass);
capptr::Alloc<void> p =
fl.take(freelist::Object::key_root, domesticate).as_void();
SNMALLOC_ASSERT(is_start_of_object(
sizeclass_t::from_small_class(sizeclass), address_cast(p)));
// If clear_meta is requested, we should also walk the free list to clear
// it.

View File

@@ -143,7 +143,7 @@ extern "C"
}
if (f.should_zero())
{
*ptr = alloc<ZeroMem::YesZero>(size);
*ptr = alloc<Zero>(size);
}
else
{
@@ -187,7 +187,7 @@ extern "C"
asize = f.aligned_size(size + extra);
}
void* p = f.should_zero() ? alloc<YesZero>(asize) : alloc(asize);
void* p = f.should_zero() ? alloc<Zero>(asize) : alloc(asize);
if (SNMALLOC_LIKELY(p != nullptr))
{
sz = bits::min(asize, sz);
@@ -255,7 +255,7 @@ extern "C"
size = f.aligned_size(size);
if (f.should_zero())
{
return alloc<ZeroMem::YesZero>(size);
return alloc<Zero>(size);
}
return alloc(size);
}
@@ -289,7 +289,7 @@ extern "C"
// allocations, because we get zeroed memory from the PAL and don't zero it
// twice. This is not profiled and so should be considered for refactoring
// if anyone cares about the performance of these APIs.
void* p = f.should_zero() ? alloc<YesZero>(size) : alloc(size);
void* p = f.should_zero() ? alloc<Zero>(size) : alloc(size);
if (SNMALLOC_LIKELY(p != nullptr))
{
sz = bits::min(size, sz);
@@ -376,7 +376,7 @@ extern "C"
}
// Include size 0 in the first sizeclass.
sz = ((sz - 1) >> (bits::BITS - 1)) + sz;
return get_scoped_allocator()->alloc<ZeroMem::YesZero>(sz);
return get_scoped_allocator()->alloc<Zero>(sz);
}
void __je_bootstrap_free(void* ptr)

View File

@@ -24,7 +24,7 @@ SNMALLOC_NAME_MANGLE(rust_alloc)(size_t alignment, size_t size)
extern "C" SNMALLOC_EXPORT void*
SNMALLOC_NAME_MANGLE(rust_alloc_zeroed)(size_t alignment, size_t size)
{
return alloc<YesZero>(aligned_size(alignment, size));
return alloc<Zero>(aligned_size(alignment, size));
}
extern "C" SNMALLOC_EXPORT void

View File

@@ -592,8 +592,6 @@ namespace snmalloc
void* ret = VirtualAlloc2FromApp(
nullptr, nullptr, size, flags, PAGE_READWRITE, &param, 1);
if (ret == nullptr)
errno = ENOMEM;
reservations.push_back(ret);
@@ -604,8 +602,6 @@ namespace snmalloc
void* PALWindows::reserve(size_t size) noexcept
{
void* ret = VirtualAlloc(nullptr, size, MEM_RESERVE, PAGE_READWRITE);
if (ret == nullptr)
errno = ENOMEM;
reservations.push_back(ret);

View File

@@ -91,7 +91,7 @@ int main()
{
static const size_t sz = 1024 * 1024;
errno = 0;
void* olarge = alloc->alloc<YesZero>(sz);
void* olarge = alloc->alloc<Zero>(sz);
int err = errno;
SNMALLOC_CHECK(alarge == address_cast(olarge));

View File

@@ -47,14 +47,14 @@ void check_calloc(void* p, size_t size)
void calloc1(size_t size)
{
void* r = snmalloc::alloc<snmalloc::ZeroMem::YesZero>(size);
void* r = snmalloc::alloc<snmalloc::Zero>(size);
check_calloc(r, size);
snmalloc::dealloc(r);
}
void calloc2(size_t size)
{
void* r = snmalloc::alloc<snmalloc::ZeroMem::YesZero>(size);
void* r = snmalloc::alloc<snmalloc::Zero>(size);
check_calloc(r, size);
snmalloc::dealloc(r, size);
}

View File

@@ -173,10 +173,12 @@ namespace
size * 2);
EXPECT(
ptr[size] == 0,
"Memory not zero initialised for {} byte reallocation from {} "
"Memory not zero initialised for {} byte reallocation from {} with "
"align {}"
"byte allocation",
size * 2,
size);
size,
align);
// The second time we run this test, we if we're allocating from a free
// list then we will reuse this, so make sure it requires explicit
// zeroing.

View File

@@ -178,7 +178,7 @@ void test_calloc()
memset(p, 0xFF, size);
snmalloc::dealloc(p, size);
p = snmalloc::alloc<YesZero>(size);
p = snmalloc::alloc<Zero>(size);
for (size_t i = 0; i < size; i++)
{
@@ -417,7 +417,7 @@ void test_calloc_16M()
// sizes >= 16M use large_alloc
const size_t size = 16'000'000;
void* p1 = snmalloc::alloc<YesZero>(size);
void* p1 = snmalloc::alloc<Zero>(size);
SNMALLOC_CHECK(snmalloc::alloc_size(snmalloc::external_pointer(p1)) >= size);
snmalloc::dealloc(p1);
}
@@ -430,7 +430,7 @@ void test_calloc_large_bug()
// not a multiple of page size.
const size_t size = (MAX_SMALL_SIZECLASS_SIZE << 3) - 7;
void* p1 = snmalloc::alloc<YesZero>(size);
void* p1 = snmalloc::alloc<Zero>(size);
SNMALLOC_CHECK(snmalloc::alloc_size(snmalloc::external_pointer(p1)) >= size);
snmalloc::dealloc(p1);
}

View File

@@ -59,7 +59,7 @@ void check_calloc(void* p, size_t size)
void calloc1(size_t size)
{
trigger_teardown();
void* r = snmalloc::alloc<snmalloc::ZeroMem::YesZero>(size);
void* r = snmalloc::alloc<snmalloc::Zero>(size);
check_calloc(r, size);
snmalloc::dealloc(r);
}
@@ -67,7 +67,7 @@ void calloc1(size_t size)
void calloc2(size_t size)
{
trigger_teardown();
void* r = snmalloc::alloc<snmalloc::ZeroMem::YesZero>(size);
void* r = snmalloc::alloc<snmalloc::Zero>(size);
check_calloc(r, size);
snmalloc::dealloc(r, size);
}

View File

@@ -5,20 +5,21 @@
using namespace snmalloc;
template<ZeroMem zero_mem>
template<typename Conts>
void test_alloc_dealloc(size_t count, size_t size, bool write)
{
{
MeasureTime m;
m << "Count: " << std::setw(6) << count << ", Size: " << std::setw(6)
<< size << ", ZeroMem: " << (zero_mem == YesZero) << ", Write: " << write;
<< size
<< ", ZeroMem: " << std::is_same_v<Conts, Zero> << ", Write: " << write;
std::unordered_set<void*> set;
// alloc 1.5x objects
for (size_t i = 0; i < ((count * 3) / 2); i++)
{
void* p = snmalloc::alloc<zero_mem>(size);
void* p = snmalloc::alloc<Conts>(size);
SNMALLOC_CHECK(set.find(p) == set.end());
if (write)
@@ -40,7 +41,7 @@ void test_alloc_dealloc(size_t count, size_t size, bool write)
// alloc 1x objects
for (size_t i = 0; i < count; i++)
{
void* p = snmalloc::alloc<zero_mem>(size);
void* p = snmalloc::alloc<Conts>(size);
SNMALLOC_CHECK(set.find(p) == set.end());
if (write)
@@ -67,18 +68,18 @@ int main(int, char**)
for (size_t size = 16; size <= 128; size <<= 1)
{
test_alloc_dealloc<NoZero>(1 << 15, size, false);
test_alloc_dealloc<NoZero>(1 << 15, size, true);
test_alloc_dealloc<YesZero>(1 << 15, size, false);
test_alloc_dealloc<YesZero>(1 << 15, size, true);
test_alloc_dealloc<Uninit>(1 << 15, size, false);
test_alloc_dealloc<Uninit>(1 << 15, size, true);
test_alloc_dealloc<Zero>(1 << 15, size, false);
test_alloc_dealloc<Zero>(1 << 15, size, true);
}
for (size_t size = 1 << 12; size <= 1 << 17; size <<= 1)
{
test_alloc_dealloc<NoZero>(1 << 10, size, false);
test_alloc_dealloc<NoZero>(1 << 10, size, true);
test_alloc_dealloc<YesZero>(1 << 10, size, false);
test_alloc_dealloc<YesZero>(1 << 10, size, true);
test_alloc_dealloc<Uninit>(1 << 10, size, false);
test_alloc_dealloc<Uninit>(1 << 10, size, true);
test_alloc_dealloc<Zero>(1 << 10, size, false);
test_alloc_dealloc<Zero>(1 << 10, size, true);
}
return 0;