Skip to content

Commit

Permalink
Fix races in arena fusing
Browse files Browse the repository at this point in the history 
* Add acquire/release where necessary for all atomic ops
* Add sentinel member to ensure safe publication when tsan is active; tsan will not catch the previous errors without this member.
* For all operations using relaxed memory order, comment why relaxed order is safe
* Add a test that exercises racy fuses and space allocated checks without mutexes or other memory barriers from the test harness

PiperOrigin-RevId: 708346752
  • Loading branch information
@protobuf-github-bot @copybara-github
protobuf-github-bot authored and copybara-github committed Dec 23, 2024
1 parent 1223341 commit eba3650
Show file tree
Hide file tree
Showing 6 changed files with 162 additions and 93 deletions.
1 change: 1 addition & 0 deletions upb/mem/BUILD
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@ cc_test(
":mem",
"//upb:port",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/cleanup",
"@com_google_absl//absl/random",
"@com_google_absl//absl/random:distributions",
"@com_google_absl//absl/synchronization",
Expand Down
107 changes: 58 additions & 49 deletions upb/mem/arena.c
View file
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 All @@ -56,15 +55,19 @@ typedef struct upb_ArenaInternal {
// == NULL at end of list.
UPB_ATOMIC(struct upb_ArenaInternal*) next;

UPB_TSAN_PUBLISHED_MEMBER;

// The last element of the linked list. This is present only as an
// optimization, so that we do not have to iterate over all members for every
// fuse. Only significant for an arena root. In other cases it is ignored.
// == 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(size_t) space_allocated;
} upb_ArenaInternal;

// All public + private state for an arena.
Expand Down Expand Up @@ -168,6 +171,7 @@ static upb_ArenaRoot _upb_Arena_FindRoot(const upb_Arena* a) {
uintptr_t poc = upb_Atomic_Load(&ai->parent_or_count, memory_order_acquire);
while (_upb_Arena_IsTaggedPointer(poc)) {
upb_ArenaInternal* next = _upb_Arena_PointerFromTagged(poc);
UPB_TSAN_CHECK_PUBLISHED(next);
UPB_ASSERT(ai != next);
uintptr_t next_poc =
upb_Atomic_Load(&next->parent_or_count, memory_order_acquire);
Expand All @@ -193,7 +197,7 @@ static upb_ArenaRoot _upb_Arena_FindRoot(const upb_Arena* a) {
// be valid and the creation of the path carries all the memory orderings
// required.
UPB_ASSERT(ai != _upb_Arena_PointerFromTagged(next_poc));
upb_Atomic_Store(&ai->parent_or_count, next_poc, memory_order_relaxed);
upb_Atomic_Store(&ai->parent_or_count, next_poc, memory_order_release);
}
ai = next;
poc = next_poc;
Expand All @@ -207,41 +211,25 @@ 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);
}
ai = upb_Atomic_Load(&ai->next, memory_order_relaxed);
UPB_TSAN_CHECK_PUBLISHED(ai);
// Relaxed is safe - no subsequent reads depend this one
memsize += upb_Atomic_Load(&ai->space_allocated, memory_order_relaxed);

ai = upb_Atomic_Load(&ai->next, memory_order_acquire);
local_fused_count++;
}

if (fused_count) *fused_count = local_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,
// it's not worth introducing a new variant for it.
UPB_TSAN_CHECK_PUBLISHED(ai);
uintptr_t poc = upb_Atomic_Load(&ai->parent_or_count, memory_order_acquire);
while (_upb_Arena_IsTaggedPointer(poc)) {
ai = _upb_Arena_PointerFromTagged(poc);
UPB_TSAN_CHECK_PUBLISHED(ai);
poc = upb_Atomic_Load(&ai->parent_or_count, memory_order_acquire);
}
return _upb_Arena_RefCountFromTagged(poc);
Expand All @@ -251,10 +239,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 +255,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 +276,14 @@ 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. Relaxed order is safe as nothing depends on order of
// size allocated.
size_t old_space_allocated =
upb_Atomic_Load(&ai->space_allocated, memory_order_relaxed);
upb_Atomic_Store(&ai->space_allocated, old_space_allocated + block_size,
memory_order_relaxed);
UPB_ASSERT(UPB_PRIVATE(_upb_ArenaHas)(a) >= size);
return true;
}
Expand Down Expand Up @@ -316,8 +312,10 @@ 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_TSAN_INIT_PUBLISHED(&a->body);

_upb_Arena_AddBlock(&a->head, mem, n);

