Repository: tornadoweb/tornado

Issue:
WebSocketHandler.on_message should allow coroutines
I'd like to be able to use python 3.5's async/await syntax as much as possible. However, when I implement a WebSocketHandler and use async/await in the on_message handler, tornado gives me the following error:

`
/home/jaapz/Programming/mijnbaopt-websocket/.env/lib/python3.5/site-packages/tornado/websocket.py:417: RuntimeWarning: coroutine 'on_message' was never awaited
`

But when I use `tornado.gen.coroutine` as a decorator on the `on_message` method and use the `yield from` syntax inside it, everything works as expected.

The async/await syntax does work on the `prepare` method of the WebSocketHandler, so it would be nice if it'd also work on `on_message`.



Discussion:
As I explained in https://stackoverflow.com/questions/35542864/how-to-use-python-3-5-style-async-and-await-in-tornado, the reason this doesn't work is that `on_message` is not allowed to be a coroutine by the base class. This is separate from the question of `async/await` vs `@gen.coroutine` (when you say that everything works as expected with `@gen.coroutine`, do your expectations include that a second call to `on_message` can start before the previous one finishes?). I agree that `on_message` should allow coroutines, but it needs to be done the right way (i.e. supporting flow control). Until then, you should use `IOLoop.current().spawn_callback` to start coroutines from `on_message`.

I'm pretty new to both asyncio and tornado, and I didn't know that it was not allowed by the API to even maken on_message a coroutine, as it "just worked" when I tried it. I assumed that meant it was ok to do that. I'll use the spawn_callback method for now, until tornado officially supports it. Thanks for the quick reply!

+1

+1 this would be really intuitive.


Resolving patch:
diff --git a/tornado/websocket.py b/tornado/websocket.py
--- a/tornado/websocket.py
+++ b/tornado/websocket.py
@@ -30,7 +30,7 @@
 
 from tornado.concurrent import TracebackFuture
 from tornado.escape import utf8, native_str, to_unicode
-from tornado import httpclient, httputil
+from tornado import gen, httpclient, httputil
 from tornado.ioloop import IOLoop, PeriodicCallback
 from tornado.iostream import StreamClosedError
 from tornado.log import gen_log, app_log
@@ -269,6 +269,10 @@ def on_message(self, message):
         """Handle incoming messages on the WebSocket
 
         This method must be overridden.
+
+        .. versionchanged:: 4.5
+
+           ``on_message`` can be a coroutine.
         """
         raise NotImplementedError
 
@@ -442,14 +446,21 @@ def __init__(self, handler):
     def _run_callback(self, callback, *args, **kwargs):
         """Runs the given callback with exception handling.
 
-        On error, aborts the websocket connection and returns False.
+        If the callback is a coroutine, returns its Future. On error, aborts the
+        websocket connection and returns None.
         """
         try:
-            callback(*args, **kwargs)
+            result = callback(*args, **kwargs)
         except Exception:
             app_log.error("Uncaught exception in %s",
                           getattr(self.request, 'path', None), exc_info=True)
             self._abort()
+        else:
+            if result is not None:
+                self.stream.io_loop.add_future(gen.convert_yielded(result),
+                                               lambda f: f.result())
+
+            return result
 
     def on_connection_close(self):
         self._abort()
@@ -810,6 +821,8 @@ def _on_masked_frame_data(self, data):
         self._on_frame_data(_websocket_mask(self._frame_mask, data))
 
     def _on_frame_data(self, data):
+        handled_future = None
+
         self._wire_bytes_in += len(data)
         if self._frame_opcode_is_control:
             # control frames may be interleaved with a series of fragmented
@@ -842,12 +855,18 @@ def _on_frame_data(self, data):
                 self._fragmented_message_buffer = data
 
         if self._final_frame:
-            self._handle_message(opcode, data)
+            handled_future = self._handle_message(opcode, data)
 
         if not self.client_terminated:
-            self._receive_frame()
+            if handled_future:
+                # on_message is a coroutine, process more frames once it's done.
+                gen.convert_yielded(handled_future).add_done_callback(
+                    lambda future: self._receive_frame())
+            else:
+                self._receive_frame()
 
     def _handle_message(self, opcode, data):
+        """Execute on_message, returning its Future if it is a coroutine."""
         if self.client_terminated:
             return
 
@@ -862,11 +881,11 @@ def _handle_message(self, opcode, data):
             except UnicodeDecodeError:
                 self._abort()
                 return
-            self._run_callback(self.handler.on_message, decoded)
+            return self._run_callback(self.handler.on_message, decoded)
         elif opcode == 0x2:
             # Binary data
             self._message_bytes_in += len(data)
