diff --git a/src/calibre/srv/tests/web_sockets.py b/src/calibre/srv/tests/web_sockets.py index e80788f373..703fa65d1a 100644 --- a/src/calibre/srv/tests/web_sockets.py +++ b/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]) diff --git a/src/calibre/srv/web_socket.py b/src/calibre/srv/web_socket.py index 48988b6f40..603d1a3db7 100644 --- a/src/calibre/srv/web_socket.py +++ 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) - for i in xrange(len(data)): - data[i] ^= self.mask[(self.pos + i) & 3] - data = bytes(data) + data = fast_mask(data, self.mask_buf, self.pos) 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 = bytearray(payload) - for i in xrange(len(payload)): - payload[i] ^= mask[i & 3] - payload = bytes(payload) + payload = fast_mask(payload, mask) return header + payload diff --git a/src/calibre/utils/speedup.c b/src/calibre/utils/speedup.c index e11b9cd0b8..f794b8e8d2 100644 --- a/src/calibre/utils/speedup.c +++ b/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} };