Repository: aaugustin/websockets (python)

Issue:
handling closed connections in ws_handler
I recently noticed that if a connection closes abnormally (and additionally is one that `websockets` is aware of), then `ws_handler` on the server side just continues running normally. This can create a problem by leaving protocol objects on the server-side running for longer than they need to be.

I sort of expected in this case that the server would send `CancelledError` in to `ws_handler` to notify the handler that there is no longer a connection to be handled.

I'm wondering then if it would make sense for `websockets` to send `CancelledError` in this case. If not, what is the right way for a `ws_handler` to "subscribe" to or be notified of the connection closing? It would be much better for `ws_handler` e.g. to be able to register a callback on a future rather than having to, say, check that `protocol.open` is true at regular intervals.

Here is example code to illustrate the issue:

```python
def log(msg):
    print(f'****: {msg}')

async def fail_connection(protocol):
    await asyncio.sleep(5)
    log('failing connection')
    await protocol.fail_connection()

async def handle_connection(protocol, path):
    asyncio.ensure_future(fail_connection(protocol))
    try:
        while True:
            await asyncio.sleep(1)
            log('still handling connection')
    finally:
        log('exiting handle_connection()')

async def run_client():
    async with websockets.connect('ws://localhost:8765'):
        print('client connected')
        await asyncio.sleep(10)

    print('client exiting')

async def run():
    async with websockets.serve(handle_connection, 'localhost', 8765):
        await run_client()
```

And here are the logs:

    client connected
    INFO:asyncio:poll 998.616 ms took 1003.063 ms: timeout
    ****: still handling connection
    INFO:asyncio:poll 999.874 ms took 1004.412 ms: timeout
    ****: still handling connection
    INFO:asyncio:poll 999.875 ms took 1003.993 ms: timeout
    ****: still handling connection
    INFO:asyncio:poll 999.880 ms took 1004.832 ms: timeout
    ****: still handling connection
    ****: failing connection
    DEBUG:websockets.protocol:server ! failing WebSocket connection: 1011 
    DEBUG:websockets.protocol:server > Frame(fin=True, opcode=8, data=b'\x03\xf3', rsv1=False, rsv2=False, rsv3=False)
    DEBUG:asyncio:poll 15.719 ms took 0.015 ms: 1 events
    DEBUG:websockets.protocol:client < Frame(fin=True, opcode=8, data=b'\x03\xf3', rsv1=False, rsv2=False, rsv3=False)
    DEBUG:websockets.protocol:client > Frame(fin=True, opcode=8, data=b'\x03\xf3', rsv1=False, rsv2=False, rsv3=False)
    DEBUG:websockets.protocol:server < Frame(fin=True, opcode=8, data=b'\x03\xf3', rsv1=False, rsv2=False, rsv3=False)
    DEBUG:websockets.protocol:server x half-closing TCP connection
    DEBUG:asyncio:poll 14.151 ms took 0.010 ms: 1 events
    DEBUG:asyncio:<_SelectorSocketTransport fd=12 read=polling write=<idle, bufsize=0>> received EOF
    DEBUG:websockets.protocol:client - eof_received()
    DEBUG:websockets.protocol:client - connection_lost(None)
    DEBUG:asyncio:<_SelectorSocketTransport fd=13 read=polling write=<idle, bufsize=0>> received EOF
    DEBUG:websockets.protocol:server - eof_received()
    DEBUG:websockets.protocol:client x closing TCP connection
    DEBUG:websockets.protocol:server - connection_lost(None)
    DEBUG:websockets.protocol:server x closing TCP connection
    ****: still handling connection
    INFO:asyncio:poll 999.879 ms took 1002.503 ms: timeout
    ****: still handling connection
    INFO:asyncio:poll 999.879 ms took 1004.720 ms: timeout
    ****: still handling connection
    INFO:asyncio:poll 999.879 ms took 1000.126 ms: timeout
    ****: still handling connection
    INFO:asyncio:poll 999.880 ms took 1003.233 ms: timeout
    ****: still handling connection
    client exiting
    ****: exiting handle_connection()
    DEBUG:websockets.protocol:server ! failing WebSocket connection: 1001 


