Support --before-graceful-exit

docs/deployment.md

@@ -133,6 +133,8 @@ Options:
                                   buffer of an incomplete event.
   --factory                       Treat APP as an application factory, i.e. a
                                   () -> <ASGI app> callable.
+  --before-graceful-exit TEXT     A callable to be executed before the server
+                                  starts the graceful shutdown.
   --help                          Show this message and exit.
 ```
 

docs/index.md

@@ -203,6 +203,8 @@ Options:
                                   buffer of an incomplete event.
   --factory                       Treat APP as an application factory, i.e. a
                                   () -> <ASGI app> callable.
+  --before-graceful-exit TEXT     A callable to be executed before the server
+                                  starts the graceful shutdown.
   --help                          Show this message and exit.
 ```
 

docs/settings.md

@@ -121,3 +121,35 @@ To understand more about the SSL context options, please refer to the [Python do
 
 * `--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.
+
+## Hooks
+
+### Before graceful exit
+
+* `--before-graceful-exit <str>` - A callable object to be executed before the server graceful exit.
+
+```python
+from starlette.applications import Starlette
+
+app = Starlette()
+app.state.should_exit = False
+
+
+@app.websocket_route('/ws')
+async def ws(websocket):
+    await websocket.accept()
+
+    while not app.state.should_exit:
+        # How to interrupt this while loop on the shutdown event?
+        await asyncio.sleep(0.1)
+
+    await websocket.close()
+
+
+def before_graceful_exit():
+    app.state.should_exit = True
+
+
+if __name__ == "__main__":
+    uvicorn.run("example:app", before_graceful_exit_hook="example:before_graceful_exit")
+```

tests/test_config.py

@@ -589,3 +589,15 @@ def test_warn_when_using_reload_and_workers(caplog: pytest.LogCaptureFixture) ->
         '"workers" flag is ignored when reloading is enabled.'
         in caplog.records[0].message
     )
+
+
+def test_import_before_graceful_exit_hook() -> None:
+    config = Config(app=asgi_app, before_graceful_exit_hook=lambda: None)
+    config.load()
+    assert callable(config.before_graceful_exit_hook)
+
+    config = Config(
+        app=asgi_app, before_graceful_exit_hook="tests.test_config:asgi_app"
+    )
+    config.load()
+    assert callable(config.before_graceful_exit_hook)

tests/test_main.py

@@ -123,3 +123,12 @@ async def test_exit_on_create_server_with_invalid_host() -> None:
         server = Server(config=config)
         await server.serve()
     assert exc_info.value.code == 1
+
+
+@pytest.mark.anyio
+async def test_before_graceful_exit_hook() -> None:
+    config = Config(
+        app=app, before_graceful_exit_hook=lambda: print("Before graceful exit")
+    )
+    async with run_server(config):
+        pass

uvicorn/config.py

@@ -226,6 +226,10 @@ def __init__(
         headers: list[tuple[str, str]] | None = None,
         factory: bool = False,
         h11_max_incomplete_event_size: int | None = None,
+        before_graceful_exit_hook: Callable[[], Any]
+        | Callable[[], Awaitable[Any]]
+        | str
+        | None = None,
     ):
         self.app = app
         self.host = host
@@ -271,6 +275,7 @@ def __init__(
         self.encoded_headers: list[tuple[bytes, bytes]] = []
         self.factory = factory
         self.h11_max_incomplete_event_size = h11_max_incomplete_event_size
+        self.before_graceful_exit_hook = before_graceful_exit_hook
 
         self.loaded = False
         self.configure_logging()
@@ -473,6 +478,11 @@ def load(self) -> None:
                     "but please consider setting the --factory flag explicitly."
                 )
 
+        if isinstance(self.before_graceful_exit_hook, str):
+            self.before_graceful_exit_hook = import_from_string(
+                self.before_graceful_exit_hook
+            )
+
         if self.interface == "auto":
             if inspect.isclass(self.loaded_app):
                 use_asgi_3 = hasattr(self.loaded_app, "__await__")

uvicorn/main.py

@@ -6,7 +6,7 @@
 import platform
 import ssl
 import sys
-from typing import Any, Callable
+from typing import Any, Awaitable, Callable
 
 import click
 
@@ -366,6 +366,12 @@ def print_version(ctx: click.Context, param: click.Parameter, value: bool) -> No
     help="Treat APP as an application factory, i.e. a () -> <ASGI app> callable.",
     show_default=True,
 )
+@click.option(
+    "--before-graceful-exit",
+    type=str,
+    default=None,
+    help="A callable to be executed before the server starts the graceful shutdown.",
+)
 def main(
     app: str,
     host: str,
@@ -414,6 +420,7 @@ def main(
     app_dir: str,
     h11_max_incomplete_event_size: int | None,
     factory: bool,
+    before_graceful_exit: str,
 ) -> None:
     run(
         app,
@@ -463,6 +470,7 @@ def main(
         factory=factory,
         app_dir=app_dir,
         h11_max_incomplete_event_size=h11_max_incomplete_event_size,
+        before_graceful_exit_hook=before_graceful_exit,
     )
 
 
@@ -515,6 +523,10 @@ def run(
     app_dir: str | None = None,
     factory: bool = False,
     h11_max_incomplete_event_size: int | None = None,
+    before_graceful_exit_hook: Callable[[], Any]
+    | Callable[[], Awaitable[Any]]
+    | str
+    | None = None,
 ) -> None:
     if app_dir is not None:
         sys.path.insert(0, app_dir)
@@ -566,6 +578,7 @@ def run(
         use_colors=use_colors,
         factory=factory,
         h11_max_incomplete_event_size=h11_max_incomplete_event_size,
+        before_graceful_exit_hook=before_graceful_exit_hook,
     )
     server = Server(config=config)
 

uvicorn/server.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import inspect
 import logging
 import os
 import platform
@@ -270,6 +271,11 @@ async def shutdown(self, sockets: list[socket.socket] | None = None) -> None:
         for sock in sockets or []:
             sock.close()
 
+        if callable(self.config.before_graceful_exit_hook):
+            f = self.config.before_graceful_exit_hook()
+            if inspect.isawaitable(f):
+                await f
+
         # Request shutdown on all existing connections.
         for connection in list(self.server_state.connections):
             connection.shutdown()