Add Proto Arena TypedBlock
PiperOrigin-RevId: 495003404
diff --git a/src/file_lists.cmake b/src/file_lists.cmake
index 1e326bc..b7367ef 100644
--- a/src/file_lists.cmake
+++ b/src/file_lists.cmake
@@ -103,7 +103,6 @@
${protobuf_SOURCE_DIR}/src/google/protobuf/arena_allocation_policy.h
${protobuf_SOURCE_DIR}/src/google/protobuf/arena_cleanup.h
${protobuf_SOURCE_DIR}/src/google/protobuf/arena_config.h
- ${protobuf_SOURCE_DIR}/src/google/protobuf/arena_impl.h
${protobuf_SOURCE_DIR}/src/google/protobuf/arenastring.h
${protobuf_SOURCE_DIR}/src/google/protobuf/arenaz_sampler.h
${protobuf_SOURCE_DIR}/src/google/protobuf/compiler/importer.h
@@ -170,6 +169,7 @@
${protobuf_SOURCE_DIR}/src/google/protobuf/repeated_field.h
${protobuf_SOURCE_DIR}/src/google/protobuf/repeated_ptr_field.h
${protobuf_SOURCE_DIR}/src/google/protobuf/service.h
+ ${protobuf_SOURCE_DIR}/src/google/protobuf/serial_arena.h
${protobuf_SOURCE_DIR}/src/google/protobuf/stubs/callback.h
${protobuf_SOURCE_DIR}/src/google/protobuf/stubs/common.h
${protobuf_SOURCE_DIR}/src/google/protobuf/stubs/logging.h
@@ -178,6 +178,7 @@
${protobuf_SOURCE_DIR}/src/google/protobuf/stubs/port.h
${protobuf_SOURCE_DIR}/src/google/protobuf/stubs/status_macros.h
${protobuf_SOURCE_DIR}/src/google/protobuf/text_format.h
+ ${protobuf_SOURCE_DIR}/src/google/protobuf/thread_safe_arena.h
${protobuf_SOURCE_DIR}/src/google/protobuf/unknown_field_set.h
${protobuf_SOURCE_DIR}/src/google/protobuf/util/delimited_message_util.h
${protobuf_SOURCE_DIR}/src/google/protobuf/util/field_comparator.h
@@ -227,7 +228,6 @@
${protobuf_SOURCE_DIR}/src/google/protobuf/arena_allocation_policy.h
${protobuf_SOURCE_DIR}/src/google/protobuf/arena_cleanup.h
${protobuf_SOURCE_DIR}/src/google/protobuf/arena_config.h
- ${protobuf_SOURCE_DIR}/src/google/protobuf/arena_impl.h
${protobuf_SOURCE_DIR}/src/google/protobuf/arenastring.h
${protobuf_SOURCE_DIR}/src/google/protobuf/arenaz_sampler.h
${protobuf_SOURCE_DIR}/src/google/protobuf/endian.h
@@ -258,6 +258,7 @@
${protobuf_SOURCE_DIR}/src/google/protobuf/port_undef.inc
${protobuf_SOURCE_DIR}/src/google/protobuf/repeated_field.h
${protobuf_SOURCE_DIR}/src/google/protobuf/repeated_ptr_field.h
+ ${protobuf_SOURCE_DIR}/src/google/protobuf/serial_arena.h
${protobuf_SOURCE_DIR}/src/google/protobuf/stubs/callback.h
${protobuf_SOURCE_DIR}/src/google/protobuf/stubs/common.h
${protobuf_SOURCE_DIR}/src/google/protobuf/stubs/logging.h
@@ -265,6 +266,7 @@
${protobuf_SOURCE_DIR}/src/google/protobuf/stubs/platform_macros.h
${protobuf_SOURCE_DIR}/src/google/protobuf/stubs/port.h
${protobuf_SOURCE_DIR}/src/google/protobuf/stubs/status_macros.h
+ ${protobuf_SOURCE_DIR}/src/google/protobuf/thread_safe_arena.h
${protobuf_SOURCE_DIR}/src/google/protobuf/wire_format_lite.h
)
diff --git a/src/google/protobuf/BUILD.bazel b/src/google/protobuf/BUILD.bazel
index 4902ce2..9581b91 100644
--- a/src/google/protobuf/BUILD.bazel
+++ b/src/google/protobuf/BUILD.bazel
@@ -244,8 +244,9 @@
hdrs = [
"arena.h",
"arena_config.h",
- "arena_impl.h",
"arenaz_sampler.h",
+ "serial_arena.h",
+ "thread_safe_arena.h",
],
include_prefix = "google/protobuf",
visibility = [
@@ -284,7 +285,6 @@
hdrs = [
"any.h",
"arena.h",
- "arena_impl.h",
"arenastring.h",
"arenaz_sampler.h",
"endian.h",
@@ -308,6 +308,8 @@
"port.h",
"repeated_field.h",
"repeated_ptr_field.h",
+ "serial_arena.h",
+ "thread_safe_arena.h",
"wire_format_lite.h",
],
copts = COPTS + select({
@@ -326,6 +328,7 @@
# In Bazel 6.0+, these will be `interface_deps`:
deps = [
":arena",
+ ":arena_align",
":arena_config",
"//src/google/protobuf/io",
"//src/google/protobuf/stubs:lite",
diff --git a/src/google/protobuf/arena.cc b/src/google/protobuf/arena.cc
index 2a62414..b208882 100644
--- a/src/google/protobuf/arena.cc
+++ b/src/google/protobuf/arena.cc
@@ -35,14 +35,19 @@
#include <cstddef>
#include <cstdint>
#include <limits>
+#include <string>
#include <typeinfo>
#include "absl/base/attributes.h"
+#include "google/protobuf/stubs/logging.h"
+#include "google/protobuf/stubs/logging.h"
#include "absl/synchronization/mutex.h"
+#include "google/protobuf/arena_align.h"
#include "google/protobuf/arena_allocation_policy.h"
-#include "google/protobuf/arena_impl.h"
#include "google/protobuf/arenaz_sampler.h"
#include "google/protobuf/port.h"
+#include "google/protobuf/serial_arena.h"
+#include "google/protobuf/thread_safe_arena.h"
#ifdef ADDRESS_SANITIZER
@@ -161,6 +166,7 @@
set_ptr(b->Pointer(offset));
limit_ = b->Limit();
head_.store(b, std::memory_order_relaxed);
+ strings_ = StringBlock::sentinel();
space_used_.store(0, std::memory_order_relaxed);
space_allocated_.store(b->size, std::memory_order_relaxed);
cached_block_length_ = 0;
@@ -195,21 +201,45 @@
return AllocateFromExisting(n);
}
-PROTOBUF_NOINLINE
-void* SerialArena::AllocateAlignedWithCleanupFallback(
- size_t n, size_t align, void (*destructor)(void*)) {
- size_t required = AlignUpTo(n, align) + cleanup::Size(destructor);
- AllocateNewBlock(required);
- return AllocateFromExistingWithCleanupFallback(n, align, destructor);
+void* SerialArena::AllocateEmbeddedFallback(size_t size, cleanup::Tag tag) {
+ AllocateNewBlock(size);
+ return BlindlyAllocateEmbedded(size, tag);
}
-PROTOBUF_NOINLINE
-void SerialArena::AddCleanupFallback(void* elem, void (*destructor)(void*)) {
- size_t required = cleanup::Size(destructor);
- AllocateNewBlock(required);
- AddCleanupFromExisting(elem, destructor);
+// https://godbolt.org/z/EcGPxbMdW
+std::string* SerialArena::AllocateStringFallback() {
+ constexpr size_t kMin = 10;
+ constexpr size_t kMax = 170;
+ size_t n = std::min(std::max(strings_->capacity() * 2, kMin), kMax);
+ strings_ = StringBlock::Create(n, strings_);
+ return strings_->Allocate();
}
+template <typename Align, typename TagOrDtor>
+void* SerialArena::AllocateWithCleanupFallback(size_t size, Align align,
+ TagOrDtor cleanup) {
+ AllocateNewBlock(align.Padded(size) + cleanup::CleanupSize(cleanup));
+ return BlindlyAllocateWithCleanup(size, align, cleanup);
+}
+
+template void* SerialArena::AllocateWithCleanupFallback(size_t,
+ ArenaAlignDefault,
+ cleanup::Tag);
+template void* SerialArena::AllocateWithCleanupFallback(size_t,
+ ArenaAlignDefault,
+ void (*)(void*));
+template void* SerialArena::AllocateWithCleanupFallback(size_t, ArenaAlign,
+ void (*)(void*));
+
+template <typename TagOrCleanup>
+void SerialArena::AddCleanupFallback(void* elem, TagOrCleanup cleanup) {
+ AllocateNewBlock(cleanup::CleanupSize(cleanup));
+ BlindlyAddCleanup(elem, cleanup);
+}
+
+template void SerialArena::AddCleanupFallback(void*, cleanup::Tag);
+template void SerialArena::AddCleanupFallback(void*, void (*)(void*));
+
void SerialArena::AllocateNewBlock(size_t n) {
size_t used = 0;
size_t wasted = 0;
@@ -271,6 +301,13 @@
}
void SerialArena::CleanupList() {
+ StringBlock* block = strings_;
+ while (StringBlock* next = block->next()) {
+ block->DestroyAll();
+ StringBlock::Delete(block);
+ block = next;
+ }
+
ArenaBlock* b = head();
if (b->IsSentry()) return;
@@ -694,29 +731,31 @@
return space_allocated;
}
-void* ThreadSafeArena::AllocateAlignedWithCleanup(size_t n, size_t align,
- void (*destructor)(void*)) {
- SerialArena* arena;
- if (PROTOBUF_PREDICT_TRUE(GetSerialArenaFast(&arena))) {
- return arena->AllocateAlignedWithCleanup(n, align, destructor);
- } else {
- return AllocateAlignedWithCleanupFallback(n, align, destructor);
- }
+template <typename Align, typename TagOrDtor>
+void* ThreadSafeArena::AllocateWithCleanupFallback(size_t size, Align align,
+ TagOrDtor cleanup) {
+ const size_t n = align.Padded(size) + cleanup::CleanupSize(cleanup);
+ SerialArena* arena = GetSerialArenaFallback(n);
+ return arena->AllocateWithCleanup(size, align, cleanup);
}
-void ThreadSafeArena::AddCleanup(void* elem, void (*cleanup)(void*)) {
- SerialArena* arena;
- if (PROTOBUF_PREDICT_FALSE(!GetSerialArenaFast(&arena))) {
- arena = GetSerialArenaFallback(kMaxCleanupNodeSize);
- }
- arena->AddCleanup(elem, cleanup);
+template void* ThreadSafeArena::AllocateWithCleanupFallback(size_t,
+ ArenaAlignDefault,
+ cleanup::Tag);
+template void* ThreadSafeArena::AllocateWithCleanupFallback(size_t,
+ ArenaAlignDefault,
+ void (*)(void*));
+template void* ThreadSafeArena::AllocateWithCleanupFallback(size_t, ArenaAlign,
+ void (*)(void*));
+
+void* ThreadSafeArena::AllocateEmbeddedFallback(size_t size, cleanup::Tag tag) {
+ SerialArena* arena = GetSerialArenaFallback(size);
+ return arena->AllocateEmbedded(size, tag);
}
-PROTOBUF_NOINLINE
-void* ThreadSafeArena::AllocateAlignedWithCleanupFallback(
- size_t n, size_t align, void (*destructor)(void*)) {
- return GetSerialArenaFallback(n + kMaxCleanupNodeSize)
- ->AllocateAlignedWithCleanup(n, align, destructor);
+std::string* ThreadSafeArena::AllocateStringFallback() {
+ SerialArena* arena = GetSerialArenaFallback(0);
+ return arena->AllocateString();
}
template <typename Functor>
@@ -845,11 +884,32 @@
return impl_.AllocateAligned<internal::AllocationClient::kArray>(n);
}
-void* Arena::AllocateAlignedWithCleanup(size_t n, size_t align,
- void (*destructor)(void*)) {
- return impl_.AllocateAlignedWithCleanup(n, align, destructor);
+template <typename TagOrCleanup>
+void Arena::AddCleanup(void* object, TagOrCleanup cleanup) {
+ impl_.AddCleanup(object, cleanup);
}
+template void Arena::AddCleanup(void*, internal::cleanup::Tag);
+template void Arena::AddCleanup(void*, void (*)(void*));
+
+void* Arena::AllocateEmbedded(size_t size, internal::cleanup::Tag tag) {
+ return impl_.AllocateEmbedded(size, tag);
+}
+
+std::string* Arena::AllocateString() { return impl_.AllocateString(); }
+
+template <typename Align, typename TagOrDtor>
+void* Arena::AllocateWithCleanup(size_t size, Align align, TagOrDtor cleanup) {
+ return impl_.AllocateWithCleanup(size, align, cleanup);
+}
+
+template void* Arena::AllocateWithCleanup(size_t, internal::ArenaAlignDefault,
+ internal::cleanup::Tag);
+template void* Arena::AllocateWithCleanup(size_t, internal::ArenaAlignDefault,
+ void (*)(void*));
+template void* Arena::AllocateWithCleanup(size_t, internal::ArenaAlign,
+ void (*)(void*));
+
} // namespace protobuf
} // namespace google
diff --git a/src/google/protobuf/arena.h b/src/google/protobuf/arena.h
index 8a7d1b4..708f918 100644
--- a/src/google/protobuf/arena.h
+++ b/src/google/protobuf/arena.h
@@ -51,8 +51,9 @@
#include <type_traits>
#include "google/protobuf/arena_align.h"
#include "google/protobuf/arena_config.h"
-#include "google/protobuf/arena_impl.h"
#include "google/protobuf/port.h"
+#include "google/protobuf/serial_arena.h"
+#include "google/protobuf/thread_safe_arena.h"
// Must be included last.
#include "google/protobuf/port_def.inc"
@@ -274,11 +275,10 @@
if (PROTOBUF_PREDICT_FALSE(arena == nullptr)) {
return new T(std::forward<Args>(args)...);
}
- auto destructor =
- internal::ObjectDestructor<std::is_trivially_destructible<T>::value,
- T>::destructor;
- return new (arena->AllocateInternal(sizeof(T), alignof(T), destructor))
- T(std::forward<Args>(args)...);
+ void* mem = std::is_trivially_destructible<T>::value
+ ? arena->AllocateAligned(sizeof(T), alignof(T))
+ : arena->AllocateWithCleanup<T>();
+ return new (mem) T(std::forward<Args>(args)...);
}
// API to delete any objects not on an arena. This can be used to safely
@@ -375,9 +375,7 @@
// arena-allocated memory.
template <typename T>
PROTOBUF_ALWAYS_INLINE void OwnDestructor(T* object) {
- if (object != nullptr) {
- impl_.AddCleanup(object, &internal::cleanup::arena_destruct_object<T>);
- }
+ if (object != nullptr) AddCleanup(object);
}
// Adds a custom member function on an object to the list of destructors that
@@ -565,15 +563,8 @@
}
}
- PROTOBUF_NDEBUG_INLINE void* AllocateInternal(size_t size, size_t align,
- void (*destructor)(void*)) {
- // Monitor allocation if needed.
- if (destructor == nullptr) {
- return AllocateAligned(size, align);
- } else {
- return AllocateAlignedWithCleanup(size, align, destructor);
- }
- }
+ template <typename T, bool enable_tags = internal::cleanup::EnableTags()>
+ void* AllocateWithCleanup();
// CreateMessage<T> requires that T supports arenas, but this private method
// works whether or not T supports arenas. These are not exposed to user code
@@ -603,12 +594,10 @@
template <typename T, typename... Args>
PROTOBUF_NDEBUG_INLINE T* DoCreateMessage(Args&&... args) {
- return InternalHelper<T>::Construct(
- AllocateInternal(sizeof(T), alignof(T),
- internal::ObjectDestructor<
- InternalHelper<T>::is_destructor_skippable::value,
- T>::destructor),
- this, std::forward<Args>(args)...);
+ void* mem = InternalHelper<T>::is_destructor_skippable::value
+ ? AllocateAligned(sizeof(T), alignof(T))
+ : AllocateWithCleanup<T>();
+ return InternalHelper<T>::Construct(mem, this, std::forward<Args>(args)...);
}
// CreateInArenaStorage is used to implement map field. Without it,
@@ -670,8 +659,17 @@
void* Allocate(size_t n);
void* AllocateForArray(size_t n);
- void* AllocateAlignedWithCleanup(size_t n, size_t align,
- void (*destructor)(void*));
+
+ template <typename T, bool enable_tags = internal::cleanup::EnableTags()>
+ void AddCleanup(T* object);
+
+ template <typename TagOrCleanup>
+ void AddCleanup(void* object, TagOrCleanup cleanup);
+
+ template <typename Align, typename TagOrDtor>
+ void* AllocateWithCleanup(size_t size, Align align, TagOrDtor cleanup);
+ void* AllocateEmbedded(size_t size, internal::cleanup::Tag tag);
+ std::string* AllocateString();
template <typename Type>
friend class internal::GenericTypeHandler;
@@ -688,6 +686,50 @@
friend struct internal::ArenaTestPeer;
};
+// Default implementation returns `AddCleanup(T*object, dtor<T>)`
+template <typename T, bool enable_tags>
+inline PROTOBUF_NDEBUG_INLINE void Arena::AddCleanup(T* object) {
+ AddCleanup(object, &internal::cleanup::arena_destruct_object<T>);
+}
+
+#ifdef PROTO_USE_TAGGED_CLEANUPS
+
+// Specialization for `AddCleanup<std::string>()`
+template <>
+inline PROTOBUF_NDEBUG_INLINE void Arena::AddCleanup<std::string, true>(
+ std::string* object) {
+ AddCleanup(object, internal::cleanup::Tag::kString);
+}
+
+// Specialization for `AddCleanup<absl::Cord>()`
+template <>
+inline PROTOBUF_NDEBUG_INLINE void Arena::AddCleanup<absl::Cord, true>(
+ absl::Cord* object) {
+ AddCleanup(object, internal::cleanup::Tag::kCord);
+}
+#endif // PROTO_USE_TAGGED_CLEANUPS
+
+template <typename T, bool enable_tags>
+inline PROTOBUF_NDEBUG_INLINE void* Arena::AllocateWithCleanup() {
+ constexpr auto align = internal::ArenaAlignOf<T>();
+ return AllocateWithCleanup(align.Ceil(sizeof(T)), align,
+ &internal::cleanup::arena_destruct_object<T>);
+}
+
+template <>
+inline PROTOBUF_NDEBUG_INLINE void*
+Arena::AllocateWithCleanup<std::string, true>() {
+ return AllocateString();
+}
+
+template <>
+inline PROTOBUF_NDEBUG_INLINE void*
+Arena::AllocateWithCleanup<absl::Cord, true>() {
+ return AllocateWithCleanup(internal::cleanup::AllocationSize<absl::Cord>(),
+ internal::ArenaAlignDefault(),
+ internal::cleanup::Tag::kCord);
+}
+
} // namespace protobuf
} // namespace google
diff --git a/src/google/protobuf/arena_cleanup.h b/src/google/protobuf/arena_cleanup.h
index d48661f..00c09f6 100644
--- a/src/google/protobuf/arena_cleanup.h
+++ b/src/google/protobuf/arena_cleanup.h
@@ -31,14 +31,16 @@
#ifndef GOOGLE_PROTOBUF_ARENA_CLEANUP_H__
#define GOOGLE_PROTOBUF_ARENA_CLEANUP_H__
+#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <string>
+#include <type_traits>
#include "absl/base/attributes.h"
#include "google/protobuf/stubs/logging.h"
-#include "google/protobuf/stubs/logging.h"
#include "absl/strings/cord.h"
+#include "google/protobuf/arena_align.h"
// Must be included last.
@@ -59,9 +61,10 @@
// lowest 2 bits of the `elem` value identifying the type of node. All node
// types must start with a `uintptr_t` that stores `Tag` in its low two bits.
enum class Tag : uintptr_t {
- kDynamic = 0, // DynamicNode
- kString = 1, // TaggedNode (std::string)
- kCord = 2, // TaggedNode (absl::Cord)
+ kDynamic = 0, // DynamicNode
+ kString = 1, // TaggedNode (std::string)
+ kCord = 2, // TaggedNode (absl::Cord)
+ kEmbeddedString = 3, // TaggedNode (std::string)
};
// DynamicNode contains the object (`elem`) that needs to be
@@ -72,57 +75,112 @@
void (*destructor)(void*);
};
+// std::max is not constexpr in c11
+template <typename T>
+inline constexpr T Max(T lhs, T rhs) {
+ return lhs < rhs ? rhs : lhs;
+}
+
// TaggedNode contains a `std::string` or `absl::Cord` object (`elem`) that
// needs to be destroyed. The lowest 2 bits of `elem` contain the non-zero
-// `kString` or `kCord` tag.
-struct TaggedNode {
+// `kString` or `kCord` tag. TaggedNode must have an alignment matching at
+// least the alignment of std::string and absl::Cord.
+struct alignas(Max(alignof(std::string), alignof(absl::Cord))) TaggedNode {
uintptr_t elem;
};
-// EnableSpecializedTags() return true if the alignment of tagged objects
+// EnableTags() return true if the alignment of tagged objects
// such as std::string allow us to poke tags in the 2 LSB bits.
-inline constexpr bool EnableSpecializedTags() {
+inline constexpr bool EnableTags() {
// For now we require 2 bits
return alignof(std::string) >= 8 && alignof(absl::Cord) >= 8;
}
-// Adds a cleanup entry identified by `tag` at memory location `pos`.
-inline ABSL_ATTRIBUTE_ALWAYS_INLINE void CreateNode(Tag tag, void* pos,
- const void* elem_raw,
- void (*destructor)(void*)) {
- auto elem = reinterpret_cast<uintptr_t>(elem_raw);
- if (EnableSpecializedTags()) {
- GOOGLE_ABSL_DCHECK_EQ(elem & 3, 0ULL); // Must be aligned
- switch (tag) {
- case Tag::kString: {
- TaggedNode n = {elem | static_cast<uintptr_t>(Tag::kString)};
- memcpy(pos, &n, sizeof(n));
- return;
- }
- case Tag::kCord: {
- TaggedNode n = {elem | static_cast<uintptr_t>(Tag::kCord)};
- memcpy(pos, &n, sizeof(n));
- return;
- }
- default:
- break;
- }
- }
- DynamicNode n = {elem, destructor};
+// Adds a cleanup entry invoking 'cleanup' on `object` at memory location `pos`.
+inline ABSL_ATTRIBUTE_ALWAYS_INLINE void CreateNode(void* pos,
+ const void* object,
+ void (*cleanup)(void*)) {
+ auto elem = reinterpret_cast<uintptr_t>(object);
+ DynamicNode n = {elem, cleanup};
memcpy(pos, &n, sizeof(n));
}
+// Adds a cleanup entry invoking the destructor of an object embedded in pos.
+inline ABSL_ATTRIBUTE_ALWAYS_INLINE void CreateNode(void* pos, Tag tag) {
+ GOOGLE_ABSL_DCHECK(tag == Tag::kEmbeddedString);
+ TaggedNode n = {static_cast<uintptr_t>(tag)};
+ memcpy(pos, &n, sizeof(n));
+}
+
+// Adds a cleanup entry invoking the destructor of `object`, who's type
+// is identified by `tag` at memory location `pos`.
+inline ABSL_ATTRIBUTE_ALWAYS_INLINE void CreateNode(void* pos,
+ const void* object,
+ Tag tag) {
+ GOOGLE_ABSL_DCHECK(tag == Tag::kString || tag == Tag::kCord);
+ GOOGLE_ABSL_DCHECK_NE(object, nullptr);
+ auto elem = reinterpret_cast<uintptr_t>(object);
+ GOOGLE_ABSL_DCHECK_EQ(elem & 3U, 0u);
+ TaggedNode n = {elem + static_cast<uintptr_t>(tag)};
+ memcpy(pos, &n, sizeof(n));
+}
+
+inline ABSL_ATTRIBUTE_ALWAYS_INLINE void Prefetch(const void* address) {
+}
+
+template <typename T>
+inline constexpr size_t NodeSize() {
+ return (std::is_same<T, std::string>::value ||
+ std::is_same<T, absl::Cord>::value)
+ ? sizeof(TaggedNode)
+ : sizeof(DynamicNode);
+}
+
+inline size_t CleanupSize(void (*)(void*)) {
+ return ArenaAlignDefault::Ceil(sizeof(DynamicNode));
+}
+inline size_t CleanupSize(Tag) {
+ return ArenaAlignDefault::Ceil(sizeof(TaggedNode));
+}
+
+template <typename T>
+inline constexpr size_t AllocationSize() {
+ return ArenaAlignDefault::Ceil(NodeSize<T>() + sizeof(T));
+}
+
+template <typename Node, typename T>
+T* NodePointerToObjectPointer(void* pos) {
+ return reinterpret_cast<T*>(reinterpret_cast<Node*>(pos) + 1);
+}
+
// Optimization: performs a prefetch on `elem_address`.
// Returns the size of the cleanup (meta) data at this address, allowing the
// caller to advance cleanup iterators without needing to examine or know
// anything about the underlying cleanup node or cleanup meta data / tags.
inline ABSL_ATTRIBUTE_ALWAYS_INLINE size_t
PrefetchNode(const void* elem_address) {
- if (EnableSpecializedTags()) {
+ if (EnableTags()) {
uintptr_t elem;
memcpy(&elem, elem_address, sizeof(elem));
- if (static_cast<Tag>(elem & 3) != Tag::kDynamic) {
- return sizeof(TaggedNode);
+ uintptr_t tag = elem & 3U;
+ elem -= tag;
+ switch (static_cast<Tag>(tag)) {
+ case Tag::kEmbeddedString:
+ GOOGLE_ABSL_DCHECK_EQ(elem, 0U);
+ return AllocationSize<std::string>();
+
+ case Tag::kString:
+ GOOGLE_ABSL_DCHECK_NE(elem, 0U);
+ Prefetch(reinterpret_cast<const void*>(elem));
+ return ArenaAlignDefault::Ceil(sizeof(TaggedNode));
+
+ case Tag::kCord:
+ GOOGLE_ABSL_DCHECK_NE(elem, 0U);
+ Prefetch(reinterpret_cast<const void*>(elem));
+ return ArenaAlignDefault::Ceil(sizeof(TaggedNode));
+
+ case Tag::kDynamic:
+ break;
}
}
return sizeof(DynamicNode);
@@ -132,23 +190,34 @@
// Returns the size of the cleanup (meta) data at this address, allowing the
// caller to advance cleanup iterators without needing to examine or know
// anything about the underlying cleanup node or cleanup meta data / tags.
-inline ABSL_ATTRIBUTE_ALWAYS_INLINE size_t DestroyNode(const void* pos) {
+inline ABSL_ATTRIBUTE_ALWAYS_INLINE size_t DestroyNode(void* pos) {
uintptr_t elem;
memcpy(&elem, pos, sizeof(elem));
- if (EnableSpecializedTags()) {
- switch (static_cast<Tag>(elem & 3)) {
- case Tag::kString: {
- // Some compilers don't like fully qualified explicit dtor calls,
- // so use an alias to avoid having to type `::`.
+ if (EnableTags()) {
+ uintptr_t tag = elem & 3U;
+ elem -= tag;
+ switch (static_cast<Tag>(tag)) {
+ case Tag::kEmbeddedString: {
using T = std::string;
- reinterpret_cast<T*>(elem - static_cast<uintptr_t>(Tag::kString))->~T();
- return sizeof(TaggedNode);
+ GOOGLE_ABSL_DCHECK_EQ(elem, 0U);
+ NodePointerToObjectPointer<TaggedNode, T>(pos)->~T();
+ return AllocationSize<std::string>();
}
+
+ case Tag::kString: {
+ using T = std::string;
+ GOOGLE_ABSL_DCHECK_NE(elem, 0U);
+ reinterpret_cast<T*>(elem)->~T();
+ return ArenaAlignDefault::Ceil(sizeof(TaggedNode));
+ }
+
case Tag::kCord: {
using T = absl::Cord;
- reinterpret_cast<T*>(elem - static_cast<uintptr_t>(Tag::kCord))->~T();
- return sizeof(TaggedNode);
+ GOOGLE_ABSL_DCHECK_NE(elem, 0U);
+ reinterpret_cast<T*>(elem)->~T();
+ return ArenaAlignDefault::Ceil(sizeof(TaggedNode));
}
+
default:
break;
}
@@ -158,62 +227,6 @@
return sizeof(DynamicNode);
}
-// Returns the `tag` identifying the type of object for `destructor` or
-// kDynamic if `destructor` does not identify a well know object type.
-inline ABSL_ATTRIBUTE_ALWAYS_INLINE Tag Type(void (*destructor)(void*)) {
- if (EnableSpecializedTags()) {
- if (destructor == &arena_destruct_object<std::string>) {
- return Tag::kString;
- }
- if (destructor == &arena_destruct_object<absl::Cord>) {
- return Tag::kCord;
- }
- }
- return Tag::kDynamic;
-}
-
-// Returns the `tag` identifying the type of object stored at memory location
-// `elem`, which represents the first uintptr_t value in the node.
-inline ABSL_ATTRIBUTE_ALWAYS_INLINE Tag Type(void* raw) {
- if (!EnableSpecializedTags()) return Tag::kDynamic;
-
- uintptr_t elem;
- memcpy(&elem, raw, sizeof(elem));
- switch (static_cast<Tag>(elem & 0x7ULL)) {
- case Tag::kDynamic:
- return Tag::kDynamic;
- case Tag::kString:
- return Tag::kString;
- case Tag::kCord:
- return Tag::kCord;
- default:
- GOOGLE_ABSL_LOG(FATAL) << "Corrupted cleanup tag: " << (elem & 0x7ULL);
- return Tag::kDynamic;
- }
-}
-
-// Returns the required size in bytes off the node type identified by `tag`.
-inline ABSL_ATTRIBUTE_ALWAYS_INLINE size_t Size(Tag tag) {
- if (!EnableSpecializedTags()) return sizeof(DynamicNode);
-
- switch (tag) {
- case Tag::kDynamic:
- return sizeof(DynamicNode);
- case Tag::kString:
- return sizeof(TaggedNode);
- case Tag::kCord:
- return sizeof(TaggedNode);
- default:
- GOOGLE_ABSL_LOG(FATAL) << "Corrupted cleanup tag: " << static_cast<int>(tag);
- return sizeof(DynamicNode);
- }
-}
-
-// Returns the required size in bytes off the node type for `destructor`.
-inline ABSL_ATTRIBUTE_ALWAYS_INLINE size_t Size(void (*destructor)(void*)) {
- return destructor == nullptr ? 0 : Size(Type(destructor));
-}
-
} // namespace cleanup
} // namespace internal
} // namespace protobuf
diff --git a/src/google/protobuf/serial_arena.h b/src/google/protobuf/serial_arena.h
new file mode 100644
index 0000000..90c85b3
--- /dev/null
+++ b/src/google/protobuf/serial_arena.h
@@ -0,0 +1,532 @@
+// Protocol Buffers - Google's data interchange format
+// Copyright 2022 Google Inc. All rights reserved.
+// https://developers.google.com/protocol-buffers/
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are
+// met:
+//
+// * Redistributions of source code must retain the above copyright
+// notice, this list of conditions and the following disclaimer.
+// * Redistributions in binary form must reproduce the above
+// copyright notice, this list of conditions and the following disclaimer
+// in the documentation and/or other materials provided with the
+// distribution.
+// * Neither the name of Google Inc. nor the names of its
+// contributors may be used to endorse or promote products derived from
+// this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+//
+// This file defines the internal class SerialArena
+
+#ifndef GOOGLE_PROTOBUF_SERIAL_ARENA_H__
+#define GOOGLE_PROTOBUF_SERIAL_ARENA_H__
+
+#include <algorithm>
+#include <atomic>
+#include <memory>
+#include <string>
+#include <type_traits>
+#include <typeinfo>
+#include <utility>
+
+#include "google/protobuf/stubs/logging.h"
+#include "google/protobuf/stubs/common.h"
+#include "google/protobuf/stubs/logging.h"
+#include "absl/numeric/bits.h"
+#include "google/protobuf/arena_align.h"
+#include "google/protobuf/arena_cleanup.h"
+#include "google/protobuf/arena_config.h"
+#include "google/protobuf/arenaz_sampler.h"
+#include "google/protobuf/port.h"
+
+
+// Must be included last.
+#include "google/protobuf/port_def.inc"
+
+namespace google {
+namespace protobuf {
+namespace internal {
+
+template <typename T>
+class TypedBlock {
+ public:
+ static constexpr size_t AllocSize(size_t count) {
+ return sizeof(TypedBlock) + sizeof(T) * count;
+ }
+
+ static TypedBlock* Emplace(void* mem, size_t size, TypedBlock* next) {
+ return new (mem) TypedBlock((size - sizeof(TypedBlock)) / sizeof(T), next);
+ }
+
+ static TypedBlock* Create(size_t count, TypedBlock* next) {
+ size_t sz = AllocSize(count);
+ return Emplace(::operator new(sz), sz, next);
+ }
+
+ static void Delete(TypedBlock* block) {
+ internal::SizedDelete(block, AllocSize(block->capacity_));
+ }
+
+ void DestroyAll() {
+ for (T& t : *this) t.~T();
+ }
+
+ static TypedBlock* sentinel();
+
+ TypedBlock* next() const { return next_; }
+ size_t capacity() const { return capacity_; }
+
+ T* begin() { return reinterpret_cast<T*>(this + 1); }
+ T* end() { return reinterpret_cast<T*>(this + 1) + count_; }
+
+ inline std::string* TryAllocate() {
+ if (ABSL_PREDICT_TRUE(count_ < capacity_)) {
+ return reinterpret_cast<T*>(this + 1) + count_++;
+ }
+ return nullptr;
+ }
+
+ inline std::string* Allocate() {
+ GOOGLE_ABSL_DCHECK_LT(count_, capacity_);
+ return reinterpret_cast<T*>(this + 1) + count_++;
+ }
+
+ private:
+ TypedBlock() = default;
+ ~TypedBlock() = default;
+ TypedBlock(size_t capacity, TypedBlock* next)
+ : next_(next), capacity_(static_cast<uint32_t>(capacity)) {
+ GOOGLE_ABSL_DCHECK_LE(capacity, std::numeric_limits<uint32_t>::max());
+ }
+
+ struct Sentinel;
+
+ struct alignas(T) {
+ TypedBlock* const next_ = nullptr;
+ const uint32_t capacity_ = 0;
+ uint32_t count_ = 0;
+ };
+};
+
+template <typename T>
+struct TypedBlock<T>::Sentinel {
+ static constexpr TypedBlock<T> kSentinel;
+};
+
+template <typename T>
+inline TypedBlock<T>* TypedBlock<T>::sentinel() {
+ return const_cast<TypedBlock<T>*>(&Sentinel::kSentinel);
+}
+
+
+// Arena blocks are variable length malloc-ed objects. The following structure
+// describes the common header for all blocks.
+struct ArenaBlock {
+ // For the sentry block with zero-size where ptr_, limit_, cleanup_nodes all
+ // point to "this".
+ constexpr ArenaBlock()
+ : next(nullptr), cleanup_nodes(this), size(0) {}
+
+ ArenaBlock(ArenaBlock* next, size_t size)
+ : next(next), cleanup_nodes(nullptr), size(size) {
+ GOOGLE_DCHECK_GT(size, sizeof(ArenaBlock));
+ }
+
+ char* Pointer(size_t n) {
+ GOOGLE_DCHECK_LE(n, size);
+ return reinterpret_cast<char*>(this) + n;
+ }
+ char* Limit() { return Pointer(size & static_cast<size_t>(-8)); }
+
+ bool IsSentry() const { return size == 0; }
+
+ ArenaBlock* const next;
+ void* cleanup_nodes;
+ const size_t size;
+ // data follows
+};
+
+enum class AllocationClient { kDefault, kArray };
+
+class ThreadSafeArena;
+
+// Tag type used to invoke the constructor of the first SerialArena.
+struct FirstSerialArena {
+ explicit FirstSerialArena() = default;
+};
+
+// A simple arena allocator. Calls to allocate functions must be properly
+// serialized by the caller, hence this class cannot be used as a general
+// purpose allocator in a multi-threaded program. It serves as a building block
+// for ThreadSafeArena, which provides a thread-safe arena allocator.
+//
+// This class manages
+// 1) Arena bump allocation + owning memory blocks.
+// 2) Maintaining a cleanup list.
+// It delegates the actual memory allocation back to ThreadSafeArena, which
+// contains the information on block growth policy and backing memory allocation
+// used.
+class PROTOBUF_EXPORT SerialArena {
+ public:
+ struct Memory {
+ void* ptr;
+ size_t size;
+ };
+
+ void CleanupList();
+ uint64_t SpaceAllocated() const {
+ return space_allocated_.load(std::memory_order_relaxed);
+ }
+ uint64_t SpaceUsed() const;
+
+ bool HasSpace(size_t n) const {
+ return n <= static_cast<size_t>(limit_ - ptr());
+ }
+
+ // See comments on `cached_blocks_` member for details.
+ PROTOBUF_ALWAYS_INLINE void* TryAllocateFromCachedBlock(size_t size) {
+ if (PROTOBUF_PREDICT_FALSE(size < 16)) return nullptr;
+ // We round up to the next larger block in case the memory doesn't match
+ // the pattern we are looking for.
+ const size_t index = absl::bit_width(size - 1) - 4;
+
+ if (index >= cached_block_length_) return nullptr;
+ auto& cached_head = cached_blocks_[index];
+ if (cached_head == nullptr) return nullptr;
+
+ void* ret = cached_head;
+ PROTOBUF_UNPOISON_MEMORY_REGION(ret, size);
+ cached_head = cached_head->next;
+ return ret;
+ }
+
+ // In kArray mode we look through cached blocks.
+ // We do not do this by default because most non-array allocations will not
+ // have the right size and will fail to find an appropriate cached block.
+ //
+ // TODO(sbenza): Evaluate if we should use cached blocks for message types of
+ // the right size. We can statically know if the allocation size can benefit
+ // from it.
+ template <AllocationClient alloc_client = AllocationClient::kDefault>
+ void* AllocateAligned(size_t n) {
+ GOOGLE_DCHECK(internal::ArenaAlignDefault::IsAligned(n));
+ GOOGLE_DCHECK_GE(limit_, ptr());
+
+ if (alloc_client == AllocationClient::kArray) {
+ if (void* res = TryAllocateFromCachedBlock(n)) {
+ return res;
+ }
+ }
+
+ if (PROTOBUF_PREDICT_FALSE(!HasSpace(n))) {
+ return AllocateAlignedFallback(n);
+ }
+ return AllocateFromExisting(n);
+ }
+
+ template <typename Align, typename TagOrDtor>
+ void* AllocateWithCleanup(size_t size, Align align, TagOrDtor cleanup);
+
+ template <typename Align, typename TagOrDtor>
+ void* TryAllocateWithCleanup(size_t size, Align align, TagOrDtor cleanup);
+
+ std::string* AllocateString();
+ std::string* TryAllocateString();
+
+ void* AllocateEmbedded(size_t size, cleanup::Tag tag);
+ void* TryAllocateEmbedded(size_t size, cleanup::Tag tag);
+
+ private:
+ void* AllocateFromExisting(size_t n) {
+ PROTOBUF_UNPOISON_MEMORY_REGION(ptr(), n);
+ void* ret = ptr();
+ set_ptr(static_cast<char*>(ret) + n);
+ return ret;
+ }
+
+ // See comments on `cached_blocks_` member for details.
+ void ReturnArrayMemory(void* p, size_t size) {
+ // We only need to check for 32-bit platforms.
+ // In 64-bit platforms the minimum allocation size from Repeated*Field will
+ // be 16 guaranteed.
+ if (sizeof(void*) < 8) {
+ if (PROTOBUF_PREDICT_FALSE(size < 16)) return;
+ } else {
+ PROTOBUF_ASSUME(size >= 16);
+ }
+
+ // We round down to the next smaller block in case the memory doesn't match
+ // the pattern we are looking for. eg, someone might have called Reserve()
+ // on the repeated field.
+ const size_t index = absl::bit_width(size) - 5;
+
+ if (PROTOBUF_PREDICT_FALSE(index >= cached_block_length_)) {
+ // We can't put this object on the freelist so make this object the
+ // freelist. It is guaranteed it is larger than the one we have, and
+ // large enough to hold another allocation of `size`.
+ CachedBlock** new_list = static_cast<CachedBlock**>(p);
+ size_t new_size = size / sizeof(CachedBlock*);
+
+ std::copy(cached_blocks_, cached_blocks_ + cached_block_length_,
+ new_list);
+
+ // We need to unpoison this memory before filling it in case it has been
+ // poisoned by another santizer client.
+ PROTOBUF_UNPOISON_MEMORY_REGION(
+ new_list + cached_block_length_,
+ (new_size - cached_block_length_) * sizeof(CachedBlock*));
+
+ std::fill(new_list + cached_block_length_, new_list + new_size, nullptr);
+
+ cached_blocks_ = new_list;
+ // Make the size fit in uint8_t. This is the power of two, so we don't
+ // need anything larger.
+ cached_block_length_ =
+ static_cast<uint8_t>(std::min(size_t{64}, new_size));
+
+ return;
+ }
+
+ auto& cached_head = cached_blocks_[index];
+ auto* new_node = static_cast<CachedBlock*>(p);
+ new_node->next = cached_head;
+ cached_head = new_node;
+ PROTOBUF_POISON_MEMORY_REGION(p, size);
+ }
+
+ public:
+ // Allocate space if the current region provides enough space.
+ bool MaybeAllocateAligned(size_t n, void** out) {
+ GOOGLE_DCHECK(internal::ArenaAlignDefault::IsAligned(n));
+ GOOGLE_DCHECK_GE(limit_, ptr());
+ if (PROTOBUF_PREDICT_FALSE(!HasSpace(n))) return false;
+ *out = AllocateFromExisting(n);
+ return true;
+ }
+
+ // If there is enough space in the current block, allocate space for one `T`
+ // object and register for destruction. The object has not been constructed
+ // and the memory returned is uninitialized.
+ template <typename T>
+ void* MaybeAllocateWithCleanup();
+
+ template <typename TagOrCleanup>
+ void AddCleanup(void* elem, TagOrCleanup cleanup);
+
+ private:
+ using StringBlock = TypedBlock<std::string>;
+ friend class ThreadSafeArena;
+
+ // Creates a new SerialArena inside mem using the remaining memory as for
+ // future allocations.
+ // The `parent` arena must outlive the serial arena, which is guaranteed
+ // because the parent manages the lifetime of the serial arenas.
+ static SerialArena* New(SerialArena::Memory mem, ThreadSafeArena& parent);
+ // Free SerialArena returning the memory passed in to New
+ template <typename Deallocator>
+ Memory Free(Deallocator deallocator);
+
+ // Members are declared here to track sizeof(SerialArena) and hotness
+ // centrally. They are (roughly) laid out in descending order of hotness.
+
+ // Next pointer to allocate from. Always 8-byte aligned. Points inside
+ // head_ (and head_->pos will always be non-canonical). We keep these
+ // here to reduce indirection.
+ std::atomic<char*> ptr_{nullptr};
+ // Limiting address up to which memory can be allocated from the head block.
+ char* limit_ = nullptr;
+ StringBlock* strings_ = StringBlock::sentinel();
+
+ std::atomic<ArenaBlock*> head_{nullptr}; // Head of linked list of blocks.
+ std::atomic<size_t> space_used_{0}; // Necessary for metrics.
+ std::atomic<size_t> space_allocated_{0};
+ ThreadSafeArena& parent_;
+
+ // Repeated*Field and Arena play together to reduce memory consumption by
+ // reusing blocks. Currently, natural growth of the repeated field types makes
+ // them allocate blocks of size `8 + 2^N, N>=3`.
+ // When the repeated field grows returns the previous block and we put it in
+ // this free list.
+ // `cached_blocks_[i]` points to the free list for blocks of size `8+2^(i+3)`.
+ // The array of freelists is grown when needed in `ReturnArrayMemory()`.
+ struct CachedBlock {
+ // Simple linked list.
+ CachedBlock* next;
+ };
+ uint8_t cached_block_length_ = 0;
+ CachedBlock** cached_blocks_ = nullptr;
+
+ // Helper getters/setters to handle relaxed operations on atomic variables.
+ ArenaBlock* head() { return head_.load(std::memory_order_relaxed); }
+ const ArenaBlock* head() const {
+ return head_.load(std::memory_order_relaxed);
+ }
+
+ char* ptr() { return ptr_.load(std::memory_order_relaxed); }
+ const char* ptr() const { return ptr_.load(std::memory_order_relaxed); }
+ void set_ptr(char* ptr) { return ptr_.store(ptr, std::memory_order_relaxed); }
+
+ // Constructor is private as only New() should be used.
+ inline SerialArena(ArenaBlock* b, ThreadSafeArena& parent);
+
+ // Constructors to handle the first SerialArena.
+ inline explicit SerialArena(ThreadSafeArena& parent);
+ inline SerialArena(FirstSerialArena, ArenaBlock* b, ThreadSafeArena& parent);
+
+ void* AllocateAlignedFallback(size_t n);
+
+ template <typename TagOrCleanup>
+ void BlindlyAddCleanup(void* elem, TagOrCleanup cleanup);
+
+ template <typename TagOrCleanup>
+ void AddCleanupFallback(void* elem, TagOrCleanup cleanup);
+
+ template <typename Align, typename TagOrDtor>
+ void* BlindlyAllocateWithCleanup(size_t size, Align align, TagOrDtor cleanup);
+
+ template <typename Align, typename TagOrDtor>
+ void* AllocateWithCleanupFallback(size_t size, Align align,
+ TagOrDtor cleanup);
+
+ void* BlindlyAllocateEmbedded(size_t size, cleanup::Tag tag);
+ void* AllocateEmbeddedFallback(size_t size, cleanup::Tag tag);
+
+ std::string* AllocateStringFallback();
+
+ inline void AllocateNewBlock(size_t n);
+ inline void Init(ArenaBlock* b, size_t offset);
+
+ public:
+ public:
+ static constexpr size_t kBlockHeaderSize =
+ ArenaAlignDefault::Ceil(sizeof(ArenaBlock));
+};
+
+template <typename TagOrCleanup>
+inline PROTOBUF_ALWAYS_INLINE void SerialArena::AddCleanup(
+ void* elem, TagOrCleanup cleanup) {
+ const size_t n = cleanup::CleanupSize(cleanup);
+ if (PROTOBUF_PREDICT_FALSE(!HasSpace(n))) {
+ return AddCleanupFallback(elem, cleanup);
+ }
+ BlindlyAddCleanup(elem, cleanup);
+}
+
+template <typename TagOrCleanup>
+inline PROTOBUF_ALWAYS_INLINE void SerialArena::BlindlyAddCleanup(
+ void* elem, TagOrCleanup cleanup) {
+ const size_t n = cleanup::CleanupSize(cleanup);
+ GOOGLE_DCHECK(HasSpace(n));
+ limit_ -= n;
+ PROTOBUF_UNPOISON_MEMORY_REGION(limit_, n);
+ cleanup::CreateNode(limit_, elem, cleanup);
+}
+
+template <typename Align, typename TagOrDtor>
+inline PROTOBUF_ALWAYS_INLINE void* SerialArena::BlindlyAllocateWithCleanup(
+ size_t size, Align align, TagOrDtor cleanup) {
+ GOOGLE_DCHECK(align.IsAligned(size));
+ char* ptr = align.CeilDefaultAligned(this->ptr());
+ PROTOBUF_UNPOISON_MEMORY_REGION(ptr, size);
+ BlindlyAddCleanup(ptr, cleanup);
+ GOOGLE_DCHECK_LE(ptr + size, limit_);
+ set_ptr(ptr + size);
+ return ptr;
+}
+
+template <typename Align, typename TagOrDtor>
+inline PROTOBUF_NDEBUG_INLINE void* SerialArena::AllocateWithCleanup(
+ size_t size, Align align, TagOrDtor cleanup) {
+ const size_t n = align.Padded(size) + cleanup::CleanupSize(cleanup);
+ if (PROTOBUF_PREDICT_TRUE(HasSpace(n))) {
+ return BlindlyAllocateWithCleanup(size, align, cleanup);
+ }
+ return AllocateWithCleanupFallback(size, align, cleanup);
+}
+
+template <typename Align, typename TagOrDtor>
+inline PROTOBUF_NDEBUG_INLINE void* SerialArena::TryAllocateWithCleanup(
+ size_t size, Align align, TagOrDtor cleanup) {
+ const size_t n = align.Padded(size) + cleanup::CleanupSize(cleanup);
+ if (PROTOBUF_PREDICT_FALSE(!HasSpace(n))) return nullptr;
+ void* ptr = BlindlyAllocateWithCleanup(size, align, cleanup);
+ PROTOBUF_ASSUME(ptr != nullptr);
+ return ptr;
+}
+
+template <typename T>
+inline PROTOBUF_ALWAYS_INLINE void* SerialArena::MaybeAllocateWithCleanup() {
+ static_assert(!std::is_trivially_destructible<T>::value, "");
+ constexpr auto align = internal::ArenaAlignOf<T>();
+ return TryAllocateWithCleanup(align.Ceil(sizeof(T)), align,
+ cleanup::arena_destruct_object<T>);
+}
+
+template <>
+inline PROTOBUF_ALWAYS_INLINE void*
+SerialArena::MaybeAllocateWithCleanup<std::string>() {
+ return TryAllocateString();
+}
+
+template <>
+inline PROTOBUF_ALWAYS_INLINE void*
+SerialArena::MaybeAllocateWithCleanup<absl::Cord>() {
+ return TryAllocateWithCleanup(cleanup::AllocationSize<absl::Cord>(),
+ ArenaAlignDefault(), cleanup::Tag::kCord);
+}
+
+inline PROTOBUF_NDEBUG_INLINE void* SerialArena::BlindlyAllocateEmbedded(
+ size_t size, cleanup::Tag tag) {
+ GOOGLE_DCHECK(HasSpace(size));
+ limit_ -= size;
+ PROTOBUF_UNPOISON_MEMORY_REGION(limit_, size);
+ cleanup::CreateNode(limit_, tag);
+ return limit_ + sizeof(cleanup::TaggedNode);
+}
+
+inline PROTOBUF_NDEBUG_INLINE void* SerialArena::AllocateEmbedded(
+ size_t size, cleanup::Tag tag) {
+ if (PROTOBUF_PREDICT_TRUE(HasSpace(size))) {
+ return BlindlyAllocateEmbedded(size, tag);
+ }
+ return AllocateEmbeddedFallback(size, tag);
+}
+
+inline PROTOBUF_NDEBUG_INLINE void* SerialArena::TryAllocateEmbedded(
+ size_t size, cleanup::Tag tag) {
+ const size_t n = sizeof(cleanup::TaggedNode) + size;
+ if (PROTOBUF_PREDICT_FALSE(!HasSpace(n))) return nullptr;
+ void* ptr = BlindlyAllocateEmbedded(size, tag);
+ PROTOBUF_ASSUME(ptr != nullptr);
+ return ptr;
+}
+
+inline PROTOBUF_NDEBUG_INLINE std::string* SerialArena::AllocateString() {
+ std::string* s = strings_->TryAllocate();
+ if (PROTOBUF_PREDICT_TRUE(s != nullptr)) return s;
+ return AllocateStringFallback();
+}
+
+inline PROTOBUF_NDEBUG_INLINE std::string* SerialArena::TryAllocateString() {
+ return strings_->TryAllocate();
+}
+
+} // namespace internal
+} // namespace protobuf
+} // namespace google
+
+#include "google/protobuf/port_undef.inc"
+
+#endif // GOOGLE_PROTOBUF_SERIAL_ARENA_H__
diff --git a/src/google/protobuf/thread_safe_arena.h b/src/google/protobuf/thread_safe_arena.h
new file mode 100644
index 0000000..757763c
--- /dev/null
+++ b/src/google/protobuf/thread_safe_arena.h
@@ -0,0 +1,320 @@
+// Protocol Buffers - Google's data interchange format
+// Copyright 2022 Google Inc. All rights reserved.
+// https://developers.google.com/protocol-buffers/
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are
+// met:
+//
+// * Redistributions of source code must retain the above copyright
+// notice, this list of conditions and the following disclaimer.
+// * Redistributions in binary form must reproduce the above
+// copyright notice, this list of conditions and the following disclaimer
+// in the documentation and/or other materials provided with the
+// distribution.
+// * Neither the name of Google Inc. nor the names of its
+// contributors may be used to endorse or promote products derived from
+// this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+//
+// This file defines the internal class ThreadSafeArena
+
+#ifndef GOOGLE_PROTOBUF_THREAD_SAFE_ARENA_H__
+#define GOOGLE_PROTOBUF_THREAD_SAFE_ARENA_H__
+
+#include <algorithm>
+#include <atomic>
+#include <string>
+#include <type_traits>
+#include <utility>
+
+#include "absl/synchronization/mutex.h"
+#include "google/protobuf/arena_align.h"
+#include "google/protobuf/arena_allocation_policy.h"
+#include "google/protobuf/arena_cleanup.h"
+#include "google/protobuf/arena_config.h"
+#include "google/protobuf/arenaz_sampler.h"
+#include "google/protobuf/port.h"
+#include "google/protobuf/serial_arena.h"
+
+// Must be included last.
+#include "google/protobuf/port_def.inc"
+
+namespace google {
+namespace protobuf {
+namespace internal {
+
+// use #ifdef the select the best implementation based on hardware / OS.
+class PROTOBUF_EXPORT ThreadSafeArena {
+ public:
+ ThreadSafeArena();
+
+ ThreadSafeArena(char* mem, size_t size);
+
+ explicit ThreadSafeArena(void* mem, size_t size,
+ const AllocationPolicy& policy);
+
+ // All protos have pointers back to the arena hence Arena must have
+ // pointer stability.
+ ThreadSafeArena(const ThreadSafeArena&) = delete;
+ ThreadSafeArena& operator=(const ThreadSafeArena&) = delete;
+ ThreadSafeArena(ThreadSafeArena&&) = delete;
+ ThreadSafeArena& operator=(ThreadSafeArena&&) = delete;
+
+ // Destructor deletes all owned heap allocated objects, and destructs objects
+ // that have non-trivial destructors, except for proto2 message objects whose
+ // destructors can be skipped. Also, frees all blocks except the initial block
+ // if it was passed in.
+ ~ThreadSafeArena();
+
+ uint64_t Reset();
+
+ uint64_t SpaceAllocated() const;
+ uint64_t SpaceUsed() const;
+
+ template <AllocationClient alloc_client = AllocationClient::kDefault>
+ void* AllocateAligned(size_t n) {
+ SerialArena* arena;
+ if (PROTOBUF_PREDICT_TRUE(GetSerialArenaFast(&arena))) {
+ return arena->AllocateAligned<alloc_client>(n);
+ } else {
+ return AllocateAlignedFallback<alloc_client>(n);
+ }
+ }
+
+ void ReturnArrayMemory(void* p, size_t size) {
+ SerialArena* arena;
+ if (PROTOBUF_PREDICT_TRUE(GetSerialArenaFast(&arena))) {
+ arena->ReturnArrayMemory(p, size);
+ }
+ }
+
+ // This function allocates n bytes if the common happy case is true and
+ // returns true. Otherwise does nothing and returns false. This strange
+ // semantics is necessary to allow callers to program functions that only
+ // have fallback function calls in tail position. This substantially improves
+ // code for the happy path.
+ PROTOBUF_NDEBUG_INLINE bool MaybeAllocateAligned(size_t n, void** out) {
+ SerialArena* arena;
+ if (PROTOBUF_PREDICT_TRUE(GetSerialArenaFast(&arena))) {
+ return arena->MaybeAllocateAligned(n, out);
+ }
+ return false;
+ }
+
+ // Add object pointer and cleanup function pointer to the list.
+ template <typename TagOrCleanup>
+ void AddCleanup(void* elem, TagOrCleanup cleanup);
+
+ template <typename Align, typename TagOrDtor>
+ void* AllocateWithCleanup(size_t size, Align align, TagOrDtor cleanup) {
+ SerialArena* arena;
+ if (PROTOBUF_PREDICT_TRUE(GetSerialArenaFast(&arena))) {
+ return arena->AllocateWithCleanup(size, align, cleanup);
+ }
+ return AllocateWithCleanupFallback(size, align, cleanup);
+ }
+
+ void* AllocateEmbedded(size_t size, cleanup::Tag tag) {
+ SerialArena* arena;
+ if (PROTOBUF_PREDICT_TRUE(GetSerialArenaFast(&arena))) {
+ return arena->AllocateEmbedded(size, tag);
+ }
+ return AllocateEmbeddedFallback(size, tag);
+ }
+
+ std::string* AllocateString() {
+ SerialArena* arena;
+ if (PROTOBUF_PREDICT_TRUE(GetSerialArenaFast(&arena))) {
+ return arena->AllocateString();
+ }
+ return AllocateStringFallback();
+ }
+
+ private:
+ friend class ArenaBenchmark;
+ friend class TcParser;
+ friend class SerialArena;
+ friend struct SerialArenaChunkHeader;
+ static uint64_t GetNextLifeCycleId();
+
+ class SerialArenaChunk;
+
+ // Returns a new SerialArenaChunk that has {id, serial} at slot 0. It may
+ // grow based on "prev_num_slots".
+ static SerialArenaChunk* NewSerialArenaChunk(uint32_t prev_capacity, void* id,
+ SerialArena* serial);
+ static SerialArenaChunk* SentrySerialArenaChunk();
+
+ // Returns the first ArenaBlock* for the first SerialArena. If users provide
+ // one, use it if it's acceptable. Otherwise returns a sentry block.
+ ArenaBlock* FirstBlock(void* buf, size_t size);
+ // Same as the above but returns a valid block if "policy" is not default.
+ ArenaBlock* FirstBlock(void* buf, size_t size,
+ const AllocationPolicy& policy);
+
+ // Adds SerialArena to the chunked list. May create a new chunk.
+ void AddSerialArena(void* id, SerialArena* serial);
+
+ // Members are declared here to track sizeof(ThreadSafeArena) and hotness
+ // centrally.
+
+ // Unique for each arena. Changes on Reset().
+ uint64_t tag_and_id_ = 0;
+
+ TaggedAllocationPolicyPtr alloc_policy_; // Tagged pointer to AllocPolicy.
+ ThreadSafeArenaStatsHandle arena_stats_;
+
+ // Adding a new chunk to head_ must be protected by mutex_.
+ absl::Mutex mutex_;
+ // Pointer to a linked list of SerialArenaChunk.
+ std::atomic<SerialArenaChunk*> head_{nullptr};
+
+ void* first_owner_;
+ // Must be declared after alloc_policy_; otherwise, it may lose info on
+ // user-provided initial block.
+ SerialArena first_arena_;
+
+ static_assert(std::is_trivially_destructible<SerialArena>{},
+ "SerialArena needs to be trivially destructible.");
+
+ const AllocationPolicy* AllocPolicy() const { return alloc_policy_.get(); }
+ void InitializeWithPolicy(const AllocationPolicy& policy);
+
+ void Init();
+
+ // Delete or Destruct all objects owned by the arena.
+ void CleanupList();
+
+ inline void CacheSerialArena(SerialArena* serial) {
+ thread_cache().last_serial_arena = serial;
+ thread_cache().last_lifecycle_id_seen = tag_and_id_;
+ }
+
+ PROTOBUF_NDEBUG_INLINE bool GetSerialArenaFast(SerialArena** arena) {
+ // If this thread already owns a block in this arena then try to use that.
+ // This fast path optimizes the case where multiple threads allocate from
+ // the same arena.
+ ThreadCache* tc = &thread_cache();
+ if (PROTOBUF_PREDICT_TRUE(tc->last_lifecycle_id_seen == tag_and_id_)) {
+ *arena = tc->last_serial_arena;
+ return true;
+ }
+ return false;
+ }
+
+ // Finds SerialArena or creates one if not found. When creating a new one,
+ // create a big enough block to accommodate n bytes.
+ SerialArena* GetSerialArenaFallback(size_t n);
+
+ template <AllocationClient alloc_client = AllocationClient::kDefault>
+ void* AllocateAlignedFallback(size_t n);
+
+ template <typename Align, typename TagOrDtor>
+ void* AllocateWithCleanupFallback(size_t size, Align align,
+ TagOrDtor cleanup);
+
+ void* AllocateEmbeddedFallback(size_t size, cleanup::Tag tag);
+ std::string* AllocateStringFallback();
+
+ // Executes callback function over SerialArenaChunk. Passes const
+ // SerialArenaChunk*.
+ template <typename Functor>
+ void WalkConstSerialArenaChunk(Functor fn) const;
+
+ // Executes callback function over SerialArenaChunk.
+ template <typename Functor>
+ void WalkSerialArenaChunk(Functor fn);
+
+ // Executes callback function over SerialArena in chunked list in reverse
+ // chronological order. Passes const SerialArena*.
+ template <typename Functor>
+ void PerConstSerialArenaInChunk(Functor 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
+ // deleting.
+ SerialArena::Memory Free(size_t* space_allocated);
+
+#ifdef _MSC_VER
+#pragma warning(disable : 4324)
+#endif
+ struct alignas(kCacheAlignment) ThreadCache {
+ // Number of per-thread lifecycle IDs to reserve. Must be power of two.
+ // To reduce contention on a global atomic, each thread reserves a batch of
+ // IDs. The following number is calculated based on a stress test with
+ // ~6500 threads all frequently allocating a new arena.
+ static constexpr size_t kPerThreadIds = 256;
+ // Next lifecycle ID available to this thread. We need to reserve a new
+ // batch, if `next_lifecycle_id & (kPerThreadIds - 1) == 0`.
+ uint64_t next_lifecycle_id{0};
+ // The ThreadCache is considered valid as long as this matches the
+ // lifecycle_id of the arena being used.
+ uint64_t last_lifecycle_id_seen{static_cast<uint64_t>(-1)};
+ SerialArena* last_serial_arena{nullptr};
+ };
+
+ // Lifecycle_id can be highly contended variable in a situation of lots of
+ // arena creation. Make sure that other global variables are not sharing the
+ // cacheline.
+#ifdef _MSC_VER
+#pragma warning(disable : 4324)
+#endif
+ using LifecycleId = uint64_t;
+ ABSL_CONST_INIT alignas(
+ kCacheAlignment) static std::atomic<LifecycleId> lifecycle_id_;
+#if defined(PROTOBUF_NO_THREADLOCAL)
+ // iOS does not support __thread keyword so we use a custom thread local
+ // storage class we implemented.
+ static ThreadCache& thread_cache();
+#elif defined(PROTOBUF_USE_DLLS)
+ // Thread local variables cannot be exposed through DLL interface but we can
+ // wrap them in static functions.
+ static ThreadCache& thread_cache();
+#else
+ static PROTOBUF_THREAD_LOCAL ThreadCache thread_cache_;
+ static ThreadCache& thread_cache() { return thread_cache_; }
+#endif
+
+ public:
+ // kBlockHeaderSize is sizeof(ArenaBlock), aligned up to the default alignment
+ // to protect the invariant that `pos` is always default aligned.
+ static constexpr size_t kBlockHeaderSize = SerialArena::kBlockHeaderSize;
+ static constexpr size_t kSerialArenaSize =
+ ArenaAlignDefault::Ceil(sizeof(SerialArena));
+ static constexpr size_t kAllocPolicySize =
+ ArenaAlignDefault::Ceil(sizeof(AllocationPolicy));
+ static constexpr size_t kMaxCleanupNodeSize = 16;
+ static_assert(ArenaAlignDefault::IsAligned(kBlockHeaderSize),
+ "kBlockHeaderSize must be default aligned.");
+ static_assert(ArenaAlignDefault::IsAligned(kSerialArenaSize),
+ "kSerialArenaSize must be default aligned.");
+};
+
+template <typename TagOrCleanup>
+inline void ThreadSafeArena::AddCleanup(void* elem, TagOrCleanup cleanup) {
+ SerialArena* arena;
+ if (PROTOBUF_PREDICT_FALSE(!GetSerialArenaFast(&arena))) {
+ arena = GetSerialArenaFallback(kMaxCleanupNodeSize);
+ }
+ arena->AddCleanup(elem, cleanup);
+}
+
+} // namespace internal
+} // namespace protobuf
+} // namespace google
+
+#include "google/protobuf/port_undef.inc"
+
+#endif // GOOGLE_PROTOBUF_THREAD_SAFE_ARENA_H__
