From 693e0d0af61f8816857f5dff818c7463072e0145 Mon Sep 17 00:00:00 2001 From: chalmer lowe Date: Mon, 20 Jul 2026 10:50:37 -0400 Subject: [PATCH 1/4] fix(proto-plus): make Marshal thread-safe and handle race conditions - Use Double-Checked Locking and Copy-on-Write in Marshal.__new__ - Use getattr safely in BaseMarshal.get_rule - Add tests for concurrency scenarios Fixes #15100 --- packages/proto-plus/proto/marshal/marshal.py | 43 ++++++++----- .../tests/test_marshal_thread_safety.py | 61 +++++++++++++++++++ 2 files changed, 88 insertions(+), 16 deletions(-) create mode 100644 packages/proto-plus/tests/test_marshal_thread_safety.py diff --git a/packages/proto-plus/proto/marshal/marshal.py b/packages/proto-plus/proto/marshal/marshal.py index 81f1a473bf8f..bb0cd8f01688 100644 --- a/packages/proto-plus/proto/marshal/marshal.py +++ b/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 diff --git a/packages/proto-plus/tests/test_marshal_thread_safety.py b/packages/proto-plus/tests/test_marshal_thread_safety.py new file mode 100644 index 000000000000..2f521afcc2b1 --- /dev/null +++ b/packages/proto-plus/tests/test_marshal_thread_safety.py @@ -0,0 +1,61 @@ +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 + 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 + + # Clean up + del Marshal._instances["fake_uninitialized"] From 412b48264b9fc2b222e1b4372ff5f12c22c98b82 Mon Sep 17 00:00:00 2001 From: chalmer lowe Date: Mon, 20 Jul 2026 12:30:48 -0400 Subject: [PATCH 2/4] test(proto-plus): use patch.dict for safe dictionary manipulation in tests Follows reviewer suggestion to avoid potential test pollution. --- .../tests/test_marshal_thread_safety.py | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/packages/proto-plus/tests/test_marshal_thread_safety.py b/packages/proto-plus/tests/test_marshal_thread_safety.py index 2f521afcc2b1..a64e47efb80f 100644 --- a/packages/proto-plus/tests/test_marshal_thread_safety.py +++ b/packages/proto-plus/tests/test_marshal_thread_safety.py @@ -47,15 +47,12 @@ class FakeMarshal: m = Marshal(name="default") - # Inject FakeMarshal into Marshal._instances - Marshal._instances["fake_uninitialized"] = FakeMarshal() + # 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 + class DummyType: + pass - # Clean up - del Marshal._instances["fake_uninitialized"] + # This should not raise AttributeError because of getattr safety + rule = m.get_rule(DummyType) + assert rule == m._noop From 896ec8e8d68177c7b06bafd9e1eb783721b71f92 Mon Sep 17 00:00:00 2001 From: Chalmer Lowe Date: Mon, 20 Jul 2026 16:48:51 -0400 Subject: [PATCH 3/4] update name for clarity --- packages/proto-plus/proto/marshal/marshal.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/proto-plus/proto/marshal/marshal.py b/packages/proto-plus/proto/marshal/marshal.py index bb0cd8f01688..c4fc4dc8799f 100644 --- a/packages/proto-plus/proto/marshal/marshal.py +++ b/packages/proto-plus/proto/marshal/marshal.py @@ -253,7 +253,7 @@ class Marshal(BaseMarshal): """ _instances = {} - _lock = threading.Lock() + _instance_creation_lock = threading.Lock() def __new__(cls, *, name: str): """Create a marshal instance. From 48542710d18aa982338acda795042c77c4c7a3c8 Mon Sep 17 00:00:00 2001 From: Chalmer Lowe Date: Mon, 20 Jul 2026 16:49:01 -0400 Subject: [PATCH 4/4] Update packages/proto-plus/proto/marshal/marshal.py --- packages/proto-plus/proto/marshal/marshal.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/proto-plus/proto/marshal/marshal.py b/packages/proto-plus/proto/marshal/marshal.py index c4fc4dc8799f..522307f5b552 100644 --- a/packages/proto-plus/proto/marshal/marshal.py +++ b/packages/proto-plus/proto/marshal/marshal.py @@ -265,7 +265,7 @@ def __new__(cls, *, name: str): """ klass = cls._instances.get(name) if klass is None: - with cls._lock: + with cls._instance_creation_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)