Repository: tornadoweb/tornado

Issue:
websocket: Allow open() to be a coroutine
According to this PR, the websocket handler's `open` method cannot be a coroutine: https://github.com/tornadoweb/tornado/issues/1650

Is this still true?


Discussion:
Yes, it's still true.
Actually, I think you could make it a coroutine but it may not do what you expect: blocking in an `open()` coroutine doesn't block incoming messages (in the way that `on_message` coroutines can be used for flow control). 
@bdarnell Thanks for the response. Are there plans to allow `open` to block `on_message`? I have to call some async methods/functions as part of my WebSocket setup and would like to be able to do this cleanly without the need for a Queue.
I don't personally have any immediate plans to work on this. It shouldn't be too hard, but it would need some refactoring of the websocket code. 

As a quick fix, you could create a tornado.locks.Event in `initialize()`, call `set()` on it at the end of `open()` and `await self.open_event.wait()` at the beginning of `on_message()`. 
Thanks. I'll keep an eye out for this change to be coming down the line.

Resolving patch:
diff --git a/tornado/websocket.py b/tornado/websocket.py
--- a/tornado/websocket.py
+++ b/tornado/websocket.py
@@ -256,18 +256,38 @@ def write_message(self, message, binary=False):
         return self.ws_connection.write_message(message, binary=binary)
 
     def select_subprotocol(self, subprotocols):
-        """Invoked when a new WebSocket requests specific subprotocols.
+        """Override to implement subprotocol negotiation.
 
         ``subprotocols`` is a list of strings identifying the
         subprotocols proposed by the client.  This method may be
         overridden to return one of those strings to select it, or
-        ``None`` to not select a subprotocol.  Failure to select a
-        subprotocol does not automatically abort the connection,
-        although clients may close the connection if none of their
-        proposed subprotocols was selected.
+        ``None`` to not select a subprotocol.
+
+        Failure to select a subprotocol does not automatically abort
+        the connection, although clients may close the connection if
+        none of their proposed subprotocols was selected.
+
+        The list may be empty, in which case this method must return
+        None. This method is always called exactly once even if no
+        subprotocols were proposed so that the handler can be advised
+        of this fact.
+
+        .. versionchanged:: 5.1
+
+           Previously, this method was called with a list containing
+           an empty string instead of an empty list if no subprotocols
+           were proposed by the client.
         """
         return None
 
+    @property
+    def selected_subprotocol(self):
+        """The subprotocol returned by `select_subprotocol`.
+
+        .. versionadded:: 5.1
+        """
+        return self.ws_connection.selected_subprotocol
+
     def get_compression_options(self):
         """Override to return compression options for the connection.
 
@@ -298,6 +318,13 @@ def open(self, *args, **kwargs):
         The arguments to `open` are extracted from the `tornado.web.URLSpec`
         regular expression, just like the arguments to
         `tornado.web.RequestHandler.get`.
