Add a graceful shutdown period to allow tasks to complete.

CHANGES/7188.feature

@@ -0,0 +1 @@
+Added a graceful shutdown period which allows pending tasks to complete before the application's cleanup is called. The period can be adjusted with the ``shutdown_timeout`` parameter. -- by :user:`Dreamsorcerer`

aiohttp/http_writer.py

@@ -72,6 +72,7 @@ def _write(self, chunk: bytes) -> None:
         self.buffer_size += size
         self.output_size += size
 
+        print("FOO", self._transport.__class__, self._transport.is_closing())
         if self._transport is None or self._transport.is_closing():
             raise ConnectionResetError("Cannot write to closing transport")
         self._transport.write(chunk)

aiohttp/web_runner.py

@@ -2,6 +2,7 @@
 import signal
 import socket
 from abc import ABC, abstractmethod
+from contextlib import suppress
 from typing import Any, List, Optional, Set, Type
 
 from yarl import URL
@@ -80,11 +81,21 @@ async def stop(self) -> None:
         # named pipes do not have wait_closed property
         if hasattr(self._server, "wait_closed"):
             await self._server.wait_closed()
+
+        # Wait for pending tasks for a given time limit.
+        with suppress(asyncio.TimeoutError):
+            await asyncio.wait_for(self._wait(), timeout=self._shutdown_timeout)
+
         await self._runner.shutdown()
         assert self._runner.server
         await self._runner.server.shutdown(self._shutdown_timeout)
         self._runner._unreg_site(self)
 
+    async def _wait(self) -> None:
+        exclude = self._runner.starting_tasks | {asyncio.current_task()}
+        while tasks := asyncio.all_tasks() - exclude:
+            await asyncio.wait(tasks)
+
 
 class TCPSite(BaseSite):
     __slots__ = ("_host", "_port", "_reuse_address", "_reuse_port")
@@ -247,7 +258,7 @@ async def start(self) -> None:
 
 
 class BaseRunner(ABC):
-    __slots__ = ("_handle_signals", "_kwargs", "_server", "_sites")
+    __slots__ = ("starting_tasks", "_handle_signals", "_kwargs", "_server", "_sites")
 
     def __init__(self, *, handle_signals: bool = False, **kwargs: Any) -> None:
         self._handle_signals = handle_signals
@@ -287,6 +298,11 @@ async def setup(self) -> None:
                 pass
 
         self._server = await self._make_server()
+        # On shutdown we want to avoid waiting on tasks which run forever.
+        # It's very likely that all tasks which run forever will have been created by
+        # the time we have completed the application startup (in self._make_server()),
+        # so we just record all running tasks here and exclude them later.
+        self.starting_tasks = asyncio.all_tasks()
 
     @abstractmethod
     async def shutdown(self) -> None:

tests/test_run_app.py

@@ -9,14 +9,15 @@
 import ssl
 import subprocess
 import sys
-from typing import Any
+import time
+from typing import Any, Callable, NoReturn
 from unittest import mock
 from uuid import uuid4
 
 import pytest
 from conftest import needs_unix
 
-from aiohttp import web
+from aiohttp import ClientConnectorError, ClientSession, web
 from aiohttp.test_utils import make_mocked_coro
 from aiohttp.web_runner import BaseRunner
 
@@ -926,3 +927,201 @@ async def init():
 
     web.run_app(init(), print=stopper(patched_loop), loop=patched_loop)
     assert count == 3
