diff --git a/src/google/protobuf/arena.cc b/src/google/protobuf/arena.cc index 170d004c96b1b..bb19c946d673e 100644 --- a/src/google/protobuf/arena.cc +++ b/src/google/protobuf/arena.cc @@ -793,8 +793,8 @@ void* ThreadSafeArena::AllocateFromStringBlock() { return GetSerialArena()->AllocateFromStringBlock(); } -template -void ThreadSafeArena::WalkConstSerialArenaChunk(Functor fn) const { +template +void ThreadSafeArena::WalkConstSerialArenaChunk(Callback fn) const { const SerialArenaChunk* chunk = head_.load(std::memory_order_acquire); for (; !chunk->IsSentry(); chunk = chunk->next_chunk()) { @@ -804,8 +804,8 @@ void ThreadSafeArena::WalkConstSerialArenaChunk(Functor fn) const { } } -template -void ThreadSafeArena::WalkSerialArenaChunk(Functor fn) { +template +void ThreadSafeArena::WalkSerialArenaChunk(Callback fn) { // By omitting an Acquire barrier we help the sanitizer that any user code // that doesn't properly synchronize Reset() or the destructor will throw a // TSAN warning. @@ -821,8 +821,12 @@ void ThreadSafeArena::WalkSerialArenaChunk(Functor fn) { } } -template -void ThreadSafeArena::PerConstSerialArenaInChunk(Functor fn) const { +template +void ThreadSafeArena::VisitSerialArena(Callback fn) const { + // In most cases, arenas are single-threaded and "first_arena_" should be + // sufficient. + fn(&first_arena_); + WalkConstSerialArenaChunk([&fn](const SerialArenaChunk* chunk) { for (const auto& each : chunk->arenas()) { const SerialArena* serial = each.load(std::memory_order_acquire); @@ -835,18 +839,17 @@ void ThreadSafeArena::PerConstSerialArenaInChunk(Functor fn) const { } uint64_t ThreadSafeArena::SpaceAllocated() const { - uint64_t space_allocated = first_arena_.SpaceAllocated(); - PerConstSerialArenaInChunk([&space_allocated](const SerialArena* serial) { + uint64_t space_allocated = 0; + VisitSerialArena([&space_allocated](const SerialArena* serial) { space_allocated += serial->SpaceAllocated(); }); return space_allocated; } uint64_t ThreadSafeArena::SpaceUsed() const { - // First arena is inlined to ThreadSafeArena and the first block's overhead is - // smaller than others that contain SerialArena. - uint64_t space_used = first_arena_.SpaceUsed(); - PerConstSerialArenaInChunk([&space_used](const SerialArena* serial) { + // `first_arena_` doesn't have kSerialArenaSize overhead, so adjust it here. + uint64_t space_used = kSerialArenaSize; + VisitSerialArena([&space_used](const SerialArena* serial) { // SerialArena on chunks directly allocated from the block and needs to be // subtracted from SpaceUsed. space_used += serial->SpaceUsed() - kSerialArenaSize; diff --git a/src/google/protobuf/thread_safe_arena.h b/src/google/protobuf/thread_safe_arena.h index b721229fe6a23..21fb2c20aae68 100644 --- a/src/google/protobuf/thread_safe_arena.h +++ b/src/google/protobuf/thread_safe_arena.h @@ -189,17 +189,18 @@ class PROTOBUF_EXPORT ThreadSafeArena { // Executes callback function over SerialArenaChunk. Passes const // SerialArenaChunk*. - template - void WalkConstSerialArenaChunk(Functor fn) const; + template + void WalkConstSerialArenaChunk(Callback fn) const; // Executes callback function over SerialArenaChunk. - template - void WalkSerialArenaChunk(Functor fn); - - // Executes callback function over SerialArena in chunked list in reverse - // chronological order. Passes const SerialArena*. - template - void PerConstSerialArenaInChunk(Functor fn) const; + template + void WalkSerialArenaChunk(Callback fn); + + // Visits SerialArena and calls "fn", including "first_arena" and ones on + // chunks. Do not rely on the order of visit. The callback function should + // accept `const SerialArena*`. + template + void VisitSerialArena(Callback fn) const; // Releases all memory except the first block which it returns. The first // block might be owned by the user and thus need some extra checks before