Merge pull request #6110 from minrk/binarycomm
support binary buffers in comm messages
commit
d269912958
@ -0,0 +1,114 @@
|
||||
// Copyright (c) IPython Development Team.
|
||||
// Distributed under the terms of the Modified BSD License.
|
||||
|
||||
define([
|
||||
'underscore',
|
||||
], function (_) {
|
||||
"use strict";
|
||||
|
||||
var _deserialize_array_buffer = function (buf) {
|
||||
var data = new DataView(buf);
|
||||
// read the header: 1 + nbufs 32b integers
|
||||
var nbufs = data.getUint32(0);
|
||||
var offsets = [];
|
||||
var i;
|
||||
for (i = 1; i <= nbufs; i++) {
|
||||
offsets.push(data.getUint32(i * 4));
|
||||
}
|
||||
var json_bytes = new Uint8Array(buf.slice(offsets[0], offsets[1]));
|
||||
var msg = JSON.parse(
|
||||
(new TextDecoder('utf8')).decode(json_bytes)
|
||||
);
|
||||
// the remaining chunks are stored as DataViews in msg.buffers
|
||||
msg.buffers = [];
|
||||
var start, stop;
|
||||
for (i = 1; i < nbufs; i++) {
|
||||
start = offsets[i];
|
||||
stop = offsets[i+1] || buf.byteLength;
|
||||
msg.buffers.push(new DataView(buf.slice(start, stop)));
|
||||
}
|
||||
return msg;
|
||||
};
|
||||
|
||||
var _deserialize_binary = function(data, callback) {
|
||||
// deserialize the binary message format
|
||||
// callback will be called with a message whose buffers attribute
|
||||
// will be an array of DataViews.
|
||||
if (data instanceof Blob) {
|
||||
// data is Blob, have to deserialize from ArrayBuffer in reader callback
|
||||
var reader = new FileReader();
|
||||
reader.onload = function () {
|
||||
var msg = _deserialize_array_buffer(this.result);
|
||||
callback(msg);
|
||||
};
|
||||
reader.readAsArrayBuffer(data);
|
||||
} else {
|
||||
// data is ArrayBuffer, can deserialize directly
|
||||
var msg = _deserialize_array_buffer(data);
|
||||
callback(msg);
|
||||
}
|
||||
};
|
||||
|
||||
var deserialize = function (data, callback) {
|
||||
// deserialize a message and pass the unpacked message object to callback
|
||||
if (typeof data === "string") {
|
||||
// text JSON message
|
||||
callback(JSON.parse(data));
|
||||
} else {
|
||||
// binary message
|
||||
_deserialize_binary(data, callback);
|
||||
}
|
||||
};
|
||||
|
||||
var _serialize_binary = function (msg) {
|
||||
// implement the binary serialization protocol
|
||||
// serializes JSON message to ArrayBuffer
|
||||
msg = _.clone(msg);
|
||||
var offsets = [];
|
||||
var buffers = [];
|
||||
msg.buffers.map(function (buf) {
|
||||
buffers.push(buf);
|
||||
});
|
||||
delete msg.buffers;
|
||||
var json_utf8 = (new TextEncoder('utf8')).encode(JSON.stringify(msg));
|
||||
buffers.unshift(json_utf8);
|
||||
var nbufs = buffers.length;
|
||||
offsets.push(4 * (nbufs + 1));
|
||||
var i;
|
||||
for (i = 0; i + 1 < buffers.length; i++) {
|
||||
offsets.push(offsets[offsets.length-1] + buffers[i].byteLength);
|
||||
}
|
||||
var msg_buf = new Uint8Array(
|
||||
offsets[offsets.length-1] + buffers[buffers.length-1].byteLength
|
||||
);
|
||||
// use DataView.setUint32 for network byte-order
|
||||
var view = new DataView(msg_buf.buffer);
|
||||
// write nbufs to first 4 bytes
|
||||
view.setUint32(0, nbufs);
|
||||
// write offsets to next 4 * nbufs bytes
|
||||
for (i = 0; i < offsets.length; i++) {
|
||||
view.setUint32(4 * (i+1), offsets[i]);
|
||||
}
|
||||
// write all the buffers at their respective offsets
|
||||
for (i = 0; i < buffers.length; i++) {
|
||||
msg_buf.set(new Uint8Array(buffers[i].buffer), offsets[i]);
|
||||
}
|
||||
|
||||
// return raw ArrayBuffer
|
||||
return msg_buf.buffer;
|
||||
};
|
||||
|
||||
var serialize = function (msg) {
|
||||
if (msg.buffers && msg.buffers.length) {
|
||||
return _serialize_binary(msg);
|
||||
} else {
|
||||
return JSON.stringify(msg);
|
||||
}
|
||||
};
|
||||
|
||||
var exports = {
|
||||
deserialize : deserialize,
|
||||
serialize: serialize
|
||||
};
|
||||
return exports;
|
||||
});
|
||||
@ -0,0 +1,113 @@
|
||||
//
|
||||
// Test binary messages on websockets.
|
||||
// Only works on slimer for now, due to old websocket impl in phantomjs.
|
||||
//
|
||||
|
||||
casper.notebook_test(function () {
|
||||
if (!this.slimerjs) {
|
||||
console.log("Can't test binary websockets on phantomjs.");
|
||||
return;
|
||||
}
|
||||
// create EchoBuffers target on js-side.
|
||||
// it just captures and echos comm messages.
|
||||
this.then(function () {
|
||||
var success = this.evaluate(function () {
|
||||
IPython._msgs = [];
|
||||
|
||||
var EchoBuffers = function(comm) {
|
||||
this.comm = comm;
|
||||
this.comm.on_msg($.proxy(this.on_msg, this));
|
||||
};
|
||||
|
||||
EchoBuffers.prototype.on_msg = function (msg) {
|
||||
IPython._msgs.push(msg);
|
||||
this.comm.send(msg.content.data, {}, {}, msg.buffers);
|
||||
};
|
||||
|
||||
IPython.notebook.kernel.comm_manager.register_target("echo", function (comm) {
|
||||
return new EchoBuffers(comm);
|
||||
});
|
||||
|
||||
return true;
|
||||
});
|
||||
this.test.assertEquals(success, true, "Created echo comm target");
|
||||
});
|
||||
|
||||
// Create a similar comm that captures messages Python-side
|
||||
this.then(function () {
|
||||
var index = this.append_cell([
|
||||
"import os",
|
||||
"from IPython.kernel.comm import Comm",
|
||||
"comm = Comm(target_name='echo')",
|
||||
"msgs = []",
|
||||
"def on_msg(msg):",
|
||||
" msgs.append(msg)",
|
||||
"comm.on_msg(on_msg)"
|
||||
].join('\n'), 'code');
|
||||
this.execute_cell(index);
|
||||
});
|
||||
|
||||
// send a message with binary data
|
||||
this.then(function () {
|
||||
var index = this.append_cell([
|
||||
"buffers = [b'\\xFF\\x00', b'\\x00\\x01\\x02']",
|
||||
"comm.send(data='hi', buffers=buffers)"
|
||||
].join('\n'), 'code');
|
||||
this.execute_cell(index);
|
||||
});
|
||||
|
||||
// wait for capture
|
||||
this.waitFor(function () {
|
||||
return this.evaluate(function () {
|
||||
return IPython._msgs.length > 0;
|
||||
});
|
||||
});
|
||||
|
||||
// validate captured buffers js-side
|
||||
this.then(function () {
|
||||
var msgs = this.evaluate(function () {
|
||||
return IPython._msgs;
|
||||
});
|
||||
this.test.assertEquals(msgs.length, 1, "Captured comm message");
|
||||
var buffers = msgs[0].buffers;
|
||||
this.test.assertEquals(buffers.length, 2, "comm message has buffers");
|
||||
|
||||
// extract attributes to test in evaluate,
|
||||
// because the raw DataViews can't be passed across
|
||||
var buf_info = function (index) {
|
||||
var buf = IPython._msgs[0].buffers[index];
|
||||
var data = {};
|
||||
data.byteLength = buf.byteLength;
|
||||
data.bytes = [];
|
||||
for (var i = 0; i < data.byteLength; i++) {
|
||||
data.bytes.push(buf.getUint8(i));
|
||||
}
|
||||
return data;
|
||||
};
|
||||
|
||||
buf0 = this.evaluate(buf_info, 0);
|
||||
buf1 = this.evaluate(buf_info, 1);
|
||||
this.test.assertEquals(buf0.byteLength, 2, 'buf[0] has correct size');
|
||||
this.test.assertEquals(buf0.bytes, [255, 0], 'buf[0] has correct bytes');
|
||||
this.test.assertEquals(buf1.byteLength, 3, 'buf[1] has correct size');
|
||||
this.test.assertEquals(buf1.bytes, [0, 1, 2], 'buf[1] has correct bytes');
|
||||
});
|
||||
|
||||
// validate captured buffers Python-side
|
||||
this.then(function () {
|
||||
var index = this.append_cell([
|
||||
"assert len(msgs) == 1, len(msgs)",
|
||||
"bufs = msgs[0]['buffers']",
|
||||
"assert len(bufs) == len(buffers), bufs",
|
||||
"assert bufs[0].bytes == buffers[0], bufs[0].bytes",
|
||||
"assert bufs[1].bytes == buffers[1], bufs[1].bytes",
|
||||
"1",
|
||||
].join('\n'), 'code');
|
||||
this.execute_cell(index);
|
||||
this.wait_for_output(index);
|
||||
this.then(function () {
|
||||
var out = this.get_output_cell(index);
|
||||
this.test.assertEquals(out['text/plain'], '1', "Python received buffers");
|
||||
});
|
||||
});
|
||||
});
|
||||
@ -0,0 +1,26 @@
|
||||
"""Test serialize/deserialize messages with buffers"""
|
||||
|
||||
import os
|
||||
|
||||
import nose.tools as nt
|
||||
|
||||
from IPython.kernel.zmq.session import Session
|
||||
from ..base.zmqhandlers import (
|
||||
serialize_binary_message,
|
||||
deserialize_binary_message,
|
||||
)
|
||||
|
||||
def test_serialize_binary():
|
||||
s = Session()
|
||||
msg = s.msg('data_pub', content={'a': 'b'})
|
||||
msg['buffers'] = [ os.urandom(3) for i in range(3) ]
|
||||
bmsg = serialize_binary_message(msg)
|
||||
nt.assert_is_instance(bmsg, bytes)
|
||||
|
||||
def test_deserialize_binary():
|
||||
s = Session()
|
||||
msg = s.msg('data_pub', content={'a': 'b'})
|
||||
msg['buffers'] = [ os.urandom(2) for i in range(3) ]
|
||||
bmsg = serialize_binary_message(msg)
|
||||
msg2 = deserialize_binary_message(bmsg)
|
||||
nt.assert_equal(msg2, msg)
|
||||
Loading…
Reference in new issue