Pull request description:
Simplify connection termination on server shutdown.
Previously, there were two scenarios where a connection handler on the
server side would receive an unexpected connection termination:

1. When the remote endpoint closed the connection or it dropped for any
   reason: this was signalled with a ConnectionClosed exception when
   interacting with the connection, usually thrown by recv() since most
   connection handlers spend the majority of their time awaiting recv()

2. When the server shut down: this caused the connection handler to be
   cancelled, meaning that a CancelledError exception was injected at
   an arbitrary location.

Now, when the server shuts down, the scenario 1 also applies. This
removes the need to prepare for CancelledError in connection handlers.

A good way to understand this change is to consider that "shutting down
the server" means "shutting down all connections".

Given that WebSocket connections are long-running, I find it sensible
to require connection handlers to detect when the connection they're
managing drops and to terminate.

I expect this to make it easier for users to write connection handlers
that behave correctly when the server shuts down, for two reasons:

- There's only one scenario to implement, which reduces the amount of
  exception handling required. In addition, this avoids the specially
  messy situation where the clean-up process for ConnectionClosed gets
  interrupted by CancelledError.
- Users are much more likely to check what happens when the connection
  drops than when the server shuts down. Making the latter behave like
  the former gives them proper behavior for free.

Another argument for this change is that cancellation wasn't properly
used there. Cancellation should only be used for terminating the
execution of a task that one started when one no longer cares about its
result. Since websockets runs a connection handler provided by the user,
it can't make the decision that whatever code comes next doesn't matter,
especially finalization code.

The counter-argument against this change is that shutdown will be much
slower for a server that reads events from an external source and sends
them over the WebSocket connection. Previously, such a server would
close immediately. Now, unless the connection handler checks for
termination explicitly, it won't notice that the server is closing
until it attempts to send the next event, which could take some time.

Fix #338.

Revert #337 and #392.

Rework #469 for consistency.

Patch:
diff --git a/docs/changelog.rst b/docs/changelog.rst
index f2f9412..f1bfd0b 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -10,26 +10,44 @@ Changelog
 
 .. warning::
 
-  **Version 7.0 renames the** ``timeout`` **argument of**
-  :func:`~server.serve()` **and** :func:`~client.connect()` **to**
-  ``close_timeout`` **.**
+    **Version 7.0 renames the** ``timeout`` **argument of**
+    :func:`~server.serve()` **and** :func:`~client.connect()` **to**
+    ``close_timeout`` **.**
 
-  This prevents confusion with ``ping_timeout``.
+    This prevents confusion with ``ping_timeout``.
 
-  For backwards compatibility, ``timeout`` is still supported.
+    For backwards compatibility, ``timeout`` is still supported.
 
 .. warning::
 
-  **Version 7.0 changes how a** :meth:`~protocol.WebSocketCommonProtocol.ping`
-  **that hasn't received a pong yet behaves when the connection is closed.**
+    **Version 7.0 changes how a server terminates connections when it's
+    closed with :meth:`~websockets.server.WebSocketServer.close`.**
 
-  The ping — as in ``ping = await websocket.ping()`` — used to be canceled
-  when the connection is closed, so that ``await ping`` raised
-  :exc:`~concurrent.futures.CancelledError`. Now ``await ping`` raises
-  :exc:`~exceptions.ConnectionClosed` like other public APIs.
+    Previously, connections handlers were canceled. Now, connections are
+    closed with close code 1001 (going away). From the perspective of the
+    connection handler, this is the same as if the remote endpoint was
+    disconnecting. This removes the need to prepare for
+    :exc:`~asyncio.CancelledError` in connection handlers.
+
+    You can restore the previous behavior by adding the following line at the
+    beginning of connection handlers::
+
+        def handler(ws, path):
+            ws.wait_closed().add_done_callback(asyncio.current_task().cancel)
+
+.. warning::
+
+    **Version 7.0 changes how a** :meth:`~protocol.WebSocketCommonProtocol.ping`
+    **that hasn't received a pong yet behaves when the connection is closed.**
+
+    The ping — as in ``ping = await websocket.ping()`` — used to be canceled
+    when the connection is closed, so that ``await ping`` raised
+    :exc:`~asyncio.CancelledError`. Now ``await ping`` raises
+    :exc:`~exceptions.ConnectionClosed` like other public APIs.
 
 * websockets sends Ping frames at regular intervals and closes the connection