+
+        `open` may be a coroutine. `on_message` will not be called until
+        `open` has returned.
+
+        .. versionchanged:: 5.1
+
+           ``open`` may be a coroutine.
         """
         pass
 
@@ -674,13 +701,17 @@ def _challenge_response(self):
         return WebSocketProtocol13.compute_accept_value(
             self.request.headers.get("Sec-Websocket-Key"))
 
+    @gen.coroutine
     def _accept_connection(self):
-        subprotocols = [s.strip() for s in self.request.headers.get_list("Sec-WebSocket-Protocol")]
-        if subprotocols:
-            selected = self.handler.select_subprotocol(subprotocols)
-            if selected:
-                assert selected in subprotocols
-                self.handler.set_header("Sec-WebSocket-Protocol", selected)
+        subprotocol_header = self.request.headers.get("Sec-WebSocket-Protocol")
+        if subprotocol_header:
+            subprotocols = [s.strip() for s in subprotocol_header.split(',')]
+        else:
+            subprotocols = []
+        self.selected_subprotocol = self.handler.select_subprotocol(subprotocols)
+        if self.selected_subprotocol:
+            assert self.selected_subprotocol in subprotocols
+            self.handler.set_header("Sec-WebSocket-Protocol", self.selected_subprotocol)
 
         extensions = self._parse_extensions_header(self.request.headers)
         for ext in extensions:
@@ -710,9 +741,11 @@ def _accept_connection(self):
         self.stream = self.handler.stream
 
         self.start_pinging()
-        self._run_callback(self.handler.open, *self.handler.open_args,
-                           **self.handler.open_kwargs)
-        IOLoop.current().add_callback(self._receive_frame_loop)
+        open_result = self._run_callback(self.handler.open, *self.handler.open_args,
+                                         **self.handler.open_kwargs)
+        if open_result is not None:
+            yield open_result
+        yield self._receive_frame_loop()
 
     def _parse_extensions_header(self, headers):
         extensions = headers.get("Sec-WebSocket-Extensions", '')
@@ -739,6 +772,8 @@ def _process_server_headers(self, key, headers):
             else:
                 raise ValueError("unsupported extension %r", ext)
 
+        self.selected_subprotocol = headers.get('Sec-WebSocket-Protocol', None)
+
     def _get_compressor_options(self, side, agreed_parameters, compression_options=None):
         """Converts a websocket agreed_parameters set to keyword arguments
         for our compressor objects.
@@ -1056,7 +1091,7 @@ class WebSocketClientConnection(simple_httpclient._HTTPConnection):
     """
     def __init__(self, request, on_message_callback=None,
                  compression_options=None, ping_interval=None, ping_timeout=None,
-                 max_message_size=None):
+                 max_message_size=None, subprotocols=[]):
         self.compression_options = compression_options
         self.connect_future = Future()
         self.protocol = None
@@ -1077,6 +1112,8 @@ def __init__(self, request, on_message_callback=None,
             'Sec-WebSocket-Key': self.key,
             'Sec-WebSocket-Version': '13',
         })
+        if subprotocols is not None:
+            request.headers['Sec-WebSocket-Protocol'] = ','.join(subprotocols)
         if self.compression_options is not None:
             # Always offer to let the server set our max_wbits (and even though
             # we don't offer it, we will accept a client_no_context_takeover
@@ -1211,11 +1248,19 @@ def get_websocket_protocol(self):
         return WebSocketProtocol13(self, mask_outgoing=True,
                                    compression_options=self.compression_options)
 
+    @property
+    def selected_subprotocol(self):
+        """The subprotocol selected by the server.
+
+        .. versionadded:: 5.1
+        """
+        return self.protocol.selected_subprotocol
+
 
 def websocket_connect(url, callback=None, connect_timeout=None,
                       on_message_callback=None, compression_options=None,
                       ping_interval=None, ping_timeout=None,
-                      max_message_size=None):
+                      max_message_size=None, subprotocols=None):
     """Client-side websocket support.
 
     Takes a url and returns a Future whose result is a
