Enable immediate garbage collection of connections.

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

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;

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"):

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.

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.

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

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):

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.
 

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):

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):