-  if it doesn't receive a matching Pong frame. See :class:`~protocol.WebSocketCommonProtocol` for details.
+  if it doesn't receive a matching Pong frame. See
+  :class:`~protocol.WebSocketCommonProtocol` for details.
 
 * Added support for sending fragmented messages.
 
@@ -107,7 +125,7 @@ Also:
   credentials.
 
 * Iterating on incoming messages no longer raises an exception when the
-  connection terminates with code 1001 (going away).
+  connection terminates with close code 1001 (going away).
 
 * A plain HTTP request now receives a 426 Upgrade Required response and
   doesn't log a stack trace.
diff --git a/docs/cheatsheet.rst b/docs/cheatsheet.rst
index 2f6a47e..3b8993a 100644
--- a/docs/cheatsheet.rst
+++ b/docs/cheatsheet.rst
@@ -13,6 +13,15 @@ Server
     :meth:`~protocol.WebSocketCommonProtocol.send` to receive and send
     messages at any time.
 
+  * When :meth:`~protocol.WebSocketCommonProtocol.recv` or
+    :meth:`~protocol.WebSocketCommonProtocol.send` raises
+    :exc:`~exceptions.ConnectionClosed`, clean up and exit. If you started
+    other :class:`asyncio.Task`, terminate them before exiting.
+
+  * If you aren't awaiting :meth:`~protocol.WebSocketCommonProtocol.recv`,
+    consider awaiting :meth:`~protocol.WebSocketCommonProtocol.wait_closed`
+    to detect quickly when the connection is closed.
+
   * You may :meth:`~protocol.WebSocketCommonProtocol.ping` or
     :meth:`~protocol.WebSocketCommonProtocol.pong` if you wish but it isn't
     needed in general.
diff --git a/docs/design.rst b/docs/design.rst
index 9386973..244d233 100644
--- a/docs/design.rst
+++ b/docs/design.rst
@@ -77,7 +77,7 @@ two tasks:
 - :attr:`~protocol.WebSocketCommonProtocol.keepalive_ping_task` runs
   :meth:`~protocol.WebSocketCommonProtocol.keepalive_ping()` which sends Ping
   frames at regular intervals and ensures that corresponding Pong frames are
-  received. It is cancelled when the connection terminates. It never exits
+  received. It is canceled when the connection terminates. It never exits
   with an exception other than :exc:`~asyncio.CancelledError`.
 
 - :attr:`~protocol.WebSocketCommonProtocol.close_connection_task` runs
@@ -340,22 +340,44 @@ sending a close frame if appropriate.
 the TCP connection.
 
 
+.. _server-shutdown:
+
+Server shutdown
+---------------
+
+:class:`~websockets.server.WebSocketServer` closes asynchronously like
+:class:`asyncio.Server`. The shutdown happen in two steps:
+
+1. Stop listening and accepting new connections;
+2. Close established connections with close code 1001 (going away) or, if
+   the opening handshake is still in progress, with HTTP status code 503
+   (Service Unavailable).
+
+The first call to :class:`~websockets.server.WebSocketServer.close` starts a
+task that performs this sequence. Further calls are ignored. This is the
+easiest way to make :class:`~websockets.server.WebSocketServer.close` and
+:class:`~websockets.server.WebSocketServer.wait_closed` idempotent.
+
+
 .. _cancellation:
 
 Cancellation
 ------------
 
 Most :doc:`public APIs <api>` of ``websockets`` are coroutines. They may be
-canceled. ``websockets`` must handle this situation.
+canceled, for example if the user starts a task that calls these coroutines
+and cancels the task later. ``websockets`` must handle this situation.
 
 Cancellation during the opening handshake is handled like any other exception:
