fix(proto-plus): make Marshal thread-safe and handle race conditions

packages/proto-plus/proto/marshal/marshal.py

@@ -13,24 +13,20 @@
 # limitations under the License.
 
 import abc
+import threading
 
-from google.protobuf import duration_pb2
-from google.protobuf import timestamp_pb2
-from google.protobuf import field_mask_pb2
-from google.protobuf import struct_pb2
-from google.protobuf import wrappers_pb2
+from google.protobuf import (
+    duration_pb2,
+    field_mask_pb2,
+    struct_pb2,
+    timestamp_pb2,
+    wrappers_pb2,
+)
 
 from proto.marshal import compat
-from proto.marshal.collections import MapComposite
-from proto.marshal.collections import Repeated
-from proto.marshal.collections import RepeatedComposite
-
+from proto.marshal.collections import MapComposite, Repeated, RepeatedComposite
 from proto.marshal.rules import bytes as pb_bytes
-from proto.marshal.rules import stringy_numbers
-from proto.marshal.rules import dates
-from proto.marshal.rules import struct
-from proto.marshal.rules import wrappers
-from proto.marshal.rules import field_mask
+from proto.marshal.rules import dates, field_mask, stringy_numbers, struct, wrappers
 from proto.primitives import ProtoType
 
 
@@ -168,7 +164,10 @@ def get_rule(self, proto_type):
         # See https://github.com/googleapis/proto-plus-python/issues/349
         if rule == self._noop and hasattr(self, "_instances"):
             for _, instance in self._instances.items():
-                rule = instance._rules.get(proto_type, self._noop)
+                # Avoid race condition where instance is added to _instances
+                # but __init__ hasn't run yet.
+                rules = getattr(instance, "_rules", {})
+                rule = rules.get(proto_type, self._noop)
                 if rule != self._noop:
                     break
         return rule
@@ -254,6 +253,7 @@ class Marshal(BaseMarshal):
     """
 
     _instances = {}
+    _lock = threading.Lock()
 
     def __new__(cls, *, name: str):
         """Create a marshal instance.
@@ -265,7 +265,18 @@ def __new__(cls, *, name: str):
         """
         klass = cls._instances.get(name)
         if klass is None:
-            klass = cls._instances[name] = super().__new__(cls)
+            with cls._lock:
+                # Double check inside lock to confirm another thread hasn't
+                # created the instance while we were waiting for the lock.
+                klass = cls._instances.get(name)
+                if klass is None:
+                    # Use Copy-on-Write to avoid 'RuntimeError: dictionary changed size during iteration'
+                    # in BaseMarshal.get_rule. This allows other threads to iterate over the old
+                    # dictionary safely while we replace it with a new one atomically.
+                    new_instances = cls._instances.copy()
+                    klass = super().__new__(cls)
+                    new_instances[name] = klass
+                    cls._instances = new_instances
 
         return klass
 

packages/proto-plus/tests/test_marshal_thread_safety.py

@@ -0,0 +1,58 @@
+from unittest.mock import patch
+
+from proto.marshal.marshal import Marshal
+
+
+def test_marshal_identity():
+    m1 = Marshal(name="foo")
+    m2 = Marshal(name="foo")
+    assert m1 is m2
+
+
+def test_marshal_different_names():
+    m1 = Marshal(name="foo")
+    m2 = Marshal(name="bar")
+    assert m1 is not m2
+
+
+def test_marshal_new_race_condition():
+    # Test the case where klass is None at line 266,
+    # but NOT None at line 271 (another thread created it).
+
+    from unittest.mock import MagicMock
+
+    mock_instances = MagicMock()
+
+    call_count = 0
+
+    def get_side_effect(name, default=None):
+        nonlocal call_count
+        call_count += 1
+        if call_count == 1:
+            return None  # First check returns None
+        # Simulate another thread having created it
+        return "fake_instance"
+
+    mock_instances.get.side_effect = get_side_effect
+
+    with patch.object(Marshal, "_instances", mock_instances):
+        instance = Marshal(name="race_test")
+        assert instance == "fake_instance"
+
+
+def test_get_rule_uninitialized_instance():
+    class FakeMarshal:
+        # No _rules attribute
+        pass
+
+    m = Marshal(name="default")
+
+    # Inject FakeMarshal into Marshal._instances safely using patch.dict
+    with patch.dict(Marshal._instances, {"fake_uninitialized": FakeMarshal()}):
+
+        class DummyType:
+            pass
+
+        # This should not raise AttributeError because of getattr safety
+        rule = m.get_rule(DummyType)
+        assert rule == m._noop