Repository: aaugustin/websockets (python)

Issue:
Stop listening for new connections for graceful stopping of server
Hello, 

We're busy moving from tornado to websockets, and it looks like we can do almost everything, but much simpler.  Great!

There is one thing I can't seem to get going: make the server stop listening for new connections, but keep current connections open until the last client exists, then finally exit altogether.  I can keep track of current clients, but I can't seem to be able to prevent new clients from connecting.  

So my idea for super graceful shutdown is:
 - stop accepting new connections
 - keep handling existing connections from clients, until they exit
 - if there are no more clients, shutdown and exit

I've experimented with the server's `close()`, which seems to disconnect existing clients, and `await wait_closed()`, which, in my case does nothing, but I suppose it is meant to be used after a normal server `close()`.  I can't find anything like tornado's `stop()`. 

Any ideas how I could accomplish this?

Thanks! 


Pull request description:
Add option keep connections open when closing server.
Fix #1174.

Patch:
diff --git a/docs/faq/server.rst b/docs/faq/server.rst
index 68490d7..22a9d3a 100644
--- a/docs/faq/server.rst
+++ b/docs/faq/server.rst
@@ -260,6 +260,21 @@ Here's an example that terminates cleanly when it receives SIGTERM on Unix:
 .. literalinclude:: ../../example/shutdown_server.py
     :emphasize-lines: 12-15,18
 
+How do I stop a server while keeping existing connections open?
+---------------------------------------------------------------
+
+Call the server's :meth:`~server.WebSocketServer.close` method with
+``close_connections=False``.
+
+Here's how to adapt the example just above::
+
+    async def server():
+        ...
+
+        server = await websockets.serve(echo, "localhost", 8765)
+        await stop
+        await server.close(close_connections=False)
+
 How do I run HTTP and WebSocket servers on the same port?
 ---------------------------------------------------------
 
diff --git a/docs/project/changelog.rst b/docs/project/changelog.rst
index 61136a3..7851db8 100644
--- a/docs/project/changelog.rst
+++ b/docs/project/changelog.rst
@@ -30,6 +30,11 @@ They may change at any time.
 
 *In development*
 
+New features
+............
+
+* Made it possible to close a server without closing existing connections.
+
 10.4
 ----
 
diff --git a/src/websockets/legacy/server.py b/src/websockets/legacy/server.py
index c32c856..89e5322 100644
--- a/src/websockets/legacy/server.py
+++ b/src/websockets/legacy/server.py
@@ -730,26 +730,30 @@ class WebSocketServer:
         """
         self.websockets.remove(protocol)
 
-    def close(self) -> None:
+    def close(self, close_connections: bool = True) -> None:
         """
         Close the server.
 
-        This method:
+        * Close the underlying :class:`~asyncio.Server`.
+        * When ``close_connections`` is :obj:`True`, which is the default,
+          close existing connections. Specifically:
 
-        * closes the underlying :class:`~asyncio.Server`;
-        * rejects new WebSocket connections with an HTTP 503 (service
-          unavailable) error; this happens when the server accepted the TCP
-          connection but didn't complete the WebSocket opening handshake prior
-          to closing;
-        * closes open WebSocket connections with close code 1001 (going away).
+          * Reject opening WebSocket connections with an HTTP 503 (service
+            unavailable) error. This happens when the server accepted the TCP
+            connection but didn't complete the opening handshake before closing.
+          * Close open WebSocket connections with close code 1001 (going away).
+
+        * Wait until all connection handlers terminate.
 
         :meth:`close` is idempotent.
 
         """
         if self.close_task is None:
-            self.close_task = self.get_loop().create_task(self._close())
+            self.close_task = self.get_loop().create_task(
+                self._close(close_connections)
+            )
 
-    async def _close(self) -> None:
+    async def _close(self, close_connections: bool) -> None:
         """
         Implementation of :meth:`close`.
 