-the TCP connection is closed and the exception is re-raised or logged.
+the TCP connection is closed and the exception is re-raised. This can only
+happen on the client side. On the server side, the opening handshake is
+managed by ``websockets`` and nothing results in a cancellation.
 
 Once the WebSocket connection is established,
 :attr:`~protocol.WebSocketCommonProtocol.transfer_data_task` and
 :attr:`~protocol.WebSocketCommonProtocol.close_connection_task` mustn't get
-accidentally canceled if a coroutine that awaits them is canceled. They must
-be shielded from cancellation.
+accidentally canceled if a coroutine that awaits them is canceled. In other
+words, they must be shielded from cancellation.
 
 :meth:`~protocol.WebSocketCommonProtocol.recv()` waits for the next message in
 the queue or for :attr:`~protocol.WebSocketCommonProtocol.transfer_data_task`
@@ -388,6 +410,10 @@ Since :attr:`~protocol.WebSocketCommonProtocol.transfer_data_task` handles
 :exc:`~asyncio.CancelledError`, cancellation doesn't propagate to
 :attr:`~protocol.WebSocketCommonProtocol.close_connnection_task`.
 
+Conversely, ``websockets`` never injects :exc:`~asyncio.CancelledError` into
+user code. It doesn't cancel connection handler coroutines. Instead it expects
+them to detect when the connection is closed and to exit.
+
 
 .. _backpressure:
 
