Optimize upb_Arena memory allocation for repeated growths

upb/hash/common.c

@@ -117,7 +117,7 @@ static bool init(upb_table* t, uint8_t size_lg2, upb_Arena* a) {
   }
   size_t bytes = upb_table_size(t) * sizeof(upb_tabent);
   if (bytes > 0) {
-    t->entries = upb_Arena_Malloc(a, bytes);
+    t->entries = upb_Arena_AllocPool(a, bytes);
     if (!t->entries) return false;
     memset(t->entries, 0, bytes);
   } else {
@@ -585,6 +585,10 @@ bool upb_strtable_resize(upb_strtable* t, size_t size_lg2, upb_Arena* a) {
     uint32_t hash = _upb_Hash_NoSeed(sv.data, sv.size);
     insert(&new_table.t, lookupkey, tabkey, val, hash, &strhash, &streql);
   }
+  if (t->t.entries) {
+    size_t old_bytes = upb_table_size(&t->t) * sizeof(upb_tabent);
+    upb_Arena_FreePool(a, t->t.entries, old_bytes);
+  }
   *t = new_table;
   return true;
 }
@@ -723,6 +727,10 @@ bool upb_exttable_resize(upb_exttable* t, size_t size_lg2, upb_Arena* a) {
     insert(&new_table.t, lookupkey, e->key, e->val, hash, &exthash, &exteql);
   }
 
+  if (t->t.entries) {
+    size_t old_bytes = upb_table_size(&t->t) * sizeof(upb_tabent);
+    upb_Arena_FreePool(a, t->t.entries, old_bytes);
+  }
   *t = new_table;
   return true;
 }
@@ -830,6 +838,10 @@ bool upb_inttable_insert(upb_inttable* t, uintptr_t key, upb_value val,
 
     UPB_ASSERT(t->t.count == new_table.count);
 
+    if (t->t.entries) {
+      size_t old_bytes = upb_table_size(&t->t) * sizeof(upb_tabent);
+      upb_Arena_FreePool(a, t->t.entries, old_bytes);
+    }
     t->t = new_table;
   }
   upb_key tabkey = {.num = key};

upb/mem/BUILD

@@ -27,6 +27,7 @@ cc_library(
     visibility = ["//visibility:public"],
     deps = [
         ":internal",
+        "//upb/base:internal",
         "//upb/port",
     ],
 )

upb/mem/arena.c

@@ -19,6 +19,7 @@
 #include <stdint.h>
 #include <string.h>
 
+#include "upb/base/internal/log2.h"
 #include "upb/mem/alloc.h"
 #include "upb/mem/internal/arena.h"
 #include "upb/port/atomic.h"
@@ -52,6 +53,12 @@ typedef struct upb_ArenaRef {
 #endif
 } upb_ArenaRef;
 
