Repository: aaugustin/websockets (python)

Issue:
Wait for connections to terminate when closing a server
Currently `WebSocketServer._close` calls `fail_connection(1001)` on all established connections. This sends a close frame with code 1001 and stops processing data immediately.

It would be nicer to close connections properly with `close(1001)`and wait for a close frame from the remote endpoint.

This is a follow-up on #394.

Pull request description:
Close connections properly in WebSocketServer.close.
Fix #541.

Patch:
diff --git a/docs/changelog.rst b/docs/changelog.rst
index 2556d70..b99d3d0 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -68,6 +68,9 @@ Also:
 
   If you were using ``ping_timeout=None`` as a workaround, you can remove it.
 
+* Changed :meth:`~server.WebSocketServer.close` to perform a proper closing
+  handshake instead of failing the connection.
+
 * Avoided a crash of a ``extra_headers`` callable returns ``None``.
 
 * Improved error messages when HTTP parsing fails.
diff --git a/src/websockets/protocol.py b/src/websockets/protocol.py
index fdadb93..acc45e8 100644
--- a/src/websockets/protocol.py
+++ b/src/websockets/protocol.py
@@ -635,7 +635,6 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
         try:
             # If close() is canceled during the wait, self.transfer_data_task
             # is canceled before the timeout elapses.
-            # This helps closing connections when shutting down a server.
             await asyncio.wait_for(
                 self.transfer_data_task, self.close_timeout, loop=self.loop
             )
diff --git a/src/websockets/server.py b/src/websockets/server.py
index 8e1db9b..4248748 100644
--- a/src/websockets/server.py
+++ b/src/websockets/server.py
@@ -46,7 +46,7 @@ from .headers import (
     parse_subprotocol,
 )
 from .http import USER_AGENT, Headers, HeadersLike, MultipleValuesError, read_request
-from .protocol import State, WebSocketCommonProtocol
+from .protocol import WebSocketCommonProtocol
 from .typing import Origin, Subprotocol
 
 
@@ -692,26 +692,22 @@ class WebSocketServer:
         # register(). See https://bugs.python.org/issue34852 for details.
         await asyncio.sleep(0)
 
-        # Close open connections. fail_connection() will cancel the transfer
-        # data task, which is expected to cause the handler task to terminate.
-        for websocket in self.websockets:
-            if websocket.state is State.OPEN:
-                websocket.fail_connection(1001)
+        # Close OPEN connections with status code 1001. Since the server was
+        # closed, handshake() closes OPENING conections with a HTTP 503 error.
+        # Wait until all connections are closed.
+
+        # asyncio.wait doesn't accept an empty first argument
+        if self.websockets:
+            await asyncio.wait(
+                [websocket.close(1001) for websocket in self.websockets], loop=self.loop
+            )
+
+        # Wait until all connection handlers are complete.
 
         # asyncio.wait doesn't accept an empty first argument.
         if self.websockets:
-            # The connection handler can terminate before or after the
-            # connection closes. Wait until both are done to avoid leaking
-            # running tasks.
-            # TODO: it would be nicer to wait only for the connection handler
-            # and let the handler wait for the connection to close.
             await asyncio.wait(
-                [websocket.handler_task for websocket in self.websockets]
-                + [
-                    websocket.close_connection_task
-                    for websocket in self.websockets
-                    if websocket.state is State.OPEN
-                ],
+                [websocket.handler_task for websocket in self.websockets],
                 loop=self.loop,
             )
 


Test patch:
diff --git a/tests/test_client_server.py b/tests/test_client_server.py
index 7281ec6..aa4bebd 100644
--- a/tests/test_client_server.py
+++ b/tests/test_client_server.py
@@ -1172,12 +1172,14 @@ class CommonClientServerTests:
     @with_server()
     def test_server_shuts_down_during_connection_handling(self):
         with self.temp_client():
+            server_ws = next(iter(self.server.websockets))
             self.server.close()
             with self.assertRaises(ConnectionClosed):
                 self.loop.run_until_complete(self.client.recv())
 
-        # Websocket connection terminates with 1001 Going Away.
+        # Websocket connection closes properly with 1001 Going Away.
         self.assertEqual(self.client.close_code, 1001)
+        self.assertEqual(server_ws.close_code, 1001)
 
     @with_server()
     @unittest.mock.patch("websockets.server.WebSocketServerProtocol.close")