From 9a8494d2701fed46055bdd9c1a70b059bb47c185 Mon Sep 17 00:00:00 2001 From: Protobuf Team Bot Date: Tue, 3 Dec 2024 15:09:19 -0800 Subject: [PATCH] Add a upb_alloc cleanup function pointer to upb_Arena. The allocator cleanup should be called last, after all the blocks have been freed. PiperOrigin-RevId: 702483612 --- upb/mem/arena.c | 16 ++++++++++++++++ upb/mem/arena.h | 7 +++++++ upb/mem/arena_test.cc | 40 ++++++++++++++++++++++++++++++++++------ upb/mem/internal/arena.h | 2 +- 4 files changed, 58 insertions(+), 7 deletions(-) diff --git a/upb/mem/arena.c b/upb/mem/arena.c index e313d605be1a..79d5e70088eb 100644 --- a/upb/mem/arena.c +++ b/upb/mem/arena.c @@ -37,6 +37,10 @@ typedef struct upb_ArenaInternal { // block. uintptr_t block_alloc; + // The cleanup for the allocator. This is called after all the blocks are + // freed in an arena. + upb_AllocCleanupFunc* upb_alloc_cleanup; + // When multiple arenas are fused together, each arena points to a parent // arena (root points to itself). The root tracks how many live arenas // reference it. @@ -304,6 +308,7 @@ static upb_Arena* _upb_Arena_InitSlow(upb_alloc* alloc) { upb_Atomic_Init(&a->body.next, NULL); upb_Atomic_Init(&a->body.tail, &a->body); upb_Atomic_Init(&a->body.blocks, NULL); + a->body.upb_alloc_cleanup = NULL; _upb_Arena_AddBlock(&a->head, mem, n); @@ -342,6 +347,7 @@ upb_Arena* upb_Arena_Init(void* mem, size_t n, upb_alloc* alloc) { upb_Atomic_Init(&a->body.next, NULL); upb_Atomic_Init(&a->body.tail, &a->body); upb_Atomic_Init(&a->body.blocks, NULL); + a->body.upb_alloc_cleanup = NULL; a->body.block_alloc = _upb_Arena_MakeBlockAlloc(alloc, 1); a->head.UPB_PRIVATE(ptr) = mem; @@ -360,6 +366,7 @@ static void _upb_Arena_DoFree(upb_ArenaInternal* ai) { (upb_ArenaInternal*)upb_Atomic_Load(&ai->next, memory_order_acquire); upb_alloc* block_alloc = _upb_ArenaInternal_BlockAlloc(ai); upb_MemBlock* block = upb_Atomic_Load(&ai->blocks, memory_order_acquire); + upb_AllocCleanupFunc* alloc_cleanup = *ai->upb_alloc_cleanup; while (block != NULL) { // Load first since we are deleting block. upb_MemBlock* next_block = @@ -367,6 +374,9 @@ static void _upb_Arena_DoFree(upb_ArenaInternal* ai) { upb_free(block_alloc, block); block = next_block; } + if (alloc_cleanup != NULL) { + alloc_cleanup(block_alloc); + } ai = next_arena; } } @@ -431,6 +441,12 @@ static void _upb_Arena_DoFuseArenaLists(upb_ArenaInternal* const parent, upb_Atomic_Store(&parent->tail, parent_tail, memory_order_relaxed); } +void upb_Arena_SetAllocCleanup(upb_Arena* a, upb_AllocCleanupFunc* func) { + upb_ArenaInternal* ai = upb_Arena_Internal(a); + UPB_ASSERT(ai->upb_alloc_cleanup == NULL); + ai->upb_alloc_cleanup = func; +} + static upb_ArenaInternal* _upb_Arena_DoFuse(const upb_Arena* a1, const upb_Arena* a2, uintptr_t* ref_delta) { diff --git a/upb/mem/arena.h b/upb/mem/arena.h index f0d9d2d6f0e2..3d9e9eaaee43 100644 --- a/upb/mem/arena.h +++ b/upb/mem/arena.h @@ -31,6 +31,8 @@ typedef struct upb_Arena upb_Arena; +typedef void upb_AllocCleanupFunc(upb_alloc* alloc); + #ifdef __cplusplus extern "C" { #endif @@ -41,6 +43,11 @@ extern "C" { UPB_API upb_Arena* upb_Arena_Init(void* mem, size_t n, upb_alloc* alloc); UPB_API void upb_Arena_Free(upb_Arena* a); +// Sets the cleanup function for the upb_alloc used by the arena. Only one +// cleanup function can be set, which will be called after all blocks are +// freed. +UPB_API void upb_Arena_SetAllocCleanup(upb_Arena* a, + upb_AllocCleanupFunc* func); UPB_API bool upb_Arena_Fuse(const upb_Arena* a, const upb_Arena* b); UPB_API bool upb_Arena_IsFused(const upb_Arena* a, const upb_Arena* b); diff --git a/upb/mem/arena_test.cc b/upb/mem/arena_test.cc index 1dce05f28bdd..292dcf340ca9 100644 --- a/upb/mem/arena_test.cc +++ b/upb/mem/arena_test.cc @@ -14,6 +14,7 @@ #include #include +#include #include #include "absl/base/thread_annotations.h" #include "absl/random/distributions.h" @@ -30,6 +31,39 @@ namespace { +struct CustomAlloc { + upb_alloc alloc; + int counter; + bool ran_cleanup; +}; + +void* CustomAllocFunc(upb_alloc* alloc, void* ptr, size_t oldsize, + size_t size) { + CustomAlloc* custom_alloc = reinterpret_cast(alloc); + if (size == 0) { + custom_alloc->counter--; + } else { + custom_alloc->counter++; + } + return upb_alloc_global.func(alloc, ptr, oldsize, size); +} + +void CustomAllocCleanup(upb_alloc* alloc) { + CustomAlloc* custom_alloc = reinterpret_cast(alloc); + EXPECT_THAT(custom_alloc->counter, 0); + custom_alloc->ran_cleanup = true; +} + +TEST(ArenaTest, ArenaWithAllocCleanup) { + CustomAlloc alloc = {{&CustomAllocFunc}, 0, false}; + upb_Arena* arena = + upb_Arena_Init(nullptr, 0, reinterpret_cast(&alloc)); + EXPECT_EQ(alloc.counter, 1); + upb_Arena_SetAllocCleanup(arena, CustomAllocCleanup); + upb_Arena_Free(arena); + EXPECT_TRUE(alloc.ran_cleanup); +} + TEST(ArenaTest, ArenaFuse) { upb_Arena* arena1 = upb_Arena_New(); upb_Arena* arena2 = upb_Arena_New(); @@ -40,12 +74,6 @@ TEST(ArenaTest, ArenaFuse) { upb_Arena_Free(arena2); } -// Do-nothing allocator for testing. -extern "C" void* TestAllocFunc(upb_alloc* alloc, void* ptr, size_t oldsize, - size_t size) { - return upb_alloc_global.func(alloc, ptr, oldsize, size); -} - TEST(ArenaTest, FuseWithInitialBlock) { char buf1[1024]; char buf2[1024]; diff --git a/upb/mem/internal/arena.h b/upb/mem/internal/arena.h index 9a6469edc16f..08aacde1eac8 100644 --- a/upb/mem/internal/arena.h +++ b/upb/mem/internal/arena.h @@ -20,7 +20,7 @@ // // We need this because the decoder inlines a upb_Arena for performance but // the full struct is not visible outside of arena.c. Yes, I know, it's awful. -#define UPB_ARENA_SIZE_HACK 7 +#define UPB_ARENA_SIZE_HACK 8 // LINT.IfChange(upb_Arena)