From 64ac479f397e7f805be66f25e5c6f51b339592b4 Mon Sep 17 00:00:00 2001 From: Peter Alm Date: Fri, 21 Aug 2026 10:25:31 +0200 Subject: [PATCH 1/2] Break reference cycles keeping closed connections alive. Closed connections could only be freed by the cyclic garbage collector, never by reference counting, because reference cycles referenced the connection: * The LoggerAdapter that injects the websocket attribute into log records held a strong reference to the connection, in all implementations: connection -> protocol -> logger -> extra dict -> connection. * In the asyncio implementation, the keepalive task retained the CancelledError raised when connection_lost() cancelled it, whose traceback references the keepalive() frame and thus the connection: connection -> task -> exception -> traceback -> frame -> connection. * In the trio implementation, the assembler held bound methods of the connection for flow control: connection -> recv_messages -> pause() and resume() callbacks -> connection. On servers handling many connections with high connect/disconnect churn, closed connections accumulated until a full collection of the oldest generation, which can lag far behind and pause the event loop for several seconds on large heaps, especially since CPython 3.13 collects the oldest generation incrementally. ConnectionLoggerAdapter now holds a weak reference to the connection and injects it into log records only while the connection is alive; records created while the connection is alive keep their own strong reference, so logging filters and handlers are unaffected. connection_lost() dereferences the keepalive task after cancelling it. The assembler drops the flow control callbacks when it's closed. Garbage collection tests are skipped on Python 3.12, where the traceback of an exception raised while closing a connection keeps frames of connection methods alive via their f_back attribute. Python 3.13 fixed that behavior. --- docs/project/changelog.rst | 7 ++++ src/websockets/asyncio/connection.py | 17 +++++--- src/websockets/sync/connection.py | 9 ++--- src/websockets/trio/connection.py | 9 ++--- src/websockets/trio/messages.py | 5 +++ src/websockets/utils.py | 30 +++++++++++++- tests/asyncio/test_connection.py | 58 +++++++++++++++++++++++++--- tests/asyncio/test_server.py | 24 ++++++++++++ tests/sync/test_connection.py | 30 +++++++++++++- tests/sync/test_server.py | 29 ++++++++++++++ tests/test_utils.py | 34 ++++++++++++++++ tests/trio/test_connection.py | 27 ++++++++++++- tests/trio/test_server.py | 27 +++++++++++++ tests/trio/utils.py | 2 + tests/utils.py | 12 ++++++ 15 files changed, 294 insertions(+), 26 deletions(-) diff --git a/docs/project/changelog.rst b/docs/project/changelog.rst index fa2540c1..bbdae504 100644 --- a/docs/project/changelog.rst +++ b/docs/project/changelog.rst @@ -32,6 +32,13 @@ notice. *In development* +Bug fixes +......... + +* Fixed reference cycles that delayed garbage collection of closed connections + until a pass of the cyclic garbage collector. Closed connections are now + freed immediately by reference counting. + .. _17.0.1: 17.0.1 diff --git a/src/websockets/asyncio/connection.py b/src/websockets/asyncio/connection.py index 74ab5e2a..97936150 100644 --- a/src/websockets/asyncio/connection.py +++ b/src/websockets/asyncio/connection.py @@ -3,7 +3,6 @@ import asyncio import collections import contextlib -import logging import random import struct import traceback @@ -22,6 +21,7 @@ from ..http11 import Request, Response from ..protocol import CLOSED, OPEN, Event, Protocol, State from ..typing import BytesLike, Data, DataLike, LoggerLike, Subprotocol +from ..utils import ConnectionLoggerAdapter from .messages import Assembler @@ -65,10 +65,9 @@ def __init__( self.write_limit_high, self.write_limit_low = write_limit # Inject reference to this instance in the protocol's logger. - self.protocol.logger = logging.LoggerAdapter( - self.protocol.logger, - {"websocket": self}, - ) + # ConnectionLoggerAdapter holds a weak reference in order to + # keep the connection garbage-collectable by reference counting. + self.protocol.logger = ConnectionLoggerAdapter(self.protocol.logger, self) # Copy attributes from the protocol for convenience. self.id: uuid.UUID = self.protocol.id @@ -111,7 +110,8 @@ def __init__( send Ping frames and measure latency with :meth:`ping`. """ - # Task that sends keepalive pings. None when ping_interval is None. + # Task that sends keepalive pings. None when ping_interval is None + # and after the connection is lost. self.keepalive_task: asyncio.Task[None] | None = None # Exception raised while reading from the connection, to be chained to @@ -1025,6 +1025,11 @@ def connection_lost(self, exc: Exception | None) -> None: if self.keepalive_task is not None: self.keepalive_task.cancel() + # Dereference the task. Else, the traceback of CancelledError + # would keep the connection in a reference cycle. This doesn't + # prevent the cancellation: the event loop keeps a strong + # reference to the task until it delivers CancelledError. + self.keepalive_task = None # If self.connection_lost_waiter isn't pending, that's a bug, because: # - it's set only here in connection_lost() which is called only once; diff --git a/src/websockets/sync/connection.py b/src/websockets/sync/connection.py index 74ac997a..a8bd8546 100644 --- a/src/websockets/sync/connection.py +++ b/src/websockets/sync/connection.py @@ -2,7 +2,6 @@ import concurrent.futures import contextlib -import logging import random import socket import struct @@ -24,6 +23,7 @@ from ..http11 import Request, Response from ..protocol import CLOSED, CONNECTING, OPEN, Event, Protocol, State from ..typing import BytesLike, Data, DataLike, LoggerLike, Subprotocol +from ..utils import ConnectionLoggerAdapter from .messages import Assembler from .utils import Deadline @@ -67,10 +67,9 @@ def __init__( max_queue_high, max_queue_low = max_queue # Inject reference to this instance in the protocol's logger. - self.protocol.logger = logging.LoggerAdapter( - self.protocol.logger, - {"websocket": self}, - ) + # ConnectionLoggerAdapter holds a weak reference in order to + # keep the connection garbage-collectable by reference counting. + self.protocol.logger = ConnectionLoggerAdapter(self.protocol.logger, self) # Copy attributes from the protocol for convenience. self.id: uuid.UUID = self.protocol.id diff --git a/src/websockets/trio/connection.py b/src/websockets/trio/connection.py index 606653ad..99c2e186 100644 --- a/src/websockets/trio/connection.py +++ b/src/websockets/trio/connection.py @@ -1,7 +1,6 @@ from __future__ import annotations import contextlib -import logging import random import struct import traceback @@ -23,6 +22,7 @@ from ..http11 import Request, Response from ..protocol import CLOSED, OPEN, Event, Protocol, State from ..typing import BytesLike, Data, DataLike, LoggerLike, Subprotocol +from ..utils import ConnectionLoggerAdapter from .messages import Assembler @@ -65,10 +65,9 @@ def __init__( max_queue_high, max_queue_low = max_queue # Inject reference to this instance in the protocol's logger. - self.protocol.logger = logging.LoggerAdapter( - self.protocol.logger, - {"websocket": self}, - ) + # ConnectionLoggerAdapter holds a weak reference in order to + # keep the connection garbage-collectable by reference counting. + self.protocol.logger = ConnectionLoggerAdapter(self.protocol.logger, self) # Copy attributes from the protocol for convenience. self.id: uuid.UUID = self.protocol.id diff --git a/src/websockets/trio/messages.py b/src/websockets/trio/messages.py index a578d81c..6ef5d588 100644 --- a/src/websockets/trio/messages.py +++ b/src/websockets/trio/messages.py @@ -282,5 +282,10 @@ def close(self) -> None: self.closed = True + # Drop the pause() and resume() callbacks. As bound methods of the + # connection, they would keep it in a reference cycle. Flow control + # isn't needed anymore. + self.pause = self.resume = lambda: None + # Unblock get() or get_iter(). self.send_frames.close() diff --git a/src/websockets/utils.py b/src/websockets/utils.py index 5a73226a..df8f7527 100644 --- a/src/websockets/utils.py +++ b/src/websockets/utils.py @@ -2,11 +2,15 @@ import base64 import hashlib +import logging import secrets import socket import sys +import weakref +from collections.abc import MutableMapping +from typing import Any -from .typing import BytesLike +from .typing import BytesLike, LoggerLike __all__ = ["accept_key", "apply_mask", "get_socket_name"] @@ -58,6 +62,30 @@ def apply_mask(data: BytesLike, mask: bytes | bytearray) -> bytes: return (data_int ^ mask_int).to_bytes(len(data), sys.byteorder) +class ConnectionLoggerAdapter(logging.LoggerAdapter[LoggerLike]): + """ + Logger adapter that adds a ``websocket`` attribute to log records. + + It holds only a weak reference to the connection. A strong reference + would create a reference cycle between the connection and its logger, + delaying garbage collection of closed connections until a full garbage + collection pass, as reference counting cannot free reference cycles. + + """ + + def __init__(self, logger: LoggerLike, websocket: object) -> None: + super().__init__(logger) + self.websocket_ref = weakref.ref(websocket) + + def process( + self, msg: Any, kwargs: MutableMapping[str, Any] + ) -> tuple[Any, MutableMapping[str, Any]]: + websocket = self.websocket_ref() + if websocket is not None: + kwargs["extra"] = {"websocket": websocket} + return msg, kwargs + + def get_socket_name(sock: socket.socket) -> str: """ Return a string representation of :meth:`~socket.socket.getsockname()`. diff --git a/tests/asyncio/test_connection.py b/tests/asyncio/test_connection.py index 5cf61940..74d45dd1 100644 --- a/tests/asyncio/test_connection.py +++ b/tests/asyncio/test_connection.py @@ -1,10 +1,12 @@ import asyncio import contextlib +import gc import itertools import logging import socket import unittest import uuid +import weakref from unittest.mock import Mock, patch from websockets.asyncio.connection import * @@ -18,7 +20,7 @@ from websockets.protocol import CLIENT, CLOSED, OPEN, SERVER, Protocol from ..protocol import RecordingProtocol -from ..utils import MS, LoggingTestCase, alist +from ..utils import MS, LoggingTestCase, alist, skip_unless_reference_counting_collects from .connection import InterceptingConnection @@ -46,7 +48,8 @@ async def asyncSetUp(self): async def asyncTearDown(self): await self.remote_connection.close() - await self.connection.close() + if hasattr(self, "connection"): # garbage collection tests delete it + await self.connection.close() # Test helpers built upon RecordingProtocol and InterceptingConnection. @@ -1113,9 +1116,10 @@ async def test_keepalive_terminates_while_sleeping(self): self.connection.ping_interval = 3 * MS self.connection.start_keepalive() await asyncio.sleep(MS) - self.assertFalse(self.connection.keepalive_task.done()) + keepalive_task = self.connection.keepalive_task + self.assertFalse(keepalive_task.done()) await self.connection.close() - self.assertTrue(self.connection.keepalive_task.done()) + self.assertTrue(keepalive_task.done()) # test_keepalive_terminates_when_sending_ping_fails is not implemented # because sending a ping cannot fail in the asyncio implementation. @@ -1129,9 +1133,10 @@ async def test_keepalive_terminates_while_waiting_for_pong(self): # 1 ms: keepalive() sends a ping frame. # 1.x ms: a pong frame is dropped. await asyncio.sleep(2 * MS) + keepalive_task = self.connection.keepalive_task # 2 ms: close the connection before ping_timeout elapses. await self.connection.close() - self.assertTrue(self.connection.keepalive_task.done()) + self.assertTrue(keepalive_task.done()) async def test_keepalive_reports_errors(self): """keepalive reports unexpected errors in logs.""" @@ -1483,6 +1488,49 @@ async def test_broadcast_type_error(self): with self.assertRaises(TypeError): broadcast([self.connection], ["⏳", "⌛️"]) + # Test garbage collection of closed connections. + + @skip_unless_reference_counting_collects + async def test_garbage_collection_after_close(self): + """Connection is freed by reference counting after a closing handshake.""" + self.connection.start_keepalive() + # Let the keepalive task start before closing the connection. + await asyncio.sleep(0) + + connection_ref = weakref.ref(self.connection) + gc.disable() + self.addCleanup(gc.enable) + + await self.connection.close() + del self.connection, self.transport + + # Let the event loop cancel the keepalive task. + for _ in range(3): + await asyncio.sleep(0) + + self.assertIsNone(connection_ref()) + + @skip_unless_reference_counting_collects + async def test_garbage_collection_after_abort(self): + """Connection is freed by reference counting after aborting the connection.""" + self.connection.start_keepalive() + # Let the keepalive task start before closing the connection. + await asyncio.sleep(0) + + connection_ref = weakref.ref(self.connection) + gc.disable() + self.addCleanup(gc.enable) + + self.connection.transport.abort() + await asyncio.shield(self.connection.connection_lost_waiter) + del self.connection, self.transport + + # Let the event loop cancel the keepalive task. + for _ in range(3): + await asyncio.sleep(0) + + self.assertIsNone(connection_ref()) + class ServerConnectionTests(ClientConnectionTests): LOCAL = SERVER diff --git a/tests/asyncio/test_server.py b/tests/asyncio/test_server.py index a687337f..18db4150 100644 --- a/tests/asyncio/test_server.py +++ b/tests/asyncio/test_server.py @@ -1,10 +1,12 @@ import asyncio import dataclasses +import gc import hmac import http import logging import socket import unittest +import weakref from websockets.asyncio.client import connect, unix_connect from websockets.asyncio.server import * @@ -22,6 +24,7 @@ MS, SERVER_CONTEXT, LoggingTestCase, + skip_unless_reference_counting_collects, temp_unix_socket_path, ) from .server import ( @@ -62,6 +65,27 @@ async def test_connection_handler_raises_exception(self): "received 1011 (internal error); then sent 1011 (internal error)", ) + @skip_unless_reference_counting_collects + async def test_garbage_collection_after_close(self): + """Closed connection is freed by reference counting in a server context.""" + async with serve(*args) as server: + async with connect(get_uri(server)) as client: + await self.assertEval(client, "ws.protocol.state.name", "OPEN") + + [handler_task] = server.handler_tasks + [connection] = server.all_connections + connection_ref = weakref.ref(connection) + gc.disable() + self.addCleanup(gc.enable) + del connection + + # Closing the client causes the connection handler to return, + # which closes the connection on the server side. + await handler_task + del handler_task + + self.assertIsNone(connection_ref()) + async def test_existing_socket(self): """Server receives connection using a pre-existing socket.""" with socket.create_server(("localhost", 0)) as sock: diff --git a/tests/sync/test_connection.py b/tests/sync/test_connection.py index cc142caf..65d0caf2 100644 --- a/tests/sync/test_connection.py +++ b/tests/sync/test_connection.py @@ -1,10 +1,12 @@ import contextlib +import gc import itertools import logging import socket import threading import time import uuid +import weakref from unittest.mock import Mock, patch from websockets.exceptions import ( @@ -18,7 +20,7 @@ from websockets.sync.connection import broadcast from ..protocol import RecordingProtocol -from ..utils import MS, LoggingTestCase +from ..utils import MS, LoggingTestCase, skip_unless_reference_counting_collects from .connection import InterceptingConnection from .utils import ThreadTestCase @@ -40,7 +42,8 @@ def setUp(self): def tearDown(self): self.remote_connection.close() - self.connection.close() + if hasattr(self, "connection"): # garbage collection tests delete it + self.connection.close() # Test helpers built upon RecordingProtocol and InterceptingConnection. @@ -1204,6 +1207,29 @@ def test_broadcast_type_error(self): with self.assertRaises(TypeError): broadcast([self.connection], ["⏳", "⌛️"]) + # Test garbage collection of closed connections. + + @skip_unless_reference_counting_collects + def test_garbage_collection_after_close(self): + """Connection is freed by reference counting after a closing handshake.""" + self.connection.start_keepalive() + + connection_ref = weakref.ref(self.connection) + gc.disable() + self.addCleanup(gc.enable) + + self.connection.close() + recv_events_thread = self.connection.recv_events_thread + keepalive_thread = self.connection.keepalive_thread + del self.connection + + # Wait until the threads that keep a reference to the connection in + # their stack frames terminate. + recv_events_thread.join() + keepalive_thread.join() + + self.assertIsNone(connection_ref()) + class ServerConnectionTests(ClientConnectionTests): LOCAL = SERVER diff --git a/tests/sync/test_server.py b/tests/sync/test_server.py index ee01a5ee..b8a2ffcc 100644 --- a/tests/sync/test_server.py +++ b/tests/sync/test_server.py @@ -1,4 +1,5 @@ import dataclasses +import gc import hmac import http import logging @@ -6,6 +7,7 @@ import threading import time import unittest +import weakref from unittest.mock import patch from websockets.exceptions import ( @@ -25,6 +27,7 @@ SERVER_CONTEXT, DeprecationTestCase, LoggingTestCase, + skip_unless_reference_counting_collects, temp_unix_socket_path, ) from .server import ( @@ -65,6 +68,32 @@ def test_connection_handler_raises_exception(self): "received 1011 (internal error); then sent 1011 (internal error)", ) + @skip_unless_reference_counting_collects + def test_garbage_collection_after_close(self): + """Closed connection is freed by reference counting in a server context.""" + with run_server() as server: + with connect(get_uri(server)) as client: + self.assertEval(client, "ws.protocol.state.name", "OPEN") + + [handler_thread] = server.handler_threads + [connection] = server.all_connections + recv_events_thread = connection.recv_events_thread + keepalive_thread = connection.keepalive_thread + connection_ref = weakref.ref(connection) + gc.disable() + self.addCleanup(gc.enable) + del connection + + # Closing the client causes the connection handler to return, + # which closes the connection on the server side. + handler_thread.join() + # Wait until the threads that keep a reference to the connection + # in their stack frames terminate. + recv_events_thread.join() + keepalive_thread.join() + + self.assertIsNone(connection_ref()) + def test_existing_socket(self): """Server receives connection using a pre-existing socket.""" with socket.create_server(("localhost", 0)) as sock: diff --git a/tests/test_utils.py b/tests/test_utils.py index cbd777e2..a59e221b 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,10 +1,13 @@ import base64 +import gc import itertools +import logging import platform import socket import unittest from websockets.utils import ( + ConnectionLoggerAdapter, accept_key, apply_mask as py_apply_mask, generate_key, @@ -111,6 +114,37 @@ def apply_mask(*args, **kwargs): raise +class FakeConnection: + """Object standing in for a connection; plain objects aren't weakrefable.""" + + +class ConnectionLoggerAdapterTests(unittest.TestCase): + def setUp(self): + self.websocket = FakeConnection() + self.adapter = ConnectionLoggerAdapter( + logging.getLogger("websockets.test"), + self.websocket, + ) + + def test_process_adds_websocket_to_extra(self): + """process makes the connection available in the extra dict.""" + msg, kwargs = self.adapter.process("message", {}) + self.assertIs(kwargs["extra"]["websocket"], self.websocket) + + def test_log_records_have_websocket_attribute(self): + """Log records have a websocket attribute referencing the connection.""" + with self.assertLogs("websockets.test", logging.INFO) as logs: + self.adapter.info("message") + self.assertIs(logs.records[0].websocket, self.websocket) + + def test_adapter_does_not_keep_websocket_alive(self): + """Adapter doesn't prevent garbage collection of the connection.""" + del self.websocket + gc.collect() + msg, kwargs = self.adapter.process("message", {}) + self.assertNotIn("extra", kwargs) + + class GetSocketNameAsStrTests(unittest.TestCase): def test_af_inet(self): with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: diff --git a/tests/trio/test_connection.py b/tests/trio/test_connection.py index 4ab19066..468d3ebd 100644 --- a/tests/trio/test_connection.py +++ b/tests/trio/test_connection.py @@ -1,7 +1,9 @@ import contextlib +import gc import itertools import logging import uuid +import weakref from unittest.mock import patch import trio.testing @@ -17,7 +19,7 @@ from websockets.trio.connection import broadcast from ..protocol import RecordingProtocol -from ..utils import MS, LoggingTestCase, alist +from ..utils import MS, LoggingTestCase, alist, skip_unless_reference_counting_collects from .connection import InterceptingConnection from .utils import IsolatedTrioTestCase @@ -48,7 +50,8 @@ async def asyncSetUp(self): async def asyncTearDown(self): await self.remote_connection.aclose() - await self.connection.aclose() + if hasattr(self, "connection"): # garbage collection tests delete it + await self.connection.aclose() # Test helpers built upon RecordingProtocol and InterceptingConnection. @@ -1435,6 +1438,26 @@ async def test_broadcast_type_error(self): with self.assertRaises(TypeError): await broadcast([self.connection], ["⏳", "⌛️"]) + # Test garbage collection of closed connections. + + @skip_unless_reference_counting_collects + async def test_garbage_collection_after_close(self): + """Connection is freed by reference counting after a closing handshake.""" + self.connection.start_keepalive() + await trio.testing.wait_all_tasks_blocked() + + connection_ref = weakref.ref(self.connection) + gc.disable() + self.addCleanup(gc.enable) + + await self.connection.aclose() + del self.connection + + # Let the tasks running recv_events() and keepalive() terminate. + await trio.testing.wait_all_tasks_blocked() + + self.assertIsNone(connection_ref()) + class ServerConnectionTests(ClientConnectionTests): LOCAL = SERVER diff --git a/tests/trio/test_server.py b/tests/trio/test_server.py index 8018175e..8cf010a6 100644 --- a/tests/trio/test_server.py +++ b/tests/trio/test_server.py @@ -1,9 +1,12 @@ import dataclasses +import gc import hmac import http import logging +import weakref import trio +import trio.testing from websockets.exceptions import ( ConnectionClosedError, @@ -21,6 +24,7 @@ MS, SERVER_CONTEXT, LoggingTestCase, + skip_unless_reference_counting_collects, ) from .server import ( EvalShellMixin, @@ -61,6 +65,29 @@ async def test_connection_handler_raises_exception(self): "received 1011 (internal error); then sent 1011 (internal error)", ) + @skip_unless_reference_counting_collects + async def test_garbage_collection_after_close(self): + """Closed connection is freed by reference counting in a server context.""" + async with run_server() as server: + async with connect(get_uri(server)) as client: + await self.assertEval(client, "ws.protocol.state.name", "OPEN") + + [connection] = server.all_connections + connection_ref = weakref.ref(connection) + gc.disable() + self.addCleanup(gc.enable) + del connection + + # Closing the client causes the connection handler to return, + # which closes the connection on the server side. + while server.all_connections: + await trio.sleep(MS) + # Let the tasks running the connection handler, recv_events(), + # and keepalive() terminate. + await trio.testing.wait_all_tasks_blocked() + + self.assertIsNone(connection_ref()) + async def test_existing_listeners(self): """Server receives connection using pre-existing listeners.""" listeners = await trio.open_tcp_listeners(0, host="localhost") diff --git a/tests/trio/utils.py b/tests/trio/utils.py index 19708524..51981efa 100644 --- a/tests/trio/utils.py +++ b/tests/trio/utils.py @@ -22,6 +22,8 @@ def __init_subclass__(cls, **kwargs): test = getattr(cls, name) if getattr(test, "converted_to_trio", False): # pragma: no cover return + if getattr(test, "__unittest_skip__", False): + continue assert inspect.iscoroutinefunction(test) setattr(cls, name, cls.convert_to_trio(test)) diff --git a/tests/utils.py b/tests/utils.py index 80a7b23e..df631f34 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -4,6 +4,7 @@ import pathlib import platform import ssl +import sys import tempfile import time import unittest @@ -67,6 +68,17 @@ MS = max(MS, 2.5 * time.get_clock_info("monotonic").resolution) +# Garbage collection tests check that closed connections are freed by reference +# counting alone. Skip them on PyPy, which doesn't implement reference counting, +# and on Python 3.12, where the traceback of an exception raised while closing +# a connection keeps frames of connection methods alive via their f_back +# attribute. Python 3.13 fixed that behavior. +skip_unless_reference_counting_collects = unittest.skipIf( + platform.python_implementation() != "CPython" or sys.version_info[:2] == (3, 12), + "test requires that reference counting frees closed connections", +) + + class GeneratorTestCase(unittest.TestCase): """ Base class for testing generator-based coroutines. From e5ebdf492f1ebdda5330b4c9faa713af3cd9dc17 Mon Sep 17 00:00:00 2001 From: Peter Alm Date: Fri, 21 Aug 2026 10:25:55 +0200 Subject: [PATCH 2/2] Break reference cycles created by the traceback of recv_exc. When recv_exc is raised in a method of the connection, its traceback references frames of connection methods, which reference the connection: connection -> recv_exc -> traceback -> frame -> connection. This prevents reference counting from freeing closed connections. It happens on any read error in the sync and trio implementations, whose recv_events() is a method of the connection, and on write errors while responding to incoming frames in all implementations. Clear the frames of recv_exc's traceback once they finished running: after connection_lost() in the asyncio implementation and after recv_events() terminates in the sync and trio implementations. Clearing frames doesn't affect formatting the traceback. In the sync implementation, the thread running recv_events() must also dereference the connection when it terminates because the traceback of recv_exc keeps a reference to its frame via the f_back attribute of the recv_events() frame. --- src/websockets/asyncio/connection.py | 7 +++++ src/websockets/sync/connection.py | 38 +++++++++++++++++++++++----- src/websockets/trio/connection.py | 16 +++++++++++- tests/asyncio/test_connection.py | 18 +++++++++++++ tests/sync/test_connection.py | 19 ++++++++++++++ tests/trio/test_connection.py | 21 +++++++++++++++ 6 files changed, 111 insertions(+), 8 deletions(-) diff --git a/src/websockets/asyncio/connection.py b/src/websockets/asyncio/connection.py index 97936150..6c25fc35 100644 --- a/src/websockets/asyncio/connection.py +++ b/src/websockets/asyncio/connection.py @@ -1019,6 +1019,13 @@ def connection_lost(self, exc: Exception | None) -> None: self.set_recv_exc(exc) + # Clear the frames of recv_exc's traceback. Else, when recv_exc was + # raised in a method of this class e.g. data_received(), its traceback + # would keep the connection in a reference cycle. Frames finished + # running, so they may be cleared; formatting isn't affected. + if self.recv_exc is not None: + traceback.clear_frames(self.recv_exc.__traceback__) + # Abort recv() and pending pings with a ConnectionClosed exception. self.recv_messages.close() self.terminate_pending_pings() diff --git a/src/websockets/sync/connection.py b/src/websockets/sync/connection.py index a8bd8546..1cfae0c6 100644 --- a/src/websockets/sync/connection.py +++ b/src/websockets/sync/connection.py @@ -31,6 +31,35 @@ __all__ = ["Connection"] +class RecvEventsThread(threading.Thread): + """ + Thread running :meth:`Connection.recv_events`. + + This thread is marked as daemon to allow creating a connection in a + non-daemon thread and using it in a daemon thread. This mustn't prevent + the interpreter from exiting. + + """ + + def __init__(self, connection: Connection) -> None: + super().__init__(daemon=True) + self.connection = connection + + def run(self) -> None: + try: + self.connection.recv_events() + finally: + recv_exc = self.connection.recv_exc + # Dereference the connection, like Thread.run() dereferences + # self._target, so that this thread doesn't keep it alive. + del self.connection + # Clear the frames of recv_exc's traceback. Else, it would keep + # the connection in a reference cycle. Frames may be cleared only + # after recv_events() terminates. + if recv_exc is not None: + traceback.clear_frames(recv_exc.__traceback__) + + class Connection: """ :mod:`threading` implementation of a WebSocket connection. @@ -127,13 +156,8 @@ def __init__( # ConnectionClosed in order to show why the TCP connection dropped. self.recv_exc: BaseException | None = None - # Receiving events from the socket. This thread is marked as daemon to - # allow creating a connection in a non-daemon thread and using it in a - # daemon thread. This mustn't prevent the interpreter from exiting. - self.recv_events_thread = threading.Thread( - target=self.recv_events, - daemon=True, - ) + # Receiving events from the socket. + self.recv_events_thread = RecvEventsThread(self) # Start recv_events only after all attributes are initialized. self.recv_events_thread.start() diff --git a/src/websockets/trio/connection.py b/src/websockets/trio/connection.py index 99c2e186..aeeaff2e 100644 --- a/src/websockets/trio/connection.py +++ b/src/websockets/trio/connection.py @@ -127,7 +127,7 @@ def __init__( self.stream_closed: trio.Event = trio.Event() # Start recv_events only after all attributes are initialized. - self.nursery.start_soon(self.recv_events) + self.nursery.start_soon(self.run_recv_events) # Public attributes @@ -868,6 +868,20 @@ def start_keepalive(self) -> None: if self.ping_interval is not None: self.nursery.start_soon(self.keepalive) + async def run_recv_events(self) -> None: + """ + Run :meth:`recv_events` then clear the frames of ``recv_exc``'s traceback. + + Else, the traceback would keep the connection in a reference cycle. + Frames may be cleared only after :meth:`recv_events` terminates. + + """ + try: + await self.recv_events() + finally: + if self.recv_exc is not None: + traceback.clear_frames(self.recv_exc.__traceback__) + async def recv_events(self) -> None: """ Read incoming data from the stream and process events. diff --git a/tests/asyncio/test_connection.py b/tests/asyncio/test_connection.py index 74d45dd1..f41b5095 100644 --- a/tests/asyncio/test_connection.py +++ b/tests/asyncio/test_connection.py @@ -1531,6 +1531,24 @@ async def test_garbage_collection_after_abort(self): self.assertIsNone(connection_ref()) + @skip_unless_reference_counting_collects + async def test_garbage_collection_after_network_error(self): + """Connection is freed by reference counting after a network error.""" + # Inject a fault by shutting down the transport for writing. + # Responding to the incoming ping will fail and set recv_exc. + self.transport.write_eof() + await self.remote_connection.ping() + with self.assertRaises(ConnectionClosedError): + await self.connection.recv() + + connection_ref = weakref.ref(self.connection) + gc.disable() + self.addCleanup(gc.enable) + + del self.connection, self.transport + + self.assertIsNone(connection_ref()) + class ServerConnectionTests(ClientConnectionTests): LOCAL = SERVER diff --git a/tests/sync/test_connection.py b/tests/sync/test_connection.py index 65d0caf2..e4393ef0 100644 --- a/tests/sync/test_connection.py +++ b/tests/sync/test_connection.py @@ -1230,6 +1230,25 @@ def test_garbage_collection_after_close(self): self.assertIsNone(connection_ref()) + @skip_unless_reference_counting_collects + def test_garbage_collection_after_network_error(self): + """Connection is freed by reference counting after a network error.""" + # Inject a fault by making sendall() fail. Responding to the + # incoming ping will fail and set recv_exc. + self.connection.socket = Mock(wraps=self.connection.socket) + self.connection.socket.sendall.side_effect = BrokenPipeError + self.remote_connection.ping() + with self.assertRaises(ConnectionClosedError): + self.connection.recv() + + connection_ref = weakref.ref(self.connection) + gc.disable() + self.addCleanup(gc.enable) + + del self.connection + + self.assertIsNone(connection_ref()) + class ServerConnectionTests(ClientConnectionTests): LOCAL = SERVER diff --git a/tests/trio/test_connection.py b/tests/trio/test_connection.py index 468d3ebd..ce917762 100644 --- a/tests/trio/test_connection.py +++ b/tests/trio/test_connection.py @@ -1458,6 +1458,27 @@ async def test_garbage_collection_after_close(self): self.assertIsNone(connection_ref()) + @skip_unless_reference_counting_collects + async def test_garbage_collection_after_network_error(self): + """Connection is freed by reference counting after a network error.""" + # Inject a fault by closing the stream for writing. Responding to + # the incoming ping will fail and set recv_exc. + self.connection.stream.send_stream.close() + await self.remote_connection.ping() + with self.assertRaises(ConnectionClosedError): + await self.connection.recv() + + connection_ref = weakref.ref(self.connection) + gc.disable() + self.addCleanup(gc.enable) + + del self.connection + + # Let the task running recv_events() terminate. + await trio.testing.wait_all_tasks_blocked() + + self.assertIsNone(connection_ref()) + class ServerConnectionTests(ClientConnectionTests): LOCAL = SERVER