diff --git a/upb/mem/arena.c b/upb/mem/arena.c index fc65f69013..544179228d 100644 --- a/upb/mem/arena.c +++ b/upb/mem/arena.c @@ -21,6 +21,10 @@ // Must be last. #include "upb/port/def.inc" +static UPB_ATOMIC(size_t) max_block_size = 32 << 10; + +void upb_Arena_SetMaxBlockSize(size_t max) { max_block_size = max; } + typedef struct upb_MemBlock { // Atomic only for the benefit of SpaceAllocated(). UPB_ATOMIC(struct upb_MemBlock*) next; @@ -258,7 +262,14 @@ static bool _upb_Arena_AllocBlock(upb_Arena* a, size_t size) { if (!ai->block_alloc) return false; upb_MemBlock* last_block = upb_Atomic_Load(&ai->blocks, memory_order_acquire); size_t last_size = last_block != NULL ? last_block->size : 128; - size_t block_size = UPB_MAX(size, last_size * 2) + kUpb_MemblockReserve; + + // Don't naturally grow beyond the max block size. + size_t clamped_size = UPB_MIN(last_size * 2, max_block_size); + + // We may need to exceed the max block size if the user requested a large + // allocation. + size_t block_size = UPB_MAX(size, clamped_size) + kUpb_MemblockReserve; + upb_MemBlock* block = upb_malloc(_upb_ArenaInternal_BlockAlloc(ai), block_size); diff --git a/upb/mem/arena.h b/upb/mem/arena.h index 83f5781f83..2ad4f5d2fa 100644 --- a/upb/mem/arena.h +++ b/upb/mem/arena.h @@ -58,6 +58,15 @@ UPB_API_INLINE void* upb_Arena_Malloc(struct upb_Arena* a, size_t size); UPB_API_INLINE void* upb_Arena_Realloc(upb_Arena* a, void* ptr, size_t oldsize, size_t size); +// Sets the maximum block size for all arenas. This is a global configuration +// setting that will affect all existing and future arenas. If +// upb_Arena_Malloc() is called with a size larger than this, we will exceed +// this size and allocate a larger block. +// +// This API is meant for experimentation only. It will likely be removed in +// the future. +void upb_Arena_SetMaxBlockSize(size_t max); + // Shrinks the last alloc from arena. // REQUIRES: (ptr, oldsize) was the last malloc/realloc from this arena. // We could also add a upb_Arena_TryShrinkLast() which is simply a no-op if diff --git a/upb/mem/arena_test.cc b/upb/mem/arena_test.cc index a7889a069b..2f0647ff27 100644 --- a/upb/mem/arena_test.cc +++ b/upb/mem/arena_test.cc @@ -171,6 +171,32 @@ TEST(ArenaTest, Contains) { upb_Arena_Free(arena2); } +TEST(ArenaTest, LargeAlloc) { + // Tests an allocation larger than the max block size. + upb_Arena* arena = upb_Arena_New(); + size_t size = 100000; + char* mem = static_cast(upb_Arena_Malloc(arena, size)); + EXPECT_NE(mem, nullptr); + for (size_t i = 0; i < size; ++i) { + mem[i] = static_cast(i); + } + for (size_t i = 0; i < size; ++i) { + EXPECT_EQ(mem[i], static_cast(i)); + } + upb_Arena_Free(arena); +} + +TEST(ArenaTest, MaxBlockSize) { + upb_Arena* arena = upb_Arena_New(); + // Perform 600 1k allocations (600k total) and ensure that the amount of + // memory allocated does not exceed 700k. + for (int i = 0; i < 600; ++i) { + upb_Arena_Malloc(arena, 1024); + } + EXPECT_LE(upb_Arena_SpaceAllocated(arena, nullptr), 700 * 1024); + upb_Arena_Free(arena); +} + #ifdef UPB_USE_C11_ATOMICS TEST(ArenaTest, FuzzFuseFreeRace) {