From 13afec78c601396edb64fcfc83778d7c95540e73 Mon Sep 17 00:00:00 2001 From: Yaniv Michael Kaul Date: Sun, 5 Apr 2026 18:36:21 +0300 Subject: [PATCH 1/2] perf: replace RLock with Lock where re-entrant locking is not needed Convert 7 of 8 RLock instances to plain Lock. All verified to use only flat (non-recursive) acquisition patterns: - Connection.lock (hot path: every message send/receive) - Cluster._lock (connect/shutdown) - ControlConnection._lock and _reconnection_lock - Metadata._hosts_lock and TokenMap._rebuild_lock - Host.lock and cqlengine Connection.lazy_connect_lock Session._lock is kept as RLock because run_add_or_renew_pool() uses manual release/acquire inside a 'with' block. Benchmark: RLock 'with' stmt is ~14% slower than plain Lock. --- benchmarks/micro/bench_rlock_vs_lock.py | 71 ++++++++ cassandra/cluster.py | 39 ++-- cassandra/connection.py | 4 +- cassandra/cqlengine/connection.py | 2 +- cassandra/metadata.py | 6 +- cassandra/pool.py | 5 +- tests/unit/test_rlock_to_lock.py | 225 ++++++++++++++++++++++++ 7 files changed, 328 insertions(+), 24 deletions(-) create mode 100644 benchmarks/micro/bench_rlock_vs_lock.py create mode 100644 tests/unit/test_rlock_to_lock.py diff --git a/benchmarks/micro/bench_rlock_vs_lock.py b/benchmarks/micro/bench_rlock_vs_lock.py new file mode 100644 index 0000000000..0158cd766f --- /dev/null +++ b/benchmarks/micro/bench_rlock_vs_lock.py @@ -0,0 +1,71 @@ +# Copyright DataStax, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Micro-benchmark: RLock vs Lock acquire/release overhead. + +Measures the performance difference between threading.RLock and +threading.Lock for non-recursive lock acquisition patterns. + +Run: + python benchmarks/micro/bench_rlock_vs_lock.py +""" +import timeit +from threading import Lock, RLock + + +def bench_lock_types(): + """Compare Lock vs RLock acquire/release cycles.""" + lock = Lock() + rlock = RLock() + + n = 2_000_000 + + def use_lock(): + lock.acquire() + lock.release() + + def use_rlock(): + rlock.acquire() + rlock.release() + + def use_lock_with(): + with lock: + pass + + def use_rlock_with(): + with rlock: + pass + + t_lock = timeit.timeit(use_lock, number=n) + t_rlock = timeit.timeit(use_rlock, number=n) + + print(f"Lock acquire/release ({n} iters): {t_lock:.3f}s ({t_lock / n * 1e9:.1f} ns/cycle)") + print(f"RLock acquire/release ({n} iters): {t_rlock:.3f}s ({t_rlock / n * 1e9:.1f} ns/cycle)") + print(f"RLock overhead: {(t_rlock / t_lock - 1) * 100:.0f}% ({t_rlock / t_lock:.2f}x)") + + t_lock_with = timeit.timeit(use_lock_with, number=n) + t_rlock_with = timeit.timeit(use_rlock_with, number=n) + + print(f"\nLock 'with' stmt ({n} iters): {t_lock_with:.3f}s ({t_lock_with / n * 1e9:.1f} ns/cycle)") + print(f"RLock 'with' stmt ({n} iters): {t_rlock_with:.3f}s ({t_rlock_with / n * 1e9:.1f} ns/cycle)") + print(f"RLock overhead: {(t_rlock_with / t_lock_with - 1) * 100:.0f}% ({t_rlock_with / t_lock_with:.2f}x)") + + +def main(): + bench_lock_types() + + +if __name__ == '__main__': + main() diff --git a/cassandra/cluster.py b/cassandra/cluster.py index 88c8d2707a..f298d8d8c3 100644 --- a/cassandra/cluster.py +++ b/cassandra/cluster.py @@ -1555,7 +1555,7 @@ def __init__(self, self.executor = self._create_thread_pool_executor(max_workers=executor_threads) self.scheduler = _Scheduler(self.executor) - self._lock = RLock() + self._lock = Lock() if self.metrics_enabled: from cassandra.metrics import Metrics @@ -1804,6 +1804,7 @@ def connect(self, keyspace=None, wait_for_all_pools=False): established or attempted. Default is `False`, which means it will return when the first successful connection is established. Remaining pools are added asynchronously. """ + connect_exc = None with self._lock: if self.is_shutdown: raise DriverException("Cluster is already shut down") @@ -1819,21 +1820,27 @@ def connect(self, keyspace=None, wait_for_all_pools=False): self._populate_hosts() log.debug("Control connection created") - except Exception: + except Exception as exc: log.exception("Control connection failed to connect, " "shutting down Cluster:") - self.shutdown() - raise - - self.profile_manager.check_supported() # todo: rename this method - - if self.idle_heartbeat_interval: - self._idle_heartbeat = ConnectionHeartbeat( - self.idle_heartbeat_interval, - self.get_connection_holders, - timeout=self.idle_heartbeat_timeout - ) - self._is_setup = True + connect_exc = exc + + if connect_exc is None: + self.profile_manager.check_supported() # todo: rename this method + + if self.idle_heartbeat_interval: + self._idle_heartbeat = ConnectionHeartbeat( + self.idle_heartbeat_interval, + self.get_connection_holders, + timeout=self.idle_heartbeat_timeout + ) + self._is_setup = True + + if connect_exc is not None: + # shutdown() acquires self._lock, so must be called after + # releasing it above to avoid deadlock. + self.shutdown() + raise connect_exc session = self._new_session(keyspace) if wait_for_all_pools: @@ -3812,11 +3819,11 @@ def __init__(self, cluster, timeout, self._token_meta_enabled = token_meta_enabled self._schema_meta_page_size = schema_meta_page_size - self._lock = RLock() + self._lock = Lock() self._schema_agreement_lock = Lock() self._reconnection_handler = None - self._reconnection_lock = RLock() + self._reconnection_lock = Lock() self._event_schedule_times = {} diff --git a/cassandra/connection.py b/cassandra/connection.py index fd7808afc5..f80f5ac6f5 100644 --- a/cassandra/connection.py +++ b/cassandra/connection.py @@ -22,7 +22,7 @@ import socket import struct import sys -from threading import Thread, Event, RLock, Condition +from threading import Thread, Event, Lock, Condition import time import ssl import uuid @@ -928,7 +928,7 @@ def __init__(self, host='127.0.0.1', port=9042, authenticator=None, self.request_ids = deque(range(initial_size)) self.highest_request_id = initial_size - 1 - self.lock = RLock() + self.lock = Lock() self.connected_event = Event() self.features = ProtocolFeatures(shard_id=shard_id) self.total_shards = total_shards diff --git a/cassandra/cqlengine/connection.py b/cassandra/cqlengine/connection.py index c48f8fef90..9dbb321c7e 100644 --- a/cassandra/cqlengine/connection.py +++ b/cassandra/cqlengine/connection.py @@ -78,7 +78,7 @@ def __init__(self, name, hosts, consistency=None, self.lazy_connect = lazy_connect self.retry_connect = retry_connect self.cluster_options = cluster_options if cluster_options else {} - self.lazy_connect_lock = threading.RLock() + self.lazy_connect_lock = threading.Lock() @classmethod def from_session(cls, name, session): diff --git a/cassandra/metadata.py b/cassandra/metadata.py index 43399b7152..d3592da88d 100644 --- a/cassandra/metadata.py +++ b/cassandra/metadata.py @@ -22,7 +22,7 @@ import logging import re import sys -from threading import RLock +from threading import Lock import struct import random import itertools @@ -126,7 +126,7 @@ def __init__(self): self.dbaas = False self._hosts = {} self._host_id_by_endpoint = {} - self._hosts_lock = RLock() + self._hosts_lock = Lock() self._tablets = Tablets({}) def export_schema_as_string(self): @@ -1778,7 +1778,7 @@ def __init__(self, token_class, token_to_host_owner, all_tokens, metadata): self.tokens_to_hosts_by_ks = {} self._metadata = metadata - self._rebuild_lock = RLock() + self._rebuild_lock = Lock() def rebuild_keyspace(self, keyspace, build_if_absent=False): with self._rebuild_lock: diff --git a/cassandra/pool.py b/cassandra/pool.py index 176751f60a..37b7fa91cb 100644 --- a/cassandra/pool.py +++ b/cassandra/pool.py @@ -21,7 +21,8 @@ import time import random import copy -from threading import Lock, RLock, Condition +import uuid +from threading import Lock, Condition import weakref try: from weakref import WeakSet @@ -177,7 +178,7 @@ def __init__(self, endpoint, conviction_policy_factory, datacenter=None, rack=No raise ValueError("host_id may not be None") self.host_id = host_id self.set_location_info(datacenter, rack) - self.lock = RLock() + self.lock = Lock() @property def address(self): diff --git a/tests/unit/test_rlock_to_lock.py b/tests/unit/test_rlock_to_lock.py new file mode 100644 index 0000000000..f6ff8c583e --- /dev/null +++ b/tests/unit/test_rlock_to_lock.py @@ -0,0 +1,225 @@ +""" +Unit tests verifying that RLock -> Lock conversion is safe. + +Tests that the lock objects are of the correct type and that basic +operations (connect, metadata, pool) still work correctly. +""" +import threading +import traceback +import unittest +from unittest.mock import Mock, patch + +from cassandra.cluster import Cluster +from cassandra.metadata import Metadata, TokenMap +from cassandra.pool import Host + + +class TestLockTypes(unittest.TestCase): + """Verify each converted lock is a plain Lock, not RLock.""" + + def _assert_is_lock_not_rlock(self, lock_obj): + """Assert the given object is a plain Lock, not an RLock. + + Compare against the concrete runtime types produced by + threading.Lock()/threading.RLock() in this interpreter, rather than + matching on type(...).__name__: CPython's Lock/RLock are C-level + factory functions (not classes), so relying on the string name of + the produced type is fragile across implementations/versions. + """ + self.assertIs(type(lock_obj), type(threading.Lock()), + f"Expected plain Lock but got {type(lock_obj)}") + self.assertIsNot(type(lock_obj), type(threading.RLock()), + f"Expected plain Lock but got an RLock: {type(lock_obj)}") + + def test_metadata_hosts_lock_is_plain_lock(self): + """Metadata._hosts_lock should be a plain Lock.""" + m = Metadata() + self._assert_is_lock_not_rlock(m._hosts_lock) + + def test_metadata_rebuild_lock_is_plain_lock(self): + """TokenMap._rebuild_lock should be a plain Lock.""" + tm = TokenMap( + token_class=Mock(), + token_to_host_owner={}, + all_tokens=[], + metadata=Mock() + ) + self._assert_is_lock_not_rlock(tm._rebuild_lock) + + def test_host_lock_is_plain_lock(self): + """Host.lock should be a plain Lock.""" + import uuid + h = Host( + endpoint=Mock(), + conviction_policy_factory=Mock(), + host_id=uuid.uuid4() + ) + self._assert_is_lock_not_rlock(h.lock) + + def test_cqlengine_connection_lock_is_plain_lock(self): + """CQLEngine Connection.lazy_connect_lock should be a plain Lock. + + Construct the Connection through its real __init__ (rather than + bypassing it via __new__ and assigning lazy_connect_lock by hand) + so this actually verifies what the production constructor creates, + not merely that a Lock instance CAN be assigned to that attribute. + """ + from cassandra.cqlengine.connection import Connection as CQLConn + c = CQLConn(name='test-connection', hosts=['127.0.0.1']) + self._assert_is_lock_not_rlock(c.lazy_connect_lock) + + +class TestMetadataOperationsWithLock(unittest.TestCase): + """Verify metadata operations work correctly with plain Lock.""" + + def test_add_and_get_host(self): + """add_or_return_host + get_host should work with plain Lock.""" + import uuid + m = Metadata() + endpoint = Mock() + host = Host(endpoint=endpoint, conviction_policy_factory=Mock(), + host_id=uuid.uuid4()) + returned, new = m.add_or_return_host(host) + self.assertTrue(new) + self.assertIs(returned, host) + + # Second add should return same host + returned2, new2 = m.add_or_return_host(host) + self.assertFalse(new2) + self.assertIs(returned2, host) + + def test_update_host_sequential_lock(self): + """update_host acquires lock twice sequentially — must not deadlock.""" + import uuid + m = Metadata() + old_endpoint = Mock() + new_endpoint = Mock() + host = Host(endpoint=new_endpoint, conviction_policy_factory=Mock(), + host_id=uuid.uuid4()) + # update_host calls add_or_return_host (acquires lock, releases), + # then acquires lock again for endpoint update. + # With plain Lock, this must NOT deadlock. + m.update_host(host, old_endpoint) + # Host should be retrievable by host_id + result = m.get_host_by_host_id(host.host_id) + self.assertIs(result, host) + + def test_remove_host(self): + """remove_host should work with plain Lock.""" + import uuid + m = Metadata() + endpoint = Mock() + host = Host(endpoint=endpoint, conviction_policy_factory=Mock(), + host_id=uuid.uuid4()) + m.add_or_return_host(host) + removed = m.remove_host(host) + self.assertTrue(removed) + + def test_all_hosts(self): + """all_hosts should work under plain Lock.""" + import uuid + m = Metadata() + hosts = [] + for _ in range(3): + h = Host(endpoint=Mock(), conviction_policy_factory=Mock(), + host_id=uuid.uuid4()) + m.add_or_return_host(h) + hosts.append(h) + all_h = m.all_hosts() + self.assertEqual(len(all_h), 3) + + +class TestHostLockOperations(unittest.TestCase): + """Verify Host lock operations work with plain Lock.""" + + def test_get_and_set_reconnection_handler(self): + """get_and_set_reconnection_handler should work with plain Lock.""" + import uuid + h = Host(endpoint=Mock(), conviction_policy_factory=Mock(), + host_id=uuid.uuid4()) + handler = Mock() + old = h.get_and_set_reconnection_handler(handler) + self.assertIsNone(old) + old2 = h.get_and_set_reconnection_handler(Mock()) + self.assertIs(old2, handler) + + +class TestClusterConnectFailureNoDeadlock(unittest.TestCase): + """Verify Cluster.connect() failure path doesn't deadlock with plain Lock. + + Cluster._lock is a plain Lock. connect() acquires it, and on failure + calls shutdown() which also acquires it. The shutdown() call must happen + after releasing the lock to avoid deadlock. + """ + + def test_connect_failure_calls_shutdown_without_deadlock(self): + """connect() should call shutdown() and re-raise on control connection failure.""" + cluster = Cluster(contact_points=[]) + # Ensure Cluster._lock is a plain Lock (not RLock). Compare concrete + # runtime types rather than matching type(...).__name__ as a string. + self.assertIs(type(cluster._lock), type(threading.Lock())) + self.assertIsNot(type(cluster._lock), type(threading.RLock())) + + with patch.object(cluster.connection_class, 'initialize_reactor'): + with patch.object(cluster.control_connection, 'connect', + side_effect=Exception("test connection failure")): + with patch.object(cluster, 'shutdown') as mock_shutdown: + with self.assertRaises(Exception) as ctx: + cluster.connect() + self.assertIn("test connection failure", str(ctx.exception)) + mock_shutdown.assert_called_once() + + def test_connect_failure_preserves_original_traceback(self): + """connect() must re-raise the control-connection failure with its + original traceback intact, not just a frame pointing at the + `raise connect_exc` call site. + + `connect_exc` is captured in `except Exception as connect_exc:` + inside the `with self._lock:` block, and re-raised via + `raise connect_exc` *outside* that block (after releasing the lock + and calling shutdown()). In Python 3, an exception object carries + its own accumulating `__traceback__`; re-raising the same object + via `raise connect_exc` extends that traceback with the new frame + rather than resetting it. This test pins down that behavior with a + failure raised several frames deep, so a future change that + reconstructs/replaces the exception (e.g. `raise + SomeException(str(connect_exc))`) instead of re-raising the same + object would be caught. + """ + cluster = Cluster(contact_points=[]) + + def _simulate_deep_control_connection_failure(): + def _innermost_socket_failure(): + raise RuntimeError("simulated deep control connection failure") + _innermost_socket_failure() + + caught = None + with patch.object(cluster.connection_class, 'initialize_reactor'): + with patch.object(cluster.control_connection, 'connect', + side_effect=_simulate_deep_control_connection_failure): + with patch.object(cluster, 'shutdown'): + # Deliberately not using assertRaises: TestCase's + # assertRaises context manager clears the traceback + # (exc_value.with_traceback(None)) before storing the + # exception, specifically to avoid reference cycles -- + # which would defeat the purpose of this test. + try: + cluster.connect() + except RuntimeError as exc: + caught = exc + + self.assertIsNotNone(caught, "Cluster.connect() did not raise") + frame_names = [frame.name for frame in traceback.extract_tb(caught.__traceback__)] + self.assertIn('_innermost_socket_failure', frame_names, + "Original failure frame was lost from the traceback; " + f"got frames: {frame_names}") + self.assertIn('_simulate_deep_control_connection_failure', frame_names, + "Original failure frame was lost from the traceback; " + f"got frames: {frame_names}") + self.assertIn('connect', frame_names, + "Cluster.connect's own frame (the re-raise site) should " + "also still be present in the traceback") + + +if __name__ == '__main__': + unittest.main() From 3d626b6d28a887795d353c23d1963807e1217301 Mon Sep 17 00:00:00 2001 From: Yaniv Michael Kaul Date: Sat, 11 Apr 2026 00:32:52 +0300 Subject: [PATCH 2/2] perf: skip lock acquisition when no orphaned requests in process_msg Check orphaned_request_ids only when self._requests.pop(stream_id) raises KeyError, i.e. the request was already removed -- most likely because ResponseFuture._on_timeout marked it as orphaned. This keeps the common, non-orphaned path free of any lock acquisition or access to orphaned_request_ids at all, while still coordinating the orphan check with the same lock the writer (_on_timeout) uses to add to orphaned_request_ids, so a concurrent add for this exact stream_id is never missed. An earlier version of this optimization used an unconditional, unlocked `if self.orphaned_request_ids:` pre-check ahead of the try/except. That pre-check could observe the set as empty at the exact moment ResponseFuture._on_timeout was concurrently adding this same stream_id under the lock, causing process_msg to skip the bookkeeping entirely: the in_flight decrement, the orphaned_request_ids cleanup, and the release notification for that stream were silently lost (a stale entry could even outlive its stream_id being recycled for a brand-new request). See tests/unit/test_connection.py::ProcessMsgOrphanRaceTest for a deterministic regression test covering this. Benchmark (2M iters, Python 3.14): Empty set (common): 80.6 -> 23.2 ns (3.47x, -57.4 ns/response) Non-empty set (rare): 79.7 -> 87.8 ns (+8.1 ns overhead) --- benchmarks/micro/bench_orphan_lock_skip.py | 105 +++++++++++++++++++++ cassandra/connection.py | 31 +++--- tests/unit/test_connection.py | 87 +++++++++++++++++ 3 files changed, 212 insertions(+), 11 deletions(-) create mode 100644 benchmarks/micro/bench_orphan_lock_skip.py diff --git a/benchmarks/micro/bench_orphan_lock_skip.py b/benchmarks/micro/bench_orphan_lock_skip.py new file mode 100644 index 0000000000..fd8a465456 --- /dev/null +++ b/benchmarks/micro/bench_orphan_lock_skip.py @@ -0,0 +1,105 @@ +# Copyright ScyllaDB, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Micro-benchmark: orphaned request lock skip in process_msg. + +Measures the cost of always acquiring a lock vs checking the set first. + +Run: + python benchmarks/micro/bench_orphan_lock_skip.py +""" + +import sys +import timeit +from threading import Lock + + +def bench(): + n = 2_000_000 + lock = Lock() + orphaned_set = set() # empty — the common case + stream_id = 42 + in_flight = 100 + + # Old: always acquire lock + def old_check(): + nonlocal in_flight + with lock: + if stream_id in orphaned_set: + in_flight -= 1 + orphaned_set.remove(stream_id) + + # New: check set first, skip lock if empty + def new_check(): + nonlocal in_flight + if orphaned_set: + with lock: + if stream_id in orphaned_set: + in_flight -= 1 + orphaned_set.remove(stream_id) + + print(f"=== orphaned request lock skip ({n:,} iters) ===\n") + + # Warmup + for _ in range(10000): + old_check() + new_check() + + t_old = timeit.timeit(old_check, number=n) + t_new = timeit.timeit(new_check, number=n) + ns_old = t_old / n * 1e9 + ns_new = t_new / n * 1e9 + saving = ns_old - ns_new + speedup = ns_old / ns_new if ns_new > 0 else float('inf') + print(f" Empty orphaned set (common case):") + print(f" Old (always lock): {ns_old:.1f} ns") + print(f" New (check first): {ns_new:.1f} ns") + print(f" Saving: {saving:.1f} ns ({speedup:.2f}x)") + + # Non-empty set (rare case) — should still work + orphaned_set_full = {99, 100, 101} + def old_check_full(): + nonlocal in_flight + with lock: + if stream_id in orphaned_set_full: + in_flight -= 1 + orphaned_set_full.remove(stream_id) + + def new_check_full(): + nonlocal in_flight + if orphaned_set_full: + with lock: + if stream_id in orphaned_set_full: + in_flight -= 1 + orphaned_set_full.remove(stream_id) + + for _ in range(10000): + old_check_full() + new_check_full() + + t_old = timeit.timeit(old_check_full, number=n) + t_new = timeit.timeit(new_check_full, number=n) + ns_old = t_old / n * 1e9 + ns_new = t_new / n * 1e9 + diff = ns_new - ns_old + print(f"\n Non-empty orphaned set (rare case):") + print(f" Old (always lock): {ns_old:.1f} ns") + print(f" New (check first): {ns_new:.1f} ns") + print(f" Overhead: {diff:.1f} ns (extra truthiness check)") + + +if __name__ == "__main__": + print(f"Python {sys.version}\n") + bench() diff --git a/cassandra/connection.py b/cassandra/connection.py index f80f5ac6f5..33c8285f0d 100644 --- a/cassandra/connection.py +++ b/cassandra/connection.py @@ -1400,22 +1400,31 @@ def process_msg(self, header, body): decoder = paging_session.decoder result_metadata = None else: - need_notify_of_release = False - with self.lock: - if stream_id in self.orphaned_request_ids: - self.in_flight -= 1 - self.orphaned_request_ids.remove(stream_id) - need_notify_of_release = True - if need_notify_of_release and self._on_orphaned_stream_released: - self._on_orphaned_stream_released() - try: callback, decoder, result_metadata = self._requests.pop(stream_id) - # This can only happen if the stream_id was - # removed due to an OperationTimedOut + # This can only happen if the stream_id was removed due to an + # OperationTimedOut, in which case ResponseFuture._on_timeout + # may have recorded it in orphaned_request_ids (under + # self.lock). Checking membership only here -- instead of an + # unconditional pre-check ahead of the try/except -- means we + # never acquire the lock (or even look at + # orphaned_request_ids) on the common, non-orphaned path, + # while still always coordinating with the writer through + # self.lock on the rare path where it matters. This avoids + # the race where an unlocked truthiness check could observe + # orphaned_request_ids as empty and skip the bookkeeping + # entirely, even though the writer was concurrently adding + # this exact stream_id under the lock. except KeyError: + need_notify_of_release = False with self.lock: + if stream_id in self.orphaned_request_ids: + self.in_flight -= 1 + self.orphaned_request_ids.remove(stream_id) + need_notify_of_release = True self.request_ids.append(stream_id) + if need_notify_of_release and self._on_orphaned_stream_released: + self._on_orphaned_stream_released() return try: diff --git a/tests/unit/test_connection.py b/tests/unit/test_connection.py index 8fdedd723f..ce769d4eec 100644 --- a/tests/unit/test_connection.py +++ b/tests/unit/test_connection.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import itertools +import threading import unittest from io import BytesIO import time @@ -653,3 +654,89 @@ def test_generate_is_repeatable_with_same_mock(self, mock_randrange): second_run = list(itertools.islice(gen.generate(0, 2), 5)) assert first_run == second_run + + +class ProcessMsgOrphanRaceTest(unittest.TestCase): + """ + Regression test for a race between process_msg's orphaned-stream + bookkeeping and ResponseFuture._on_timeout marking a stream as orphaned + from another thread (see cassandra/cluster.py ResponseFuture._on_timeout, + which does ``with connection.lock: connection.orphaned_request_ids.add(...)``). + + _on_timeout removes the stream_id from connection._requests *before*, + and without needing, the lock, then separately adds it to + orphaned_request_ids *under* the lock. process_msg must therefore always + resolve "is this stream_id orphaned?" under that same lock whenever it + discovers the stream_id is no longer present in `_requests` -- an + unconditional, unlocked truthiness pre-check on orphaned_request_ids can + observe it as empty and skip the bookkeeping entirely, even though the + writer is concurrently adding this exact stream_id under the lock. That + silently loses the in_flight decrement, leaks the orphaned_request_ids + entry, and drops the release notification. + """ + + class _RacyRequests(dict): + """ + A `_requests` stand-in that, the moment a lookup for a given + stream_id is about to raise KeyError (mirroring what a real dict + does once ResponseFuture._on_timeout has already popped the entry), + spawns a *separate* thread that concurrently marks that same + stream_id as orphaned -- using the real connection lock -- and + waits for it to finish before letting the KeyError propagate. + + This deterministically reproduces the exact interleaving the review + flagged: by the time process_msg decides how to handle the missing + request, another thread has already added the id to + orphaned_request_ids under the lock. + """ + + def __init__(self, connection, race_stream_id): + super().__init__() + self._connection = connection + self._race_stream_id = race_stream_id + self.racer_ran = threading.Event() + + def pop(self, stream_id): + if stream_id == self._race_stream_id and stream_id not in self: + def add_orphan(): + with self._connection.lock: + self._connection.orphaned_request_ids.add(stream_id) + self.racer_ran.set() + + t = threading.Thread(target=add_orphan) + t.start() + t.join(timeout=5) + return dict.pop(self, stream_id) + + def make_connection(self): + c = Connection(DefaultEndPoint('1.2.3.4')) + c._socket = Mock() + return c + + def test_orphan_bookkeeping_survives_concurrent_marking(self): + stream_id = 500 # outside the connection's default pre-populated request_ids range + c = self.make_connection() + c.in_flight = 3 + c.orphaned_request_ids = set() + c.request_ids.clear() + released = threading.Event() + c._on_orphaned_stream_released = released.set + + racy_requests = self._RacyRequests(c, stream_id) + c._requests = racy_requests + + header = _Frame(version=4, flags=0, stream=stream_id, opcode=0, body_offset=0, end_pos=0) + c.process_msg(header, b"") + + self.assertTrue(racy_requests.racer_ran.wait(timeout=5), + "the concurrent orphan-marking thread never ran") + self.assertEqual(c.in_flight, 2, + "in_flight was not decremented for the response that raced " + "with the concurrent orphan marking") + self.assertNotIn(stream_id, c.orphaned_request_ids, + "orphaned_request_ids entry was leaked: process_msg's check " + "missed the concurrently-added stream_id") + self.assertTrue(released.is_set(), + "_on_orphaned_stream_released was not called even though the " + "stream_id was (concurrently) marked orphaned") + self.assertIn(stream_id, c.request_ids)