diff --git a/mitmproxy/addons/clientplayback.py b/mitmproxy/addons/clientplayback.py index b956bc600..64d90e5a9 100644 --- a/mitmproxy/addons/clientplayback.py +++ b/mitmproxy/addons/clientplayback.py @@ -158,7 +158,9 @@ class ClientPlayback: def running(self): self.options = ctx.options self.playback_task = asyncio_utils.create_task( - self.playback(), name="client playback" + self.playback(), + name="client playback", + keep_ref=False, ) async def done(self): @@ -177,7 +179,9 @@ class ClientPlayback: h = ReplayHandler(self.inflight, self.options) if ctx.options.client_replay_concurrency == -1: t = asyncio_utils.create_task( - h.replay(), name="client playback awaiting response" + h.replay(), + name="client playback awaiting response", + keep_ref=False, ) # keep a reference so this is not garbage collected self.replay_tasks.add(t) diff --git a/mitmproxy/addons/keepserving.py b/mitmproxy/addons/keepserving.py index 04554c9f3..183000347 100644 --- a/mitmproxy/addons/keepserving.py +++ b/mitmproxy/addons/keepserving.py @@ -7,8 +7,6 @@ from mitmproxy.utils import asyncio_utils class KeepServing: - _watch_task: asyncio.Task | None = None - def load(self, loader): loader.add_option( "keepserving", @@ -45,6 +43,8 @@ class KeepServing: ctx.options.rfile, ] if any(opts) and not ctx.options.keepserving: - self._watch_task = asyncio_utils.create_task( - self.watch(), name="keepserving" + asyncio_utils.create_task( + self.watch(), + name="keepserving", + keep_ref=True, ) diff --git a/mitmproxy/addons/proxyserver.py b/mitmproxy/addons/proxyserver.py index 5b76e38d0..247b0bd45 100644 --- a/mitmproxy/addons/proxyserver.py +++ b/mitmproxy/addons/proxyserver.py @@ -120,14 +120,11 @@ class Proxyserver(ServerManager): is_running: bool _connect_addr: Address | None = None - _update_task: asyncio.Task | None = None - _inject_tasks: set[asyncio.Task] def __init__(self): self.connections = {} self.servers = Servers(self) self.is_running = False - self._inject_tasks = set() def __repr__(self): return f"Proxyserver({len(self.connections)} active conns)" @@ -289,8 +286,10 @@ class Proxyserver(ServerManager): ) if self.is_running: - self._update_task = asyncio_utils.create_task( - self.servers.update(modes), name="update servers" + asyncio_utils.create_task( + self.servers.update(modes), + name="update servers", + keep_ref=True, ) async def setup_servers(self) -> bool: @@ -315,14 +314,12 @@ class Proxyserver(ServerManager): if connection_id not in self.connections: raise ValueError("Flow is not from a live connection.") - t = asyncio_utils.create_task( + asyncio_utils.create_task( self.connections[connection_id].server_event(event), name=f"inject_event", + keep_ref=True, client=event.flow.client_conn.peername, ) - # Python 3.11 Use TaskGroup instead. - self._inject_tasks.add(t) - t.add_done_callback(self._inject_tasks.remove) @command.command("inject.websocket") def inject_websocket( diff --git a/mitmproxy/addons/readfile.py b/mitmproxy/addons/readfile.py index 1a702a91f..b3248cfbb 100644 --- a/mitmproxy/addons/readfile.py +++ b/mitmproxy/addons/readfile.py @@ -10,6 +10,7 @@ from mitmproxy import ctx from mitmproxy import exceptions from mitmproxy import flowfilter from mitmproxy import io +from mitmproxy.utils import asyncio_utils logger = logging.getLogger(__name__) @@ -74,7 +75,11 @@ class ReadFile: def running(self): if ctx.options.rfile: - self._read_task = asyncio.create_task(self.doread(ctx.options.rfile)) + self._read_task = asyncio_utils.create_task( + self.doread(ctx.options.rfile), + name="readfile", + keep_ref=False, + ) @command.command("readfile.reading") def reading(self) -> bool: diff --git a/mitmproxy/addons/script.py b/mitmproxy/addons/script.py index 69a7bc805..bc5bc0e43 100644 --- a/mitmproxy/addons/script.py +++ b/mitmproxy/addons/script.py @@ -92,6 +92,7 @@ class Script: self.reloadtask = asyncio_utils.create_task( self.watcher(), name=f"script watcher for {path}", + keep_ref=False, ) else: self.loadscript() diff --git a/mitmproxy/master.py b/mitmproxy/master.py index 7ed6fde5c..c5d118c44 100644 --- a/mitmproxy/master.py +++ b/mitmproxy/master.py @@ -65,8 +65,12 @@ class Master: # This may block for some proxy modes, so we also monitor should_exit. await asyncio.wait( [ - asyncio.create_task(ps.setup_servers()), - asyncio.create_task(self.should_exit.wait()), + asyncio_utils.create_task( + ps.setup_servers(), name="setup_servers", keep_ref=False + ), + asyncio_utils.create_task( + self.should_exit.wait(), name="should_exit", keep_ref=False + ), ], return_when=asyncio.FIRST_COMPLETED, ) diff --git a/mitmproxy/proxy/server.py b/mitmproxy/proxy/server.py index 1071138d7..9a516bbfb 100644 --- a/mitmproxy/proxy/server.py +++ b/mitmproxy/proxy/server.py @@ -103,14 +103,12 @@ class ConnectionHandler(metaclass=abc.ABCMeta): max_conns: collections.defaultdict[Address, asyncio.Semaphore] layer: "layer.Layer" wakeup_timer: set[asyncio.Task] - hook_tasks: set[asyncio.Task] def __init__(self, context: Context) -> None: self.client = context.client self.transports = {} self.max_conns = collections.defaultdict(lambda: asyncio.Semaphore(5)) self.wakeup_timer = set() - self.hook_tasks = set() # Ask for the first layer right away. # In a reverse proxy scenario, this is necessary as we would otherwise hang @@ -135,6 +133,7 @@ class ConnectionHandler(metaclass=abc.ABCMeta): watch = asyncio_utils.create_task( self.timeout_watchdog.watch(), name="timeout watchdog", + keep_ref=False, client=self.client.peername, ) @@ -150,6 +149,7 @@ class ConnectionHandler(metaclass=abc.ABCMeta): handler = asyncio_utils.create_task( self.handle_connection(self.client), name=f"client connection handler", + keep_ref=False, client=self.client.peername, ) self.transports[self.client].handler = handler @@ -394,6 +394,7 @@ class ConnectionHandler(metaclass=abc.ABCMeta): handler = asyncio_utils.create_task( self.open_connection(command), name=f"server connection handler {command.connection.address}", + keep_ref=False, client=self.client.peername, ) self.transports[command.connection] = ConnectionIO( @@ -403,6 +404,7 @@ class ConnectionHandler(metaclass=abc.ABCMeta): task = asyncio_utils.create_task( self.wakeup(command), name=f"wakeup timer ({command.delay:.1f}s)", + keep_ref=False, client=self.client.peername, ) assert task is not None @@ -422,14 +424,12 @@ class ConnectionHandler(metaclass=abc.ABCMeta): elif isinstance(command, commands.CloseConnection): self.close_connection(command.connection, False) elif isinstance(command, commands.StartHook): - t = asyncio_utils.create_task( + asyncio_utils.create_task( self.hook_task(command), name=f"handle_hook({command.name})", + keep_ref=True, client=self.client.peername, ) - # Python 3.11 Use TaskGroup instead. - self.hook_tasks.add(t) - t.add_done_callback(self.hook_tasks.remove) elif isinstance(command, commands.Log): self.log(command.message, command.level) else: diff --git a/mitmproxy/tools/web/app.py b/mitmproxy/tools/web/app.py index 46e2c8d46..386be6446 100644 --- a/mitmproxy/tools/web/app.py +++ b/mitmproxy/tools/web/app.py @@ -1,6 +1,5 @@ from __future__ import annotations -import asyncio import hashlib import json import logging @@ -35,6 +34,7 @@ from mitmproxy.tcp import TCPFlow from mitmproxy.tcp import TCPMessage from mitmproxy.udp import UDPFlow from mitmproxy.udp import UDPMessage +from mitmproxy.utils import asyncio_utils from mitmproxy.utils.emoji import emoji from mitmproxy.utils.strutils import always_str from mitmproxy.websocket import WebSocketMessage @@ -290,7 +290,6 @@ class FilterHelp(RequestHandler): class WebSocketEventBroadcaster(tornado.websocket.WebSocketHandler): # raise an error if inherited class doesn't specify its own instance. connections: ClassVar[set[WebSocketEventBroadcaster]] - _send_tasks: ClassVar[set[asyncio.Task]] = set() def open(self, *args, **kwargs): self.connections.add(self) @@ -306,9 +305,11 @@ class WebSocketEventBroadcaster(tornado.websocket.WebSocketHandler): except tornado.websocket.WebSocketClosedError: cls.connections.discard(conn) - t = asyncio.create_task(wrapper()) - cls._send_tasks.add(t) - t.add_done_callback(cls._send_tasks.remove) + asyncio_utils.create_task( + wrapper(), + name="WebSocketEventBroadcaster", + keep_ref=True, + ) @classmethod def broadcast(cls, **kwargs): diff --git a/mitmproxy/utils/asyncio_utils.py b/mitmproxy/utils/asyncio_utils.py index c545a13d6..dffb7c292 100644 --- a/mitmproxy/utils/asyncio_utils.py +++ b/mitmproxy/utils/asyncio_utils.py @@ -8,18 +8,30 @@ from contextlib import contextmanager from mitmproxy.utils import human +_KEEP_ALIVE = set() + def create_task( coro: Coroutine, *, name: str, + keep_ref: bool, client: tuple | None = None, ) -> asyncio.Task: """ - Like asyncio.create_task, but also store some debug info on the task object. + Wrapper around `asyncio.create_task`. + + - Use `keep_ref` to keep an internal reference. + This ensures that the task is not garbage collected mid-execution if no other reference is kept. + - Use `client` to pass the client address as additional debug info on the task. """ - t = asyncio.create_task(coro) + t = asyncio.create_task(coro) # noqa: TID251 set_task_debug_info(t, name=name, client=client) + if keep_ref and not t.done(): + # The event loop only keeps weak references to tasks. + # A task that isn’t referenced elsewhere may get garbage collected at any time, even before it’s done. + _KEEP_ALIVE.add(t) + t.add_done_callback(_KEEP_ALIVE.discard) return t diff --git a/pyproject.toml b/pyproject.toml index 707a06746..99b97d152 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -251,9 +251,16 @@ ignore_errors = true extend-exclude = ["mitmproxy/contrib/"] [tool.ruff.lint] -select = ["E", "F", "I"] +select = ["E", "F", "I", "TID251"] ignore = ["F541", "E501"] +[tool.ruff.lint.per-file-ignores] +"examples/**" = ["TID251"] +"test/**" = ["TID251"] + +[tool.ruff.lint.flake8-tidy-imports.banned-api] +"asyncio.create_task".msg = "Use mitmproxy.utils.asyncio_utils.create_task instead to avoid GC footgun." + [tool.ruff.lint.isort] # these rules are a bit weird, but they mimic our existing reorder_python_imports style. diff --git a/release/selftest.py b/release/selftest.py index 4ad7acb4f..97fa819e4 100644 --- a/release/selftest.py +++ b/release/selftest.py @@ -11,6 +11,7 @@ import sys from pathlib import Path from mitmproxy import ctx +from mitmproxy.utils import asyncio_utils def load(_): @@ -23,8 +24,11 @@ def load(_): def running(): - # attach is somewhere so that it's not collected. - ctx.task = asyncio.create_task(make_request()) # type: ignore + asyncio_utils.create_task( + make_request(), + name="selftest", + keep_ref=True, + ) async def make_request(): diff --git a/test/mitmproxy/utils/test_asyncio_utils.py b/test/mitmproxy/utils/test_asyncio_utils.py index 05891221b..a45135794 100644 --- a/test/mitmproxy/utils/test_asyncio_utils.py +++ b/test/mitmproxy/utils/test_asyncio_utils.py @@ -8,18 +8,22 @@ from mitmproxy.utils import asyncio_utils async def ttask(): + await asyncio.sleep(0) asyncio_utils.set_current_task_debug_info(name="newname") await asyncio.sleep(999) async def test_simple(monkeypatch): monkeypatch.setenv("PYTEST_CURRENT_TEST", "test_foo") - task = asyncio_utils.create_task(ttask(), name="ttask", client=("127.0.0.1", 42313)) + task = asyncio_utils.create_task( + ttask(), name="ttask", keep_ref=True, client=("127.0.0.1", 42313) + ) assert ( asyncio_utils.task_repr(task) == "127.0.0.1:42313: ttask [created in test_foo] (age: 0s)" ) await asyncio.sleep(0) + await asyncio.sleep(0) assert "newname" in asyncio_utils.task_repr(task) delattr(task, "created") assert asyncio_utils.task_repr(task)