Repository: encode/uvicorn

Issue:
How do I cancel a streaming response on shutdown?
The code below simulates a server sent event with uvicorn 0.9.

When run, if the HTTP request is *not* initiated the lifespan startup and shutdown events get called correctly on `^C`.

If the HTTP request *is* made the first `^C` will **not stop the responses**. The second `^C` stops the responses, but **the shutdown event is not called**.

What do I need to do to handle this situation?

```python
import asyncio
from datetime import datetime

import uvicorn

class App:

    def __call__(self, scope):
        if scope['type'] == 'lifespan':
            return self.lifespan
        elif scope['type'] == 'http':
            return self.http

    async def lifespan(self, receive, send):
        """Handle lifespan messages"""

        message = await receive()
        assert message["type"] == "lifespan.startup"
        print('startup')
        await send({"type": "lifespan.startup.complete"})

        message = await receive()
        assert message["type"] == "lifespan.shutdown"
        print('shutdown')
        await send({"type": "lifespan.shutdown.complete"})
        
    async def http(self, receive, send):
        """Handle http messages"""

        message = await receive()
        assert message['type'] == 'http.request'

        while message['more_body']:
            message = await receive()

        await send({
            'type': 'http.response.start',
            'status': 200,
            'headers': [
                (b'cache-control', b'no-cache'),
                (b'content-type', b'text/event-stream'),
                (b'connection', b'keep-alive')
            ]
        })

        while True:
            await send({
                'type': 'http.response.body',
                'body': f'data: {datetime.now().isoformat()}\n\n'.encode('utf-8'),
                'more_body': True
            })

            try:
                # Check the receive, timing out after a second to send more data
                receive_task = asyncio.create_task(receive())
                await asyncio.wait_for(receive_task, 1)
                message = receive_task.result()
                if message['type'] == 'http.disconnect':
                    print('disconnect')
                    break
            except asyncio.TimeoutError:
                print('timeout')
                receive_task.cancel()

if __name__ == '__main__':
    uvicorn.run(App(), port=9009)
```

Discussion:

Patch:
diff --git a/docs/deployment.md b/docs/deployment.md
index b65c431..65163b5 100644
--- a/docs/deployment.md
+++ b/docs/deployment.md
@@ -104,6 +104,9 @@ Options:
   --timeout-keep-alive INTEGER    Close Keep-Alive connections if no new data
                                   is received within this timeout.  [default:
                                   5]
+  --timeout-graceful-shutdown INTEGER
+                                  Maximum number of seconds to wait for
+                                  graceful shutdown.
   --ssl-keyfile TEXT              SSL key file
   --ssl-certfile TEXT             SSL certificate file
   --ssl-keyfile-password TEXT     SSL keyfile password
diff --git a/docs/index.md b/docs/index.md
index 17efc23..300a318 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -171,6 +171,9 @@ Options:
   --timeout-keep-alive INTEGER    Close Keep-Alive connections if no new data
                                   is received within this timeout.  [default:
                                   5]
+  --timeout-graceful-shutdown INTEGER
+                                  Maximum number of seconds to wait for
+                                  graceful shutdown.
   --ssl-keyfile TEXT              SSL key file
   --ssl-certfile TEXT             SSL certificate file
   --ssl-keyfile-password TEXT     SSL keyfile password
diff --git a/docs/settings.md b/docs/settings.md
index 3be8c89..fbc1a5f 100644
--- a/docs/settings.md
+++ b/docs/settings.md
@@ -115,3 +115,4 @@ connecting IPs in the `forwarded-allow-ips` configuration.
 ## Timeouts
 
 * `--timeout-keep-alive <int>` - Close Keep-Alive connections if no new data is received within this timeout. **Default:** *5*.
+* `--timeout-graceful-shutdown <int>` - Maximum number of seconds to wait for graceful shutdown. After this timeout, the server will start terminating requests. 
\ No newline at end of file
diff --git a/setup.cfg b/setup.cfg
index 2037895..fc0a641 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -2,7 +2,7 @@
 disallow_untyped_defs = True
 ignore_missing_imports = True
 follow_imports = silent
