Use SerialArena PiperOrigin-RevId: 488996168
diff --git a/python/google/protobuf/pyext/message.cc b/python/google/protobuf/pyext/message.cc index 108dd52..afd0007 100644 --- a/python/google/protobuf/pyext/message.cc +++ b/python/google/protobuf/pyext/message.cc
@@ -1886,7 +1886,7 @@ : io::CodedInputStream::GetDefaultRecursionLimit(); const char* ptr; internal::ParseContext ctx( - depth, false, &ptr, + depth, false, &ptr, self->message->GetArenaForAllocation(), absl::string_view(static_cast<const char*>(data.buf), data.len)); PyBuffer_Release(&data); ctx.data().pool = factory->pool->pool;
diff --git a/src/google/protobuf/arena.h b/src/google/protobuf/arena.h index abce219..d212407 100644 --- a/src/google/protobuf/arena.h +++ b/src/google/protobuf/arena.h
@@ -88,6 +88,7 @@ class EpsCopyInputStream; // defined in parse_context.h class RepeatedPtrFieldBase; // defined in repeated_ptr_field.h class TcParser; // defined in generated_message_tctable_impl.h +class ParseContext; template <typename Type> class GenericTypeHandler; // defined in repeated_field.h @@ -701,6 +702,7 @@ friend class internal::LazyField; // For CreateMaybeMessage. friend class internal::EpsCopyInputStream; // For parser performance friend class internal::TcParser; // For parser performance + friend class internal::ParseContext; friend class MessageLite; template <typename Key, typename T> friend class Map; @@ -710,6 +712,8 @@ friend struct internal::ArenaTestPeer; }; +// static_assert(std::is_standard_layout_v<Arena>, ""); + } // namespace protobuf } // namespace google
diff --git a/src/google/protobuf/arena_impl.h b/src/google/protobuf/arena_impl.h index 9d2de93..d02a697 100644 --- a/src/google/protobuf/arena_impl.h +++ b/src/google/protobuf/arena_impl.h
@@ -56,7 +56,10 @@ namespace google { namespace protobuf { +class Arena; + namespace internal { +class RepeatedPtrFieldBase; // To prevent sharing cache lines between threads #ifdef __cpp_aligned_new @@ -296,6 +299,8 @@ return true; } + Arena* parent() const { return reinterpret_cast<Arena*>(&parent_); } + // 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. @@ -336,6 +341,14 @@ AddCleanupFromExisting(elem, destructor); } + template <typename T, typename... Args> + T* MakeWithCleanup(Args&&... args) { + static_assert(!std::is_trivially_destructible_v<T>, ""); + return new (AllocateAlignedWithCleanup(sizeof(T), alignof(T), + cleanup::arena_destruct_object<T>)) + T(std::forward<Args>(args)...); + } + private: void* AllocateFromExistingWithCleanupFallback(size_t n, size_t align, void (*destructor)(void*)) { @@ -361,6 +374,7 @@ private: friend class ThreadSafeArena; + friend class internal::RepeatedPtrFieldBase; // For ReturnArrayMemory // Creates a new SerialArena inside mem using the remaining memory as for // future allocations. @@ -598,6 +612,14 @@ template <AllocationClient alloc_client = AllocationClient::kDefault> void* AllocateAlignedFallback(size_t n); + public: + SerialArena* GetSerialArena() { + SerialArena* arena; + if (PROTOBUF_PREDICT_TRUE(GetSerialArenaFast(&arena))) return arena; + return GetSerialArenaFallback(0); + } + + private: // Executes callback function over SerialArenaChunk. Passes const // SerialArenaChunk*. template <typename Functor>
diff --git a/src/google/protobuf/arenastring.cc b/src/google/protobuf/arenastring.cc index 8c21b1d..813f838 100644 --- a/src/google/protobuf/arenastring.cc +++ b/src/google/protobuf/arenastring.cc
@@ -297,6 +297,21 @@ return ptr; } +const char* EpsCopyInputStream::ReadArenaString(const char* ptr, + ArenaStringPtr* s, + SerialArena* arena) { + ScopedCheckPtrInvariants check(&s->tagged_ptr_); + GOOGLE_DCHECK(arena != nullptr); + + int size = ReadSize(&ptr); + if (!ptr) return nullptr; + + auto* str = s->NewString(arena); + ptr = ReadString(ptr, size, str); + GOOGLE_PROTOBUF_PARSER_ASSERT(ptr); + return ptr; +} + } // namespace internal } // namespace protobuf } // namespace google
diff --git a/src/google/protobuf/arenastring.h b/src/google/protobuf/arenastring.h index c6d61d7..ab9cb83 100644 --- a/src/google/protobuf/arenastring.h +++ b/src/google/protobuf/arenastring.h
@@ -369,6 +369,18 @@ } } + template <typename... Args> + inline std::string* NewString(SerialArena* arena, Args&&... args) { + if (arena == nullptr) { + auto* s = new std::string(std::forward<Args>(args)...); + return tagged_ptr_.SetAllocated(s); + } else { + auto* s = + arena->MakeWithCleanup<std::string>(std::forward<Args>(args)...); + return tagged_ptr_.SetMutableArena(s); + } + } + TaggedStringPtr tagged_ptr_; bool IsFixedSizeArena() const { return false; }
diff --git a/src/google/protobuf/extension_set_inl.h b/src/google/protobuf/extension_set_inl.h index 8256455..04db91a 100644 --- a/src/google/protobuf/extension_set_inl.h +++ b/src/google/protobuf/extension_set_inl.h
@@ -240,7 +240,7 @@ const char* p; // We can't use regular parse from string as we have to track // proper recursion depth and descriptor pools. - ParseContext tmp_ctx(ctx->depth(), false, &p, payload); + ParseContext tmp_ctx(ctx->depth(), false, &p, GetArena(), payload); tmp_ctx.data().pool = ctx->data().pool; tmp_ctx.data().factory = ctx->data().factory; GOOGLE_PROTOBUF_PARSER_ASSERT(value->_InternalParse(p, &tmp_ctx) &&
diff --git a/src/google/protobuf/generated_message_tctable_impl.h b/src/google/protobuf/generated_message_tctable_impl.h index dc2e976..1cc4a33 100644 --- a/src/google/protobuf/generated_message_tctable_impl.h +++ b/src/google/protobuf/generated_message_tctable_impl.h
@@ -654,8 +654,8 @@ static inline const char* RepeatedString(PROTOBUF_TC_PARAM_DECL); static inline const char* ParseRepeatedStringOnce( - const char* ptr, Arena* arena, SerialArena* serial_arena, - ParseContext* ctx, RepeatedPtrField<std::string>& field); + const char* ptr, SerialArena* serial_arena, ParseContext* ctx, + RepeatedPtrField<std::string>& field); static void UnknownPackedEnum(MessageLite* msg, ParseContext* ctx, const TcParseTableBase* table, uint32_t tag,
diff --git a/src/google/protobuf/generated_message_tctable_lite.cc b/src/google/protobuf/generated_message_tctable_lite.cc index bbb6144..f21ace1 100644 --- a/src/google/protobuf/generated_message_tctable_lite.cc +++ b/src/google/protobuf/generated_message_tctable_lite.cc
@@ -1518,7 +1518,7 @@ PROTOBUF_ALWAYS_INLINE inline const char* ReadStringIntoArena( MessageLite* /*msg*/, const char* ptr, ParseContext* ctx, uint32_t /*aux_idx*/, const TcParseTableBase* /*table*/, - ArenaStringPtr& field, Arena* arena) { + ArenaStringPtr& field, SerialArena* arena) { return ctx->ReadArenaString(ptr, &field, arena); } @@ -1549,7 +1549,7 @@ ptr += sizeof(TagType); hasbits |= (uint64_t{1} << data.hasbit_idx()); auto& field = RefAt<FieldType>(msg, data.offset()); - auto arena = msg->GetArenaForAllocation(); + auto* arena = ctx->data().serial_arena; if (arena) { ptr = ReadStringIntoArena(msg, ptr, ctx, data.aux_idx(), table, field, arena); @@ -1668,14 +1668,12 @@ } }; - auto* arena = field.GetOwningArena(); - SerialArena* serial_arena; - if (PROTOBUF_PREDICT_TRUE(arena != nullptr && - arena->impl_.GetSerialArenaFast(&serial_arena) && - field.PrepareForParse())) { + SerialArena* serial_arena = ctx->data().serial_arena; + if (PROTOBUF_PREDICT_TRUE(serial_arena != nullptr && + field.PrepareForParse(serial_arena))) { do { ptr += sizeof(TagType); - ptr = ParseRepeatedStringOnce(ptr, arena, serial_arena, ctx, field); + ptr = ParseRepeatedStringOnce(ptr, serial_arena, ctx, field); if (PROTOBUF_PREDICT_FALSE(ptr == nullptr || !validate_last_string())) { return Error(PROTOBUF_TC_PARAM_PASS); @@ -1785,7 +1783,7 @@ case field_layout::kRepMessage: case field_layout::kRepGroup: { auto& field = RefAt<MessageLite*>(msg, current_entry->offset); - if (!msg->GetArenaForAllocation()) { + if (ctx->data().serial_arena == nullptr) { delete field; } break; @@ -1828,9 +1826,9 @@ if (split == default_split) { // Allocate split instance when needed. uint32_t size = GetSizeofSplit(table); - Arena* arena = msg->GetArenaForAllocation(); + auto* arena = ctx->data().serial_arena; split = (arena == nullptr) ? ::operator new(size) - : arena->AllocateAligned(size); + : arena->AllocateAligned(AlignUpTo8(size)); memcpy(split, default_split, size); } out = split; @@ -2202,9 +2200,9 @@ case field_layout::kRepAString: { auto& field = RefAt<ArenaStringPtr>(base, entry.offset); if (need_init) field.InitDefault(); - Arena* arena = msg->GetArenaForAllocation(); - if (arena) { - ptr = ctx->ReadArenaString(ptr, &field, arena); + SerialArena* serial_arena = ctx->data().serial_arena; + if (serial_arena) { + ptr = ctx->ReadArenaString(ptr, &field, serial_arena); } else { std::string* str = field.MutableNoCopy(nullptr); ptr = InlineGreedyStringParser(str, ptr, ctx); @@ -2226,12 +2224,12 @@ } PROTOBUF_ALWAYS_INLINE const char* TcParser::ParseRepeatedStringOnce( - const char* ptr, Arena* arena, SerialArena* serial_arena, ParseContext* ctx, + const char* ptr, SerialArena* serial_arena, ParseContext* ctx, RepeatedPtrField<std::string>& field) { int size = ReadSize(&ptr); if (PROTOBUF_PREDICT_FALSE(!ptr)) return {}; - auto* str = Arena::Create<std::string>(arena); - field.AddAllocatedForParse(str); + auto* str = serial_arena->MakeWithCleanup<std::string>(); + field.AddAllocatedForParse(str, serial_arena); ptr = ctx->ReadString(ptr, size, str); if (PROTOBUF_PREDICT_FALSE(!ptr)) return {}; PROTOBUF_ASSUME(ptr != nullptr); @@ -2257,15 +2255,12 @@ const char* ptr2 = ptr; uint32_t next_tag; - auto* arena = field.GetOwningArena(); - SerialArena* serial_arena; - if (PROTOBUF_PREDICT_TRUE( - arena != nullptr && - arena->impl_.GetSerialArenaFast(&serial_arena) && - field.PrepareForParse())) { + SerialArena* serial_arena = ctx->data().serial_arena; + if (PROTOBUF_PREDICT_TRUE(serial_arena != nullptr && + field.PrepareForParse(serial_arena))) { do { ptr = ptr2; - ptr = ParseRepeatedStringOnce(ptr, arena, serial_arena, ctx, field); + ptr = ParseRepeatedStringOnce(ptr, serial_arena, ctx, field); if (PROTOBUF_PREDICT_FALSE(ptr == nullptr || !MpVerifyUtf8(field[field.size() - 1], table, entry, xform_val))) {
diff --git a/src/google/protobuf/inlined_string_field.cc b/src/google/protobuf/inlined_string_field.cc index 34b73d2..1ae31ee8 100644 --- a/src/google/protobuf/inlined_string_field.cc +++ b/src/google/protobuf/inlined_string_field.cc
@@ -75,6 +75,14 @@ return UnsafeMutablePointer(); } +std::string* InlinedStringField::MutableSlow( + ::google::protobuf::internal::SerialArena* arena, bool donated, + uint32_t* donating_states, uint32_t mask, MessageLite* msg) { + (void)mask; + (void)msg; + return UnsafeMutablePointer(); +} + void InlinedStringField::SetAllocated(const std::string* default_value, std::string* value, Arena* arena, bool donated, uint32_t* donating_states,
diff --git a/src/google/protobuf/inlined_string_field.h b/src/google/protobuf/inlined_string_field.h index fe52a4d..b1d89b6 100644 --- a/src/google/protobuf/inlined_string_field.h +++ b/src/google/protobuf/inlined_string_field.h
@@ -365,6 +365,9 @@ alignas(std::string) char value_[sizeof(std::string)]; + std::string* MutableSlow(::google::protobuf::internal::SerialArena* arena, bool donated, + uint32_t* donating_states, uint32_t mask, + MessageLite* msg); std::string* MutableSlow(::google::protobuf::Arena* arena, bool donated, uint32_t* donating_states, uint32_t mask, MessageLite* msg);
diff --git a/src/google/protobuf/lazy_field_lite.cc b/src/google/protobuf/lazy_field_lite.cc index 1830523..78ceb47 100644 --- a/src/google/protobuf/lazy_field_lite.cc +++ b/src/google/protobuf/lazy_field_lite.cc
@@ -65,7 +65,8 @@ const char* ptr; ParseContext local_ctx = ctx != nullptr ? ctx->Spawn(&ptr, input) - : ParseContext(GetInitDepth(option), false, &ptr, input); + : ParseContext(GetInitDepth(option), false, &ptr, + message->GetArenaForAllocation(), input); if (ctx == nullptr || ctx->lazy_parse_mode() == ParseContext::LazyParseMode::kEagerVerify) {
diff --git a/src/google/protobuf/message_lite.cc b/src/google/protobuf/message_lite.cc index 161758d..021cf9b 100644 --- a/src/google/protobuf/message_lite.cc +++ b/src/google/protobuf/message_lite.cc
@@ -149,7 +149,8 @@ MessageLite::ParseFlags parse_flags) { const char* ptr; internal::ParseContext ctx(io::CodedInputStream::GetDefaultRecursionLimit(), - aliasing, &ptr, input); + aliasing, &ptr, msg->GetArenaForAllocation(), + input); ptr = msg->_InternalParse(ptr, &ctx); // ctx has an explicit limit set (length of string_view). if (PROTOBUF_PREDICT_TRUE(ptr && ctx.EndedAtLimit())) { @@ -163,7 +164,8 @@ MessageLite::ParseFlags parse_flags) { const char* ptr; internal::ParseContext ctx(io::CodedInputStream::GetDefaultRecursionLimit(), - aliasing, &ptr, input); + aliasing, &ptr, msg->GetArenaForAllocation(), + input); ptr = msg->_InternalParse(ptr, &ctx); // ctx has no explicit limit (hence we end on end of stream) if (PROTOBUF_PREDICT_TRUE(ptr && ctx.EndedAtEndOfStream())) { @@ -177,7 +179,8 @@ MessageLite::ParseFlags parse_flags) { const char* ptr; internal::ParseContext ctx(io::CodedInputStream::GetDefaultRecursionLimit(), - aliasing, &ptr, input.zcis, input.limit); + aliasing, &ptr, msg->GetArenaForAllocation(), + input.zcis, input.limit); ptr = msg->_InternalParse(ptr, &ctx); if (PROTOBUF_PREDICT_FALSE(!ptr)) return false; ctx.BackUp(ptr); @@ -227,7 +230,7 @@ ZeroCopyCodedInputStream zcis(input); const char* ptr; internal::ParseContext ctx(input->RecursionBudget(), zcis.aliasing_enabled(), - &ptr, &zcis); + &ptr, GetArenaForAllocation(), &zcis); // MergePartialFromCodedStream allows terminating the wireformat by 0 or // end-group tag. Leaving it up to the caller to verify correct ending by // calling LastTagWas on input. We need to maintain this behavior.
diff --git a/src/google/protobuf/message_lite.h b/src/google/protobuf/message_lite.h index 38e903e..888a0ae 100644 --- a/src/google/protobuf/message_lite.h +++ b/src/google/protobuf/message_lite.h
@@ -440,8 +440,11 @@ // Returns the arena, used for allocating internal objects(e.g., child // messages, etc), or owning incoming objects (e.g., set allocated). + public: + // XXX Fix this Arena* GetArenaForAllocation() const { return _internal_metadata_.arena(); } + protected: // Returns true if this message is enabled for message-owned arena (MOA) // trials. No lite messages are eligible for MOA. static bool InMoaTrial() { return false; }
diff --git a/src/google/protobuf/parse_context.h b/src/google/protobuf/parse_context.h index e1ff16c..b94a578 100644 --- a/src/google/protobuf/parse_context.h +++ b/src/google/protobuf/parse_context.h
@@ -182,6 +182,9 @@ PROTOBUF_NODISCARD const char* ReadArenaString(const char* ptr, ArenaStringPtr* s, Arena* arena); + PROTOBUF_NODISCARD const char* ReadArenaString(const char* ptr, + ArenaStringPtr* s, + SerialArena* arena); template <typename Tag, typename T> PROTOBUF_NODISCARD const char* ReadRepeatedFixed(const char* ptr, @@ -401,11 +404,14 @@ struct Data { const DescriptorPool* pool = nullptr; MessageFactory* factory = nullptr; + internal::SerialArena* serial_arena = nullptr; }; template <typename... T> - ParseContext(int depth, bool aliasing, const char** start, T&&... args) + ParseContext(int depth, bool aliasing, const char** start, Arena* arena, + T&&... args) : EpsCopyInputStream(aliasing), depth_(depth) { + if (arena) data().serial_arena = arena->impl_.GetSerialArena(); *start = InitFrom(std::forward<T>(args)...); } @@ -426,7 +432,8 @@ // The spawned context always disables aliasing (different input). template <typename... T> ParseContext Spawn(const char** start, T&&... args) { - ParseContext spawned(depth_, false, start, std::forward<T>(args)...); + ParseContext spawned(depth_, false, start, nullptr, + std::forward<T>(args)...); // Transfer key context states. spawned.data_ = data_; return spawned;
diff --git a/src/google/protobuf/repeated_ptr_field.cc b/src/google/protobuf/repeated_ptr_field.cc index 1661474..318a24b 100644 --- a/src/google/protobuf/repeated_ptr_field.cc +++ b/src/google/protobuf/repeated_ptr_field.cc
@@ -92,6 +92,39 @@ return &rep_->elements[current_size_]; } +void RepeatedPtrFieldBase::InternalExtendForParse(SerialArena* serial_arena) { + int new_size = current_size_ + 1; + if (total_size_ >= new_size) { + return; + } + Rep* old_rep = rep_; + new_size = internal::CalculateReserveSize<void*, kRepHeaderSize>(total_size_, + new_size); + GOOGLE_CHECK_LE(static_cast<int64_t>(new_size), + static_cast<int64_t>( + (std::numeric_limits<size_t>::max() - kRepHeaderSize) / + sizeof(old_rep->elements[0]))) + << "Requested size is too large to fit into size_t."; + size_t bytes = kRepHeaderSize + sizeof(old_rep->elements[0]) * new_size; + rep_ = reinterpret_cast<Rep*>( + serial_arena->AllocateAligned<AllocationClient::kArray>(bytes)); + const int old_total_size = total_size_; + total_size_ = new_size; + if (old_rep) { + if (old_rep->allocated_size > 0) { + memcpy(rep_->elements, old_rep->elements, + old_rep->allocated_size * sizeof(rep_->elements[0])); + } + rep_->allocated_size = old_rep->allocated_size; + + const size_t old_size = + old_total_size * sizeof(rep_->elements[0]) + kRepHeaderSize; + serial_arena->ReturnArrayMemory(old_rep, old_size); + } else { + rep_->allocated_size = 0; + } +} + void RepeatedPtrFieldBase::Reserve(int new_size) { if (new_size > current_size_) { InternalExtend(new_size - current_size_);
diff --git a/src/google/protobuf/repeated_ptr_field.h b/src/google/protobuf/repeated_ptr_field.h index 8d3a310..c0600e9 100644 --- a/src/google/protobuf/repeated_ptr_field.h +++ b/src/google/protobuf/repeated_ptr_field.h
@@ -322,9 +322,9 @@ // - there are no preallocated elements. // Returns true if the invariants hold and `AddAllocatedForParse` can be // used. - bool PrepareForParse() { + bool PrepareForParse(SerialArena* arena) { if (current_size_ == total_size_) { - InternalExtend(1); + InternalExtendForParse(arena); } return rep_->allocated_size == current_size_; } @@ -333,12 +333,13 @@ // Can only be invoked after a call to `PrepareForParse` that returned `true`, // or other calls to `AddAllocatedForParse`. template <typename TypeHandler> - void AddAllocatedForParse(typename TypeHandler::Type* value) { + void AddAllocatedForParse(typename TypeHandler::Type* value, + SerialArena* arena) { PROTOBUF_ASSUME(rep_ != nullptr); PROTOBUF_ASSUME(current_size_ == rep_->allocated_size); if (current_size_ == total_size_) { // The array is completely full with no cleared objects, so grow it. - InternalExtend(1); + InternalExtendForParse(arena); } rep_->elements[current_size_++] = value; ++rep_->allocated_size; @@ -767,6 +768,7 @@ // common behavior from Reserve() and MergeFrom() to reduce code size. // |extend_amount| must be > 0. void** InternalExtend(int extend_amount); + void InternalExtendForParse(SerialArena* serial_arena); // Internal helper for Add: adds "obj" as the next element in the // array, including potentially resizing the array with Reserve if @@ -1198,8 +1200,8 @@ void ExtractSubrangeInternal(int start, int num, Element** elements, std::false_type); - void AddAllocatedForParse(Element* p) { - return RepeatedPtrFieldBase::AddAllocatedForParse<TypeHandler>(p); + void AddAllocatedForParse(Element* p, internal::SerialArena* arena) { + return RepeatedPtrFieldBase::AddAllocatedForParse<TypeHandler>(p, arena); } friend class Arena;
diff --git a/src/google/protobuf/wire_format.cc b/src/google/protobuf/wire_format.cc index ff542b7..8c1e5f7 100644 --- a/src/google/protobuf/wire_format.cc +++ b/src/google/protobuf/wire_format.cc
@@ -696,7 +696,8 @@ const char* p; // We can't use regular parse from string as we have to track // proper recursion depth and descriptor pools. - ParseContext tmp_ctx(ctx->depth(), false, &p, payload); + ParseContext tmp_ctx(ctx->depth(), false, &p, + value->GetArenaForAllocation(), payload); tmp_ctx.data().pool = ctx->data().pool; tmp_ctx.data().factory = ctx->data().factory; GOOGLE_PROTOBUF_PARSER_ASSERT(value->_InternalParse(p, &tmp_ctx) &&