diff --git a/docs/project/changelog.rst b/docs/project/changelog.rst index fa2540c1..d38e8a4a 100644 --- a/docs/project/changelog.rst +++ b/docs/project/changelog.rst @@ -32,6 +32,11 @@ notice. *In development* +Improvements +............ + +* Connections are now garbage collected immediately once closed. + .. _17.0.1: 17.0.1 diff --git a/src/websockets/asyncio/connection.py b/src/websockets/asyncio/connection.py index 74ab5e2a..87d96ab5 100644 --- a/src/websockets/asyncio/connection.py +++ b/src/websockets/asyncio/connection.py @@ -8,6 +8,7 @@ import struct import traceback import uuid +import weakref from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Iterable, Mapping from types import TracebackType from typing import Any, Literal, Self, cast, overload @@ -67,7 +68,7 @@ def __init__( # Inject reference to this instance in the protocol's logger. self.protocol.logger = logging.LoggerAdapter( self.protocol.logger, - {"websocket": self}, + {"websocket": weakref.proxy(self)}, ) # Copy attributes from the protocol for convenience. @@ -1025,6 +1026,8 @@ def connection_lost(self, exc: Exception | None) -> None: if self.keepalive_task is not None: self.keepalive_task.cancel() + # Break reference cycle to allow immediate garbage collection. + 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/legacy/protocol.py b/src/websockets/legacy/protocol.py index 0f86d3d9..cb97e481 100644 --- a/src/websockets/legacy/protocol.py +++ b/src/websockets/legacy/protocol.py @@ -11,6 +11,7 @@ import traceback import uuid import warnings +import weakref from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Iterable, Mapping from typing import Any, Callable, Deque, cast @@ -205,7 +206,10 @@ def __init__( # Logger or LoggerAdapter for this connection. if logger is None: logger = logging.getLogger("websockets.protocol") - self.logger: LoggerLike = logging.LoggerAdapter(logger, {"websocket": self}) + self.logger: LoggerLike = logging.LoggerAdapter( + logger, + {"websocket": weakref.proxy(self)}, + ) """Logger for this connection.""" # Track if DEBUG is enabled. Shortcut logging calls if it isn't. @@ -1278,6 +1282,8 @@ async def close_connection(self) -> None: # Cancel the keepalive ping task. if hasattr(self, "keepalive_ping_task"): self.keepalive_ping_task.cancel() + # Break reference cycle to allow immediate garbage collection. + del self.keepalive_ping_task # A client should wait for a TCP close from the server. if self.is_client and hasattr(self, "transfer_data_task"): diff --git a/src/websockets/sync/connection.py b/src/websockets/sync/connection.py index 74ac997a..c687d78f 100644 --- a/src/websockets/sync/connection.py +++ b/src/websockets/sync/connection.py @@ -10,6 +10,7 @@ import time import traceback import uuid +import weakref from collections.abc import Iterable, Iterator, Mapping from types import TracebackType from typing import Any, Literal, Self, overload @@ -69,7 +70,7 @@ def __init__( # Inject reference to this instance in the protocol's logger. self.protocol.logger = logging.LoggerAdapter( self.protocol.logger, - {"websocket": self}, + {"websocket": weakref.proxy(self)}, ) # Copy attributes from the protocol for convenience. diff --git a/src/websockets/trio/connection.py b/src/websockets/trio/connection.py index 606653ad..26e03210 100644 --- a/src/websockets/trio/connection.py +++ b/src/websockets/trio/connection.py @@ -6,6 +6,7 @@ import struct import traceback import uuid +import weakref from collections.abc import AsyncIterable, AsyncIterator, Iterable, Mapping from types import TracebackType from typing import Any, Literal, Self, overload @@ -67,7 +68,7 @@ def __init__( # Inject reference to this instance in the protocol's logger. self.protocol.logger = logging.LoggerAdapter( self.protocol.logger, - {"websocket": self}, + {"websocket": weakref.proxy(self)}, ) # Copy attributes from the protocol for convenience. diff --git a/src/websockets/trio/messages.py b/src/websockets/trio/messages.py index a578d81c..a1b348f6 100644 --- a/src/websockets/trio/messages.py +++ b/src/websockets/trio/messages.py @@ -284,3 +284,8 @@ def close(self) -> None: # Unblock get() or get_iter(). self.send_frames.close() + + # Break reference cycle to allow immediate garbage collection: + # Connection -> Assembler -> bound methods -> Connection. + self.pause = lambda: None + self.resume = lambda: None diff --git a/tests/asyncio/test_connection.py b/tests/asyncio/test_connection.py index 5cf61940..e0b4c7be 100644 --- a/tests/asyncio/test_connection.py +++ b/tests/asyncio/test_connection.py @@ -5,6 +5,7 @@ import socket import unittest import uuid +import weakref from unittest.mock import Mock, patch from websockets.asyncio.connection import * @@ -46,7 +47,8 @@ async def asyncSetUp(self): async def asyncTearDown(self): await self.remote_connection.close() - await self.connection.close() + if hasattr(self, "connection"): + await self.connection.close() # Test helpers built upon RecordingProtocol and InterceptingConnection. @@ -1112,10 +1114,13 @@ async def test_keepalive_terminates_while_sleeping(self): """keepalive task terminates while waiting to send a ping.""" self.connection.ping_interval = 3 * MS self.connection.start_keepalive() + # Keep a reference to self.connection.keepalive_task + # because it's set to None once cancelled. + keepalive_task = self.connection.keepalive_task await asyncio.sleep(MS) - self.assertFalse(self.connection.keepalive_task.done()) + 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. @@ -1126,12 +1131,15 @@ async def test_keepalive_terminates_while_waiting_for_pong(self): self.connection.ping_timeout = 4 * MS async with self.drop_frames_rcvd(): self.connection.start_keepalive() + # Keep a reference to self.connection.keepalive_task + # because it's set to None once cancelled. + keepalive_task = self.connection.keepalive_task # 1 ms: keepalive() sends a ping frame. # 1.x ms: a pong frame is dropped. await asyncio.sleep(2 * MS) # 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.""" @@ -1329,6 +1337,18 @@ async def test_unexpected_failure_in_send_context(self, send_text): await self.connection.send("😀") self.assertIsInstance(raised.exception.__cause__, AssertionError) + # Test garbage collection. + + async def test_no_reference_cycle(self): + """Connection is garbage collected immediately after deletion.""" + self.connection.start_keepalive() + await asyncio.sleep(0) # let the event loop start the keepalive task + await self.connection.close() + + connection_ref = weakref.ref(self.connection) + del self.connection + self.assertIsNone(connection_ref(), "still alive after deletion") + # Test broadcast. async def test_broadcast_text(self): diff --git a/tests/legacy/test_protocol.py b/tests/legacy/test_protocol.py index 79d97ced..d628cd86 100644 --- a/tests/legacy/test_protocol.py +++ b/tests/legacy/test_protocol.py @@ -1208,6 +1208,8 @@ async def create_protocol(): self.protocol.is_client = initial_protocol.is_client self.protocol.side = initial_protocol.side + self.keepalive_ping_task = self.protocol.keepalive_ping_task + def test_keepalive_ping(self): self.restart_protocol_with_keepalive_ping() @@ -1223,7 +1225,7 @@ def test_keepalive_ping(self): self.assertOneFrameSent(True, OP_PING, ping_2) # The keepalive ping task goes on. - self.assertFalse(self.protocol.keepalive_ping_task.done()) + self.assertFalse(self.keepalive_ping_task.done()) def test_keepalive_ping_not_acknowledged_closes_connection(self): self.restart_protocol_with_keepalive_ping() @@ -1242,7 +1244,7 @@ def test_keepalive_ping_not_acknowledged_closes_connection(self): ) # The keepalive ping task is complete. - self.assertEqual(self.protocol.keepalive_ping_task.result(), None) + self.assertEqual(self.keepalive_ping_task.result(), None) def test_keepalive_ping_stops_when_connection_closing(self): self.restart_protocol_with_keepalive_ping() @@ -1253,7 +1255,7 @@ def test_keepalive_ping_stops_when_connection_closing(self): self.assertNoFrameSent() # The keepalive ping task terminated. - self.assertTrue(self.protocol.keepalive_ping_task.cancelled()) + self.assertTrue(self.keepalive_ping_task.cancelled()) self.loop.run_until_complete(close_task) # cleanup @@ -1262,7 +1264,7 @@ def test_keepalive_ping_stops_when_connection_closed(self): self.close_connection() # The keepalive ping task terminated. - self.assertTrue(self.protocol.keepalive_ping_task.cancelled()) + self.assertTrue(self.keepalive_ping_task.cancelled()) def test_keepalive_ping_does_not_crash_when_connection_lost(self): self.restart_protocol_with_keepalive_ping() @@ -1283,7 +1285,7 @@ def test_keepalive_ping_does_not_crash_when_connection_lost(self): with self.assertRaises(ConnectionClosed): pong_waiter.result() # The keepalive ping task terminated properly. - self.assertIsNone(self.protocol.keepalive_ping_task.result()) + self.assertIsNone(self.keepalive_ping_task.result()) # Unclog incoming queue to terminate the test quickly. self.loop.run_until_complete(self.protocol.recv()) @@ -1311,7 +1313,7 @@ def test_keepalive_ping_with_no_ping_timeout(self): self.assertOneFrameSent(True, OP_PING, ping_2) # The keepalive ping task goes on. - self.assertFalse(self.protocol.keepalive_ping_task.done()) + self.assertFalse(self.keepalive_ping_task.done()) def test_keepalive_ping_unexpected_error(self): self.restart_protocol_with_keepalive_ping() @@ -1326,7 +1328,7 @@ async def ping(): # The keepalive ping task is complete. # It logs and swallows the exception. - self.assertEqual(self.protocol.keepalive_ping_task.result(), None) + self.assertEqual(self.keepalive_ping_task.result(), None) # Test the protocol logic for closing the connection. diff --git a/tests/sync/test_connection.py b/tests/sync/test_connection.py index cc142caf..3d40dd8b 100644 --- a/tests/sync/test_connection.py +++ b/tests/sync/test_connection.py @@ -5,6 +5,7 @@ import threading import time import uuid +import weakref from unittest.mock import Mock, patch from websockets.exceptions import ( @@ -40,7 +41,8 @@ def setUp(self): def tearDown(self): self.remote_connection.close() - self.connection.close() + if hasattr(self, "connection"): + self.connection.close() # Test helpers built upon RecordingProtocol and InterceptingConnection. @@ -1049,6 +1051,19 @@ def test_close_expected_during_connecting(self): connection.close_socket() connection.recv_events_thread.join() + # Test garbage collection. + + def test_no_reference_cycle(self): + """Connection is garbage collected immediately after deletion.""" + self.connection.start_keepalive() + time.sleep(0) # let the keepalive thread start + self.connection.close() + time.sleep(MS) # let the recv_events thread terminate + + connection_ref = weakref.ref(self.connection) + del self.connection + self.assertIsNone(connection_ref(), "still alive after deletion") + # Test broadcast. def test_broadcast_text(self): diff --git a/tests/trio/test_connection.py b/tests/trio/test_connection.py index 4ab19066..e804ef1c 100644 --- a/tests/trio/test_connection.py +++ b/tests/trio/test_connection.py @@ -2,6 +2,7 @@ import itertools import logging import uuid +import weakref from unittest.mock import patch import trio.testing @@ -48,7 +49,8 @@ async def asyncSetUp(self): async def asyncTearDown(self): await self.remote_connection.aclose() - await self.connection.aclose() + if hasattr(self, "connection"): + await self.connection.aclose() # Test helpers built upon RecordingProtocol and InterceptingConnection. @@ -1285,6 +1287,20 @@ async def test_unexpected_failure_in_send_context(self, send_text): await self.connection.send("😀") self.assertIsInstance(raised.exception.__cause__, AssertionError) + # Test garbage collection. + + async def test_no_reference_cycle(self): + """Connection is garbage collected immediately after deletion.""" + self.connection.start_keepalive() + await trio.testing.wait_all_tasks_blocked() + await self.connection.aclose() + # Wait for recv_events(), which runs in self.nursery, to terminate. + await trio.testing.wait_all_tasks_blocked() + + connection_ref = weakref.ref(self.connection) + del self.connection + self.assertIsNone(connection_ref(), "still alive after deletion") + # Test broadcast. async def test_broadcast_text(self):