Permalink
Browse files

Add cython-based speedup for websocket mask function.

This optimization is currently activated only if Cython is present
when Tornado is installed.
  • Loading branch information...
1 parent 0d0f583 commit e8dc5e427d2e244c519c4211fb1796e747b77ed1 @bdarnell bdarnell committed Oct 25, 2013
Showing with 86 additions and 16 deletions.
  1. +1 −0 MANIFEST.in
  2. +11 −0 setup.py
  3. +16 −0 tornado/speedups.pyx
  4. +26 −1 tornado/test/websocket_test.py
  5. +28 −15 tornado/websocket.py
  6. +4 −0 tox.ini
View
@@ -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
View
@@ -23,17 +23,28 @@
except ImportError:
pass
+try:
+ from Cython.Build import cythonize
+except ImportError:
+ cythonize = None
+
kwargs = {}
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
View
@@ -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 = <char*> 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 <bytes>(buf[:data_len])
+ finally:
+ PyMem_Free(buf)
@@ -2,8 +2,14 @@
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 @@ def test_websocket_headers(self):
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)
View
@@ -586,7 +586,7 @@ def _write_frame(self, fin, opcode, data):
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 @@ def _on_masking_key(self, data):
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
View
@@ -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]

0 comments on commit e8dc5e4

Please sign in to comment.