diff --git a/README.rst b/README.rst index 1016c168f..2c4f1411b 100644 --- a/README.rst +++ b/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()) diff --git a/compliance/asyncio/server.py b/compliance/asyncio/server.py index 84deb9727..dbc6d909d 100644 --- a/compliance/asyncio/server.py +++ b/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 diff --git a/compliance/sync/server.py b/compliance/sync/server.py index 494f56a44..0de2f6707 100644 --- a/compliance/sync/server.py +++ b/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: diff --git a/docs/intro/tutorial1.rst b/docs/intro/tutorial1.rst index 88640e660..734617433 100644 --- a/docs/intro/tutorial1.rst +++ b/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. diff --git a/docs/project/changelog.rst b/docs/project/changelog.rst index 934f84155..40910e56b 100644 --- a/docs/project/changelog.rst +++ b/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. diff --git a/example/asyncio/echo.py b/example/asyncio/echo.py index 28d877be7..0d210fc64 100755 --- a/example/asyncio/echo.py +++ b/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__": diff --git a/example/asyncio/server.py b/example/asyncio/server.py index 574e053bf..72ac12a0f 100755 --- a/example/asyncio/server.py +++ b/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__": diff --git a/example/django/authentication.py b/example/django/authentication.py index e61d70432..ec7d21370 100644 --- a/example/django/authentication.py +++ b/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__": diff --git a/example/faq/health_check_server.py b/example/faq/health_check_server.py index 3fdffb501..9c500caf5 100755 --- a/example/faq/health_check_server.py +++ b/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()) diff --git a/example/quick/counter.py b/example/quick/counter.py index b31345ce2..130ec7bfd 100755 --- a/example/quick/counter.py +++ b/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()) diff --git a/example/quick/server.py b/example/quick/server.py index a01f91703..7e21e7d47 100755 --- a/example/quick/server.py +++ b/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()) diff --git a/example/quick/show_time.py b/example/quick/show_time.py index b56aada7b..3c62b7043 100755 --- a/example/quick/show_time.py +++ b/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()) diff --git a/example/tls/server.py b/example/tls/server.py index 92c6629b5..19805c30d 100755 --- a/example/tls/server.py +++ b/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()) diff --git a/example/tutorial/step1/app.py b/example/tutorial/step1/app.py index bc8f02484..ecf29ac84 100644 --- a/example/tutorial/step1/app.py +++ b/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__": diff --git a/example/tutorial/step2/app.py b/example/tutorial/step2/app.py index fe50fb3af..afbc94fe4 100644 --- a/example/tutorial/step2/app.py +++ b/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__": diff --git a/experiments/authentication/app.py b/experiments/authentication/app.py index 0bdd7fd2f..2aaf248ba 100644 --- a/experiments/authentication/app.py +++ b/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") diff --git a/experiments/routing.py b/experiments/routing.py index 7fc4ad4b3..2a7220587 100644 --- a/experiments/routing.py +++ b/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__": diff --git a/src/websockets/asyncio/router.py b/src/websockets/asyncio/router.py index ecf8cd678..9b85ce4e8 100644 --- a/src/websockets/asyncio/router.py +++ b/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`, diff --git a/src/websockets/asyncio/server.py b/src/websockets/asyncio/server.py index 7ccc5c042..6ce776243 100644 --- a/src/websockets/asyncio/server.py +++ b/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,7 +560,7 @@ 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__( @@ -543,7 +568,7 @@ async def __aexit__( 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. diff --git a/src/websockets/sync/server.py b/src/websockets/sync/server.py index 368050e55..0c2e34183 100644 --- a/src/websockets/sync/server.py +++ b/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. diff --git a/src/websockets/trio/server.py b/src/websockets/trio/server.py index 91f741b33..4c237e47d 100644 --- a/src/websockets/trio/server.py +++ b/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. diff --git a/tests/asyncio/test_server.py b/tests/asyncio/test_server.py index bd7cac0a0..30753a7f7 100644 --- a/tests/asyncio/test_server.py +++ b/tests/asyncio/test_server.py @@ -567,14 +567,91 @@ async def test_close_keeps_handlers_running(self): async with asyncio.timeout(5 * MS): await server.wait_closed() + async def test_serve_context_manager_closes_server(self): + """Server closes when exiting serve() used as a context manager.""" + async with serve(handler, "localhost", 0) as server: + self.assertTrue(server.is_serving()) + + self.assertFalse(server.is_serving()) + async def test_context_manager_closes_server(self): """Server closes when exiting context manager.""" - with self.assertLogs("websockets", logging.INFO) as logs: - async with serve(handler, "localhost", 0) as server: - async with connect(get_uri(server)) as client: - await self.assertEval(client, "ws.protocol.state.name", "OPEN") + server = await serve(handler, "localhost", 0) + async with server: + self.assertTrue(server.is_serving()) + + self.assertFalse(server.is_serving()) + + async def test_serve_forever(self): + """Server runs until serve_forever() is canceled.""" + server = await serve(*args) + serve_forever_task = asyncio.get_running_loop().create_task( + server.serve_forever() + ) + await asyncio.sleep(0) + + self.assertTrue(server.is_serving()) + + serve_forever_task.cancel() + with self.assertRaises(asyncio.CancelledError): + await serve_forever_task + + self.assertFalse(server.is_serving()) + + async def test_close_cancels_serve_forever(self): + """Closing the server cancels serve_forever().""" + server = await serve(*args) + serve_forever_task = asyncio.get_running_loop().create_task( + server.serve_forever() + ) + await asyncio.sleep(0) + + server.close() + with self.assertRaises(asyncio.CancelledError): + await serve_forever_task + + async def test_serve_forever_closes_open_connections(self): + """Canceling serve_forever() closes open connections with code 1001.""" + server = await serve(*args) + serve_forever_task = asyncio.get_running_loop().create_task( + server.serve_forever() + ) + await asyncio.sleep(0) + + async with connect(get_uri(server)) as client: + serve_forever_task.cancel() + with self.assertRaises(ConnectionClosedOK) as raised: + await client.recv() + self.assertEqual( + str(raised.exception), + "received 1001 (going away); then sent 1001 (going away)", + ) + + with self.assertRaises(asyncio.CancelledError): + await serve_forever_task + + async def test_serve_forever_twice(self): + """Server rejects awaiting serve_forever() concurrently.""" + server = await serve(*args) + serve_forever_task = asyncio.get_running_loop().create_task( + server.serve_forever() + ) + await asyncio.sleep(0) + + with self.assertRaises(RuntimeError): + await server.serve_forever() + + serve_forever_task.cancel() + with self.assertRaises(asyncio.CancelledError): + await serve_forever_task + + async def test_serve_forever_after_close(self): + """Server rejects serve_forever() after close().""" + async with serve(*args) as server: + pass - self.assertEqual(logs.records[-1].getMessage(), "server closed") + with self.assertRaises(RuntimeError): + await server.serve_forever() async def test_sockets(self): """Server provides a sockets property.""" diff --git a/tests/proxy.py b/tests/proxy.py index d5839ed05..68f2eec6c 100644 --- a/tests/proxy.py +++ b/tests/proxy.py @@ -83,7 +83,7 @@ class ProxyMixin: @classmethod async def run_proxy(cls): cls.proxy_loop = loop = asyncio.get_event_loop() - cls.proxy_stop = stop = loop.create_future() + cls.proxy_stop = stop = asyncio.Event() cls.proxy_options = options = Options( mode=[cls.proxy_mode], @@ -103,7 +103,7 @@ async def run_proxy(cls): ) task = loop.create_task(cls.proxy_master.run()) - await stop + await stop.wait() for server in master.addons.get("proxyserver").servers: await server.stop() @@ -145,6 +145,6 @@ def tearDown(self): @classmethod def tearDownClass(cls): - cls.proxy_loop.call_soon_threadsafe(cls.proxy_stop.set_result, None) + cls.proxy_loop.call_soon_threadsafe(cls.proxy_stop.set) cls.proxy_thread.join() super().tearDownClass() diff --git a/tests/sync/test_server.py b/tests/sync/test_server.py index d54f94a6e..44bc61d53 100644 --- a/tests/sync/test_server.py +++ b/tests/sync/test_server.py @@ -477,15 +477,13 @@ def test_shutdown_keeps_handlers_running(self): def test_context_manager_closes_server(self): """Server closes when exiting context manager.""" - with self.assertLogs("websockets", logging.INFO) as logs: - with serve(handler, "localhost", 0) as server: - thread = threading.Thread(target=server.serve_forever) - thread.start() - self.addCleanup(thread.join) - with connect(get_uri(server)) as client: - self.assertEval(client, "ws.protocol.state.name", "OPEN") + with serve(handler, "localhost", 0) as server: + thread = threading.Thread(target=server.serve_forever) + thread.start() + self.addCleanup(thread.join) + self.assertFalse(server.socket_closed.is_set()) - self.assertEqual(logs.records[-1].getMessage(), "server closed") + self.assertTrue(server.socket_closed.is_set()) def test_fileno(self): """Server provides a fileno method.""" diff --git a/tests/trio/test_server.py b/tests/trio/test_server.py index 5b0e0506a..8018175e4 100644 --- a/tests/trio/test_server.py +++ b/tests/trio/test_server.py @@ -585,13 +585,11 @@ async def close_server(): async def test_context_manager_closes_server(self): """Server closes when exiting context manager.""" - with self.assertLogs("websockets", logging.INFO) as logs: - server = await self.nursery.start(serve, handler, 0) - async with server: - async with connect(get_uri(server)) as client: - await self.assertEval(client, "ws.protocol.state.name", "OPEN") + server = await self.nursery.start(serve, handler, 0) + async with server: + self.assertFalse(server.handlers_waiter.is_set()) - self.assertEqual(logs.records[-1].getMessage(), "server closed") + self.assertTrue(server.handlers_waiter.is_set()) async def test_listeners(self): """Server provides a listeners attribute."""