diff --git a/upb/hash/common.c b/upb/hash/common.c index 17b68235e116d..18b950d92ca5f 100644 --- a/upb/hash/common.c +++ b/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}; diff --git a/upb/mem/BUILD b/upb/mem/BUILD index 206b7a43b2c6e..11156d4ccb859 100644 --- a/upb/mem/BUILD +++ b/upb/mem/BUILD @@ -27,6 +27,7 @@ cc_library( visibility = ["//visibility:public"], deps = [ ":internal", + "//upb/base:internal", "//upb/port", ], ) diff --git a/upb/mem/arena.c b/upb/mem/arena.c index 4c9eeca0ed4d8..b0890ca429865 100644 --- a/upb/mem/arena.c +++ b/upb/mem/arena.c @@ -19,6 +19,7 @@ #include #include +#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; + } +} diff --git a/upb/mem/arena.h b/upb/mem/arena.h index 5a9377bec91ac..f161c6da7df3f 100644 --- a/upb/mem/arena.h +++ b/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; diff --git a/upb/mem/arena_test.cc b/upb/mem/arena_test.cc index c56f052742836..9c0e72b25890c 100644 --- a/upb/mem/arena_test.cc +++ b/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 diff --git a/upb/mem/internal/arena.h b/upb/mem/internal/arena.h index 3319f91990add..cbcf4ed4e243d 100644 --- a/upb/mem/internal/arena.h +++ b/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 \ diff --git a/upb/message/array.c b/upb/message/array.c index 19446c9cbbad7..b620ba01c55f6 100644 --- a/upb/message/array.c +++ b/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; } diff --git a/upb/port/sanitizers.h b/upb/port/sanitizers.h index 850f962da1ff6..3963abfa1dc72 100644 --- a/upb/port/sanitizers.h +++ b/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