Rebuild serve_forever in the asyncio implementation.

README.rst

@@ -45,8 +45,8 @@ Here's an echo server with the ``asyncio`` API:
             await websocket.send(message)
 
     async def main():
-        async with serve(echo, "localhost", 8765) as server:
-            await server.serve_forever()
+        server = await serve(echo, "localhost", 8765)
+        await server.serve_forever()
 
     asyncio.run(main())
 

compliance/asyncio/server.py

@@ -23,12 +23,12 @@ async def main():
         echo,
         HOST,
         PORT,
-        server_header="websockets.sync",
+        server_header="websockets.asyncio",
         max_size=2**25,
     ) as server:
         try:
             await server.serve_forever()
-        except KeyboardInterrupt:
+        except asyncio.CancelledError:
             pass
 
 

compliance/sync/server.py

@@ -22,7 +22,7 @@ def main():
         echo,
         HOST,
         PORT,
-        server_header="websockets.asyncio",
+        server_header="websockets.sync",
         max_size=2**25,
     ) as server:
         try:

docs/intro/tutorial1.rst

@@ -196,8 +196,8 @@ Create an ``app.py`` file next to ``connect4.py`` with this content:
 
 
     async def main():
-        async with serve(handler, "", 8001) as server:
-            await server.serve_forever()
+        server = await serve(handler, "", 8001)
+        await server.serve_forever()
 
 
     if __name__ == "__main__":
@@ -218,9 +218,8 @@ arguments:
   on the same local network can connect.
 * The third argument is the port on which the server listens.
 
-Invoking :func:`~asyncio.server.serve` as an asynchronous context manager, in an
-``async with`` block, ensures that the server shuts down properly when
-terminating the program.
+:meth:`~asyncio.server.Server.serve_forever` doesn't return. When stopping the
+program with Ctrl-C, it gets cancelled and it shuts down the server properly.
 
 For each connection, the ``handler()`` coroutine runs an infinite loop that
 receives messages from the browser and prints them.

docs/project/changelog.rst

@@ -136,6 +136,9 @@ Bug fixes
 
 * Restored compatibility of the ``websockets`` CLI with Windows.
 
+* Fixed :meth:`~asyncio.server.Server.serve_forever` in the :mod:`asyncio`
+  implementation so that canceling it always closes connections gracefully.
+
 * Fixed a bug that could delay or block the client in the :mod:`threading`
   implementation on macOS when the opening handshake fails.
 

example/asyncio/echo.py

@@ -12,8 +12,8 @@ async def echo(websocket):
 
 
 async def main():
-    async with serve(echo, "localhost", 8765) as server:
-        await server.serve_forever()
+    server = await serve(echo, "localhost", 8765)
+    await server.serve_forever()
 
 
 if __name__ == "__main__":

example/asyncio/server.py

@@ -17,8 +17,8 @@ async def hello(websocket):
 
 
 async def main():
-    async with serve(hello, "localhost", 8765) as server:
-        await server.serve_forever()
+    server = await serve(hello, "localhost", 8765)
+    await server.serve_forever()
 
 
 if __name__ == "__main__":

example/django/authentication.py

@@ -22,8 +22,8 @@ async def handler(websocket):
 
 
 async def main():
-    async with serve(handler, "localhost", 8888) as server:
-        await server.serve_forever()
+    server = await serve(handler, "localhost", 8888)
+    await server.serve_forever()
 
 
 if __name__ == "__main__":

example/faq/health_check_server.py

@@ -13,7 +13,7 @@ async def echo(websocket):
         await websocket.send(message)
 
 async def main():
-    async with serve(echo, "localhost", 8765, process_request=health_check) as server:
-        await server.serve_forever()
+    server = await serve(echo, "localhost", 8765, process_request=health_check)
+    await server.serve_forever()
 
 asyncio.run(main())

example/quick/counter.py

@@ -43,8 +43,8 @@ async def counter(websocket):
         broadcast(USERS, users_event())
 
 async def main():
-    async with serve(counter, "localhost", 6789) as server:
-        await server.serve_forever()
+    server = await serve(counter, "localhost", 6789)
+    await server.serve_forever()
 
 if __name__ == "__main__":
     asyncio.run(main())

example/quick/server.py

@@ -14,8 +14,8 @@ async def hello(websocket):
     print(f">>> {greeting}")
 
 async def main():
-    async with serve(hello, "localhost", 8765) as server:
-        await server.serve_forever()
+    server = await serve(hello, "localhost", 8765)
+    await server.serve_forever()
 
 if __name__ == "__main__":
     asyncio.run(main())