-            self._run_callback(self.handler.on_message, data)
+            return self._run_callback(self.handler.on_message, data)
         elif opcode == 0x8:
             # Close
             self.client_terminated = True
@@ -883,7 +902,7 @@ def _handle_message(self, opcode, data):
         elif opcode == 0xA:
             # Pong
             self.last_pong = IOLoop.current().time()
-            self._run_callback(self.handler.on_pong, data)
+            return self._run_callback(self.handler.on_pong, data)
         else:
             self._abort()
 


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
@@ -1,5 +1,6 @@
 from __future__ import absolute_import, division, print_function, with_statement
 
+import sys
 import traceback
 
 from tornado.concurrent import Future
@@ -7,7 +8,7 @@
 from tornado.httpclient import HTTPError, HTTPRequest
 from tornado.log import gen_log, app_log
 from tornado.testing import AsyncHTTPTestCase, gen_test, bind_unused_port, ExpectLog
-from tornado.test.util import unittest
+from tornado.test.util import unittest, skipBefore35, exec_test
 from tornado.web import Application, RequestHandler
 
 try:
@@ -92,6 +93,22 @@ def open(self, arg):
         self.write_message(arg)
 
 
+class CoroutineOnMessageHandler(TestWebSocketHandler):
+    def initialize(self, close_future, compression_options=None):
+        super(CoroutineOnMessageHandler, self).initialize(close_future,
+                                                          compression_options)
+        self.sleeping = 0
+
+    @gen.coroutine
+    def on_message(self, message):
+        if self.sleeping > 0:
+            self.write_message('another coroutine is already sleeping')
+        self.sleeping += 1
+        yield gen.sleep(0.01)
+        self.sleeping -= 1
+        self.write_message(message)
+
+
 class WebSocketBaseTestCase(AsyncHTTPTestCase):
     @gen.coroutine
     def ws_connect(self, path, **kwargs):
@@ -126,6 +143,8 @@ def get_app(self):
              dict(close_future=self.close_future)),
             ('/path_args/(.*)', PathArgsHandler,
              dict(close_future=self.close_future)),
+            ('/coroutine', CoroutineOnMessageHandler,
+             dict(close_future=self.close_future)),
         ])
 
     def test_http_request(self):
@@ -259,6 +278,17 @@ def test_path_args(self):
         res = yield ws.read_message()
         self.assertEqual(res, 'hello')
 
+    @gen_test
+    def test_coroutine(self):
+        ws = yield self.ws_connect('/coroutine')
+        # Send both messages immediately, coroutine must process one at a time.
+        yield ws.write_message('hello1')
+        yield ws.write_message('hello2')
+        res = yield ws.read_message()
+        self.assertEqual(res, 'hello1')
+        res = yield ws.read_message()
+        self.assertEqual(res, 'hello2')
+
     @gen_test
     def test_check_origin_valid_no_path(self):
         port = self.get_http_port()
@@ -330,6 +360,42 @@ def test_check_origin_invalid_subdomains(self):
         self.assertEqual(cm.exception.code, 403)
 
 
+if sys.version_info >= (3, 5):
+    NativeCoroutineOnMessageHandler = exec_test(globals(), locals(), """
+class NativeCoroutineOnMessageHandler(TestWebSocketHandler):
+    def initialize(self, close_future, compression_options=None):
+        super().initialize(close_future, compression_options)
+        self.sleeping = 0
+
+    async def on_message(self, message):
+        if self.sleeping > 0:
+            self.write_message('another coroutine is already sleeping')
+        self.sleeping += 1
+        await gen.sleep(0.01)
+        self.sleeping -= 1
+        self.write_message(message)""")['NativeCoroutineOnMessageHandler']
+
+
+class WebSocketNativeCoroutineTest(WebSocketBaseTestCase):
+    def get_app(self):
+        self.close_future = Future()
+        return Application([
+            ('/native', NativeCoroutineOnMessageHandler,
+             dict(close_future=self.close_future))])
+
+    @skipBefore35
+    @gen_test
+    def test_native_coroutine(self):
+        ws = yield self.ws_connect('/native')
+        # Send both messages immediately, coroutine must process one at a time.
+        yield ws.write_message('hello1')
+        yield ws.write_message('hello2')
+        res = yield ws.read_message()
+        self.assertEqual(res, 'hello1')
+        res = yield ws.read_message()
+        self.assertEqual(res, 'hello2')
+
+
 class CompressionTestMixin(object):
     MESSAGE = 'Hello world. Testing 123 123'