Add --timeout-graceful-shutdown parameter (#1950)

Co-authored-by: Marcelo Trylesinski <[email protected]>

Author: JonasKs

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

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

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. 

setup.cfg

@@ -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

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}")

uvicorn/config.py

@@ -230,6 +230,7 @@ def __init__(
         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 @@ def __init__(
         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

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,

uvicorn/server.py

@@ -277,6 +277,28 @@ async def shutdown(self, sockets: Optional[List[socket.socket]] = None) -> None:
             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 @@ async def shutdown(self, sockets: Optional[List[socket.socket]] = None) -> None:
             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.