example/quick/show_time.py

@@ -13,8 +13,8 @@ async def show_time(websocket):
         await asyncio.sleep(random.random() * 2 + 1)
 
 async def main():
-    async with serve(show_time, "localhost", 5678) as server:
-        await server.serve_forever()
+    server = await serve(show_time, "localhost", 5678)
+    await server.serve_forever()
 
 if __name__ == "__main__":
     asyncio.run(main())

example/tls/server.py

@@ -20,8 +20,8 @@ async def hello(websocket):
 ssl_context.load_cert_chain(localhost_pem)
 
 async def main():
-    async with serve(hello, "localhost", 8765, ssl=ssl_context) as server:
-        await server.serve_forever()
+    server = await serve(hello, "localhost", 8765, ssl=ssl_context)
+    await server.serve_forever()
 
 if __name__ == "__main__":
     asyncio.run(main())

example/tutorial/step1/app.py

@@ -57,8 +57,8 @@ async def handler(websocket):
 
 
 async def main():
-    async with serve(handler, "", 8001) as server:
-        await server.serve_forever()
+    server = await serve(handler, "", 8001)
+    await server.serve_forever()
 
 
 if __name__ == "__main__":

example/tutorial/step2/app.py

@@ -182,8 +182,8 @@ async def handler(websocket):
 
 
 async def main():
-    async with serve(handler, "", 8001) as server:
-        await server.serve_forever()
+    server = await serve(handler, "", 8001)
+    await server.serve_forever()
 
 
 if __name__ == "__main__":

experiments/authentication/app.py

@@ -171,9 +171,9 @@ async def main():
     """Start one HTTP server and four WebSocket servers."""
     # Set the stop condition when receiving SIGINT or SIGTERM.
     loop = asyncio.get_running_loop()
-    stop = loop.create_future()
-    loop.add_signal_handler(signal.SIGINT, stop.set_result, None)
-    loop.add_signal_handler(signal.SIGTERM, stop.set_result, None)
+    stop = asyncio.Event()
+    loop.add_signal_handler(signal.SIGINT, stop.set)
+    loop.add_signal_handler(signal.SIGTERM, stop.set)
 
     async with (
         serve(handler, host="", port=8000, process_request=serve_html),
@@ -183,7 +183,7 @@ async def main():
         serve(handler, host="", port=8004, process_request=basic_auth),
     ):
         print("Running on http://localhost:8000/")
-        await stop
+        await stop.wait()
         print("\rExiting")
 
 

experiments/routing.py

@@ -146,8 +146,8 @@ def format_timedelta(delta):
 
 
 async def main():
-    async with route(url_map, "localhost", 8888) as server:
-        await server.serve_forever()
+    server = await route(url_map, "localhost", 8888)
+    await server.serve_forever()
 
 
 if __name__ == "__main__":

src/websockets/asyncio/router.py

@@ -76,12 +76,11 @@ async def channel_handler(websocket, channel_id):
                 ...
             ])
 
-            # set this future to exit the server
-            stop = asyncio.get_running_loop().create_future()
+            # set this event to exit the server
+            stop = asyncio.Event()
 
             async with route(url_map, ...) as server:
-                await stop
-
+                await stop.wait()
 
         Refer to the documentation of :mod:`werkzeug.routing` for details.
 
@@ -99,6 +98,10 @@ async def channel_handler(websocket, channel_id):
         There is no need to specify ``websocket=True`` in each rule. It is added
         automatically.
 
+        Like :func:`~websockets.sync.server.serve`, :func:`route` returns a
+        :class:`~websockets.sync.server.Server` that you can also run with
+        :meth:`~websockets.sync.server.Server.serve_forever`.
+
         Args:
             url_map: Mapping of URL patterns to connection handlers.
             server_name: Name of the server as seen by clients. If :obj:`None`,

src/websockets/asyncio/server.py

@@ -227,7 +227,10 @@ class Server:
     """
     WebSocket server returned by :func:`serve`.
 
-    This class mirrors the API of :class:`asyncio.Server`.
+    This class mirrors most of the API of :class:`asyncio.Server`.
+
+    It doesn't provide ``close_clients`` or ``abort_clients``; by default,
+    :meth:`close` closes existing connections with code 1001 (going away).
 
     Args:
         handler: Connection handler. It receives the WebSocket connection,
@@ -509,7 +512,7 @@ async def start_serving(self) -> None:  # pragma: no cover
         """
         await self.server.start_serving()
 
