[Py/FreeThreading] Fixed remaining race conditions in Dealloc()

This change modifies all remaining `Dealloc()` functions to use `EraseIfEqual` if they were not already. This prevents the same race that was fixed for descriptors in cl/874084218.

PiperOrigin-RevId: 952273589

python/google/protobuf/descriptor_pool.py

@@ -34,6 +34,7 @@ directly instead of this class.
 __author__ = 'matthewtoia@google.com (Matt Toia)'
 
 import collections
+import copy
 import threading
 import warnings
 
@@ -127,6 +128,18 @@ class DescriptorPool(object):
     )
     self._edition_defaults = None
     self._feature_cache = dict()
+    self._lock = threading.RLock()
+
+  def __deepcopy__(self, memo):
+    cls = self.__class__
+    result = cls.__new__(cls)
+    memo[id(self)] = result
+    for k, v in self.__dict__.items():
+      if k == '_lock':
+        result._lock = threading.RLock()
+      else:
+        setattr(result, k, copy.deepcopy(v, memo))
+    return result
 
   def _CheckConflictRegister(self, desc, desc_name, file_name):
     """Check if the descriptor name conflicts with another of the same name.
@@ -843,7 +856,13 @@ class DescriptorPool(object):
     Returns:
       A FileDescriptor matching the passed in proto.
     """
-    if file_proto.name not in self._file_descriptors:
+    if file_proto.name in self._file_descriptors:
+      return self._file_descriptors[file_proto.name]
+
+    with self._lock:
+      if file_proto.name in self._file_descriptors:
+        return self._file_descriptors[file_proto.name]
+
       built_deps = list(self._GetDeps(file_proto.dependency))
       direct_deps = [self.FindFileByName(n) for n in file_proto.dependency]
       public_deps = [direct_deps[i] for i in file_proto.public_dependency]
@@ -938,20 +957,20 @@ class DescriptorPool(object):
 
       self._file_descriptors[file_proto.name] = file_descriptor
 
-    # Add extensions to the pool
-    def AddExtensionForNested(message_type):
-      for nested in message_type.nested_types:
-        AddExtensionForNested(nested)
-      for extension in message_type.extensions:
+      # Add extensions to the pool
+      def AddExtensionForNested(message_type):
+        for nested in message_type.nested_types:
+          AddExtensionForNested(nested)
+        for extension in message_type.extensions:
+          self._AddExtensionDescriptor(extension)
+
+      file_desc = self._file_descriptors[file_proto.name]
+      for extension in file_desc.extensions_by_name.values():
         self._AddExtensionDescriptor(extension)
+      for message_type in file_desc.message_types_by_name.values():
+        AddExtensionForNested(message_type)
 
