Improve websocket performance by masking in native code

This commit is contained in:
Kovid Goyal 2015-10-26 14:14:20 +05:30
parent 3f1b5782f0
commit 2b6021fcbd
3 changed files with 37 additions and 11 deletions

View File

@ -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])

View File

@ -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

View File

@ -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}
};