-# NOTE: If you are seing this, feel free to create a PR to cover the below files.
+# NOTE: If you are seeing this, feel free to create a PR to cover the below files.
 exclude = (?x)
   ^(
     |tests/supervisors/test_reload.py
@@ -52,3 +52,5 @@ rules =
   "sys_platform == 'darwin'": py-darwin
   "sys_version_info >= (3, 8)": py-gte-38
   "sys_version_info < (3, 8)": py-lt-38
+  "sys_version_info < (3, 9)": py-gte-39
+  "sys_version_info < (3, 9)": py-lt-39
diff --git a/uvicorn/config.py b/uvicorn/config.py
index bd041ec..55bddc0 100644
--- a/uvicorn/config.py
+++ b/uvicorn/config.py
@@ -230,6 +230,7 @@ class Config:
         backlog: int = 2048,
         timeout_keep_alive: int = 5,
         timeout_notify: int = 30,
+        timeout_graceful_shutdown: Optional[int] = None,
         callback_notify: Optional[Callable[..., Awaitable[None]]] = None,
         ssl_keyfile: Optional[str] = None,
         ssl_certfile: Optional[Union[str, os.PathLike]] = None,
@@ -272,6 +273,7 @@ class Config:
         self.backlog = backlog
         self.timeout_keep_alive = timeout_keep_alive
         self.timeout_notify = timeout_notify
+        self.timeout_graceful_shutdown = timeout_graceful_shutdown
         self.callback_notify = callback_notify
         self.ssl_keyfile = ssl_keyfile
         self.ssl_certfile = ssl_certfile
diff --git a/uvicorn/main.py b/uvicorn/main.py
index 690271d..95e62a4 100644
--- a/uvicorn/main.py
+++ b/uvicorn/main.py
@@ -273,6 +273,12 @@ def print_version(ctx: click.Context, param: click.Parameter, value: bool) -> No
     help="Close Keep-Alive connections if no new data is received within this timeout.",
     show_default=True,
 )
+@click.option(
+    "--timeout-graceful-shutdown",
+    type=int,
+    default=None,
+    help="Maximum number of seconds to wait for graceful shutdown.",
+)
 @click.option(
     "--ssl-keyfile", type=str, default=None, help="SSL key file", show_default=True
 )