Expand Down Expand Up @@ -355,12 +353,13 @@ 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);
a->head.UPB_PRIVATE(ptr) = mem;
a->head.UPB_PRIVATE(end) = UPB_PTR_AT(mem, n - sizeof(upb_ArenaState), char);
UPB_TSAN_INIT_PUBLISHED(&a->body);
#ifdef UPB_TRACING_ENABLED
upb_Arena_LogInit(&a->head, n);
#endif
Expand All @@ -370,16 +369,21 @@ upb_Arena* upb_Arena_Init(void* mem, size_t n, upb_alloc* alloc) {
static void _upb_Arena_DoFree(upb_ArenaInternal* ai) {
UPB_ASSERT(_upb_Arena_RefCountFromTagged(ai->parent_or_count) == 1);
while (ai != NULL) {
UPB_TSAN_CHECK_PUBLISHED(ai);
// Load first since arena itself is likely from one of its blocks.
upb_ArenaInternal* next_arena =
(upb_ArenaInternal*)upb_Atomic_Load(&ai->next, memory_order_acquire);
// Freeing may have memory barriers that confuse tsan, so assert immdiately
// after load here
if (next_arena) {
UPB_TSAN_CHECK_PUBLISHED(next_arena);
}
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_MemBlock* next_block = block->next;
upb_free(block_alloc, block);
block = next_block;
}
Expand All @@ -396,6 +400,7 @@ void upb_Arena_Free(upb_Arena* a) {
retry:
while (_upb_Arena_IsTaggedPointer(poc)) {
ai = _upb_Arena_PointerFromTagged(poc);
UPB_TSAN_CHECK_PUBLISHED(ai);
poc = upb_Atomic_Load(&ai->parent_or_count, memory_order_acquire);
}

Expand Down Expand Up @@ -425,29 +430,33 @@ void upb_Arena_Free(upb_Arena* a) {

static void _upb_Arena_DoFuseArenaLists(upb_ArenaInternal* const parent,
upb_ArenaInternal* child) {
UPB_TSAN_CHECK_PUBLISHED(parent);
upb_ArenaInternal* parent_tail =
upb_Atomic_Load(&parent->tail, memory_order_relaxed);
upb_Atomic_Load(&parent->tail, memory_order_acquire);

do {
UPB_TSAN_CHECK_PUBLISHED(parent_tail);
// Our tail might be stale, but it will always converge to the true tail.
upb_ArenaInternal* parent_tail_next =
upb_Atomic_Load(&parent_tail->next, memory_order_relaxed);
upb_Atomic_Load(&parent_tail->next, memory_order_acquire);
while (parent_tail_next != NULL) {
parent_tail = parent_tail_next;
UPB_TSAN_CHECK_PUBLISHED(parent_tail);
parent_tail_next =
upb_Atomic_Load(&parent_tail->next, memory_order_relaxed);
upb_Atomic_Load(&parent_tail->next, memory_order_acquire);
}

UPB_TSAN_CHECK_PUBLISHED(child);
upb_ArenaInternal* displaced =
upb_Atomic_Exchange(&parent_tail->next, child, memory_order_relaxed);
parent_tail = upb_Atomic_Load(&child->tail, memory_order_relaxed);
upb_Atomic_Exchange(&parent_tail->next, child, memory_order_acq_rel);
parent_tail = upb_Atomic_Load(&child->tail, memory_order_acquire);

// If we displaced something that got installed racily, we can simply
// reinstall it on our new tail.
child = displaced;
} while (child != NULL);

upb_Atomic_Store(&parent->tail, parent_tail, memory_order_relaxed);
upb_Atomic_Store(&parent->tail, parent_tail, memory_order_release);
}

void upb_Arena_SetAllocCleanup(upb_Arena* a, upb_AllocCleanupFunc* func) {
Expand Down Expand Up @@ -524,6 +533,8 @@ static bool _upb_Arena_FixupRefs(upb_ArenaInternal* new_root,
if (_upb_Arena_IsTaggedPointer(poc)) return false;
uintptr_t with_refs = poc - ref_delta;
UPB_ASSERT(!_upb_Arena_IsTaggedPointer(with_refs));
// Relaxed order is safe here, because the value we're setting is a refcount,
// never a tagged pointer, and the refcount is not needed for ordering.
return upb_Atomic_CompareExchangeStrong(&new_root->parent_or_count, &poc,
with_refs, memory_order_relaxed,
memory_order_relaxed);
Expand Down Expand Up @@ -601,15 +612,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;
}
Loading

0 comments on commit eba3650

Please sign in to comment.