Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Pass size to upb_alloc when freeing an arena #19771

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Pass size to upb_alloc when freeing an arena #19771

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions upb/mem/alloc.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,11 @@ UPB_INLINE void upb_free(upb_alloc* alloc, void* ptr) {
alloc->func(alloc, ptr, 0, 0);
}

UPB_INLINE void upb_free_sized(upb_alloc* alloc, void* ptr, size_t size) {
UPB_ASSERT(alloc);
alloc->func(alloc, ptr, size, 0);
}

// The global allocator used by upb. Uses the standard malloc()/free().

extern upb_alloc upb_alloc_global;
Expand Down
66 changes: 27 additions & 39 deletions upb/mem/arena.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,7 @@ void upb_Arena_SetMaxBlockSize(size_t max) {
}

typedef struct upb_MemBlock {
// Atomic only for the benefit of SpaceAllocated().
UPB_ATOMIC(struct upb_MemBlock*) next;
struct upb_MemBlock* next;
uint32_t size;
// Data follows.
} upb_MemBlock;
Expand Down Expand Up @@ -62,9 +61,11 @@ typedef struct upb_ArenaInternal {
// == self when no other list members.
UPB_ATOMIC(struct upb_ArenaInternal*) tail;

// Linked list of blocks to free/cleanup. Atomic only for the benefit of
// upb_Arena_SpaceAllocated().
UPB_ATOMIC(upb_MemBlock*) blocks;
// Linked list of blocks to free/cleanup.
upb_MemBlock* blocks;

// Total space allocated in blocks, atomic only for SpaceAllocated
UPB_ATOMIC(uint32_t) space_allocated;
} upb_ArenaInternal;