diff --git a/src/websockets/exceptions.py b/src/websockets/exceptions.py
index b34a2c0..b1618fa 100644
--- a/src/websockets/exceptions.py
+++ b/src/websockets/exceptions.py
@@ -1,5 +1,6 @@
 __all__ = [
     'AbortHandshake',
+    'CancelHandshake',
     'ConnectionClosed',
     'DuplicateParameter',
     'InvalidHandshake',
@@ -43,6 +44,13 @@ class AbortHandshake(InvalidHandshake):
         super().__init__(message)
 
 
+class CancelHandshake(InvalidHandshake):
+    """
+    Exception raised to cancel a handshake when the connection is closed.
+
+    """
+
+
 class InvalidMessage(InvalidHandshake):
     """
     Exception raised when the HTTP message in a handshake request is malformed.
diff --git a/src/websockets/framing.py b/src/websockets/framing.py
index afbc664..00a24d8 100644
--- a/src/websockets/framing.py
+++ b/src/websockets/framing.py
@@ -40,7 +40,7 @@ __all__ = [
 ]
 
 DATA_OPCODES = OP_CONT, OP_TEXT, OP_BINARY = 0x00, 0x01, 0x02
-CTRL_OPCODES = OP_CLOSE, OP_PING, OP_PONG = 0x08, 0x09, 0x0a
+CTRL_OPCODES = OP_CLOSE, OP_PING, OP_PONG = 0x08, 0x09, 0x0A
 
 # Close code that are allowed in a close frame.
 # Using a list optimizes `code in EXTERNAL_CLOSE_CODES`.
diff --git a/src/websockets/protocol.py b/src/websockets/protocol.py
index 6e15853..6147eff 100644
--- a/src/websockets/protocol.py
+++ b/src/websockets/protocol.py
@@ -76,7 +76,7 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
             await process(message)
 
     The iterator yields incoming messages. It exits normally when the
-    connection is closed with the status code 1000 (OK) or 1001 (going away).
+    connection is closed with the close code 1000 (OK) or 1001 (going away).
     It raises a :exc:`~websockets.exceptions.ConnectionClosed` exception when
     the connection is closed with any other status code.
 
@@ -352,9 +352,10 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
         """
         return self.state is State.CLOSED
 
+    @asyncio.coroutine
     def wait_closed(self):
         """
-        Return a :class:`asyncio.Future` that completes when the connection is closed.
+        Wait until the connection is closed.
 
         This is identical to :attr:`closed`, except it can be awaited.
 
@@ -362,7 +363,7 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
         of its cause, in tasks that interact with the WebSocket connection.
 
         """
-        return asyncio.shield(self.connection_lost_waiter)
+        yield from asyncio.shield(self.connection_lost_waiter)
 
     @asyncio.coroutine
     def recv(self):
@@ -941,9 +942,9 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
 
                 # ping() cannot raise ConnectionClosed, only CancelledError:
                 # - If the connection is CLOSING, keepalive_ping_task will be
-                #   cancelled by close_connection() before ping() returns.
+                #   canceled by close_connection() before ping() returns.
                 # - If the connection is CLOSED, keepalive_ping_task must be
-                #   cancelled already.
+                #   canceled already.
                 ping_waiter = yield from self.ping()
 
                 if self.ping_timeout is not None:
@@ -959,7 +960,7 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
         except asyncio.CancelledError:
             raise
 
-        except Exception as exc:
+        except Exception:
             logger.warning("Unexpected exception in keepalive ping task", exc_info=True)
 
     @asyncio.coroutine
@@ -1072,7 +1073,11 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
 
         """
         logger.debug(
-            "%s ! failing WebSocket connection: %d %s", self.side, code, reason
+            "%s ! failing WebSocket connection in the %s state: %d %s",
+            self.side,
+            self.state.name,
+            code,
+            reason or '[no reason]',
         )
 
         # Cancel transfer_data_task if the opening handshake succeeded.
@@ -1199,7 +1204,7 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
             "%s x code = %d, reason = %s",
             self.side,
             self.close_code,
-            self.close_reason or '[empty]',
+            self.close_reason or '[no reason]',
         )
         self.abort_keepalive_pings()
         # If self.connection_lost_waiter isn't pending, that's a bug, because:
diff --git a/src/websockets/server.py b/src/websockets/server.py
index 625c212..857dde5 100644
--- a/src/websockets/server.py
+++ b/src/websockets/server.py
@@ -14,13 +14,13 @@ from .compatibility import (
     BAD_REQUEST,
     FORBIDDEN,
     INTERNAL_SERVER_ERROR,
-    SERVICE_UNAVAILABLE,
     SWITCHING_PROTOCOLS,
     UPGRADE_REQUIRED,
     asyncio_ensure_future,
 )
 from .exceptions import (
     AbortHandshake,
+    CancelHandshake,
     InvalidHandshake,
     InvalidHeader,
     InvalidMessage,
@@ -32,7 +32,7 @@ from .extensions.permessage_deflate import ServerPerMessageDeflateFactory
 from .handshake import build_response, check_request
 from .headers import build_extension_list, parse_extension_list, parse_subprotocol_list
 from .http import USER_AGENT, Headers, MultipleValuesError, read_request
-from .protocol import WebSocketCommonProtocol
+from .protocol import State, WebSocketCommonProtocol
 
 
 __all__ = ['serve', 'unix_serve', 'WebSocketServerProtocol']
@@ -97,7 +97,8 @@ class WebSocketServerProtocol(WebSocketCommonProtocol):
         Handle the lifecycle of a WebSocket connection.
 
         Since this method doesn't have a caller able to handle exceptions, it
-        attemps to log relevant ones and close the connection properly.
+        attemps to log relevant ones and guarantees that the TCP connection is
+        closed before exiting.
 
         """
         try:
@@ -109,17 +110,14 @@ class WebSocketServerProtocol(WebSocketCommonProtocol):
                     available_subprotocols=self.available_subprotocols,
                     extra_headers=self.extra_headers,
                 )
-            except ConnectionError as exc:
+            except ConnectionError:
                 logger.debug("Connection error in opening handshake", exc_info=True)
                 raise
+            except CancelHandshake:
+                yield from self.fail_connection()
+                return
             except Exception as exc:
-                if self._is_server_shutting_down(exc):
-                    status, headers, body = (
-                        SERVICE_UNAVAILABLE,
-                        [],
-                        b"Server is shutting down.\n",
-                    )
-                elif isinstance(exc, AbortHandshake):
+                if isinstance(exc, AbortHandshake):
                     status, headers, body = exc.status, exc.headers, exc.body
                 elif isinstance(exc, InvalidOrigin):
                     logger.debug("Invalid origin", exc_info=True)
@@ -157,29 +155,23 @@ class WebSocketServerProtocol(WebSocketCommonProtocol):
 
                 yield from self.write_http_response(status, headers, body)
                 yield from self.fail_connection()
