Respect graceful_timeout and fix HTTP/2 shutdown cancellation

src/hypercorn/asyncio/run.py

@@ -163,18 +163,17 @@ async def _server_callback(reader: asyncio.StreamReader, writer: asyncio.StreamW
 
         for server in servers:
             server.close()
-            await server.wait_closed()
 
         try:
-            gathered_server_tasks = asyncio.gather(*server_tasks)
-            await asyncio.wait_for(gathered_server_tasks, config.graceful_timeout)
-        except asyncio.TimeoutError:
-            pass
+            if server_tasks:
+                _, pending = await asyncio.wait(server_tasks, timeout=config.graceful_timeout)
+                for task in pending:
+                    task.cancel()
+                if pending:
+                    # Bounded, so an app that swallows CancelledError
+                    # cannot keep the worker alive.
+                    await asyncio.wait(pending, timeout=config.graceful_timeout)
         finally:
-            # Retrieve the Gathered Tasks Cancelled Exception, to
-            # prevent a warning that this hasn't been done.
-            gathered_server_tasks.exception()
-
             await lifespan.wait_for_shutdown()
             lifespan_task.cancel()
             await lifespan_task

src/hypercorn/asyncio/tcp_server.py

@@ -82,7 +82,7 @@ async def protocol_send(self, event: Event) -> None:
                 try:
                     self.writer.write(event.data)
                     await self.writer.drain()
-                except (ConnectionError, RuntimeError):
+                except (OSError, RuntimeError):
                     await self.protocol.handle(Closed())
         elif isinstance(event, Closed):
             await self._close()
@@ -118,13 +118,7 @@ async def _close(self) -> None:
         try:
             self.writer.close()
             await self.writer.wait_closed()
-        except (
-            BrokenPipeError,
-            ConnectionAbortedError,
-            ConnectionResetError,
-            RuntimeError,
-            asyncio.CancelledError,
-        ):
+        except (OSError, RuntimeError, asyncio.CancelledError):
             pass  # Already closed
         finally:
             await self.idle_task.stop()

src/hypercorn/asyncio/worker_context.py

@@ -15,21 +15,15 @@ async def restart(self, task_group: TaskGroup, action: Callable) -> None:
         async with self._lock:
             if self._handle is not None:
                 self._handle.cancel()
-                try:
-                    await self._handle
-                except asyncio.CancelledError:
-                    pass
+                await asyncio.wait({self._handle})
 
             self._handle = task_group._task_group.create_task(action())  # type: ignore
 
     async def stop(self) -> None:
         async with self._lock:
             if self._handle is not None:
                 self._handle.cancel()
-                try:
-                    await self._handle
-                except asyncio.CancelledError:
-                    pass
+                await asyncio.wait({self._handle})
 
             self._handle = None
 

src/hypercorn/protocol/h2.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Awaitable, Callable
+from contextlib import suppress
 
 import h2
 import h2.connection
@@ -175,9 +176,11 @@ async def _send_data(self, stream_id: int) -> None:
         except (h2.exceptions.StreamClosedError, KeyError, h2.exceptions.ProtocolError):
             # Stream or connection has closed whilst waiting to send
             # data, not a problem - just force close it.
-            await self.stream_buffers[stream_id].close()
-            del self.stream_buffers[stream_id]
-            self.priority.remove_stream(stream_id)
+            buffer = self.stream_buffers.pop(stream_id, None)
+            if buffer is not None:
+                await buffer.close()
+            with suppress(priority.MissingStreamError):
+                self.priority.remove_stream(stream_id)
 
     async def handle(self, event: Event) -> None:
         if isinstance(event, RawData):
@@ -228,6 +231,8 @@ async def stream_send(self, event: StreamEvent) -> None:
                 if idle and self.context.terminated.is_set():
                     self.connection.close_connection()
                     await self._flush()
+                    await self.send(Closed())
+                    return
                 await self.send(Updated(idle=idle))
             elif isinstance(event, Request):
                 await self._create_server_push(event.stream_id, event.raw_path, event.headers)
@@ -256,9 +261,12 @@ async def _handle_events(self, events: list[h2.events.Event]) -> None:
                 if self.keep_alive_requests > self.config.keep_alive_max_requests:
                     self.connection.close_connection()
             elif isinstance(event, h2.events.DataReceived):
-                await self.streams[event.stream_id].handle(
-                    Body(stream_id=event.stream_id, data=event.data)
-                )
+                stream = self.streams.get(event.stream_id)
+                if stream is not None:
+                    # The stream is gone when the request was refused
+                    # as it arrived (terminated) or the response was
+                    # sent before the full request was received.
+                    await stream.handle(Body(stream_id=event.stream_id, data=event.data))
                 self.connection.acknowledge_received_data(
                     event.flow_controlled_length, event.stream_id
                 )

tests/asyncio/test_run.py

@@ -0,0 +1,248 @@
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+import h2.config
+import h2.connection
+import h2.errors
+import h2.events
+import pytest
+
+from hypercorn.app_wrappers import ASGIWrapper
+from hypercorn.asyncio.run import worker_serve
+from hypercorn.config import Config
+from hypercorn.typing import ASGIReceiveCallable, ASGISendCallable, Scope
+
+
+async def _lifespan(receive: ASGIReceiveCallable, send: ASGISendCallable) -> None:
+    while True:
+        message = await receive()
+        if message["type"] == "lifespan.startup":
+            await send({"type": "lifespan.startup.complete"})
+        elif message["type"] == "lifespan.shutdown":
+            await send({"type": "lifespan.shutdown.complete"})
+            return
+
+
+async def slow_framework(
+    scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallable
+) -> None:
+    if scope["type"] == "lifespan":
+        await _lifespan(receive, send)
+        return
+    await asyncio.sleep(100)
+
+
+async def streaming_framework(
+    scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallable
+) -> None:
+    if scope["type"] == "lifespan":
+        await _lifespan(receive, send)
+        return
+    # Bidirectional stream: respond immediately, then keep both directions open.
+    await send({"type": "http.response.start", "status": 200, "headers": []})
+    await send({"type": "http.response.body", "body": b"start", "more_body": True})
+    while True:
+        message = await receive()
+        if message["type"] == "http.disconnect":
+            return
+
+
+@pytest.mark.asyncio
+async def test_graceful_timeout_bounds_active_connections() -> None:
+    config = Config()
+    config.bind = ["127.0.0.1:0"]
+    config.graceful_timeout = 0.5
+    sockets = config.create_sockets()
+    port = sockets.insecure_sockets[0].getsockname()[1]
+
+    shutdown = asyncio.Event()
+    serving = asyncio.ensure_future(
+        worker_serve(
+            ASGIWrapper(slow_framework), config, sockets=sockets, shutdown_trigger=shutdown.wait
+        )
+    )
+    await asyncio.sleep(0.1)
+
+    reader, writer = await asyncio.open_connection("127.0.0.1", port)
+    writer.write(b"GET / HTTP/1.1\r\nHost: hypercorn\r\n\r\n")
+    await writer.drain()
+    await asyncio.sleep(0.1)
+
+    shutdown.set()
+    try:
+        # The request never completes, so the drain must be bounded by
+        # graceful_timeout rather than waiting for the connection to close.
+        await asyncio.wait_for(serving, timeout=5)
+    finally:
+        writer.close()
+
+
+@pytest.mark.asyncio
+async def test_graceful_timeout_cancels_active_h2_stream(
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    config = Config()
+    config.bind = ["127.0.0.1:0"]
+    config.graceful_timeout = 0.5
+    sockets = config.create_sockets()
+    port = sockets.insecure_sockets[0].getsockname()[1]
+
+    shutdown = asyncio.Event()
+    serving = asyncio.ensure_future(
+        worker_serve(
+            ASGIWrapper(streaming_framework),
+            config,
+            sockets=sockets,
+            shutdown_trigger=shutdown.wait,
+        )
+    )
+    await asyncio.sleep(0.1)
+
+    reader, writer = await asyncio.open_connection("127.0.0.1", port)
+    client = h2.connection.H2Connection(h2.config.H2Configuration(client_side=True))
+    client.initiate_connection()
+    client.send_headers(
+        1,
+        [(":method", "POST"), (":path", "/"), (":scheme", "http"), (":authority", "hypercorn")],
+    )
+    client.send_data(1, b"start")  # No end_stream, the request stays open
+    writer.write(client.data_to_send())
+    await writer.drain()
+    await asyncio.sleep(0.1)
+    client.receive_data(await reader.read(65536))
+    writer.write(client.data_to_send())
+    await writer.drain()
+
+    shutdown.set()
+    try:
+        # The stream never completes, so it is cancelled when graceful_timeout
+        # expires. This must not raise out of the worker.
+        await asyncio.wait_for(serving, timeout=5)
+
+        data = b""
+        while chunk := await asyncio.wait_for(reader.read(65536), timeout=1):
+            data += chunk
+        events = client.receive_data(data)
+        assert isinstance(events[-1], h2.events.ConnectionTerminated)
+        assert events[-1].error_code == h2.errors.ErrorCodes.NO_ERROR
+        assert "Unhandled exception" not in caplog.text
+    finally:
+        writer.close()
+
+
+class CancellationSwallowingFramework:
+    # Some frameworks catch the cancellation and carry on, e.g. restate_sdk's
+    # ServerInvocationContext.enter() does `except asyncio.CancelledError: pass`
+    # and then keeps awaiting the request. The server must still exit.
+    # `release` lets the test unwind the stuck request afterwards; otherwise a
+    # failing run also hangs the event loop teardown.
+
+    def __init__(self) -> None:
+        self.release = asyncio.Event()
+
+    async def __call__(
+        self, scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallable
+    ) -> None:
+        if scope["type"] == "lifespan":
+            await _lifespan(receive, send)
+            return
+        try:
+            await asyncio.sleep(100)
+        except asyncio.CancelledError:
+            pass
+        await self.release.wait()
+
+
+H2_REQUEST_HEADERS = [
+    (":method", "POST"),
+    (":path", "/"),
+    (":scheme", "http"),
+    (":authority", "hypercorn"),
+]
+
+
+async def _serve(
+    framework: Any, graceful_timeout: float
+) -> tuple[asyncio.Future, asyncio.Event, int]:
+    config = Config()
+    config.bind = ["127.0.0.1:0"]
+    config.graceful_timeout = graceful_timeout
+    sockets = config.create_sockets()
+    port = sockets.insecure_sockets[0].getsockname()[1]
+    shutdown = asyncio.Event()
+    serving = asyncio.ensure_future(
+        worker_serve(
+            ASGIWrapper(framework), config, sockets=sockets, shutdown_trigger=shutdown.wait
+        )
+    )
+    await asyncio.sleep(0.1)
+    return serving, shutdown, port
+
+
+async def _open_h2_stream(
+    port: int,
+) -> tuple[asyncio.StreamReader, asyncio.StreamWriter, h2.connection.H2Connection]:
+    reader, writer = await asyncio.open_connection("127.0.0.1", port)
+    client = h2.connection.H2Connection(h2.config.H2Configuration(client_side=True))
+    client.initiate_connection()
+    client.send_headers(1, H2_REQUEST_HEADERS)
+    client.send_data(1, b"start")  # No end_stream, the request stays open
+    writer.write(client.data_to_send())
+    await writer.drain()
+    await asyncio.sleep(0.1)
+    client.receive_data(await reader.read(65536))
+    writer.write(client.data_to_send())
+    await writer.drain()
+    return reader, writer, client
+
+
+@pytest.mark.asyncio
+async def test_graceful_timeout_bounds_app_that_swallows_cancellation() -> None:
+    framework = CancellationSwallowingFramework()
+    serving, shutdown, port = await _serve(framework, 0.5)
+    reader, writer, client = await _open_h2_stream(port)
+
+    shutdown.set()
+    try:
+        # Cancelling the connection tasks is not enough on its own: the app
+        # ignores the first cancellation, so waiting for the cancellation to
+        # complete hangs. The drain must be bounded regardless of the app.
+        # asyncio.wait rather than wait_for: wait_for would cancel `serving`
+        # and then hang on that cancellation for the same reason.
+        done, _ = await asyncio.wait({serving}, timeout=5)
+        assert serving in done, "worker_serve did not exit after graceful_timeout"
+        serving.result()
+    finally:
+        framework.release.set()
+        writer.close()
+        await asyncio.wait({serving}, timeout=5)
+
+
+@pytest.mark.asyncio
+async def test_request_during_drain_does_not_crash_worker(
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    serving, shutdown, port = await _serve(streaming_framework, 2)
+    reader, writer, client = await _open_h2_stream(port)
+
+    shutdown.set()
+    await asyncio.sleep(0.1)
+    # A client that has not seen a GOAWAY yet starts a new request. HEADERS
+    # and DATA arrive in one read. The terminated worker refuses the stream
+    # with RST_STREAM and must ignore the DATA that follows it. It must not
+    # let the connection's exception escape worker_serve: a non-zero worker
+    # exit makes the master shut down every other worker.
+    client.send_headers(3, H2_REQUEST_HEADERS)
+    client.send_data(3, b"body")
+    writer.write(client.data_to_send())
+    await writer.drain()
+
+    try:
+        done, _ = await asyncio.wait({serving}, timeout=5)
+        assert serving in done, "worker_serve did not exit after graceful_timeout"
+        serving.result()
+        assert "Unhandled exception" not in caplog.text
+    finally:
+        writer.close()