From a670b3b67a83731b3ae3f3e36e7b1ef8c802c415 Mon Sep 17 00:00:00 2001
From: Joe Cheng <joe@posit.co>
Date: Wed, 1 Nov 2023 14:38:52 -0700
Subject: [PATCH] Don't wait_closed() on servers until connections are closed

The asyncio.Server.wait_closed() method was a no-op in versions of Python
earlier than 3.12. The intention of this method was to block until all
existing connections are closed, but due to a bug in its implementation,
it wouldn't actually wait. This bug was fixed in Python 3.12, which
exposed uvicorn's dependence on the buggy behavior: the implementation of
uvicorn.Server.shutdown() called wait_closed() before asking existing
connections to close. As a result, attempting to stop a Uvicorn server
with an open connection would result in a blocked process, with further
attempts to Ctrl+C having no effect.
---
 tests/protocols/test_websocket.py | 24 ++++++++++++++++++++++++
 uvicorn/server.py                 | 16 ++++++++++++++--
 2 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/tests/protocols/test_websocket.py b/tests/protocols/test_websocket.py
index 17f2a92d1..9f98e3c7b 100644
--- a/tests/protocols/test_websocket.py
+++ b/tests/protocols/test_websocket.py
@@ -109,6 +109,30 @@ async def open_connection(url):
         assert is_open
 
 
+@pytest.mark.anyio
+@pytest.mark.parametrize("http_protocol_cls", HTTP_PROTOCOLS)
+async def test_shutdown(
+    ws_protocol_cls: "typing.Type[WSProtocol | WebSocketProtocol]",
+    http_protocol_cls,
+    unused_tcp_port: int,
+):
+    class App(WebSocketResponse):
+        async def websocket_connect(self, message):
+            await self.send({"type": "websocket.accept"})
+
+    config = Config(
+        app=App,
+        ws=ws_protocol_cls,
+        http=http_protocol_cls,
+        lifespan="off",
+        port=unused_tcp_port,
+    )
+    async with run_server(config) as server:
+        async with websockets.client.connect(f"ws://127.0.0.1:{unused_tcp_port}"):
+            # Attempt shutdown while connection is still open
+            await server.shutdown()
+
+
 @pytest.mark.anyio
 @pytest.mark.parametrize("http_protocol_cls", HTTP_PROTOCOLS)
 async def test_supports_permessage_deflate_extension(
diff --git a/uvicorn/server.py b/uvicorn/server.py
index 3e0db9d01..39075176a 100644
--- a/uvicorn/server.py
+++ b/uvicorn/server.py
@@ -268,8 +268,6 @@ async def shutdown(self, sockets: Optional[List[socket.socket]] = None) -> None:
             server.close()
         for sock in sockets or []:
             sock.close()
-        for server in self.servers:
-            await server.wait_closed()
 
         # Request shutdown on all existing connections.
         for connection in list(self.server_state.connections):
@@ -312,6 +310,20 @@ async def _wait_tasks_to_complete(self) -> None:
             while self.server_state.tasks and not self.force_exit:
                 await asyncio.sleep(0.1)
 
+        # Wait for servers to close. They won't do so until all connections are
+        # closed, which we've already waited for above, so this should be quick.
+        servers_closed = asyncio.gather(
+            *[server.wait_closed() for server in self.servers]
+        )
+        # Give the servers_closed future a chance to complete so we don't
+        # spuriously log about this operation.
+        await asyncio.sleep(0.1)
+        if not servers_closed.done() and not self.force_exit:
+            msg = "Waiting for servers to close. (CTRL+C to force quit)"
+            logger.info(msg)
+            while not servers_closed.done() and not self.force_exit:
+                await asyncio.sleep(0.1)
+
     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.