-    async def serve_forever(self) -> None:  # pragma: no cover
+    async def serve_forever(self) -> None:
         """
         See :meth:`asyncio.Server.serve_forever`.
 
@@ -525,7 +528,29 @@ async def serve_forever(self) -> None:  # pragma: no cover
         instead of exiting a :func:`serve` context.
 
         """
-        await self.server.serve_forever()
+        # This is a copy-paste of asyncio.Server.serve_forever(), with
+        # self.server instead of self, except it calls our close() and
+        # wait_closed() when cancelled to ensure a graceful shutdown.
+        if self.server._serving_forever_fut is not None:  # type: ignore[attr-defined]
+            raise RuntimeError(
+                f"server {self.server!r} is already being awaited on serve_forever()"
+            )
+        if self.server._sockets is None:  # type: ignore[attr-defined]
+            raise RuntimeError(f"server {self.server!r} is closed")
+
+        self.server._start_serving()  # type: ignore[attr-defined]
+        self.server._serving_forever_fut = self.server._loop.create_future()  # type: ignore[attr-defined]
+
+        try:
+            await self.server._serving_forever_fut  # type: ignore[attr-defined]
+        except asyncio.CancelledError:
+            try:
+                self.close()
+                await self.wait_closed()
+            finally:
+                raise
+        finally:
+            self.server._serving_forever_fut = None  # type: ignore[attr-defined]
 
     @property
     def sockets(self) -> tuple[socket.socket, ...]:
@@ -535,15 +560,15 @@ def sockets(self) -> tuple[socket.socket, ...]:
         """
         return self.server.sockets
 
-    async def __aenter__(self) -> Self:  # pragma: no cover
+    async def __aenter__(self) -> Self:
         return self
 
     async def __aexit__(
         self,
         exc_type: type[BaseException] | None,
         exc_value: BaseException | None,
         traceback: TracebackType | None,
-    ) -> None:  # pragma: no cover
+    ) -> None:
         self.close()
         await self.wait_closed()
 
@@ -577,12 +602,18 @@ async def handler(websocket):
         async with serve(handler, host, port):
             await stop.wait()
 
-    Alternatively, call :meth:`~Server.serve_forever` to serve requests and
-    cancel it to stop the server::
+    Alternatively, call :meth:`~Server.serve_forever` to serve requests, then
+    cancel it or call :meth:`~Server.close` to stop the server::
 
         server = await serve(handler, host, port)
         await server.serve_forever()
 
+    The following pattern is functional but redundant: by the time the context
+    manager exits, :meth:`~Server.serve_forever` has already closed the server::
+
+        async with serve(handler, host, port) as server:
+            await server.serve_forever()
+
     Args:
         handler: Connection handler. It receives the WebSocket connection,
             which is a :class:`ServerConnection`, in argument.
@@ -645,8 +676,8 @@ async def handler(websocket):
             to 32 KiB. You may pass a ``(high, low)`` tuple to set the
             high-water and low-water marks.
         logger: Logger for this server.
-            It defaults to ``logging.getLogger("websockets.server")``. See the
-            :doc:`logging guide <../../topics/logging>` for details.
+            It defaults to ``logging.getLogger("websockets.server")``.
+            See the :doc:`logging guide <../../topics/logging>` for details.
         create_connection: Factory for the :class:`ServerConnection` managing
             the connection. Set it to a wrapper or a subclass to customize
             connection handling.

src/websockets/sync/server.py

@@ -603,8 +603,8 @@ def handler(websocket):
             and low-water marks. If you want to disable flow control entirely,
             you may set it to ``None``, although that's a bad idea.
         logger: Logger for this server.
-            It defaults to ``logging.getLogger("websockets.server")``. See the
-            :doc:`logging guide <../../topics/logging>` for details.
+            It defaults to ``logging.getLogger("websockets.server")``.
+            See the :doc:`logging guide <../../topics/logging>` for details.
         create_connection: Factory for the :class:`ServerConnection` managing
             the connection. Set it to a wrapper or a subclass to customize
             connection handling.

src/websockets/trio/server.py

@@ -513,8 +513,8 @@ async def handler(websocket):
             and low-water marks. If you want to disable flow control entirely,
             you may set it to ``None``, although that's a bad idea.
         logger: Logger for this server.
-            It defaults to ``logging.getLogger("websockets.server")``. See the
-            :doc:`logging guide <../../topics/logging>` for details.
+            It defaults to ``logging.getLogger("websockets.server")``.
+            See the :doc:`logging guide <../../topics/logging>` for details.
         create_connection: Factory for the :class:`ServerConnection` managing
             the connection. Set it to a wrapper or a subclass to customize
             connection handling.