From d3172f5d7368a2b47baa5e996545f3de44009dac Mon Sep 17 00:00:00 2001
From: Joshua Haberman <haberman@google.com>
Date: Fri, 28 Jun 2024 14:41:48 -0700
Subject: [PATCH] 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       | 13 ++++++++++++-
 upb/mem/arena.h       |  9 +++++++++
 upb/mem/arena_test.cc | 26 ++++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/upb/mem/arena.c b/upb/mem/arena.c
index fc65f69013..544179228d 100644
--- a/upb/mem/arena.c
+++ b/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);
 
diff --git a/upb/mem/arena.h b/upb/mem/arena.h
index 83f5781f83..2ad4f5d2fa 100644
--- a/upb/mem/arena.h
+++ b/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
diff --git a/upb/mem/arena_test.cc b/upb/mem/arena_test.cc
index a7889a069b..2f0647ff27 100644
--- a/upb/mem/arena_test.cc
+++ b/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) {