websocket: Fix use of render_string in websocket handlers

PR #1917 caused websocket connections to be "finished" early, which
broke the use of render_string by setting self.ui to None.

Fixes #2010

diff --git a/tornado/test/websocket_test.py b/tornado/test/websocket_test.py
index 7bdca8773cd92b8bfc272c6cc99a9d277cb3d5d4..d47a74e651e86ab8ecd8ab76d6936c23edd06e05 100644 (file)
--- a/tornado/test/websocket_test.py
+++ b/tornado/test/websocket_test.py
@@ -8,6 +8,7 @@ from tornado.concurrent import Future
 from tornado import gen
 from tornado.httpclient import HTTPError, HTTPRequest
 from tornado.log import gen_log, app_log
+from tornado.template import DictLoader
 from tornado.testing import AsyncHTTPTestCase, gen_test, bind_unused_port, ExpectLog
 from tornado.test.util import unittest, skipBefore35, exec_test
 from tornado.web import Application, RequestHandler
@@ -130,6 +131,11 @@ class CoroutineOnMessageHandler(TestWebSocketHandler):
         self.write_message(message)
 
 
+class RenderMessageHandler(TestWebSocketHandler):
+    def on_message(self, message):
+        self.write_message(self.render_string('message.html', message=message))
+
+
 class WebSocketBaseTestCase(AsyncHTTPTestCase):
     @gen.coroutine
     def ws_connect(self, path, **kwargs):
@@ -168,7 +174,15 @@ class WebSocketTest(WebSocketBaseTestCase):
              dict(close_future=self.close_future)),
             ('/coroutine', CoroutineOnMessageHandler,
              dict(close_future=self.close_future)),
-        ])
+            ('/render', RenderMessageHandler,
+             dict(close_future=self.close_future)),
+        ], template_loader=DictLoader({
+            'message.html': '<b>{{ message }}</b>',
+        }))
+
+    def tearDown(self):
+        super(WebSocketTest, self).tearDown()
+        RequestHandler._template_loaders.clear()
 
     def test_http_request(self):
         # WS server, HTTP client.
@@ -219,6 +233,14 @@ class WebSocketTest(WebSocketBaseTestCase):
         self.assertEqual(response, u'hello \u00e9')
         yield self.close(ws)
 
+    @gen_test
+    def test_render_message(self):
+        ws = yield self.ws_connect('/render')
+        ws.write_message('hello')
+        response = yield ws.read_message()
+        self.assertEqual(response, '<b>hello</b>')
+        yield self.close(ws)
+
     @gen_test
     def test_error_in_on_message(self):
         ws = yield self.ws_connect('/error_in_on_message')

diff --git a/tornado/web.py b/tornado/web.py
index 8ff52e9ce5046b39709c40307c1874aabe1dc717..d79889fa3768df9daaa3eb18ad8de068735068a3 100644 (file)
--- a/tornado/web.py
+++ b/tornado/web.py
@@ -993,6 +993,9 @@ class RequestHandler(object):
         self._log()
         self._finished = True
         self.on_finish()
+        self._break_cycles()
+
+    def _break_cycles(self):
         # Break up a reference cycle between this handler and the
         # _ui_module closures to allow for faster GC on CPython.
         self.ui = None

diff --git a/tornado/websocket.py b/tornado/websocket.py
index ce13d262b87a1af3210a51cd7791fe6d007f13f7..69437ee4e3d9cc335b0eb1b785a3565fff14ab70 100644 (file)
--- a/tornado/websocket.py
+++ b/tornado/websocket.py
@@ -434,6 +434,16 @@ class WebSocketHandler(tornado.web.RequestHandler):
         if not self._on_close_called:
             self._on_close_called = True
             self.on_close()
+            self._break_cycles()
+
+    def _break_cycles(self):
+        # WebSocketHandlers call finish() early, but we don't want to
+        # break up reference cycles (which makes it impossible to call
+        # self.render_string) until after we've really closed the
+        # connection (if it was established in the first place,
+        # indicated by status code 101).
+        if self.get_status() != 101 or self._on_close_called:
+            super(WebSocketHandler, self)._break_cycles()
 
     def send_error(self, *args, **kwargs):
         if self.stream is None: