Break reference cycles keeping closed connections alive

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

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
@@ -1019,12 +1019,24 @@ 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()
 
         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;

src/websockets/sync/connection.py

@@ -2,7 +2,6 @@
 
 import concurrent.futures
 import contextlib
-import logging
 import random
 import socket
 import struct
@@ -24,13 +23,43 @@
 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
 
 
 __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.
@@ -67,10 +96,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
@@ -128,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()

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
@@ -128,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
 
@@ -869,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.

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

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()`.