From: Ben Darnell Date: Fri, 25 Oct 2013 19:46:13 +0000 (-0400) Subject: Add cython-based speedup for websocket mask function. X-Git-Tag: v3.2.0b1~64 X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=e8dc5e427d2e244c519c4211fb1796e747b77ed1;p=thirdparty%2Ftornado.git Add cython-based speedup for websocket mask function. This optimization is currently activated only if Cython is present when Tornado is installed. --- diff --git a/MANIFEST.in b/MANIFEST.in index b710aac98..ea526699f 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,4 +1,5 @@ recursive-include demos *.py *.yaml *.html *.css *.js *.xml *.sql README +include tornado/speedups.pyx include tornado/ca-certificates.crt include tornado/test/README include tornado/test/csv_translations/fr_FR.csv diff --git a/setup.py b/setup.py index 7e21325d3..1d019fc7c 100644 --- a/setup.py +++ b/setup.py @@ -23,6 +23,11 @@ try: except ImportError: pass +try: + from Cython.Build import cythonize +except ImportError: + cythonize = None + kwargs = {} version = "3.2.dev2" @@ -30,10 +35,16 @@ version = "3.2.dev2" with open('README.rst') as f: long_description = f.read() +if cythonize is not None: + extensions = cythonize('tornado/speedups.pyx') +else: + extensions = [] + distutils.core.setup( name="tornado", version=version, packages = ["tornado", "tornado.test", "tornado.platform"], + ext_modules = extensions, package_data = { "tornado": ["ca-certificates.crt"], # data files need to be listed both here (which determines what gets diff --git a/tornado/speedups.pyx b/tornado/speedups.pyx new file mode 100644 index 000000000..46344453b --- /dev/null +++ b/tornado/speedups.pyx @@ -0,0 +1,16 @@ +# -*- python -*- +from cpython.mem cimport PyMem_Malloc, PyMem_Free + +def websocket_mask(bytes mask_bytes, bytes data_bytes): + cdef size_t data_len = len(data_bytes) + cdef char* data = data_bytes + cdef char* mask = mask_bytes + cdef size_t i + cdef char* buf = PyMem_Malloc(data_len) + try: + for i in xrange(data_len): + buf[i] = data[i] ^ mask[i % 4] + # Is there a zero-copy equivalent of this? + return (buf[:data_len]) + finally: + PyMem_Free(buf) diff --git a/tornado/test/websocket_test.py b/tornado/test/websocket_test.py index da9d780b1..66c48e576 100644 --- a/tornado/test/websocket_test.py +++ b/tornado/test/websocket_test.py @@ -2,8 +2,14 @@ from tornado.concurrent import Future from tornado.httpclient import HTTPError, HTTPRequest from tornado.log import gen_log from tornado.testing import AsyncHTTPTestCase, gen_test, bind_unused_port, ExpectLog +from tornado.test.util import unittest from tornado.web import Application, RequestHandler -from tornado.websocket import WebSocketHandler, websocket_connect, WebSocketError +from tornado.websocket import WebSocketHandler, websocket_connect, WebSocketError, _websocket_mask_python + +try: + from tornado import speedups +except ImportError: + speedups = None class TestWebSocketHandler(WebSocketHandler): """Base class for testing handlers that exposes the on_close event. @@ -110,3 +116,22 @@ class WebSocketTest(AsyncHTTPTestCase): self.assertEqual(response, 'hello') ws.close() yield self.close_future + + +class MaskFunctionMixin(object): + # Subclasses should define self.mask(mask, data) + def test_mask(self): + self.assertEqual(self.mask(b'abcd', b''), b'') + self.assertEqual(self.mask(b'abcd', b'b'), b'\x03') + self.assertEqual(self.mask(b'abcd', b'54321'), b'TVPVP') + self.assertEqual(self.mask(b'ZXCV', b'98765432'), b'c`t`olpd') + + +class PythonMaskFunctionTest(MaskFunctionMixin, unittest.TestCase): + def mask(self, mask, data): + return _websocket_mask_python(mask, data) + +@unittest.skipIf(speedups is None, "tornado.speedups module not present") +class CythonMaskFunctionTest(MaskFunctionMixin, unittest.TestCase): + def mask(self, mask, data): + return speedups.websocket_mask(mask, data) diff --git a/tornado/websocket.py b/tornado/websocket.py index 75b1d8f23..bdfa779c7 100644 --- a/tornado/websocket.py +++ b/tornado/websocket.py @@ -586,7 +586,7 @@ class WebSocketProtocol13(WebSocketProtocol): frame += struct.pack("!BQ", 127 | mask_bit, l) if self.mask_outgoing: mask = os.urandom(4) - data = mask + self._apply_mask(mask, data) + data = mask + _websocket_mask(mask, data) frame += data self.stream.write(frame) @@ -671,21 +671,8 @@ class WebSocketProtocol13(WebSocketProtocol): except StreamClosedError: self._abort() - def _apply_mask(self, mask, data): - mask = array.array("B", mask) - unmasked = array.array("B", data) - for i in xrange(len(data)): - unmasked[i] = unmasked[i] ^ mask[i % 4] - if hasattr(unmasked, 'tobytes'): - # tostring was deprecated in py32. It hasn't been removed, - # but since we turn on deprecation warnings in our tests - # we need to use the right one. - return unmasked.tobytes() - else: - return unmasked.tostring() - def _on_masked_frame_data(self, data): - self._on_frame_data(self._apply_mask(self._frame_mask, data)) + self._on_frame_data(_websocket_mask(self._frame_mask, data)) def _on_frame_data(self, data): if self._frame_opcode_is_control: @@ -882,3 +869,29 @@ def websocket_connect(url, io_loop=None, callback=None, connect_timeout=None): if callback is not None: io_loop.add_future(conn.connect_future, callback) return conn.connect_future + +def _websocket_mask_python(mask, data): + """Websocket masking function. + + `mask` is a `bytes` object of length 4; `data` is a `bytes` object of any length. + Returns a `bytes` object of the same length as `data` with the mask applied + as specified in section 5.3 of RFC 6455. + + This pure-python implementation may be replaced by an optimized version when available. + """ + mask = array.array("B", mask) + unmasked = array.array("B", data) + for i in xrange(len(data)): + unmasked[i] = unmasked[i] ^ mask[i % 4] + if hasattr(unmasked, 'tobytes'): + # tostring was deprecated in py32. It hasn't been removed, + # but since we turn on deprecation warnings in our tests + # we need to use the right one. + return unmasked.tobytes() + else: + return unmasked.tostring() + +try: + from tornado.speedups import websocket_mask as _websocket_mask +except ImportError: + _websocket_mask = _websocket_mask_python diff --git a/tox.ini b/tox.ini index 77918ef08..2d9afaffd 100644 --- a/tox.ini +++ b/tox.ini @@ -30,6 +30,7 @@ deps = unittest2 [testenv:py26-full] basepython = python2.6 deps = + Cython futures mock pycurl @@ -39,6 +40,7 @@ deps = [testenv:py27-full] basepython = python2.7 deps = + Cython futures mock pycurl @@ -148,6 +150,7 @@ commands = python -m tornado.test.runtests --locale=zh_TW {posargs:} # there. basepython = pypy deps = + Cython futures mock @@ -168,6 +171,7 @@ setenv = LANG=en_US.utf-8 [testenv:py32-full] basepython = python3.2 deps = + Cython mock [testenv:py33]