From d6ecec884c687b04de6ab2d45fca3d115892a732 Mon Sep 17 00:00:00 2001 From: Protobuf Team Bot Date: Fri, 21 Apr 2023 08:40:33 -0700 Subject: [PATCH] Internal change. PiperOrigin-RevId: 526046192 --- src/google/protobuf/arena.cc | 40 ++++++++++++++++--- src/google/protobuf/arena.h | 9 +++++ src/google/protobuf/arena_cleanup.h | 20 ++++++++++ src/google/protobuf/arena_test_util.h | 3 ++ .../protobuf/generated_message_tctable_impl.h | 13 +++--- .../generated_message_tctable_lite.cc | 13 +++--- src/google/protobuf/map.h | 16 ++++++-- src/google/protobuf/map_test.cc | 6 +++ src/google/protobuf/map_test.inc | 11 ++++- src/google/protobuf/serial_arena.h | 3 ++ src/google/protobuf/thread_safe_arena.h | 4 ++ 11 files changed, 116 insertions(+), 22 deletions(-) diff --git a/src/google/protobuf/arena.cc b/src/google/protobuf/arena.cc index 5d25d65e1f..51afdbdafe 100644 --- a/src/google/protobuf/arena.cc +++ b/src/google/protobuf/arena.cc @@ -37,6 +37,7 @@ #include #include #include +#include #include "absl/base/attributes.h" #include "absl/synchronization/mutex.h" @@ -156,6 +157,29 @@ SerialArena::SerialArena(FirstSerialArena, ArenaBlock* b, limit_ = b->Limit(); } +std::vector SerialArena::PeekCleanupListForTesting() { + std::vector res; + + ArenaBlock* b = head(); + if (b->IsSentry()) return res; + + const auto peek_list = [&](const char* pos, const char* end) { + while (pos != end) { + pos += cleanup::PeekNode(pos, res); + } + }; + + peek_list(limit_, b->Limit()); + for (b = b->next; b; b = b->next) { + peek_list(reinterpret_cast(b->cleanup_nodes), b->Limit()); + } + return res; +} + +std::vector ThreadSafeArena::PeekCleanupListForTesting() { + return GetSerialArena()->PeekCleanupListForTesting(); +} + void SerialArena::Init(ArenaBlock* b, size_t offset) { set_ptr(b->Pointer(offset)); limit_ = b->Limit(); @@ -754,11 +778,15 @@ void* ThreadSafeArena::AllocateAlignedWithCleanup(size_t n, size_t align, } void ThreadSafeArena::AddCleanup(void* elem, void (*cleanup)(void*)) { + GetSerialArena()->AddCleanup(elem, cleanup); +} + +SerialArena* ThreadSafeArena::GetSerialArena() { SerialArena* arena; if (PROTOBUF_PREDICT_FALSE(!GetSerialArenaFast(&arena))) { arena = GetSerialArenaFallback(kMaxCleanupNodeSize); } - arena->AddCleanup(elem, cleanup); + return arena; } PROTOBUF_NOINLINE @@ -770,11 +798,7 @@ void* ThreadSafeArena::AllocateAlignedWithCleanupFallback( PROTOBUF_NOINLINE void* ThreadSafeArena::AllocateFromStringBlock() { - SerialArena* arena; - if (PROTOBUF_PREDICT_FALSE(!GetSerialArenaFast(&arena))) { - arena = GetSerialArenaFallback(0); - } - return arena->AllocateFromStringBlock(); + return GetSerialArena()->AllocateFromStringBlock(); } template @@ -908,6 +932,10 @@ void* Arena::AllocateAlignedWithCleanup(size_t n, size_t align, return impl_.AllocateAlignedWithCleanup(n, align, destructor); } +std::vector Arena::PeekCleanupListForTesting() { + return impl_.PeekCleanupListForTesting(); +} + } // namespace protobuf } // namespace google diff --git a/src/google/protobuf/arena.h b/src/google/protobuf/arena.h index 0bced3c953..47966891fb 100644 --- a/src/google/protobuf/arena.h +++ b/src/google/protobuf/arena.h @@ -36,6 +36,7 @@ #include #include #include +#include #if defined(_MSC_VER) && !defined(_LIBCPP_STD_VER) && !_HAS_EXCEPTIONS // Work around bugs in MSVC header when _HAS_EXCEPTIONS=0. #include @@ -667,6 +668,14 @@ class PROTOBUF_EXPORT PROTOBUF_ALIGNAS(8) Arena final { void* AllocateAlignedWithCleanup(size_t n, size_t align, void (*destructor)(void*)); + // Test only API. + // It returns the objects that are in the cleanup list for the current + // SerialArena. This API is meant for tests that want to see if something was + // added or not to the cleanup list. Sometimes adding something to the cleanup + // list has no visible side effect so peeking into the list is the only way to + // test. + std::vector PeekCleanupListForTesting(); + template friend class internal::GenericTypeHandler; friend class internal::InternalMetadata; // For user_arena(). diff --git a/src/google/protobuf/arena_cleanup.h b/src/google/protobuf/arena_cleanup.h index 5c7799f4a3..cd2f2cc8cb 100644 --- a/src/google/protobuf/arena_cleanup.h +++ b/src/google/protobuf/arena_cleanup.h @@ -34,6 +34,7 @@ #include #include #include +#include #include "absl/base/attributes.h" #include "absl/log/absl_check.h" @@ -158,6 +159,25 @@ inline ABSL_ATTRIBUTE_ALWAYS_INLINE size_t DestroyNode(const void* pos) { return sizeof(DynamicNode); } +// Append in `out` the pointer to the to-be-cleaned object in `pos`. +// Return the length of the cleanup node to allow the caller to advance the +// position, like `DestroyNode` does. +inline size_t PeekNode(const void* pos, std::vector& out) { + uintptr_t elem; + memcpy(&elem, pos, sizeof(elem)); + out.push_back(reinterpret_cast(elem & ~3)); + if (EnableSpecializedTags()) { + switch (static_cast(elem & 3)) { + case Tag::kString: + case Tag::kCord: + return sizeof(TaggedNode); + default: + break; + } + } + 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*)) { diff --git a/src/google/protobuf/arena_test_util.h b/src/google/protobuf/arena_test_util.h index 7a6de35c80..1065b62828 100644 --- a/src/google/protobuf/arena_test_util.h +++ b/src/google/protobuf/arena_test_util.h @@ -79,6 +79,9 @@ struct ArenaTestPeer { static void ReturnArrayMemory(Arena* arena, void* p, size_t size) { arena->ReturnArrayMemory(p, size); } + static auto PeekCleanupListForTesting(Arena* arena) { + return arena->PeekCleanupListForTesting(); + } }; class NoHeapChecker { diff --git a/src/google/protobuf/generated_message_tctable_impl.h b/src/google/protobuf/generated_message_tctable_impl.h index bf20515675..c076455a56 100644 --- a/src/google/protobuf/generated_message_tctable_impl.h +++ b/src/google/protobuf/generated_message_tctable_impl.h @@ -821,13 +821,16 @@ class PROTOBUF_EXPORT TcParser final { static void InitializeMapNodeEntry(void* obj, MapTypeCard type_card, UntypedMapBase& map, - const TcParseTableBase::FieldAux* aux); + const TcParseTableBase::FieldAux* aux, + bool is_key); static void DestroyMapNode(NodeBase* node, MapAuxInfo map_info, UntypedMapBase& map); - static const char* ParseOneMapEntry( - NodeBase* node, const char* ptr, ParseContext* ctx, - const TcParseTableBase::FieldAux* aux, const TcParseTableBase* table, - const TcParseTableBase::FieldEntry& entry); + static const char* ParseOneMapEntry(NodeBase* node, const char* ptr, + ParseContext* ctx, + const TcParseTableBase::FieldAux* aux, + const TcParseTableBase* table, + const TcParseTableBase::FieldEntry& entry, + Arena* arena); // Mini field lookup: static const TcParseTableBase::FieldEntry* FindFieldEntry( diff --git a/src/google/protobuf/generated_message_tctable_lite.cc b/src/google/protobuf/generated_message_tctable_lite.cc index 20a3233c73..f2b36ee491 100644 --- a/src/google/protobuf/generated_message_tctable_lite.cc +++ b/src/google/protobuf/generated_message_tctable_lite.cc @@ -2447,7 +2447,8 @@ void TcParser::WriteMapEntryAsUnknown(MessageLite* msg, PROTOBUF_ALWAYS_INLINE inline void TcParser::InitializeMapNodeEntry( void* obj, MapTypeCard type_card, UntypedMapBase& map, - const TcParseTableBase::FieldAux* aux) { + const TcParseTableBase::FieldAux* aux, bool is_key) { + (void)is_key; switch (type_card.cpp_type()) { case MapTypeCard::kBool: memset(obj, 0, sizeof(bool)); @@ -2497,7 +2498,7 @@ const char* ReadFixed(void* obj, const char* ptr) { const char* TcParser::ParseOneMapEntry( NodeBase* node, const char* ptr, ParseContext* ctx, const TcParseTableBase::FieldAux* aux, const TcParseTableBase* table, - const TcParseTableBase::FieldEntry& entry) { + const TcParseTableBase::FieldEntry& entry, Arena* arena) { using WFL = WireFormatLite; const auto map_info = aux[0].map_info; @@ -2631,13 +2632,13 @@ PROTOBUF_NOINLINE const char* TcParser::MpMap(PROTOBUF_TC_PARAM_DECL) { while (true) { NodeBase* node = map.AllocNode(map_info.node_size_info); - InitializeMapNodeEntry(node->GetVoidKey(), map_info.key_type_card, map, - aux); + InitializeMapNodeEntry(node->GetVoidKey(), map_info.key_type_card, map, aux, + true); InitializeMapNodeEntry(node->GetVoidValue(map_info.node_size_info), - map_info.value_type_card, map, aux); + map_info.value_type_card, map, aux, false); ptr = ctx->ParseLengthDelimitedInlined(ptr, [&](const char* ptr) { - return ParseOneMapEntry(node, ptr, ctx, aux, table, entry); + return ParseOneMapEntry(node, ptr, ctx, aux, table, entry, map.arena()); }); if (PROTOBUF_PREDICT_TRUE(ptr != nullptr)) { diff --git a/src/google/protobuf/map.h b/src/google/protobuf/map.h index 4d2466e871..869ebf1003 100644 --- a/src/google/protobuf/map.h +++ b/src/google/protobuf/map.h @@ -1070,6 +1070,12 @@ class KeyMapBase : public UntypedMapBase { } }; +template +bool InitializeMapKey(T*, K&&, Arena*) { + return false; +} + + } // namespace internal // This is the class for Map's internal value_type. @@ -1555,13 +1561,17 @@ class Map : private internal::KeyMapBase> { std::is_same::type, key_type>::value, K&&, key_type>::type; Node* node = static_cast(this->AllocNode(sizeof(Node))); + // Even when arena is nullptr, CreateInArenaStorage is still used to // ensure the arena of submessage will be consistent. Otherwise, // submessage may have its own arena when message-owned arena is enabled. // Note: This only works if `Key` is not arena constructible. - Arena::CreateInArenaStorage(const_cast(&node->kv.first), - this->alloc_.arena(), - static_cast(std::forward(k))); + if (!internal::InitializeMapKey(const_cast(&node->kv.first), + std::forward(k), this->alloc_.arena())) { + Arena::CreateInArenaStorage(const_cast(&node->kv.first), + this->alloc_.arena(), + static_cast(std::forward(k))); + } // Note: if `T` is arena constructible, `Args` needs to be empty. Arena::CreateInArenaStorage(&node->kv.second, this->alloc_.arena(), std::forward(args)...); diff --git a/src/google/protobuf/map_test.cc b/src/google/protobuf/map_test.cc index 205e57ee92..7dffe21e7b 100644 --- a/src/google/protobuf/map_test.cc +++ b/src/google/protobuf/map_test.cc @@ -29,8 +29,10 @@ // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. #include +#include #include "absl/container/flat_hash_set.h" +#include "google/protobuf/arena_test_util.h" #include "google/protobuf/map_proto2_unittest.pb.h" #include "google/protobuf/map_unittest.pb.h" #include "google/protobuf/reflection_tester.h" @@ -71,6 +73,9 @@ struct is_internal_map_value_type : std::true_type {}; namespace { +using ::testing::FieldsAre; +using ::testing::UnorderedElementsAre; + template void MapTest_Aligned() { @@ -90,6 +95,7 @@ TEST(MapTest, Aligned8) { MapTest_Aligned(); } TEST(MapTest, Aligned8OnArena) { MapTest_Aligned(); } + } // namespace } // namespace internal } // namespace protobuf diff --git a/src/google/protobuf/map_test.inc b/src/google/protobuf/map_test.inc index fd3cd86afb..c27c1cace5 100644 --- a/src/google/protobuf/map_test.inc +++ b/src/google/protobuf/map_test.inc @@ -4164,8 +4164,15 @@ TEST(ArenaTest, StringMapNoLeak) { data.append("a"); } (*message->mutable_map_string_string())[data] = data; - // We rely on heap checkers to detect memory leak for us. - ASSERT_FALSE(message == nullptr); + // Now we recreate the map via parsing + message->ParseFromString(message->SerializeAsString()); + // And we cause value strings to grow to make sure they are handled correctly. + std::string& value = (*message->mutable_map_string_string())[data]; + ASSERT_EQ(value, data); + // Consume all the capacity. + while (value.size() < value.capacity()) value.append("a"); + // And then grow the value. + value.append("a"); } TEST(ArenaTest, IsInitialized) { diff --git a/src/google/protobuf/serial_arena.h b/src/google/protobuf/serial_arena.h index 8ac5cddce1..0a38a26b01 100644 --- a/src/google/protobuf/serial_arena.h +++ b/src/google/protobuf/serial_arena.h @@ -39,6 +39,7 @@ #include #include #include +#include #include "google/protobuf/stubs/common.h" #include "absl/base/attributes.h" @@ -293,6 +294,8 @@ class PROTOBUF_EXPORT SerialArena { ABSL_ATTRIBUTE_RETURNS_NONNULL void* AllocateFromStringBlock(); + std::vector PeekCleanupListForTesting(); + private: bool MaybeAllocateString(void*& p); ABSL_ATTRIBUTE_RETURNS_NONNULL void* AllocateFromStringBlockFallback(); diff --git a/src/google/protobuf/thread_safe_arena.h b/src/google/protobuf/thread_safe_arena.h index d2a3e6e72a..5a527bc637 100644 --- a/src/google/protobuf/thread_safe_arena.h +++ b/src/google/protobuf/thread_safe_arena.h @@ -126,6 +126,8 @@ class PROTOBUF_EXPORT ThreadSafeArena { void* AllocateFromStringBlock(); + std::vector PeekCleanupListForTesting(); + private: friend class ArenaBenchmark; friend class TcParser; @@ -204,6 +206,8 @@ class PROTOBUF_EXPORT ThreadSafeArena { // create a big enough block to accommodate n bytes. SerialArena* GetSerialArenaFallback(size_t n); + SerialArena* GetSerialArena(); + template void* AllocateAlignedFallback(size_t n);