Internal change.

PiperOrigin-RevId: 526046192

src/google/protobuf/arena.cc

@@ -37,6 +37,7 @@
 #include <limits>
 #include <string>
 #include <typeinfo>
+#include <vector>
 
 #include "absl/base/attributes.h"
 #include "absl/synchronization/mutex.h"
@@ -156,6 +157,29 @@ SerialArena::SerialArena(FirstSerialArena, ArenaBlock* b,
   limit_ = b->Limit();
 }
 
+std::vector<void*> SerialArena::PeekCleanupListForTesting() {
+  std::vector<void*> 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<char*>(b->cleanup_nodes), b->Limit());
+  }
+  return res;
+}
+
+std::vector<void*> 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 <typename Functor>
@@ -908,6 +932,10 @@ void* Arena::AllocateAlignedWithCleanup(size_t n, size_t align,
   return impl_.AllocateAlignedWithCleanup(n, align, destructor);
 }
 
+std::vector<void*> Arena::PeekCleanupListForTesting() {
+  return impl_.PeekCleanupListForTesting();
+}
+
 }  // namespace protobuf
 }  // namespace google
 

src/google/protobuf/arena.h

@@ -36,6 +36,7 @@
 #include <limits>
 #include <type_traits>
 #include <utility>
+#include <vector>
 #if defined(_MSC_VER) && !defined(_LIBCPP_STD_VER) && !_HAS_EXCEPTIONS
 // Work around bugs in MSVC <typeinfo> header when _HAS_EXCEPTIONS=0.
 #include <exception>
@@ -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<void*> PeekCleanupListForTesting();
+
   template <typename Type>
   friend class internal::GenericTypeHandler;
   friend class internal::InternalMetadata;  // For user_arena().

src/google/protobuf/arena_cleanup.h

@@ -34,6 +34,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <string>
+#include <vector>
 
 #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<void*>& out) {
+  uintptr_t elem;
+  memcpy(&elem, pos, sizeof(elem));
+  out.push_back(reinterpret_cast<void*>(elem & ~3));
+  if (EnableSpecializedTags()) {
+    switch (static_cast<Tag>(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*)) {

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 {

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(

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)) {

src/google/protobuf/map.h

@@ -1070,6 +1070,12 @@ class KeyMapBase : public UntypedMapBase {
   }
 };
 
+template <typename T, typename K>
+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<internal::KeyForBase<Key>> {
         std::is_same<typename std::decay<K>::type, key_type>::value, K&&,
         key_type>::type;
     Node* node = static_cast<Node*>(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<Key*>(&node->kv.first),
-                                this->alloc_.arena(),
-                                static_cast<TypeToInit>(std::forward<K>(k)));
+    if (!internal::InitializeMapKey(const_cast<Key*>(&node->kv.first),
+                                    std::forward<K>(k), this->alloc_.arena())) {
+      Arena::CreateInArenaStorage(const_cast<Key*>(&node->kv.first),
+                                  this->alloc_.arena(),
+                                  static_cast<TypeToInit>(std::forward<K>(k)));
+    }
     // Note: if `T` is arena constructible, `Args` needs to be empty.
     Arena::CreateInArenaStorage(&node->kv.second, this->alloc_.arena(),
                                 std::forward<Args>(args)...);

src/google/protobuf/map_test.cc

@@ -29,8 +29,10 @@
 // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 #include <type_traits>
+#include <utility>
 
 #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<AlignedAs8> : std::true_type {};
 
 namespace {
 
+using ::testing::FieldsAre;
+using ::testing::UnorderedElementsAre;
+
 
 template <typename Aligned, bool on_arena = false>
 void MapTest_Aligned() {
@@ -90,6 +95,7 @@ TEST(MapTest, Aligned8) { MapTest_Aligned<AlignedAs8>(); }
 TEST(MapTest, Aligned8OnArena) { MapTest_Aligned<AlignedAs8, true>(); }
 
 
+
 }  // namespace
 }  // namespace internal
 }  // namespace protobuf

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) {

src/google/protobuf/serial_arena.h

@@ -39,6 +39,7 @@
 #include <type_traits>
 #include <typeinfo>
 #include <utility>
+#include <vector>
 
 #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<void*> PeekCleanupListForTesting();
+
  private:
   bool MaybeAllocateString(void*& p);
   ABSL_ATTRIBUTE_RETURNS_NONNULL void* AllocateFromStringBlockFallback();

src/google/protobuf/thread_safe_arena.h

@@ -126,6 +126,8 @@ class PROTOBUF_EXPORT ThreadSafeArena {
 
   void* AllocateFromStringBlock();
 
+  std::vector<void*> 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 <AllocationClient alloc_client = AllocationClient::kDefault>
   void* AllocateAlignedFallback(size_t n);