Simplify pattern for shutting down servers.

This avoids having to create manually a future. Also, it's robust to
receiving multiple times the signal.

Fix #1593.

Author: aaugustin

docs/faq/client.rst

@@ -103,7 +103,7 @@ You can close the connection.
 Here's an example that terminates cleanly when it receives SIGTERM on Unix:
 
 .. literalinclude:: ../../example/faq/shutdown_client.py
-    :emphasize-lines: 11-13
+    :emphasize-lines: 10-12
 
 How do I disable TLS/SSL certificate verification?
 --------------------------------------------------

docs/faq/server.rst

@@ -310,7 +310,7 @@ Exit the :func:`~serve` context manager.
 Here's an example that terminates cleanly when it receives SIGTERM on Unix:
 
 .. literalinclude:: ../../example/faq/shutdown_server.py
-    :emphasize-lines: 13-16,19
+    :emphasize-lines: 14-16
 
 How do I stop a server while keeping existing connections open?
 ---------------------------------------------------------------

docs/howto/haproxy.rst

@@ -15,7 +15,7 @@ Run server processes
 Save this app to ``app.py``:
 
 .. literalinclude:: ../../example/deployment/haproxy/app.py
-    :emphasize-lines: 24
+    :language: python
 
 Each server process listens on a different port by extracting an incremental
 index from an environment variable set by Supervisor.

docs/howto/nginx.rst

@@ -15,7 +15,7 @@ Run server processes
 Save this app to ``app.py``:
 
 .. literalinclude:: ../../example/deployment/nginx/app.py
-    :emphasize-lines: 21,23
+    :language: python
 
 We'd like nginx to connect to websockets servers via Unix sockets in order to
 avoid the overhead of TCP for communicating between processes running in the

docs/intro/tutorial3.rst

@@ -120,14 +120,10 @@ that runs for each request. When it returns an HTTP response, websockets sends
 that response instead of opening a WebSocket connection. Here, requests to
 ``/healthz`` return an HTTP 200 status code.
 
-To catch the ``SIGTERM`` signal, ``main()`` creates a :class:`~asyncio.Future`
-called ``stop`` and registers a signal handler that sets the result of this
-future. The value of the future doesn't matter; it's only for waiting for
-``SIGTERM``.
-
-Then, by using :func:`~asyncio.server.serve` as a context manager and exiting
-the context when ``stop`` has a result, ``main()`` ensures that the server
-closes connections cleanly and exits on ``SIGTERM``.
+``main()`` registers a signal handler that closes the server when receiving the
+``SIGTERM`` signal. Then, it waits for the server to be closed. Additionally,
+using :func:`~asyncio.server.serve` as a context manager ensures that the server
+will always be closed cleanly, even if the program crashes.
 
 Deploy the WebSocket server
 ---------------------------
@@ -157,14 +153,14 @@ Commit and push your changes:
 
     $ git add .
     $ git commit -m "Deploy to Koyeb."
-    [main ac96d65] Deploy to Koyeb.
-     3 files changed, 18 insertions(+), 2 deletions(-)
+    [main 4a4b6e9] Deploy to Koyeb.
+     3 files changed, 15 insertions(+), 2 deletions(-)
      create mode 100644 Procfile
      create mode 100644 requirements.txt
     $ git push
     ...
     To github.com:python-websockets/websockets-tutorial.git
-    + 6bd6032...ac96d65 main -> main
+    + 6bd6032...4a4b6e9 main -> main
 
 Sign up or log in to Koyeb.
 
@@ -239,7 +235,7 @@ Commit your changes:
     $ git push
     ...
     To github.com:python-websockets/websockets-tutorial.git
-    + ac96d65...0903526 main -> main
+    + 4a4b6e9...968eaaa main -> main
 
 Deploy the web application
 --------------------------

docs/topics/deployment.rst

@@ -101,8 +101,7 @@ Here's an example:
     :emphasize-lines: 13-16,19
 
 When exiting the context manager, :func:`~asyncio.server.serve` closes all
-connections
-with code 1001 (going away). As a consequence:
+connections with code 1001 (going away). As a consequence:
 
 * If the connection handler is awaiting
   :meth:`~asyncio.server.ServerConnection.recv`, it receives a

example/deployment/fly/app.py

@@ -18,18 +18,10 @@ def health_check(connection, request):
 
 
 async def main():