-
                 return
 
             try:
                 yield from self.ws_handler(self, path)
-            except Exception as exc:
-                if self._is_server_shutting_down(exc):
-                    if not self.closed:
-                        self.fail_connection(1001)
-                else:
-                    logger.error("Error in connection handler", exc_info=True)
-                    if not self.closed:
-                        self.fail_connection(1011)
+            except Exception:
+                logger.error("Error in connection handler", exc_info=True)
+                if not self.closed:
+                    self.fail_connection(1011)
                 raise
 
             try:
                 yield from self.close()
-            except ConnectionError as exc:
+            except ConnectionError:
                 logger.debug("Connection error in closing handshake", exc_info=True)
                 raise
-            except Exception as exc:
-                if not self._is_server_shutting_down(exc):
-                    logger.warning("Error in closing handshake", exc_info=True)
+            except Exception:
+                logger.warning("Error in closing handshake", exc_info=True)
                 raise
 
         except Exception:
@@ -196,13 +188,6 @@ class WebSocketServerProtocol(WebSocketCommonProtocol):
             # connections before terminating.
             self.ws_server.unregister(self)
 
-    def _is_server_shutting_down(self, exc):
-        """
-        Decide whether an exception means that the server is shutting down.
-
-        """
-        return isinstance(exc, asyncio.CancelledError) and self.ws_server.closing
-
     @asyncio.coroutine
     def read_http_request(self):
         """
@@ -467,8 +452,14 @@ class WebSocketServerProtocol(WebSocketCommonProtocol):
 
         # Hook for customizing request handling, for example checking
         # authentication or treating some paths as plain HTTP endpoints.
-
         early_response = yield from self.process_request(path, request_headers)
+
+        # Give up immediately and don't attempt to write a HTTP response if
+        # the TCP connection was closed while process_request() was running.
+        # This happens if the server shuts down and calls fail_connection().
+        if self.state != State.CONNECTING:
+            raise CancelHandshake()
+
         if early_response is not None:
             raise AbortHandshake(*early_response)
 
@@ -538,9 +529,15 @@ class WebSocketServer:
         # Store a reference to loop to avoid relying on self.server._loop.
         self.loop = loop
 
-        self.closing = False
+        # Keep track of active connections.
         self.websockets = set()
 
+        # Task responsible for closing the server and terminating connections.
+        self.close_task = None
+
+        # Completed when the server is closed and connections are terminated.
+        self.closed_waiter = asyncio.Future(loop=loop)
+
     def wrap(self, server):
         """
         Attach to a given :class:`~asyncio.Server`.
@@ -574,53 +571,65 @@ class WebSocketServer:
 
     def close(self):
         """
-        Close the underlying server, and clean up connections.
+        Close the server and terminate connections with close code 1001.
 
-        This calls :meth:`~asyncio.Server.close` on the underlying
-        :class:`~asyncio.Server` object, closes open connections with
-        status code 1001, and stops accepting new connections.
+        This method is idempotent.
 
         """
-        # Make a note that the server is shutting down. Websocket connections
-        # check this attribute to decide to send a "going away" close code.
-        self.closing = True
+        if self.close_task is None:
+            self.close_task = asyncio_ensure_future(self._close(), loop=self.loop)
 
-        # Stop accepting new connections.
-        self.server.close()
+    @asyncio.coroutine
+    def _close(self):
+        """
+        Implementation of :meth:`close`.
 
-        # Close open connections. For each connection, two tasks are running:
-        # 1. self.transfer_data_task receives incoming WebSocket messages
-        # 2. self.handler_task runs the opening handshake, the handler provided
-        #    by the user and the closing handshake
-        # In the general case, cancelling the handler task will cause the
-        # handler provided by the user to exit with a CancelledError, which
-        # will then cause the transfer data task to terminate.
-        for websocket in self.websockets:
-            websocket.handler_task.cancel()
+        This calls :meth:`~asyncio.Server.close` on the underlying
+        :class:`~asyncio.Server` object to stop accepting new connections and
+        then closes open connections with close code 1001.
 
-    @asyncio.coroutine
-    def wait_closed(self):
         """