// All public + private state for an arena.
Expand Down Expand Up @@ -207,11 +208,7 @@ size_t upb_Arena_SpaceAllocated(upb_Arena* arena, size_t* fused_count) {
size_t local_fused_count = 0;

while (ai != NULL) {
upb_MemBlock* block = upb_Atomic_Load(&ai->blocks, memory_order_relaxed);
while (block != NULL) {
memsize += sizeof(upb_MemBlock) + block->size;
block = upb_Atomic_Load(&block->next, memory_order_relaxed);
}
memsize += upb_Atomic_Load(&ai->space_allocated, memory_order_relaxed);
ai = upb_Atomic_Load(&ai->next, memory_order_relaxed);
local_fused_count++;
}
Expand All @@ -220,21 +217,6 @@ size_t upb_Arena_SpaceAllocated(upb_Arena* arena, size_t* fused_count) {
return memsize;
}

bool UPB_PRIVATE(_upb_Arena_Contains)(const upb_Arena* a, void* ptr) {
upb_ArenaInternal* ai = upb_Arena_Internal(a);
UPB_ASSERT(ai);

upb_MemBlock* block = upb_Atomic_Load(&ai->blocks, memory_order_relaxed);
while (block) {
uintptr_t beg = (uintptr_t)block;
uintptr_t end = beg + block->size;
if ((uintptr_t)ptr >= beg && (uintptr_t)ptr < end) return true;
block = upb_Atomic_Load(&block->next, memory_order_relaxed);
}

return false;
}

uint32_t upb_Arena_DebugRefCount(upb_Arena* a) {
upb_ArenaInternal* ai = upb_Arena_Internal(a);
// These loads could probably be relaxed, but given that this is debug-only,
Expand All @@ -251,10 +233,10 @@ static void _upb_Arena_AddBlock(upb_Arena* a, void* ptr, size_t size) {
upb_ArenaInternal* ai = upb_Arena_Internal(a);
upb_MemBlock* block = ptr;

// Insert into linked list.
block->size = (uint32_t)size;
upb_Atomic_Init(&block->next, ai->blocks);
upb_Atomic_Store(&ai->blocks, block, memory_order_release);
// Insert into linked list.
block->next = ai->blocks;
ai->blocks = block;

a->UPB_PRIVATE(ptr) = UPB_PTR_AT(block, kUpb_MemblockReserve, char);
a->UPB_PRIVATE(end) = UPB_PTR_AT(block, size, char);
Expand All @@ -267,7 +249,7 @@ static bool _upb_Arena_AllocBlock(upb_Arena* a, size_t size) {
upb_ArenaInternal* ai = upb_Arena_Internal(a);
if (!ai->block_alloc) return false;
size_t last_size = 128;
upb_MemBlock* last_block = upb_Atomic_Load(&ai->blocks, memory_order_relaxed);
upb_MemBlock* last_block = ai->blocks;
if (last_block) {
last_size = a->UPB_PRIVATE(end) - (char*)last_block;
}
Expand All @@ -288,6 +270,13 @@ static bool _upb_Arena_AllocBlock(upb_Arena* a, size_t size) {

if (!block) return false;
_upb_Arena_AddBlock(a, block, block_size);
// Atomic add not required here, as threads won't race allocating blocks, plus
// atomic fetch-add is slower than load/add/store on arm devices compiled
// targetting pre-v8.1.
uint32_t old_block_size =
upb_Atomic_Load(&ai->space_allocated, memory_order_relaxed);
upb_Atomic_Store(&ai->space_allocated, block_size + old_block_size,
memory_order_relaxed);
UPB_ASSERT(UPB_PRIVATE(_upb_ArenaHas)(a) >= size);
return true;
}
Expand Down Expand Up @@ -316,7 +305,8 @@ static upb_Arena* _upb_Arena_InitSlow(upb_alloc* alloc) {
upb_Atomic_Init(&a->body.parent_or_count, _upb_Arena_TaggedFromRefcount(1));
upb_Atomic_Init(&a->body.next, NULL);
upb_Atomic_Init(&a->body.tail, &a->body);
upb_Atomic_Init(&a->body.blocks, NULL);
upb_Atomic_Init(&a->body.space_allocated, n);
a->body.blocks = NULL;
a->body.upb_alloc_cleanup = NULL;

_upb_Arena_AddBlock(&a->head, mem, n);
Expand Down Expand Up @@ -355,7 +345,8 @@ upb_Arena* upb_Arena_Init(void* mem, size_t n, upb_alloc* alloc) {
upb_Atomic_Init(&a->body.parent_or_count, _upb_Arena_TaggedFromRefcount(1));
upb_Atomic_Init(&a->body.next, NULL);
upb_Atomic_Init(&a->body.tail, &a->body);
upb_Atomic_Init(&a->body.blocks, NULL);
upb_Atomic_Init(&a->body.space_allocated, 0);
a->body.blocks = NULL;
a->body.upb_alloc_cleanup = NULL;

a->body.block_alloc = _upb_Arena_MakeBlockAlloc(alloc, 1);
Expand All @@ -374,13 +365,12 @@ static void _upb_Arena_DoFree(upb_ArenaInternal* ai) {
upb_ArenaInternal* next_arena =
(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_MemBlock* block = ai->blocks;
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);
upb_MemBlock* next_block = block->next;
upb_free_sized(block_alloc, block, block->size);
block = next_block;
}
if (alloc_cleanup != NULL) {
Expand Down Expand Up @@ -601,15 +591,13 @@ void UPB_PRIVATE(_upb_Arena_SwapIn)(upb_Arena* des, const upb_Arena* src) {

*des = *src;
desi->block_alloc = srci->block_alloc;
upb_MemBlock* blocks = upb_Atomic_Load(&srci->blocks, memory_order_relaxed);
upb_Atomic_Init(&desi->blocks, blocks);
desi->blocks = srci->blocks;
}

void UPB_PRIVATE(_upb_Arena_SwapOut)(upb_Arena* des, const upb_Arena* src) {
upb_ArenaInternal* desi = upb_Arena_Internal(des);
upb_ArenaInternal* srci = upb_Arena_Internal(src);

*des = *src;
upb_MemBlock* blocks = upb_Atomic_Load(&srci->blocks, memory_order_relaxed);
upb_Atomic_Store(&desi->blocks, blocks, memory_order_relaxed);
desi->blocks = srci->blocks;
}
64 changes: 26 additions & 38 deletions upb/mem/arena_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -172,44 +172,6 @@ TEST(ArenaTest, FuzzSingleThreaded) {
}
}

TEST(ArenaTest, Contains) {
upb_Arena* arena1 = upb_Arena_New();
upb_Arena* arena2 = upb_Arena_New();
void* ptr1a = upb_Arena_Malloc(arena1, 8);
void* ptr2a = upb_Arena_Malloc(arena2, 8);

EXPECT_TRUE(UPB_PRIVATE(_upb_Arena_Contains)(arena1, ptr1a));
EXPECT_TRUE(UPB_PRIVATE(_upb_Arena_Contains)(arena2, ptr2a));
EXPECT_FALSE(UPB_PRIVATE(_upb_Arena_Contains)(arena1, ptr2a));
EXPECT_FALSE(UPB_PRIVATE(_upb_Arena_Contains)(arena2, ptr1a));

void* ptr1b = upb_Arena_Malloc(arena1, 1000000);
void* ptr2b = upb_Arena_Malloc(arena2, 1000000);

EXPECT_TRUE(UPB_PRIVATE(_upb_Arena_Contains)(arena1, ptr1a));
EXPECT_TRUE(UPB_PRIVATE(_upb_Arena_Contains)(arena1, ptr1b));
EXPECT_TRUE(UPB_PRIVATE(_upb_Arena_Contains)(arena2, ptr2a));
EXPECT_TRUE(UPB_PRIVATE(_upb_Arena_Contains)(arena2, ptr2b));
EXPECT_FALSE(UPB_PRIVATE(_upb_Arena_Contains)(arena1, ptr2a));
EXPECT_FALSE(UPB_PRIVATE(_upb_Arena_Contains)(arena1, ptr2b));
EXPECT_FALSE(UPB_PRIVATE(_upb_Arena_Contains)(arena2, ptr1a));
EXPECT_FALSE(UPB_PRIVATE(_upb_Arena_Contains)(arena2, ptr1b));

upb_Arena_Fuse(arena1, arena2);

EXPECT_TRUE(UPB_PRIVATE(_upb_Arena_Contains)(arena1, ptr1a));
EXPECT_TRUE(UPB_PRIVATE(_upb_Arena_Contains)(arena1, ptr1b));
EXPECT_TRUE(UPB_PRIVATE(_upb_Arena_Contains)(arena2, ptr2a));
EXPECT_TRUE(UPB_PRIVATE(_upb_Arena_Contains)(arena2, ptr2b));
EXPECT_FALSE(UPB_PRIVATE(_upb_Arena_Contains)(arena1, ptr2a));
EXPECT_FALSE(UPB_PRIVATE(_upb_Arena_Contains)(arena1, ptr2b));
EXPECT_FALSE(UPB_PRIVATE(_upb_Arena_Contains)(arena2, ptr1a));
EXPECT_FALSE(UPB_PRIVATE(_upb_Arena_Contains)(arena2, ptr1b));

upb_Arena_Free(arena1);
upb_Arena_Free(arena2);
}

TEST(ArenaTest, LargeAlloc) {
// Tests an allocation larger than the max block size.
upb_Arena* arena = upb_Arena_New();
Expand Down Expand Up @@ -284,6 +246,32 @@ TEST(ArenaTest, FuzzFuseFuseRace) {
for (auto& t : threads) t.join();
}

TEST(ArenaTest, FuzzAllocSpaceAllocatedRace) {
Environment env;
upb_Arena_SetMaxBlockSize(512);
upb_Arena* arena = upb_Arena_New();
absl::Notification done;
std::vector<std::thread> threads;
for (int i = 0; i < 10; ++i) {
threads.emplace_back([&]() {
while (!done.HasBeenNotified()) {
size_t count;
upb_Arena_SpaceAllocated(arena, &count);
}
});
}

absl::BitGen gen;
auto end = absl::Now() + absl::Seconds(2);
while (absl::Now() < end) {
upb_Arena_Malloc(arena, 256);
}
done.Notify();
for (auto& t : threads) t.join();
upb_Arena_Free(arena);
upb_Arena_SetMaxBlockSize(32 << 10);
}

TEST(ArenaTest, ArenaIncRef) {
upb_Arena* arena1 = upb_Arena_New();
EXPECT_EQ(upb_Arena_DebugRefCount(arena1), 1);
Expand Down
7 changes: 1 addition & 6 deletions upb/mem/internal/arena.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 8
#define UPB_ARENA_SIZE_HACK 9

// LINT.IfChange(upb_Arena)

Expand All @@ -40,11 +40,6 @@ void UPB_PRIVATE(_upb_Arena_SwapIn)(struct upb_Arena* des,
void UPB_PRIVATE(_upb_Arena_SwapOut)(struct upb_Arena* des,
const struct upb_Arena* src);

// Returns whether |ptr| was allocated directly by |a| (so care must be used
// with fused arenas).
UPB_API bool UPB_ONLYBITS(_upb_Arena_Contains)(const struct upb_Arena* a,
void* ptr);

UPB_INLINE size_t UPB_PRIVATE(_upb_ArenaHas)(const struct upb_Arena* a) {
return (size_t)(a->UPB_ONLYBITS(end) - a->UPB_ONLYBITS(ptr));
}
Expand Down
Loading