-    # Set the stop condition when receiving SIGTERM.
-    loop = asyncio.get_running_loop()
-    stop = loop.create_future()
-    loop.add_signal_handler(signal.SIGTERM, stop.set_result, None)
-
-    async with serve(
-        echo,
-        host="",
-        port=8080,
-        process_request=health_check,
-    ):
-        await stop
+    async with serve(echo, "", 8080, process_request=health_check) as server:
+        loop = asyncio.get_running_loop()
+        loop.add_signal_handler(signal.SIGTERM, server.close)
+        await server.wait_closed()
 
 
 if __name__ == "__main__":

example/deployment/haproxy/app.py

@@ -13,17 +13,11 @@ async def echo(websocket):
 
 
 async def main():
-    # Set the stop condition when receiving SIGTERM.
-    loop = asyncio.get_running_loop()
-    stop = loop.create_future()
-    loop.add_signal_handler(signal.SIGTERM, stop.set_result, None)
-
-    async with serve(
-        echo,
-        host="localhost",
-        port=8000 + int(os.environ["SUPERVISOR_PROCESS_NAME"][-2:]),
-    ):
-        await stop
+    port = 8000 + int(os.environ["SUPERVISOR_PROCESS_NAME"][-2:])
+    async with serve(echo, "localhost", port) as server:
+        loop = asyncio.get_running_loop()
+        loop.add_signal_handler(signal.SIGTERM, server.close)
+        await server.wait_closed()
 
 
 if __name__ == "__main__":

example/deployment/heroku/app.py

@@ -13,17 +13,11 @@ async def echo(websocket):
 
 
 async def main():
-    # Set the stop condition when receiving SIGTERM.
-    loop = asyncio.get_running_loop()
-    stop = loop.create_future()
-    loop.add_signal_handler(signal.SIGTERM, stop.set_result, None)
-
-    async with serve(
-        echo,
-        host="",
-        port=int(os.environ["PORT"]),
-    ):
-        await stop
+    port = int(os.environ["PORT"])
+    async with serve(echo, "localhost", port) as server:
+        loop = asyncio.get_running_loop()
+        loop.add_signal_handler(signal.SIGTERM, server.close)
+        await server.wait_closed()
 
 
 if __name__ == "__main__":

example/deployment/koyeb/app.py

@@ -19,18 +19,11 @@ def health_check(connection, request):
 
 
 async def main():
-    # Set the stop condition when receiving SIGINT.
-    loop = asyncio.get_running_loop()
-    stop = loop.create_future()
-    loop.add_signal_handler(signal.SIGINT, stop.set_result, None)
-
-    async with serve(
-        echo,
-        host="",
-        port=int(os.environ["PORT"]),
-        process_request=health_check,
-    ):
-        await stop
+    port = int(os.environ["PORT"])
+    async with serve(echo, "", port, process_request=health_check) as server:
+        loop = asyncio.get_running_loop()
+        loop.add_signal_handler(signal.SIGINT, server.close)
+        await server.wait_closed()
 
 
 if __name__ == "__main__":

example/deployment/kubernetes/app.py

@@ -31,18 +31,10 @@ def health_check(connection, request):
 
 
 async def main():
-    # Set the stop condition when receiving SIGTERM.
-    loop = asyncio.get_running_loop()
-    stop = loop.create_future()
-    loop.add_signal_handler(signal.SIGTERM, stop.set_result, None)
-
-    async with serve(
-        slow_echo,
-        host="",
-        port=80,
-        process_request=health_check,
-    ):
-        await stop
+    async with serve(slow_echo, "", 80, process_request=health_check) as server:
+        loop = asyncio.get_running_loop()
+        loop.add_signal_handler(signal.SIGTERM, server.close)
+        await server.wait_closed()
 
 
 if __name__ == "__main__":

example/deployment/nginx/app.py

@@ -13,16 +13,11 @@ async def echo(websocket):
 
 
 async def main():
-    # Set the stop condition when receiving SIGTERM.
-    loop = asyncio.get_running_loop()
-    stop = loop.create_future()
-    loop.add_signal_handler(signal.SIGTERM, stop.set_result, None)
-
-    async with unix_serve(
-        echo,
-        path=f"{os.environ['SUPERVISOR_PROCESS_NAME']}.sock",
-    ):
-        await stop
+    path = f"{os.environ['SUPERVISOR_PROCESS_NAME']}.sock"
+    async with unix_serve(echo, path) as server:
+        loop = asyncio.get_running_loop()
+        loop.add_signal_handler(signal.SIGTERM, server.close)
+        await server.wait_closed()
 
 
 if __name__ == "__main__":

example/deployment/render/app.py

@@ -18,18 +18,10 @@ def health_check(connection, request):
 
 
 async def main():
-    # Set the stop condition when receiving SIGTERM.
-    loop = asyncio.get_running_loop()
-    stop = loop.create_future()
-    loop.add_signal_handler(signal.SIGTERM, stop.set_result, None)
-
-    async with serve(
-        echo,
-        host="",
-        port=8080,
-        process_request=health_check,
-    ):
-        await stop
+    async with serve(echo, "", 8080, process_request=health_check) as server:
+        loop = asyncio.get_running_loop()
+        loop.add_signal_handler(signal.SIGTERM, server.close)
+        await server.wait_closed()
 
 
 if __name__ == "__main__":

example/deployment/supervisor/app.py

@@ -12,18 +12,10 @@ async def echo(websocket):
 
 
 async def main():
-    # Set the stop condition when receiving SIGTERM.
-    loop = asyncio.get_running_loop()
-    stop = loop.create_future()
-    loop.add_signal_handler(signal.SIGTERM, stop.set_result, None)
-
-    async with serve(
-        echo,
-        host="",
-        port=8080,
-        reuse_port=True,
-    ):
-        await stop
+    async with serve(echo, "", 8080, reuse_port=True) as server:
+        loop = asyncio.get_running_loop()
+        loop.add_signal_handler(signal.SIGTERM, server.close)
+        await server.wait_closed()
 
 
 if __name__ == "__main__":

example/faq/shutdown_client.py

@@ -6,8 +6,7 @@
 from websockets.asyncio.client import connect
 
 async def client():
-    uri = "ws://localhost:8765"
-    async with connect(uri) as websocket:
+    async with connect("ws://localhost:8765") as websocket:
         # Close the connection when receiving SIGTERM.
         loop = asyncio.get_running_loop()
         loop.add_signal_handler(signal.SIGTERM, loop.create_task, websocket.close())

example/faq/shutdown_server.py

@@ -5,17 +5,15 @@
 
 from websockets.asyncio.server import serve
 
-async def echo(websocket):
+async def handler(websocket):
     async for message in websocket:
-        await websocket.send(message)
+        ...
 
 async def server():
-    # Set the stop condition when receiving SIGTERM.
-    loop = asyncio.get_running_loop()
-    stop = loop.create_future()
-    loop.add_signal_handler(signal.SIGTERM, stop.set_result, None)
-
-    async with serve(echo, "localhost", 8765):
-        await stop
+    async with serve(handler, "localhost", 8765) as server:
+        # Close the server when receiving SIGTERM.
+        loop = asyncio.get_running_loop()
+        loop.add_signal_handler(signal.SIGTERM, server.close)
+        await server.wait_closed()
 
 asyncio.run(server())

example/tutorial/step3/app.py

@@ -190,14 +190,11 @@ def health_check(connection, request):
 
 
 async def main():
-    # Set the stop condition when receiving SIGTERM.
-    loop = asyncio.get_running_loop()
-    stop = loop.create_future()
-    loop.add_signal_handler(signal.SIGTERM, stop.set_result, None)
-
     port = int(os.environ.get("PORT", "8001"))
-    async with serve(handler, "", port, process_request=health_check):
-        await stop
+    async with serve(handler, "", port, process_request=health_check) as server:
+        loop = asyncio.get_running_loop()
+        loop.add_signal_handler(signal.SIGTERM, server.close)
+        await server.wait_closed()
 
 
 if __name__ == "__main__":

experiments/compression/server.py

@@ -35,15 +35,6 @@ async def handler(ws):
 
 
 async def server():
-    loop = asyncio.get_running_loop()
-    stop = loop.create_future()
-
-    # Set the stop condition when receiving SIGTERM.
-    print("Stop the server with:")
-    print(f"kill -TERM {os.getpid()}")
-    print()
-    loop.add_signal_handler(signal.SIGTERM, stop.set_result, None)
-
     async with serve(
         handler,
         "localhost",
@@ -55,9 +46,15 @@ async def server():
                 compress_settings={"memLevel": ML},
             )
         ],
-    ):
+    ) as server:
+        print("Stop the server with:")
+        print(f"kill -TERM {os.getpid()}")
+        print()
+        loop = asyncio.get_running_loop()
+        loop.add_signal_handler(signal.SIGTERM, server.close)
+
         tracemalloc.start()
-        await stop
+        await server.wait_closed()
 
 
 asyncio.run(server())