-    file_desc = self._file_descriptors[file_proto.name]
-    for extension in file_desc.extensions_by_name.values():
-      self._AddExtensionDescriptor(extension)
-    for message_type in file_desc.message_types_by_name.values():
-      AddExtensionForNested(message_type)
-
-    return file_desc
+      return file_desc
 
   def _ConvertMessageDescriptor(
       self, desc_proto, package=None, file_desc=None, scope=None, syntax=None

python/google/protobuf/internal/thread_safe_test.py

@@ -376,6 +376,120 @@ class FreeThreadingTest(unittest.TestCase):
     for thread in threads:
       thread.join()
 
+  @unittest.skipIf(
+      api_implementation.Type() == 'upb',
+      'Upb has not been fixed to handle this case.',
+  )
+  def testConcurrentDescriptorDeallocRace(self):
+    """Tests descriptor cache interning under concurrent deallocation."""
+    pool = descriptor_pool.DescriptorPool()
+    file_proto = descriptor_pb2.FileDescriptorProto(name='race.proto')
+    file_proto.message_type.add(name='RaceMessage')
+    pool.Add(file_proto)
+
+    barrier = threading.Barrier(10)
+    errors = []
+
+    def Worker():
+      barrier.wait()
+      for _ in range(500):
+        try:
+          d1 = pool.FindMessageTypeByName('RaceMessage')
+          d2 = pool.FindMessageTypeByName('RaceMessage')
+          if d1 is not d2:
+            errors.append('Descriptor interning broken')
+            break
+          # Explicitly delete local references to trigger concurrent tp_dealloc
+          del d1
+          del d2
+        except Exception as e:
+          errors.append(str(e))
+          break
+
+    threads = [threading.Thread(target=Worker) for _ in range(10)]
+    for t in threads:
+      t.start()
+    for t in threads:
+      t.join()
+    self.assertEqual([], errors)
+
+  @unittest.skipIf(
+      api_implementation.Type() == 'upb',
+      'Upb has not been fixed to handle this case.',
+  )
+  def testConcurrentSubmessageDeallocRace(self):
+    """Tests child submessage wrapper interning under concurrent deallocation."""
+    msg = unittest_proto3_pb2.TestAllTypes()
+    msg.repeated_nested_message.add(bb=123)
+
+    barrier = threading.Barrier(10)
+    errors = []
+
+    def Worker():
+      barrier.wait()
+      for _ in range(500):
+        try:
+          m1 = msg.repeated_nested_message[0]
+          m2 = msg.repeated_nested_message[0]
+          if m1 is not m2:
+            errors.append('Child submessage interning broken')
+            break
+          del m1
+          del m2
+        except Exception as e:
+          errors.append(str(e))
+          break
+
+    threads = [threading.Thread(target=Worker) for _ in range(10)]
+    for t in threads:
+      t.start()
+    for t in threads:
+      t.join()
+    self.assertEqual([], errors)
+
+  @unittest.skipIf(
+      api_implementation.Type() == 'upb',
+      'Upb has not been fixed to handle this case.',
+  )
+  def testConcurrentCompositeFieldDeallocRace(self):
+    """Tests composite field wrapper interning under concurrent deallocation."""
+    msg = unittest_proto3_pb2.TestAllTypes()
+
+    barrier = threading.Barrier(10)
+    errors = []
+
+    def Worker():
+      barrier.wait()
+      for _ in range(500):
+        try:
+          # Test singular composite field wrapper
+          sub1 = msg.optional_nested_message
+          sub2 = msg.optional_nested_message
+          if sub1 is not sub2:
+            errors.append('Singular composite field interning broken')
+            break
+          del sub1
+          del sub2
+
+          # Test repeated container wrapper
+          rep1 = msg.repeated_int32
+          rep2 = msg.repeated_int32
+          if rep1 is not rep2:
+            errors.append('Repeated container interning broken')
+            break
+          del rep1
+          del rep2
+        except Exception as e:
+          errors.append(str(e))
+          break
+
+    threads = [threading.Thread(target=Worker) for _ in range(10)]
+    for t in threads:
+      t.start()
+    for t in threads:
+      t.join()
+    self.assertEqual([], errors)
+
   @unittest.skipIf(not ALSO_RUN_BENCHMARKS, 'Benchmarks are disabled.')
   def testConcurrentGetOptionsBenchmark(self):
     """Benchmarks concurrent GetOptions calls."""

python/google/protobuf/pyext/message.cc

@@ -1417,12 +1417,17 @@ static void Dealloc(CMessage* self) {
     if (self->parent_field_descriptor->is_repeated()) {
       CMessage::SubMessagesMap* child_submessages =
           parent->child_submessages.TryGet();
-      if (child_submessages) child_submessages->Erase(self->message);
+      if (child_submessages) {
+        child_submessages->EraseIfEqual(self->message,
+                                        reinterpret_cast<PyObject*>(self));
+      }
     } else {
       CMessage::CompositeFieldsMap* composite_fields =
           parent->composite_fields.TryGet();
-      if (composite_fields)
-        composite_fields->Erase(self->parent_field_descriptor);
+      if (composite_fields) {
+        composite_fields->EraseIfEqual(self->parent_field_descriptor,
+                                       reinterpret_cast<PyObject*>(self));
+      }
     }
     Py_CLEAR(self->parent);
   }
@@ -1635,7 +1640,8 @@ static int InternalReparentFields(
     Py_INCREF(new_message);
     Py_DECREF(to_release->parent);
     to_release->parent = new_message;
-    self_child_submessages->Erase(to_release->message);
+    self_child_submessages->Erase(to_release->message,
+                                  to_release->AsPyObject());
     new_child_submessages->Set(to_release->message, to_release->AsPyObject());
   }
 
@@ -1646,7 +1652,8 @@ static int InternalReparentFields(
     Py_INCREF(new_message);
     Py_DECREF(to_release->parent);
     to_release->parent = new_message;
-    self_composite_fields->Erase(to_release->parent_field_descriptor);
+    self_composite_fields->Erase(to_release->parent_field_descriptor,
+                                 to_release->AsPyObject());
     new_composite_fields->Set(to_release->parent_field_descriptor,
                               to_release->AsPyObject());
   }
@@ -2886,7 +2893,7 @@ void ContainerBase::RemoveFromParentCache() {
   if (parent) {
     if (CMessage::CompositeFieldsMap* fields =
             parent->composite_fields.TryGet()) {
-      fields->Erase(this->parent_field_descriptor);
+      fields->EraseIfEqual(this->parent_field_descriptor, this->AsPyObject());
     }
     Py_CLEAR(parent);
   }
@@ -2930,7 +2937,7 @@ CMessage* CMessage::MaybeReleaseSubMessage(const Message* sub_message) {
   released->parent_field_descriptor = nullptr;
   released->state = MESSAGE_MUTABLE;
   // Delete it from the cache.
-  sub_messages->Erase(sub_message);
+  sub_messages->Erase(sub_message, released->AsPyObject());
   // child_submessages->Get returned a new reference.
   Py_DECREF(released);
   return released;

python/google/protobuf/pyext/weak_value_map.cc

@@ -58,17 +58,14 @@ bool PyWeakValueMap::TrySet(const void* key, PyObject*& value) {
   return false;
 }
 
-bool PyWeakValueMap::Erase(const void* key) {
-  absl::MutexLock lock(&mutex_);
-  return cache_.erase(key) != 0;
-}
-
-void PyWeakValueMap::EraseIfEqual(const void* key, PyObject* value) {
+bool PyWeakValueMap::EraseIfEqualImpl(const void* key, PyObject* value) {
   absl::MutexLock lock(&mutex_);
   auto it = cache_.find(key);
   if (it != cache_.end() && it->second == value) {
     cache_.erase(it);
+    return true;
   }
+  return false;
 }
 
 bool PyWeakValueMap::IsEmpty() const {
@@ -99,12 +96,12 @@ PyObject* PyWeakValueMap::Get(const void* key, const PyTypeObject* type) {
   return it->second;
 }
 
-bool PyWeakValueMap::Erase(const void* key) { return cache_.erase(key) != 0; }
-
-void PyWeakValueMap::EraseIfEqual(const void* key, PyObject* value) {
+bool PyWeakValueMap::EraseIfEqualImpl(const void* key, PyObject* value) {
   auto it = cache_.find(key);
+  // In a single-threaded build, the object is guaranteed to be in the map.
   ABSL_CHECK(it != cache_.end() && it->second == value);
   cache_.erase(it);
+  return true;
 }
 
 bool PyWeakValueMap::IsEmpty() const { return cache_.empty(); }

python/google/protobuf/pyext/weak_value_map.h

@@ -47,11 +47,22 @@ class PyWeakValueMap {
   // Returns a new reference to the cached value, or nullptr if not found.
   PyObject* Get(const void* key, const PyTypeObject* type);
 
-  // Removes the entry from the map. Returns true if the entry was found.
-  bool Erase(const void* key);
-
   // Removes the entry from the cache, but only if it matches the given value.
-  void EraseIfEqual(const void* key, PyObject* value);
+  //
+  // This is useful in Dealloc() functions, since Dealloc() can always race with
+  // other threads that insert a new value into the map.
+  void EraseIfEqual(const void* key, PyObject* value) {
+    (void)EraseIfEqualImpl(key, value);
+  }
+
+  // Removes the entry from the cache. Checks that the entry was present and
+  // matched the given value.
+  //
+  // This is useful in cases where the caller holds a strong reference to the
+  // value, and can guarantee that the value is still present in the map.
+  void Erase(const void* key, PyObject* value) {
+    ABSL_CHECK(EraseIfEqualImpl(key, value));
+  }
 
   // Returns true if the map is empty.
   bool IsEmpty() const;
@@ -73,6 +84,7 @@ class PyWeakValueMap {
   void ForEach(Func&& func);
 
  private:
+  bool EraseIfEqualImpl(const void* key, PyObject* value);
 #ifdef Py_GIL_DISABLED
   mutable absl::Mutex mutex_;
   absl::flat_hash_map<const void*, PyObject*> cache_ ABSL_GUARDED_BY(mutex_);