Improve websocket performance by masking in native code

src/calibre/srv/tests/web_sockets.py

@@ -318,3 +318,11 @@ class WebSocketTest(BaseTest):
                 simple_test([(CLOSE, struct.pack(b'!H', code))], send_close=False, close_code=code)
             for code in (0,999,1004,1005,1006,1012,1013,1014,1015,1016,1100,2000,2999):
                 simple_test([(CLOSE, struct.pack(b'!H', code))], send_close=False, close_code=PROTOCOL_ERROR)
+
+    def test_websocket_perf(self):
+        with WSTestServer(EchoHandler) as server:
+            simple_test = partial(self.simple_test, server)
+            for sz in (64, 256, 1024, 4096, 8192, 16384):
+                sz *= 1024
+                t, b = 'a'*sz, b'a'*sz
+                simple_test([t, b], [t, b])

src/calibre/srv/web_socket.py

@@ -15,9 +15,15 @@ from Queue import Queue, Empty
 from threading import Lock, Thread
 
 from calibre import as_unicode
+from calibre.constants import plugins
 from calibre.srv.loop import ServerLoop, HandleInterrupt, WRITE, READ, RDWR, Connection
 from calibre.srv.http_response import HTTPConnection, create_http_handler
 from calibre.srv.utils import DESIRED_SEND_BUFFER_SIZE
+speedup, err = plugins['speedup']
+if not speedup:
+    raise RuntimeError('Failed to load speedup module with error: ' + err)
+fast_mask = speedup.websocket_mask
+del speedup, err
 
 HANDSHAKE_STR = (
     "HTTP/1.1 101 Switching Protocols\r\n"
@@ -131,8 +137,6 @@ class ReadFrame(object):  # {{{
         self.mask_buf += data
         if len(self.mask_buf) < 4:
             return
-        self.mask = bytearray(self.mask_buf)
-        del self.mask_buf
         self.state = self.read_payload
         self.pos = 0
         self.frame_starting = True
@@ -148,10 +152,7 @@ class ReadFrame(object):  # {{{
                 return
         else:
             data = b''
-        data = bytearray(data)
+        data = fast_mask(data, self.mask_buf, self.pos)
-        for i in xrange(len(data)):
-            data[i] ^= self.mask[(self.pos + i) & 3]
-        data = bytes(data)
         self.pos += len(data)
         frame_finished = self.pos >= self.payload_length
         conn.ws_data_received(data, self.opcode, self.frame_starting, frame_finished, self.fin)
@@ -177,11 +178,7 @@ def create_frame(fin, opcode, payload, mask=None, rsv=0):
         header = bytes(bytearray((b1, b2 | 127))) + struct.pack(b'!Q', l)
     if mask is not None:
         header += mask
-        mask = bytearray(mask)
+        payload = fast_mask(payload, mask)
-        payload = bytearray(payload)
-        for i in xrange(len(payload)):
-            payload[i] ^= mask[i & 3]
-        payload = bytes(payload)
 
     return header + payload
 

src/calibre/utils/speedup.c

@@ -217,6 +217,23 @@ speedup_create_texture(PyObject *self, PyObject *args, PyObject *kw) {
     return ret;
 }
 
+static PyObject*
+speedup_websocket_mask(PyObject *self, PyObject *args) {
+	PyObject *data = NULL, *mask = NULL, *ans = NULL;
+	Py_ssize_t offset_ = 0;
+	size_t offset = 0, i = 0;
+	char *data_buf = NULL, *mask_buf = NULL, *ans_buf = NULL;
+    if(!PyArg_ParseTuple(args, "OO|n", &data, &mask, &offset)) return NULL;
+	offset = (size_t)offset_;
+	ans = PyBytes_FromStringAndSize(NULL, PyBytes_GET_SIZE(data));
+	if (ans != NULL) {
+		data_buf = PyBytes_AS_STRING(data); mask_buf = PyBytes_AS_STRING(mask); ans_buf = PyBytes_AS_STRING(ans);
+		for(i = 0; i < PyBytes_GET_SIZE(ans); i++)
+			ans_buf[i] = data_buf[i] ^ mask_buf[(i + offset) & 3];
+	}
+	return ans;
+}
+
 static PyMethodDef speedup_methods[] = {
     {"parse_date", speedup_parse_date, METH_VARARGS,
         "parse_date()\n\nParse ISO dates faster."
@@ -244,6 +261,10 @@ static PyMethodDef speedup_methods[] = {
         "fdopen(fd, name, mode [, bufsize=-1)\n\nCreate a python file object from an OS file descriptor with a name. Note that this does not do any validation of mode, so you must ensure fd already has the correct flags set."
     },
 
+    {"websocket_mask", speedup_websocket_mask, METH_VARARGS,
+        "websocket_mask(data, mask [, offset=0)\n\nXOR the data (bytestring) with the specified (must be 4-byte bytestring) mask"
+    },
+
     {NULL, NULL, 0, NULL}
 };