@@ -1238,6 +1283,11 @@ def websocket_connect(url, callback=None, connect_timeout=None,
     ``websocket_connect``. In both styles, a message of ``None``
     indicates that the connection has been closed.
 
+    ``subprotocols`` may be a list of strings specifying proposed
+    subprotocols. The selected protocol may be found on the
+    ``selected_subprotocol`` attribute of the connection object
+    when the connection is complete.
+
     .. versionchanged:: 3.2
        Also accepts ``HTTPRequest`` objects in place of urls.
 
@@ -1250,6 +1300,9 @@ def websocket_connect(url, callback=None, connect_timeout=None,
 
     .. versionchanged:: 5.0
        The ``io_loop`` argument (deprecated since version 4.1) has been removed.
+
+    .. versionchanged:: 5.1
+       Added the ``subprotocols`` argument.
     """
     if isinstance(url, httpclient.HTTPRequest):
         assert connect_timeout is None
@@ -1266,7 +1319,8 @@ def websocket_connect(url, callback=None, connect_timeout=None,
                                      compression_options=compression_options,
                                      ping_interval=ping_interval,
                                      ping_timeout=ping_timeout,
-                                     max_message_size=max_message_size)
+                                     max_message_size=max_message_size,
+                                     subprotocols=subprotocols)
     if callback is not None:
         IOLoop.current().add_future(conn.connect_future, callback)
     return conn.connect_future


Test patch:
diff --git a/tornado/test/websocket_test.py b/tornado/test/websocket_test.py
--- a/tornado/test/websocket_test.py
+++ b/tornado/test/websocket_test.py
@@ -7,6 +7,7 @@
 from tornado.concurrent import Future
 from tornado import gen
 from tornado.httpclient import HTTPError, HTTPRequest
+from tornado.locks import Event
 from tornado.log import gen_log, app_log
 from tornado.simple_httpclient import SimpleAsyncHTTPClient
 from tornado.template import DictLoader
@@ -143,6 +144,43 @@ def on_message(self, message):
         self.write_message(self.render_string('message.html', message=message))
 
 
+class SubprotocolHandler(TestWebSocketHandler):
+    def initialize(self, **kwargs):
+        super(SubprotocolHandler, self).initialize(**kwargs)
+        self.select_subprotocol_called = False
+
+    def select_subprotocol(self, subprotocols):
+        if self.select_subprotocol_called:
+            raise Exception("select_subprotocol called twice")
+        self.select_subprotocol_called = True
+        if 'goodproto' in subprotocols:
+            return 'goodproto'
+        return None
+
+    def open(self):
+        if not self.select_subprotocol_called:
+            raise Exception("select_subprotocol not called")
+        self.write_message("subprotocol=%s" % self.selected_subprotocol)
+
+
+class OpenCoroutineHandler(TestWebSocketHandler):
+    def initialize(self, test, **kwargs):
+        super(OpenCoroutineHandler, self).initialize(**kwargs)
+        self.test = test
+        self.open_finished = False
+
+    @gen.coroutine
+    def open(self):
+        yield self.test.message_sent.wait()
+        yield gen.sleep(0.010)
+        self.open_finished = True
+
+    def on_message(self, message):
+        if not self.open_finished:
+            raise Exception('on_message called before open finished')
+        self.write_message('ok')
+
+
 class WebSocketBaseTestCase(AsyncHTTPTestCase):
     @gen.coroutine
     def ws_connect(self, path, **kwargs):
@@ -183,6 +221,10 @@ def get_app(self):
              dict(close_future=self.close_future)),
             ('/render', RenderMessageHandler,
              dict(close_future=self.close_future)),
+            ('/subprotocol', SubprotocolHandler,
+             dict(close_future=self.close_future)),
+            ('/open_coroutine', OpenCoroutineHandler,
+             dict(close_future=self.close_future, test=self)),
         ], template_loader=DictLoader({
             'message.html': '<b>{{ message }}</b>',
         }))
@@ -443,6 +485,32 @@ def test_check_origin_invalid_subdomains(self):
 
         self.assertEqual(cm.exception.code, 403)
 
+    @gen_test
+    def test_subprotocols(self):
+        ws = yield self.ws_connect('/subprotocol', subprotocols=['badproto', 'goodproto'])
+        self.assertEqual(ws.selected_subprotocol, 'goodproto')
+        res = yield ws.read_message()
+        self.assertEqual(res, 'subprotocol=goodproto')
+        yield self.close(ws)
+
+    @gen_test
+    def test_subprotocols_not_offered(self):
+        ws = yield self.ws_connect('/subprotocol')
+        self.assertIs(ws.selected_subprotocol, None)
+        res = yield ws.read_message()
+        self.assertEqual(res, 'subprotocol=None')
+        yield self.close(ws)
+
+    @gen_test
+    def test_open_coroutine(self):
+        self.message_sent = Event()
+        ws = yield self.ws_connect('/open_coroutine')
+        yield ws.write_message('hello')
+        self.message_sent.set()
+        res = yield ws.read_message()
+        self.assertEqual(res, 'ok')
+        yield self.close(ws)
+
 
 if sys.version_info >= (3, 5):
     NativeCoroutineOnMessageHandler = exec_test(globals(), locals(), """