+typedef struct upb_ArenaFreeBlock {
+  struct upb_ArenaFreeBlock* next_same_size;
+  struct upb_ArenaFreeBlock* next_size_class;
+  size_t size;
+} upb_ArenaFreeBlock;
+
 typedef struct upb_ArenaInternal {
   // upb_alloc* together with a low bit which signals if there is an initial
   // block.
@@ -60,6 +67,9 @@ typedef struct upb_ArenaInternal {
   // Linked list of blocks to free/cleanup.
   upb_MemBlock* blocks;
 
+  // Pool of power-of-2 sized unused memory blocks.
+  upb_ArenaFreeBlock* freelists;
+
 #ifndef NDEBUG
   // Stack of pointers to other arenas that this arena owns.
   // Used for debug-only ref cycle checks.
@@ -532,6 +542,7 @@ static upb_Arena* _upb_Arena_InitSlow(upb_alloc* alloc, size_t first_size) {
                   _upb_Arena_TaggedFromTail(&a->body));
   upb_Atomic_Init(&a->body.space_allocated, 0);
   a->body.blocks = NULL;
+  a->body.freelists = NULL;
 #ifndef NDEBUG
   a->body.refs = NULL;
 #endif
@@ -575,6 +586,7 @@ upb_Arena* upb_Arena_Init(void* mem, size_t n, upb_alloc* alloc) {
                   _upb_Arena_TaggedFromTail(&a->body));
   upb_Atomic_Init(&a->body.space_allocated, 0);
   a->body.blocks = NULL;
+  a->body.freelists = NULL;
 #ifndef NDEBUG
   a->body.refs = NULL;
 #endif
@@ -1061,3 +1073,151 @@ void UPB_PRIVATE(_upb_Arena_FreeBlock)(upb_Arena* a, void* block) {
   upb_alloc* alloc = upb_Arena_GetUpbAlloc(a);
   upb_free_sized(alloc, b, b->size);
 }
+
+#define UPB_ARENA_POOL_MIN_SIZE sizeof(upb_ArenaFreeBlock)
+#define UPB_MAP_ENTRY_SIZE (8 + 2 * sizeof(void*))
+
+static size_t _upb_Arena_PoolBlockSize(size_t size) {
+  if (size == 0) return UPB_MAP_ENTRY_SIZE;
+
+  size_t ratio = (size + UPB_MAP_ENTRY_SIZE - 1) / UPB_MAP_ENTRY_SIZE;
+  size_t next_pow2 = upb_RoundUpToPowerOfTwo(ratio);
+  if (next_pow2 == SIZE_MAX) return SIZE_MAX;
+
+  if (next_pow2 > SIZE_MAX / UPB_MAP_ENTRY_SIZE) {
+    return SIZE_MAX;
+  }
+  return next_pow2 * UPB_MAP_ENTRY_SIZE;
+}
+
+void* upb_Arena_AllocPool(upb_Arena* a, size_t size) {
+  UPB_ASSERT(a);
+  UPB_ASSERT(size > 0);
+
+  size_t pool_size = _upb_Arena_PoolBlockSize(size);
+  if (pool_size == SIZE_MAX) return NULL;
+  UPB_ASSERT(pool_size >= UPB_ARENA_POOL_MIN_SIZE);
+  UPB_ASSERT(pool_size % UPB_MAP_ENTRY_SIZE == 0);
+  UPB_ASSERT(((pool_size / UPB_MAP_ENTRY_SIZE) &
+              ((pool_size / UPB_MAP_ENTRY_SIZE) - 1)) == 0);
+
+  upb_ArenaInternal* ai = upb_Arena_Internal(a);
+  UPB_PRIVATE(upb_Xsan_AccessReadWrite)(UPB_XSAN(ai));
+
+  upb_ArenaFreeBlock* prev = NULL;
+  upb_ArenaFreeBlock* curr = ai->freelists;
+
+  while (curr && curr->size < pool_size) {
+    UPB_ASSERT(curr->size >= UPB_ARENA_POOL_MIN_SIZE);
+    UPB_ASSERT(curr->size % UPB_MAP_ENTRY_SIZE == 0);
+    UPB_ASSERT(((curr->size / UPB_MAP_ENTRY_SIZE) &
+                ((curr->size / UPB_MAP_ENTRY_SIZE) - 1)) == 0);
+    UPB_ASSERT(UPB_PRIVATE(_upb_Arena_IsAligned)(curr));
+    if (prev) {
+      UPB_ASSERT(prev->size < curr->size);
+    }
+    prev = curr;
+    curr = curr->next_size_class;
+  }
+
+  if (curr && curr->size == pool_size) {
+    UPB_ASSERT(UPB_PRIVATE(_upb_Arena_IsAligned)(curr));
+    void* ret = curr;
+    upb_ArenaFreeBlock* next_same = curr->next_same_size;
+    if (next_same) {
+      UPB_ASSERT(next_same->size == pool_size);
+      UPB_ASSERT(UPB_PRIVATE(_upb_Arena_IsAligned)(next_same));
+      next_same->next_size_class = curr->next_size_class;
+      if (prev) {
+        prev->next_size_class = next_same;
+      } else {
+        ai->freelists = next_same;
+      }
+    } else {
+      if (prev) {
+        prev->next_size_class = curr->next_size_class;
+      } else {
+        ai->freelists = curr->next_size_class;
+      }
+    }
+    void* unpoisoned =
+        UPB_PRIVATE(upb_Xsan_NewUnpoisonedRegion)(UPB_XSAN(ai), ret, size);
+    UPB_ASSERT(UPB_PRIVATE(_upb_Arena_IsAligned)(unpoisoned));
+    return unpoisoned;
+  }
+
+  void* ret = upb_Arena_Malloc(a, pool_size);
+  if (!ret) return NULL;
+  UPB_ASSERT(UPB_PRIVATE(_upb_Arena_IsAligned)(ret));
+
+  size_t aligned_size = UPB_ALIGN_MALLOC(size);
+  if (aligned_size < pool_size) {
+    UPB_PRIVATE(upb_Xsan_PoisonRegion)((char*)ret + aligned_size,
+                                       pool_size - aligned_size);
+  }
+
+  return ret;
+}
+
+void upb_Arena_FreePool(upb_Arena* a, void* ptr, size_t size) {
+  UPB_ASSERT(a);
+  if (!ptr) return;
+  if (size < UPB_ARENA_POOL_MIN_SIZE) return;
+
+  size_t pool_size = _upb_Arena_PoolBlockSize(size);
+  if (pool_size == SIZE_MAX) return;
+  UPB_ASSERT(pool_size >= UPB_ARENA_POOL_MIN_SIZE);
+  UPB_ASSERT(pool_size % UPB_MAP_ENTRY_SIZE == 0);
+  UPB_ASSERT(((pool_size / UPB_MAP_ENTRY_SIZE) &
+              ((pool_size / UPB_MAP_ENTRY_SIZE) - 1)) == 0);
+
+  upb_ArenaInternal* ai = upb_Arena_Internal(a);
+  UPB_PRIVATE(upb_Xsan_AccessReadWrite)(UPB_XSAN(ai));
+
+  void* untagged_ptr = UPB_PRIVATE(upb_Xsan_UntagPointer)(ptr);
+  UPB_ASSERT(UPB_PRIVATE(_upb_Arena_IsAligned)(untagged_ptr));
+
+  size_t metadata_size = UPB_ALIGN_MALLOC(sizeof(upb_ArenaFreeBlock));
+  UPB_PRIVATE(_upb_Xsan_UnpoisonRegion)(untagged_ptr, metadata_size, 0);
+
+  if (pool_size > metadata_size) {
+    UPB_PRIVATE(upb_Xsan_PoisonRegion)((char*)untagged_ptr + metadata_size,
+                                       pool_size - metadata_size);
+  }
+
+  upb_ArenaFreeBlock* new_block = (upb_ArenaFreeBlock*)untagged_ptr;
+
+  upb_ArenaFreeBlock* prev = NULL;
+  upb_ArenaFreeBlock* curr = ai->freelists;
+
+  while (curr && curr->size < pool_size) {
+    UPB_ASSERT(curr->size >= UPB_ARENA_POOL_MIN_SIZE);
+    UPB_ASSERT(curr->size % UPB_MAP_ENTRY_SIZE == 0);
+    UPB_ASSERT(((curr->size / UPB_MAP_ENTRY_SIZE) &
+                ((curr->size / UPB_MAP_ENTRY_SIZE) - 1)) == 0);
+    UPB_ASSERT(UPB_PRIVATE(_upb_Arena_IsAligned)(curr));
+    if (prev) {
+      UPB_ASSERT(prev->size < curr->size);
+    }
+    prev = curr;
+    curr = curr->next_size_class;
+  }
+
+  if (curr && curr->size == pool_size) {
+    UPB_ASSERT(UPB_PRIVATE(_upb_Arena_IsAligned)(curr));
+    // Size class exists. Make new_block the new head, pushing curr onto its
+    // stack.
+    new_block->next_size_class = curr->next_size_class;
+    new_block->next_same_size = curr;
+  } else {
+    // Size class does not exist. Insert new_block as a new size class head.
+    new_block->next_size_class = curr;
+    new_block->next_same_size = NULL;
+  }
+  new_block->size = pool_size;
+  if (prev) {
+    prev->next_size_class = new_block;
+  } else {
+    ai->freelists = new_block;
+  }
+}

upb/mem/arena.h

@@ -150,6 +150,9 @@ UPB_NODISCARD UPB_API_INLINE void* upb_Arena_Realloc(upb_Arena* a, void* ptr,
                                                      size_t oldsize,
                                                      size_t size);
 
+UPB_NODISCARD UPB_API void* upb_Arena_AllocPool(upb_Arena* a, size_t size);
+UPB_API void upb_Arena_FreePool(upb_Arena* a, void* ptr, size_t size);
+
 static const size_t UPB_PRIVATE(kUpbDefaultMaxBlockSize) =
     UPB_DEFAULT_MAX_BLOCK_SIZE;
 

upb/mem/arena_test.cc

@@ -1010,4 +1010,70 @@ TEST(ArenaDeathTest, ArenaRefFuseCycle) {
 
 #endif  // UPB_SUPPRESS_MISSING_ATOMICS
 
+TEST(ArenaTest, PoolAlloc) {
+  upb_Arena* arena = upb_Arena_New();
+
+  // 1. Test basic reuse
+  void* ptr1 = upb_Arena_AllocPool(arena, 32);
+  EXPECT_NE(ptr1, nullptr);
+
+  // Free it back to the pool
+  upb_Arena_FreePool(arena, ptr1, 32);
+
+  // Allocate again, it should return the SAME block
+  void* ptr2 = upb_Arena_AllocPool(arena, 32);
+  EXPECT_TRUE(UPB_PRIVATE(upb_Xsan_PtrEq)(ptr1, ptr2));
+
+  // 2. Test multiple size classes
+  void* ptr_small = upb_Arena_AllocPool(arena, 32);
+  void* ptr_large = upb_Arena_AllocPool(arena, 64);
+  EXPECT_NE(ptr_small, ptr_large);
+
+  upb_Arena_FreePool(arena, ptr_small, 32);
+  upb_Arena_FreePool(arena, ptr_large, 64);
+
+  // Allocate 64 first, should get ptr_large
+  void* ptr_large2 = upb_Arena_AllocPool(arena, 64);
+  EXPECT_TRUE(UPB_PRIVATE(upb_Xsan_PtrEq)(ptr_large2, ptr_large));
+
+  // Allocate 32, should get ptr_small
+  void* ptr_small2 = upb_Arena_AllocPool(arena, 32);
+  EXPECT_TRUE(UPB_PRIVATE(upb_Xsan_PtrEq)(ptr_small2, ptr_small));
+
+  // 3. Test stack behavior (LIFO)
+  void* a1 = upb_Arena_AllocPool(arena, 32);
+  void* a2 = upb_Arena_AllocPool(arena, 32);
+  EXPECT_NE(a1, a2);
+
+  upb_Arena_FreePool(arena, a1, 32);
+  upb_Arena_FreePool(arena, a2, 32);
+
+  // Since it's a stack, we should get a2 first, then a1
+  void* r1 = upb_Arena_AllocPool(arena, 32);
+  void* r2 = upb_Arena_AllocPool(arena, 32);
+  EXPECT_TRUE(UPB_PRIVATE(upb_Xsan_PtrEq)(r1, a2));
+  EXPECT_TRUE(UPB_PRIVATE(upb_Xsan_PtrEq)(r2, a1));
+
+  // 4. Test rounding up
+  void* b1 = upb_Arena_AllocPool(arena, 25);
+  upb_Arena_FreePool(arena, b1, 25);
+
+  void* b2 = upb_Arena_AllocPool(arena, 30);
+  EXPECT_TRUE(UPB_PRIVATE(upb_Xsan_PtrEq)(b1, b2));
+
+  upb_Arena_FreePool(arena, b2, 30);
+
+  // 5. Test that blocks smaller than UPB_ARENA_POOL_MIN_SIZE are not pooled
+  // 8 bytes is smaller than MIN_SIZE on both 32-bit (12) and 64-bit (24).
+  void* c1 = upb_Arena_AllocPool(arena, 8);
+  upb_Arena_FreePool(arena, c1, 8);  // Should be rejected and not pooled
+
+  void* c2 = upb_Arena_AllocPool(arena, 8);
+  EXPECT_FALSE(UPB_PRIVATE(upb_Xsan_PtrEq)(c1, c2));
+
+  upb_Arena_FreePool(arena, c2, 8);
+
+  upb_Arena_Free(arena);
+}
+
 }  // namespace

upb/mem/internal/arena.h

@@ -23,9 +23,9 @@
 // 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.
 #ifndef NDEBUG
-#define UPB_ARENA_BASE_SIZE_HACK 10
+#define UPB_ARENA_BASE_SIZE_HACK 11
 #else
-#define UPB_ARENA_BASE_SIZE_HACK 9
+#define UPB_ARENA_BASE_SIZE_HACK 10
 #endif
 
 #define UPB_ARENA_SIZE_HACK                                                   \

upb/message/array.c

@@ -180,11 +180,29 @@ bool UPB_PRIVATE(_upb_Array_Realloc)(upb_Array* array, size_t min_capacity,
   if (upb_ShlOverflow(&new_bytes, lg2)) {
     return false;
   }
-  ptr = upb_Arena_Realloc(arena, ptr, old_bytes, new_bytes);
-  if (!ptr) return false;
-
-  UPB_PRIVATE(_upb_Array_SetTaggedPtr)(array, ptr, lg2);
-  array->UPB_PRIVATE(capacity) = new_capacity;
+  if (ptr && upb_Arena_TryExtend(arena, ptr, old_bytes, new_bytes)) {
+    array->UPB_PRIVATE(capacity) = new_capacity;
+  } else {
+    const size_t map_entry_size = 8 + 2 * sizeof(void*);
+    size_t ratio = (new_bytes + map_entry_size - 1) / map_entry_size;
+    size_t next_pow2 = upb_RoundUpToPowerOfTwo(ratio);
+    size_t pool_bytes = next_pow2 * map_entry_size;
+
+    void* new_ptr = upb_Arena_AllocPool(arena, pool_bytes);
+    if (!new_ptr) return false;
+
+    if (ptr && old_bytes > 0) {
+      memcpy(new_ptr, ptr, old_bytes);
+      const size_t array_size =
+          UPB_ALIGN_UP(sizeof(struct upb_Array), UPB_MALLOC_ALIGN);
+      if (ptr != UPB_PTR_AT(array, array_size, void)) {
+        upb_Arena_FreePool(arena, ptr, old_bytes);
+      }
+    }
+    ptr = new_ptr;
+    UPB_PRIVATE(_upb_Array_SetTaggedPtr)(array, ptr, lg2);
+    array->UPB_PRIVATE(capacity) = pool_bytes >> lg2;
+  }
   return true;
 }
 

upb/port/sanitizers.h

@@ -162,6 +162,14 @@ UPB_INLINE bool UPB_PRIVATE(upb_Xsan_PtrEq)(const void *a, const void *b) {
 #endif
 }
 
+UPB_INLINE void* UPB_PRIVATE(upb_Xsan_UntagPointer)(const void* ptr) {
+#if UPB_HWASAN
+  return __hwasan_tag_pointer(ptr, 0);
+#else
+  return (void*)ptr;
+#endif
+}
+
 // These annotations improve TSAN's ability to detect data races.  By
 // proactively accessing a non-atomic variable at the point where it is
 // "logically" accessed, we can trigger TSAN diagnostics that might have