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

@@ -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,13 +366,17 @@ 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 =
           upb_Atomic_Load(&block->next, memory_order_acquire);
       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) {

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);
 

upb/mem/arena_test.cc

@@ -14,6 +14,7 @@
 #include <thread>
 #include <vector>
 
+#include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #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<CustomAlloc*>(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<CustomAlloc*>(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<upb_alloc*>(&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];

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)