Add a maximum block size for arena growth.

This should significantly reduce the size of large arenas. Previously, a large arena would nearly double in size if the most recent block filled up.  This could end up wasting large amounts of memory.  After this CL, we will waste at most the max block size, which defaults to 32k.

This more or less matches the behavior of the C++ arena.

PiperOrigin-RevId: 647802280

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

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

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<char*>(upb_Arena_Malloc(arena, size));
+  EXPECT_NE(mem, nullptr);
+  for (size_t i = 0; i < size; ++i) {
+    mem[i] = static_cast<char>(i);
+  }
+  for (size_t i = 0; i < size; ++i) {
+    EXPECT_EQ(mem[i], static_cast<char>(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) {