@@ -387,6 +393,7 @@ def main(
     backlog: int,
     limit_max_requests: int,
     timeout_keep_alive: int,
+    timeout_graceful_shutdown: typing.Optional[int],
     ssl_keyfile: str,
     ssl_certfile: str,
     ssl_keyfile_password: str,
@@ -434,6 +441,7 @@ def main(
         backlog=backlog,
         limit_max_requests=limit_max_requests,
         timeout_keep_alive=timeout_keep_alive,
+        timeout_graceful_shutdown=timeout_graceful_shutdown,
         ssl_keyfile=ssl_keyfile,
         ssl_certfile=ssl_certfile,
         ssl_keyfile_password=ssl_keyfile_password,
@@ -486,6 +494,7 @@ def run(
     backlog: int = 2048,
     limit_max_requests: typing.Optional[int] = None,
     timeout_keep_alive: int = 5,
+    timeout_graceful_shutdown: typing.Optional[int] = None,
     ssl_keyfile: typing.Optional[str] = None,
     ssl_certfile: typing.Optional[typing.Union[str, os.PathLike]] = None,
     ssl_keyfile_password: typing.Optional[str] = None,
@@ -536,6 +545,7 @@ def run(
         backlog=backlog,
         limit_max_requests=limit_max_requests,
         timeout_keep_alive=timeout_keep_alive,
+        timeout_graceful_shutdown=timeout_graceful_shutdown,
         ssl_keyfile=ssl_keyfile,
         ssl_certfile=ssl_certfile,
         ssl_keyfile_password=ssl_keyfile_password,
diff --git a/uvicorn/server.py b/uvicorn/server.py
index 426fd7d..94add91 100644
--- a/uvicorn/server.py
+++ b/uvicorn/server.py
@@ -277,6 +277,28 @@ class Server:
             connection.shutdown()
         await asyncio.sleep(0.1)
 
+        # When 3.10 is not supported anymore, use `async with asyncio.timeout(...):`.
+        try:
+            await asyncio.wait_for(
+                self._wait_tasks_to_complete(),
+                timeout=self.config.timeout_graceful_shutdown,
+            )
+        except asyncio.TimeoutError:
+            logger.error(
+                "Cancel %s running task(s), timeout graceful shutdown exceeded",
+                len(self.server_state.tasks),
+            )
+            for t in self.server_state.tasks:
+                if sys.version_info < (3, 9):  # pragma: py-gte-39
+                    t.cancel()
+                else:  # pragma: py-lt-39
+                    t.cancel(msg="Task cancelled, timeout graceful shutdown exceeded")
+
+        # Send the lifespan shutdown event, and wait for application shutdown.
+        if not self.force_exit:
+            await self.lifespan.shutdown()
+
+    async def _wait_tasks_to_complete(self) -> None:
         # Wait for existing connections to finish sending responses.
         if self.server_state.connections and not self.force_exit:
             msg = "Waiting for connections to close. (CTRL+C to force quit)"
@@ -291,10 +313,6 @@ class Server:
             while self.server_state.tasks and not self.force_exit:
                 await asyncio.sleep(0.1)
 
-        # Send the lifespan shutdown event, and wait for application shutdown.
-        if not self.force_exit:
-            await self.lifespan.shutdown()
-
     def install_signal_handlers(self) -> None:
         if threading.current_thread() is not threading.main_thread():
             # Signals can only be listened to from the main thread.


Test patch:
diff --git a/tests/supervisors/test_signal.py b/tests/supervisors/test_signal.py
new file mode 100644
index 0000000..2288f0f
--- /dev/null
+++ b/tests/supervisors/test_signal.py
@@ -0,0 +1,111 @@
+import asyncio
+import signal
+from asyncio import Event
+
+import httpx
+import pytest
+
+from tests.utils import run_server
+from uvicorn import Server
+from uvicorn.config import Config
+
+
+@pytest.mark.anyio
+async def test_sigint_finish_req(unused_tcp_port: int):
+    """
+    1. Request is sent
+    2. Sigint is sent to uvicorn
+    3. Shutdown sequence start
+    4. Request is finished before timeout_graceful_shutdown=1
+
+    Result: Request should go through, even though the server was cancelled.
+    """
+
+    server_event = Event()
+
+    async def wait_app(scope, receive, send):
+        await send({"type": "http.response.start", "status": 200, "headers": []})
+        await send({"type": "http.response.body", "body": b"start", "more_body": True})
+        await server_event.wait()
+        await send({"type": "http.response.body", "body": b"end", "more_body": False})
+
+    config = Config(
+        app=wait_app, reload=False, port=unused_tcp_port, timeout_graceful_shutdown=1
+    )
+    server: Server
+    async with run_server(config) as server:
+        async with httpx.AsyncClient() as client:
+            Event()
+            req = asyncio.create_task(client.get(f"http://127.0.0.1:{unused_tcp_port}"))
+            await asyncio.sleep(0.1)  # ensure next tick
+            server.handle_exit(sig=signal.SIGINT, frame=None)  # exit
+            server_event.set()  # continue request
+            # ensure httpx has processed the response and result is complete
+            await req
+            assert req.result().status_code == 200
+
+
+@pytest.mark.anyio
+async def test_sigint_abort_req(unused_tcp_port: int, caplog):
+    """
+    1. Request is sent
+    2. Sigint is sent to uvicorn
+    3. Shutdown sequence start
+    4. Request is _NOT_ finished before timeout_graceful_shutdown=1
+
+    Result: Request is cancelled mid-execution, and httpx will raise a `RemoteProtocolError`
+    """
+
+    async def forever_app(scope, receive, send):
+        server_event = Event()
+        await send({"type": "http.response.start", "status": 200, "headers": []})
+        await send({"type": "http.response.body", "body": b"start", "more_body": True})
+        await server_event.wait()  # we never continue this one, so this request will time out
+        await send({"type": "http.response.body", "body": b"end", "more_body": False})
+
+    config = Config(
+        app=forever_app, reload=False, port=unused_tcp_port, timeout_graceful_shutdown=1
+    )
+    server: Server
+    async with run_server(config) as server:
+        async with httpx.AsyncClient() as client:
+            Event()
+            req = asyncio.create_task(client.get(f"http://127.0.0.1:{unused_tcp_port}"))
+            await asyncio.sleep(0.1)  # next tick
+            # trigger exit, this request should time out in ~1 sec
+            server.handle_exit(sig=signal.SIGINT, frame=None)
+            with pytest.raises(httpx.RemoteProtocolError):
+                await req
+
+        # req.result()
+    assert (
+        "Cancel 1 running task(s), timeout graceful shutdown exceeded"
+        in caplog.messages
+    )
+
+
+@pytest.mark.anyio
+async def test_sigint_deny_request_after_triggered(unused_tcp_port: int, caplog):
+    """
+    1. Server is started
+    2. Shutdown sequence start
+    3. Request is sent, but not accepted
+
+    Result: Request should fail, and not be able to be sent, since server is no longer accepting connections
+    """
+
+    async def app(scope, receive, send):
+        await send({"type": "http.response.start", "status": 200, "headers": []})
+        await asyncio.sleep(1)
+
+    config = Config(
+        app=app, reload=False, port=unused_tcp_port, timeout_graceful_shutdown=1
+    )
+    server: Server
+    async with run_server(config) as server:
+        # exit and ensure we do not accept more requests
+        server.handle_exit(sig=signal.SIGINT, frame=None)
+        await asyncio.sleep(0.1)  # next tick
+        async with httpx.AsyncClient() as client:
+            with pytest.raises(httpx.ConnectError):
+                await client.get(f"http://127.0.0.1:{unused_tcp_port}")