websocket: write_message consistently raises WebSocketClosedError

diff --git a/docs/releases/v5.0.0.rst b/docs/releases/v5.0.0.rst
index 33d42b67178ac7dfac3cfe862d46bc3580e6cfad..8d9b56ba6dd7acc6b5949d14eda305dcaca3d174 100644 (file)
--- a/docs/releases/v5.0.0.rst
+++ b/docs/releases/v5.0.0.rst
@@ -309,7 +309,6 @@ Other notes
   improve performance.
 - Requests with invalid websocket headers now get a response with
   status code 400 instead of a closed connection.
-- `.WebSocketHandler.write_message` now raises `.StreamClosedError` if
-  called after the connection is closed.
-  TODO(bdarnell): fix this up with WebSocketClosedError
+- `.WebSocketHandler.write_message` now raises `.WebSocketClosedError` if
+  the connection closes while the write is in progress.
 - The ``io_loop`` argument to `.websocket_connect` has been removed.

diff --git a/tornado/test/websocket_test.py b/tornado/test/websocket_test.py
index 54734d815590d96e00364e9b8c23fcfc00b2715a..f7e8016a43a55e3bbd59e27057f9c2a41c119977 100644 (file)
--- a/tornado/test/websocket_test.py
+++ b/tornado/test/websocket_test.py
@@ -7,7 +7,6 @@ import traceback
 from tornado.concurrent import Future
 from tornado import gen
 from tornado.httpclient import HTTPError, HTTPRequest
-from tornado.iostream import StreamClosedError
 from tornado.log import gen_log, app_log
 from tornado.template import DictLoader
 from tornado.testing import AsyncHTTPTestCase, gen_test, bind_unused_port, ExpectLog
@@ -25,7 +24,7 @@ except ImportError:
     traceback.print_exc()
     raise
 
-from tornado.websocket import WebSocketHandler, websocket_connect, WebSocketError
+from tornado.websocket import WebSocketHandler, websocket_connect, WebSocketError, WebSocketClosedError
 
 try:
     from tornado import speedups
@@ -50,10 +49,11 @@ class TestWebSocketHandler(WebSocketHandler):
 
 
 class EchoHandler(TestWebSocketHandler):
+    @gen.coroutine
     def on_message(self, message):
         try:
-            self.write_message(message, isinstance(message, bytes))
-        except StreamClosedError:
+            yield self.write_message(message, isinstance(message, bytes))
+        except WebSocketClosedError:
             pass
 
 
@@ -342,7 +342,7 @@ class WebSocketTest(WebSocketBaseTestCase):
         ws = yield self.ws_connect('/close_reason')
         msg = yield ws.read_message()
         self.assertIs(msg, None)
-        with self.assertRaises(StreamClosedError):
+        with self.assertRaises(WebSocketClosedError):
             ws.write_message('hello')
 
     @gen_test

diff --git a/tornado/websocket.py b/tornado/websocket.py
index 320127801743f9a3e19cd440e654c1b43ea6f396..53832d2b1245f660c06d335e27e834d619ff91a8 100644 (file)
--- a/tornado/websocket.py
+++ b/tornado/websocket.py
@@ -237,6 +237,7 @@ class WebSocketHandler(tornado.web.RequestHandler):
         is allowed.
 
         If the connection is already closed, raises `WebSocketClosedError`.
+        Returns a `.Future` which can be used for flow control.
 
         .. versionchanged:: 3.2
            `WebSocketClosedError` was added (previously a closed connection
@@ -244,6 +245,10 @@ class WebSocketHandler(tornado.web.RequestHandler):
 
         .. versionchanged:: 4.3
            Returns a `.Future` which can be used for flow control.
+
+        .. versionchanged:: 5.0
+           Consistently raises `WebSocketClosedError`. Previously could
+           sometimes raise `.StreamClosedError`.
         """
         if self.ws_connection is None:
             raise WebSocketClosedError()
@@ -788,7 +793,23 @@ class WebSocketProtocol13(WebSocketProtocol):
         if self._compressor:
             message = self._compressor.compress(message)
             flags |= self.RSV1
-        return self._write_frame(True, opcode, message, flags=flags)
+        # For historical reasons, write methods in Tornado operate in a semi-synchronous
+        # mode in which awaiting the Future they return is optional (But errors can
+        # still be raised). This requires us to go through an awkward dance here
+        # to transform the errors that may be returned while presenting the same
+        # semi-synchronous interface.
+        try:
+            fut = self._write_frame(True, opcode, message, flags=flags)
+        except StreamClosedError:
+            raise WebSocketClosedError()
+
+        @gen.coroutine
+        def wrapper():
+            try:
+                yield fut
+            except StreamClosedError:
+                raise WebSocketClosedError()
+        return wrapper()
 
     def write_ping(self, data):
         """Send ping frame."""
@@ -1144,8 +1165,16 @@ class WebSocketClientConnection(simple_httpclient._HTTPConnection):
         future_set_result_unless_cancelled(self.connect_future, self)
 
     def write_message(self, message, binary=False):
-        """Sends a message to the WebSocket server."""
-        return self.protocol.write_message(message, binary)
+        """Sends a message to the WebSocket server.
+
+        If the stream is closed, raises `WebSocketClosedError`.
+        Returns a `.Future` which can be used for flow control.
+
+        .. versionchanged:: 5.0
+           Exception raised on a closed stream changed from `.StreamClosedError`
+           to `WebSocketClosedError`.
+        """
+        return self.protocol.write_message(message, binary=binary)
 
     def read_message(self, callback=None):
         """Reads a message from the WebSocket server.