fix use of asyncio.create_task (#7443)

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)

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

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(

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:

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

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

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:

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

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
 
 

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.

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

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)