From b9cf68ed5e158f0afedc6a99d973f648a7c1d883 Mon Sep 17 00:00:00 2001 From: Protobuf Team Bot Date: Wed, 8 Jan 2025 19:38:06 -0800 Subject: [PATCH] Allow suggesting the initial block size without requiring an initial block to be provided and managed separately. This can avoid small mallocs when the approximate arena size is known. PiperOrigin-RevId: 713501871 --- upb/mem/arena.c | 25 +++++++++++++------------ upb/mem/arena.h | 15 ++++++++++++--- upb/mem/arena.hpp | 24 ++++-------------------- upb/mem/arena_test.cc | 17 +++++++++++++++++ 4 files changed, 46 insertions(+), 35 deletions(-) diff --git a/upb/mem/arena.c b/upb/mem/arena.c index f0aa39a4bab8a..274336e163a19 100644 --- a/upb/mem/arena.c +++ b/upb/mem/arena.c @@ -343,14 +343,16 @@ void* UPB_PRIVATE(_upb_Arena_SlowMalloc)(upb_Arena* a, size_t size) { return upb_Arena_Malloc(a, size - UPB_ASAN_GUARD_SIZE); } -static upb_Arena* _upb_Arena_InitSlow(upb_alloc* alloc) { +static upb_Arena* _upb_Arena_InitSlow(upb_alloc* alloc, size_t first_size) { const size_t first_block_overhead = sizeof(upb_ArenaState) + kUpb_MemblockReserve; upb_ArenaState* a; // We need to malloc the initial block. char* mem; - size_t block_size = first_block_overhead + 256; + size_t block_size = + first_block_overhead + + UPB_MAX(256, UPB_ALIGN_MALLOC(first_size) + UPB_ASAN_GUARD_SIZE); if (!alloc || !(mem = upb_malloc(alloc, block_size))) { return NULL; } @@ -377,26 +379,25 @@ upb_Arena* upb_Arena_Init(void* mem, size_t n, upb_alloc* alloc) { UPB_ASSERT(sizeof(void*) * UPB_ARENA_SIZE_HACK >= sizeof(upb_ArenaState)); upb_ArenaState* a; - if (n) { + if (mem) { /* Align initial pointer up so that we return properly-aligned pointers. */ void* aligned = (void*)UPB_ALIGN_UP((uintptr_t)mem, UPB_MALLOC_ALIGN); size_t delta = (uintptr_t)aligned - (uintptr_t)mem; n = delta <= n ? n - delta : 0; mem = aligned; + /* Round block size down to alignof(*a) since we will allocate the arena + * itself at the end. */ + n = UPB_ALIGN_DOWN(n, UPB_ALIGN_OF(upb_ArenaState)); + } else { + n = UPB_ALIGN_UP(n, UPB_ALIGN_OF(upb_ArenaState)); } - /* Round block size down to alignof(*a) since we will allocate the arena - * itself at the end. */ - n = UPB_ALIGN_DOWN(n, UPB_ALIGN_OF(upb_ArenaState)); - - if (UPB_UNLIKELY(n < sizeof(upb_ArenaState))) { + if (UPB_UNLIKELY(n < sizeof(upb_ArenaState) || !mem)) { + upb_Arena* ret = _upb_Arena_InitSlow(alloc, mem ? 0 : n); #ifdef UPB_TRACING_ENABLED - upb_Arena* ret = _upb_Arena_InitSlow(alloc); upb_Arena_LogInit(ret, n); - return ret; -#else - return _upb_Arena_InitSlow(alloc); #endif + return ret; } a = UPB_PTR_AT(mem, n - sizeof(upb_ArenaState), upb_ArenaState); diff --git a/upb/mem/arena.h b/upb/mem/arena.h index 04d11b6d14ade..7e22aec5c63fd 100644 --- a/upb/mem/arena.h +++ b/upb/mem/arena.h @@ -38,9 +38,14 @@ typedef void upb_AllocCleanupFunc(upb_alloc* alloc); extern "C" { #endif -// Creates an arena from the given initial block (if any -- n may be 0). -// Additional blocks will be allocated from |alloc|. If |alloc| is NULL, this -// is a fixed-size arena and cannot grow. +// Creates an arena from the given initial block (if any -- mem may be NULL). If +// an initial block is specified, the arena's lifetime cannot be extended by +// |upb_Arena_IncRefFor| or |upb_Arena_Fuse|. Additional blocks will be +// allocated from |alloc|. If |alloc| is NULL, this is a fixed-size arena and +// cannot grow. If an initial block is specified, |n| is its length; if there is +// no initial block, |n| is a hint of the size that should be allocated for the +// first block of the arena, such that `upb_Arena_Malloc(hint)` will not require +// another call to |alloc|. UPB_API upb_Arena* upb_Arena_Init(void* mem, size_t n, upb_alloc* alloc); UPB_API void upb_Arena_Free(upb_Arena* a); @@ -76,6 +81,10 @@ UPB_API_INLINE upb_Arena* upb_Arena_New(void) { return upb_Arena_Init(NULL, 0, &upb_alloc_global); } +UPB_API_INLINE upb_Arena* upb_Arena_NewSized(size_t size_hint) { + return upb_Arena_Init(NULL, size_hint, &upb_alloc_global); +} + 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, diff --git a/upb/mem/arena.hpp b/upb/mem/arena.hpp index 420adfa463c30..2ac0731ec6a03 100644 --- a/upb/mem/arena.hpp +++ b/upb/mem/arena.hpp @@ -8,8 +8,10 @@ #ifndef UPB_MEM_ARENA_HPP_ #define UPB_MEM_ARENA_HPP_ +#include "upb/mem/alloc.h" #ifdef __cplusplus +#include #include #include "upb/mem/arena.h" @@ -23,6 +25,8 @@ class Arena { Arena(char* initial_block, size_t size) : ptr_(upb_Arena_Init(initial_block, size, &upb_alloc_global), upb_Arena_Free) {} + explicit Arena(size_t size) + : ptr_(upb_Arena_NewSized(size), upb_Arena_Free) {} upb_Arena* ptr() const { return ptr_.get(); } @@ -35,26 +39,6 @@ class Arena { protected: std::unique_ptr ptr_; }; - -// InlinedArena seeds the arenas with a predefined amount of memory. No heap -// memory will be allocated until the initial block is exceeded. -template -class InlinedArena : public Arena { - public: - InlinedArena() : Arena(initial_block_, N) {} - ~InlinedArena() { - // Explicitly destroy the arena now so that it does not outlive - // initial_block_. - ptr_.reset(); - } - - private: - InlinedArena(const InlinedArena&) = delete; - InlinedArena& operator=(const InlinedArena&) = delete; - - char initial_block_[N]; -}; - } // namespace upb #endif // __cplusplus diff --git a/upb/mem/arena_test.cc b/upb/mem/arena_test.cc index 95f3bdfd775a1..4f916196e6a58 100644 --- a/upb/mem/arena_test.cc +++ b/upb/mem/arena_test.cc @@ -110,6 +110,23 @@ TEST(ArenaTest, SizedFree) { EXPECT_EQ(sizes.size(), 0); } +TEST(ArenaTest, SizeHint) { + absl::flat_hash_map sizes; + SizeTracker alloc; + alloc.alloc.func = size_checking_allocfunc; + alloc.delegate_alloc = &upb_alloc_global; + alloc.sizes = &sizes; + + upb_Arena* arena = upb_Arena_Init(nullptr, 2459, &alloc.alloc); + EXPECT_EQ(sizes.size(), 1); + EXPECT_NE(upb_Arena_Malloc(arena, 2459), nullptr); + EXPECT_EQ(sizes.size(), 1); + EXPECT_NE(upb_Arena_Malloc(arena, 500), nullptr); + EXPECT_EQ(sizes.size(), 2); + upb_Arena_Free(arena); + EXPECT_EQ(sizes.size(), 0); +} + TEST(ArenaTest, ArenaFuse) { upb_Arena* arena1 = upb_Arena_New(); upb_Arena* arena2 = upb_Arena_New();