-        Wait until the underlying server and all connections are closed.
+        # Stop accepting new connections.
+        self.server.close()
 
-        This calls :meth:`~asyncio.Server.wait_closed` on the underlying
-        :class:`~asyncio.Server` object and waits until closing handshakes
-        are complete and all connections are closed.
+        # Wait until self.server.close() completes.
+        yield from self.server.wait_closed()
 
-        This method must be called after :meth:`close()`.
+        # Wait until all accepted connections reach connection_made() and call
+        # register(). See https://bugs.python.org/issue34852 for details.
+        yield from 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:
+            websocket.fail_connection(1001)
 
-        """
         # asyncio.wait doesn't accept an empty first argument.
         if self.websockets:
-            # Either the handler or the connection can terminate first,
-            # depending on how the client behaves and the server is
-            # implemented.
+            # 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.
             yield from asyncio.wait(
                 [websocket.handler_task for websocket in self.websockets]
                 + [websocket.close_connection_task for websocket in self.websockets],
                 loop=self.loop,
             )
-        yield from self.server.wait_closed()
+
+        # Tell wait_closed() to return.
+        self.closed_waiter.set_result(None)
+
+    @asyncio.coroutine
+    def wait_closed(self):
+        """
+        Wait until the server is closed and all connections are terminated.
+
+        When :meth:`wait_closed()` returns, all TCP connections are closed and
+        there are no pending tasks left.
+
+        """
+        yield from asyncio.shield(self.closed_waiter)
 
     @property
     def sockets(self):
@@ -694,11 +703,11 @@ class Serve:
     delegates to the WebSocket handler. Once the handler completes, the server
     performs the closing handshake and closes the connection.
 
-    When a server is closed with
-    :meth:`~websockets.server.WebSocketServer.close`, all running WebSocket
-    handlers are canceled. They may intercept :exc:`~asyncio.CancelledError`
-    and perform cleanup actions before re-raising that exception. If a handler
-    started new tasks, it should cancel them as well in that case.
+    When a server is closed with :meth:`~WebSocketServer.close`, it closes all
+    connections with close code 1001 (going away). WebSocket handlers — which
+    are running the coroutine passed in the ``ws_handler`` — will receive a
+    :exc:`~websockets.exceptions.ConnectionClosed` exception on their current
+    or next interaction with the WebSocket connection.
 
     Since there's no useful way to propagate exceptions triggered in handlers,
     they're sent to the ``'websockets.server'`` logger instead. Debugging is


Test patch:
diff --git a/tests/test_client_server.py b/tests/test_client_server.py
index 55fee73..20a7837 100644
--- a/tests/test_client_server.py
+++ b/tests/test_client_server.py
@@ -64,11 +64,8 @@ def handler(ws, path):
     elif path == '/subprotocol':
         yield from ws.send(repr(ws.subprotocol))
     elif path == '/slow_stop':
-        try:
-            yield from asyncio.sleep(1000 * MS)
-        except asyncio.CancelledError:
-            yield from asyncio.sleep(MS)
-            raise
+        yield from ws.wait_closed()
+        yield from asyncio.sleep(2 * MS)
     else:
         yield from ws.send((yield from ws.recv()))
 
@@ -174,6 +171,12 @@ class HealthCheckServerProtocol(WebSocketServerProtocol):
             return OK, [('X-Access', 'OK')], b'status = green\n'
 
 
+class SlowServerProtocol(WebSocketServerProtocol):
+    @asyncio.coroutine
+    def process_request(self, path, request_headers):
+        yield from asyncio.sleep(10 * MS)
+
+
 class FooClientProtocol(WebSocketClientProtocol):
     pass
 
@@ -286,16 +289,6 @@ class ClientServerTests(unittest.TestCase):
         reply = self.loop.run_until_complete(self.client.recv())
         self.assertEqual(reply, "Hello!")
 
-    def test_server_close_while_client_connected(self):
-        with self.temp_server(loop=self.loop):
-            # This endpoint waits just a bit when the connection is canceled
-            # in order to test that wait_closed() really waits for completion.
-            self.start_client('/slow_stop')
-        with self.assertRaises(ConnectionClosed):
-            self.loop.run_until_complete(self.client.recv())
-        # Connection ends with 1001 going away.
-        self.assertEqual(self.client.close_code, 1001)
-
     def test_explicit_event_loop(self):
         with self.temp_server(loop=self.loop):
             with self.temp_client(loop=self.loop):
@@ -923,18 +916,12 @@ class ClientServerTests(unittest.TestCase):
         # The server should stop properly anyway. It used to hang because the
         # task handling the connection was waiting for the opening handshake.
 
-    @with_server()
-    @unittest.mock.patch('websockets.server.read_request')
-    def test_server_shuts_down_during_opening_handshake(self, _read_request):
-        _read_request.side_effect = asyncio.CancelledError
-
-        self.server.closing = True
-        with self.assertRaises(InvalidHandshake) as raised:
+    @with_server(create_protocol=SlowServerProtocol)
+    def test_server_shuts_down_during_opening_handshake(self):
+        self.loop.call_later(5 * MS, self.server.close)
+        with self.assertRaises(InvalidHandshake):
             self.start_client()
 
-        # Opening handshake fails with 503 Service Unavailable
-        self.assertEqual(str(raised.exception), "Status code not 101: 503")
-
     @with_server()
     def test_server_shuts_down_during_connection_handling(self):
         with self.temp_client():
@@ -959,6 +946,22 @@ class ClientServerTests(unittest.TestCase):
         # Websocket connection terminates abnormally.
         self.assertEqual(self.client.close_code, 1006)
 
+    @with_server()
+    def test_server_shuts_down_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')
+        server_ws = next(iter(self.server.websockets))
+
+        # Test that the handler task keeps running after close().
+        self.server.close()
+        self.loop.run_until_complete(asyncio.sleep(MS))
+        self.assertFalse(server_ws.handler_task.done())
+
+        # Test that the handler task terminates before wait_closed() returns.
+        self.loop.run_until_complete(self.server.wait_closed())
+        self.assertTrue(server_ws.handler_task.done())
+
     @with_server(create_protocol=ForbiddenServerProtocol)
     def test_invalid_status_error_during_client_connect(self):
         with self.assertRaises(InvalidStatusCode) as raised:
diff --git a/tests/test_framing.py b/tests/test_framing.py
index ee85157..9da64f1 100644
--- a/tests/test_framing.py
+++ b/tests/test_framing.py
@@ -121,13 +121,13 @@ class FramingTests(unittest.TestCase):
                     self.decode(encoded)
 
     def test_good_opcode(self):
-        for opcode in list(range(0x00, 0x03)) + list(range(0x08, 0x0b)):
+        for opcode in list(range(0x00, 0x03)) + list(range(0x08, 0x0B)):
             encoded = bytes([0x80 | opcode, 0])
             with self.subTest(encoded=encoded):
                 self.decode(encoded)  # does not raise an exception
 
     def test_bad_opcode(self):
-        for opcode in list(range(0x03, 0x08)) + list(range(0x0b, 0x10)):
+        for opcode in list(range(0x03, 0x08)) + list(range(0x0B, 0x10)):
             encoded = bytes([0x80 | opcode, 0])
             with self.subTest(encoded=encoded):
                 with self.assertRaises(WebSocketProtocolError):
diff --git a/tests/test_protocol.py b/tests/test_protocol.py
index d9aea19..cceb63e 100644
--- a/tests/test_protocol.py
+++ b/tests/test_protocol.py
@@ -371,9 +371,10 @@ class CommonTests:
         self.assertTrue(self.protocol.closed)
 
     def test_wait_closed(self):
-        self.assertFalse(self.protocol.wait_closed().done())
+        wait_closed = asyncio_ensure_future(self.protocol.wait_closed())
+        self.assertFalse(wait_closed.done())
         self.close_connection()
-        self.assertTrue(self.protocol.wait_closed().done())
+        self.assertTrue(wait_closed.done())
 
     # Test the recv coroutine.