实现和测试WebSocket服务器连接超时

时间:2014-06-01 14:38:13

标签: python tornado

我正在Tornado 3.2中实现WebSockets服务器。连接到服务器的客户端不会成为浏览器。

对于服务器和客户端之间来回通信的情况,我想添加最大值。服务器在关闭连接之前等待客户端响应的时间。

这大致是我一直在尝试的:

import datetime
import tornado

class WSHandler(WebSocketHandler):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.timeout = None

    def _close_on_timeout(self):
        if self.ws_connection:
            self.close()

    def open(self):
        initialize()

    def on_message(self, message):
        # Remove previous timeout, if one exists.
        if self.timeout:
            tornado.ioloop.IOLoop.instance().remove_timeout(self.timeout)
            self.timeout = None

        if is_last_message:
            self.write_message(message)
            self.close()
        else:
            # Add a new timeout.
            self.timeout = tornado.ioloop.IOLoop.instance().add_timeout(
                datetime.timedelta(milliseconds=1000), self._close_on_timeout)
            self.write_message(message)

我是一个klutz并且有一个更简单的方法吗?我甚至无法通过上面的add_timeout安排一个简单的打印声明。

我还需要一些测试帮助。这就是我到目前为止所做的:

from tornado.websocket import websocket_connect
from tornado.testing import AsyncHTTPTestCase, gen_test
import time

class WSTests(AsyncHTTPTestCase):

    @gen_test
    def test_long_response(self):
        ws = yield websocket_connect('ws://address', io_loop=self.io_loop)

        # First round trip.
        ws.write_message('First message.')
        result = yield ws.read_message()
        self.assertEqual(result, 'First response.')

        # Wait longer than the timeout.
        # The test is in its own IOLoop, so a blocking sleep should be okay?
        time.sleep(1.1)

        # Expect either write or read to fail because of a closed socket.
        ws.write_message('Second message.')
        result = yield ws.read_message()

        self.assertNotEqual(result, 'Second response.')

客户端在写入和读取套接字时没有问题。这可能是因为add_timeout没有被触发。

测试是否需要以某种方式产生允许服务器上的超时回调运行?我想不会因为文档说测试是在他们自己的IOLoop中运行的。

修改

根据Ben的建议,这是工作版本。

import datetime
import tornado

class WSHandler(WebSocketHandler):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.timeout = None

    def _close_on_timeout(self):
        if self.ws_connection:
            self.close()

    def open(self):
        initialize()

    def on_message(self, message):
        # Remove previous timeout, if one exists.
        if self.timeout:
            tornado.ioloop.IOLoop.current().remove_timeout(self.timeout)
            self.timeout = None

        if is_last_message:
            self.write_message(message)
            self.close()
        else:
            # Add a new timeout.
            self.timeout = tornado.ioloop.IOLoop.current().add_timeout(
                datetime.timedelta(milliseconds=1000), self._close_on_timeout)
            self.write_message(message)

测试:

from tornado.websocket import websocket_connect
from tornado.testing import AsyncHTTPTestCase, gen_test
import time

class WSTests(AsyncHTTPTestCase):

    @gen_test
    def test_long_response(self):
        ws = yield websocket_connect('ws://address', io_loop=self.io_loop)

        # First round trip.
        ws.write_message('First message.')
        result = yield ws.read_message()
        self.assertEqual(result, 'First response.')

        # Wait a little more than the timeout.
        yield gen.Task(self.io_loop.add_timeout, datetime.timedelta(seconds=1.1))

        # Expect either write or read to fail because of a closed socket.
        ws.write_message('Second message.')
        result = yield ws.read_message()
        self.assertEqual(result, None)

2 个答案:

答案 0 :(得分:2)

第一个示例中的超时处理代码对我来说是正确的。

为了进行测试,每个测试用例都有自己的IOLoop,但是测试和运行的其他任何东西都只有一个IOLoop,所以你必须在这里使用add_timeout而不是time.sleep()来避免阻塞服务器

答案 1 :(得分:1)

Ey Ben,我知道这个问题很久以前就已经解决了,但我想与任何阅读此解决方案的用户分享这个问题。 它基本上是基于你的,但它解决了外部服务的问题,可以使用组合而不是继承轻松集成到任何websocket中:

class TimeoutWebSocketService():
    _default_timeout_delta_ms = 10 * 60 * 1000  # 10 min

    def __init__(self, websocket, ioloop=None, timeout=None):
        # Timeout
        self.ioloop = ioloop or tornado.ioloop.IOLoop.current()
        self.websocket = websocket
        self._timeout = None
        self._timeout_delta_ms = timeout or TimeoutWebSocketService._default_timeout_delta_ms

    def _close_on_timeout(self):
        self._timeout = None
        if self.websocket.ws_connection:
            self.websocket.close()

    def refresh_timeout(self, timeout=None):
        timeout = timeout or self._timeout_delta_ms
        if timeout > 0:
            # Clean last timeout, if one exists
            self.clean_timeout()

            # Add a new timeout (must be None from clean).
            self._timeout = self.ioloop.add_timeout(
                datetime.timedelta(milliseconds=timeout), self._close_on_timeout)

    def clean_timeout(self):
        if self._timeout is not None:
            # Remove previous timeout, if one exists.
            self.ioloop.remove_timeout(self._timeout)
            self._timeout = None

为了使用该服务,就像创建一个新的TimeoutWebService实例一样简单(可选择以ms为单位的超时,以及应该执行它的ioloop)并调用方法''refresh_timeout''来设置第一次超时或重置已存在的超时,或''clean_timeout''停止超时服务。

class BaseWebSocketHandler(WebSocketHandler):
    def prepare(self):
        self.timeout_service = TimeoutWebSocketService(timeout=(1000*60))

        ## Optionally starts the service here 
        self.timeout_service.refresh_timeout()

        ## rest of prepare method 

    def on_message(self):
        self.timeout_service.refresh_timeout()

    def on_close(self):
        self.timeout_service.clean_timeout()

由于采用了这种方法,您可以控制何时以及在何种条件下重新启动可能因应用程序而异的超时。例如,如果用户执行X条件,或者消息是预期的消息,您可能只想刷新超时。

我希望人们享受这个解决方案!