@@ -770,21 +774,22 @@ class WebSocketServer:
         # register(). See https://bugs.python.org/issue34852 for details.
         await asyncio.sleep(0, **loop_if_py_lt_38(self.get_loop()))
 
-        # Close OPEN connections with status code 1001. Since the server was
-        # closed, handshake() closes OPENING connections with a HTTP 503
-        # error. Wait until all connections are closed.
-
-        close_tasks = [
-            asyncio.create_task(websocket.close(1001))
-            for websocket in self.websockets
-            if websocket.state is not State.CONNECTING
-        ]
-        # asyncio.wait doesn't accept an empty first argument.
-        if close_tasks:
-            await asyncio.wait(
-                close_tasks,
-                **loop_if_py_lt_38(self.get_loop()),
-            )
+        if close_connections:
+            # Close OPEN connections with status code 1001. Since the server was
+            # closed, handshake() closes OPENING connections with a HTTP 503
+            # error. Wait until all connections are closed.
+
+            close_tasks = [
+                asyncio.create_task(websocket.close(1001))
+                for websocket in self.websockets
+                if websocket.state is not State.CONNECTING
+            ]
+            # asyncio.wait doesn't accept an empty first argument.
+            if close_tasks:
+                await asyncio.wait(
+                    close_tasks,
+                    **loop_if_py_lt_38(self.get_loop()),
+                )
 
         # Wait until all connection handlers are complete.
 
@@ -903,18 +908,22 @@ class Serve:
     server performs the closing handshake and closes the connection.
 
     Awaiting :func:`serve` yields a :class:`WebSocketServer`. This object
-    provides :meth:`~WebSocketServer.close` and
-    :meth:`~WebSocketServer.wait_closed` methods for shutting down the server.
+    provides a :meth:`~WebSocketServer.close` method to shut down the server::
+
+        stop = asyncio.Future()  # set this future to exit the server
+
+        server = await serve(...)
+        await stop
+        await server.close()
 
-    :func:`serve` can be used as an asynchronous context manager::
+    :func:`serve` can be used as an asynchronous context manager. Then, the
+    server is shut down automatically when exiting the context::
 
         stop = asyncio.Future()  # set this future to exit the server
 
         async with serve(...):
             await stop
 
-    The server is shut down automatically when exiting the context.
-
     Args:
         ws_handler: connection handler. It receives the WebSocket connection,
             which is a :class:`WebSocketServerProtocol`, in argument.


Test patch:
diff --git a/tests/legacy/test_client_server.py b/tests/legacy/test_client_server.py
index d02daed..9db15c0 100644
--- a/tests/legacy/test_client_server.py
+++ b/tests/legacy/test_client_server.py
@@ -1214,14 +1214,27 @@ class CommonClientServerTests:
             server_ws = next(iter(self.server.websockets))
             self.server.close()
             with self.assertRaises(ConnectionClosed):
+                self.loop.run_until_complete(self.client.send("Hello!"))
                 self.loop.run_until_complete(self.client.recv())
 
-        # Websocket connection closes properly with 1001 Going Away.
+        # Server closed the connection with 1001 Going Away.
         self.assertEqual(self.client.close_code, 1001)
         self.assertEqual(server_ws.close_code, 1001)
 
     @with_server()
-    def test_server_shuts_down_waits_until_handlers_terminate(self):
+    def test_server_shuts_down_gracefully_during_connection_handling(self):
+        with self.temp_client():
+            server_ws = next(iter(self.server.websockets))
+            self.server.close(close_connections=False)
+            self.loop.run_until_complete(self.client.send("Hello!"))
+            self.loop.run_until_complete(self.client.recv())
+
+        # Client closed the connection with 1000 OK.
+        self.assertEqual(self.client.close_code, 1000)
+        self.assertEqual(server_ws.close_code, 1000)
+
+    @with_server()
+    def test_server_shuts_down_and_waits_until_handlers_terminate(self):
         # This handler waits a bit after the connection is closed in order
         # to test that wait_closed() really waits for handlers to complete.
         self.start_client("/slow_stop")