+
+
+class TestShutdown:
+    def raiser(self) -> NoReturn:
+        raise KeyboardInterrupt
+
+    async def stop(self, request: web.Request) -> web.Response:
+        asyncio.get_running_loop().call_soon(self.raiser)
+        return web.Response()
+
+    def run_app(
+        self, port: int, timeout: int, task, extra_test=None, cleanup=None
+    ) -> asyncio.Task:
+        async def test() -> None:
+            await asyncio.sleep(1)
+            async with ClientSession() as sess:
+                async with sess.get(f"http://localhost:{port}/") as resp:
+                    pass
+                async with sess.get(f"http://localhost:{port}/stop") as resp:
+                    pass
+
+                if extra_test:
+                    await extra_test(sess)
+
+        async def run_test(app: web.Application) -> None:
+            nonlocal test_task
+            test_task = asyncio.create_task(test())
+            yield
+            await test_task
+
+        async def handler(request: web.Request) -> web.Response:
+            nonlocal t
+            t = asyncio.create_task(task())
+            return web.Response(text="FOO")
+
+        t = test_task = None
+        app = web.Application()
+        app.cleanup_ctx.append(run_test)
+        if cleanup:
+            app.on_shutdown.append(cleanup)
+        app.router.add_get("/", handler)
+        app.router.add_get("/stop", self.stop)
+
+        web.run_app(app, port=port, shutdown_timeout=timeout)
+        assert test_task.exception() is None
+        return t
+
+    def test_shutdown_wait_for_task(
+        self, aiohttp_unused_port: Callable[[], int]
+    ) -> None:
+        port = aiohttp_unused_port()
+        finished = False
+
+        async def task():
+            nonlocal finished
+            await asyncio.sleep(2)
+            finished = True
+
+        t = self.run_app(port, 3, task)
+
+        assert finished is True
+        assert t.done()
+        assert not t.cancelled()
+
+    def test_shutdown_timeout_task(
+        self, aiohttp_unused_port: Callable[[], int]
+    ) -> None:
+        port = aiohttp_unused_port()
+        finished = False
+
+        async def task():
+            nonlocal finished
+            await asyncio.sleep(2)
+            finished = True
+
+        t = self.run_app(port, 1, task)
+
+        assert finished is False
+        assert t.done()
+        assert t.cancelled()
+
+    def test_shutdown_wait_for_spawned_task(
+        self, aiohttp_unused_port: Callable[[], int]
+    ) -> None:
+        port = aiohttp_unused_port()
+        finished = False
+        finished_sub = False
+        sub_t = None
+
+        async def sub_task():
+            nonlocal finished_sub
+            await asyncio.sleep(1.5)
+            finished_sub = True
+
+        async def task():
+            nonlocal finished, sub_t
+            await asyncio.sleep(0.5)
+            sub_t = asyncio.create_task(sub_task())
+            finished = True
+
+        t = self.run_app(port, 3, task)
+
+        assert finished is True
+        assert t.done()
+        assert not t.cancelled()
+        assert finished_sub is True
+        assert sub_t.done()
+        assert not sub_t.cancelled()
+
+    def test_shutdown_timeout_not_reached(
+        self, aiohttp_unused_port: Callable[[], int]
+    ) -> None:
+        port = aiohttp_unused_port()
+        finished = False
+
+        async def task():
+            nonlocal finished
+            await asyncio.sleep(1)
+            finished = True
+
+        start_time = time.time()
+        t = self.run_app(port, 15, task)
+
+        assert finished is True
+        assert t.done()
+        # Verify run_app has not waited for timeout.
+        assert time.time() - start_time < 10
+
+    def test_shutdown_new_conn_rejected(
+        self, aiohttp_unused_port: Callable[[], int]
+    ) -> None:
+        port = aiohttp_unused_port()
+        finished = False
+
+        async def task() -> None:
+            nonlocal finished
+            await asyncio.sleep(2)
+            finished = True
+
+        async def test(sess: ClientSession) -> None:
+            # Ensure we are in the middle of shutdown (waiting for task()).
+            await asyncio.sleep(1)
+            with pytest.raises(ClientConnectorError):
+                # Use a new session to try and open a new connection.
+                async with ClientSession() as sess:
+                    async with sess.get(f"http://localhost:{port}/") as resp:
+                        pass
+            assert finished is False
+
+        t = self.run_app(port, 3, task, test)
+
+        assert finished is True
+        assert t.done()
+
+    def test_shutdown_pending_handler_responds(
+        self, aiohttp_unused_port: Callable[[], int]
+    ) -> None:
+        port = aiohttp_unused_port()
+        finished = False
+
+        async def test() -> None:
+            async def test_resp(sess):
+                async with sess.get(f"http://localhost:{port}/") as resp:
+                    assert await resp.text() == "FOO"
+
+            await asyncio.sleep(1)
+            async with ClientSession() as sess:
+                t = asyncio.create_task(test_resp(sess))
+                await asyncio.sleep(1)
+                # Handler is in-progress while we trigger server shutdown.
+                async with sess.get(f"http://localhost:{port}/stop") as resp:
+                    pass
+
+                assert finished == False
+                # Handler should still complete and produce a response.
+                await t
+
+        async def run_test(app: web.Application) -> None:
+            nonlocal t
+            t = asyncio.create_task(test())
+            yield
+            await t
+
+        async def handler(request: web.Request) -> web.Response:
+            nonlocal finished
+            await asyncio.sleep(3)
+            finished = True
+            return web.Response(text="FOO")
+
+        t = None
+        app = web.Application()
+        app.cleanup_ctx.append(run_test)
+        app.router.add_get("/", handler)
+        app.router.add_get("/stop", self.stop)
+
+        web.run_app(app, port=port, shutdown_timeout=5)
+        assert t.exception() is None
+        assert finished is True