Skip to content
Permalink
Branch: master
Find file Copy path
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
341 lines (295 sloc) 12.3 KB
use std::convert::TryInto;
use std::io::{BufRead, Read};
use std::mem::size_of;
use capnp::message::{ReaderOptions, ScratchSpace, ScratchSpaceHeapAllocator};
use capnp::serialize::write_message;
use capnp::serialize_packed::{
read_message as read_message_packed, write_message as write_message_packed,
};
use capnp::Error;
use crate::iex::{IexMessage, IexPayload};
use crate::marketdata_capnp::message;
use crate::marketdata_capnp::{multi_message, Side};
use crate::{RunnerDeserialize, RunnerSerialize, StreamVec, Summarizer};
pub struct CapnpWriter<'a> {
// We have to be very careful with how messages are built, as running
// `init_root` and rebuilding will still accumulate garbage if using
// the standard HeapAllocator.
// https://github.com/capnproto/capnproto-rust/issues/111
_words: Vec<capnp::Word>,
scratch: ScratchSpace<'a>,
packed: bool,
}
impl<'a> CapnpWriter<'a> {
pub fn new(packed: bool) -> CapnpWriter<'a> {
// Cap'n'Proto words are 8 bytes, MTU is 1500 bytes, theoretically need only 188 words.
// In practice, let's just make sure everything fits.
let mut words = capnp::Word::allocate_zeroed_vec(1024);
let scratch = ScratchSpace::new(unsafe { std::mem::transmute(&mut words[..]) });
CapnpWriter {
_words: words,
scratch,
packed,
}
}
fn builder(&mut self) -> capnp::message::Builder<ScratchSpaceHeapAllocator<'a, 'a>> {
// Builders are only safe to use for serializing a single message. We can re-use the
// backing memory (unsafe because now both `self` and the returned builder have a
// mutable reference to `self.scratch), but Bad Things happen if we don't drop
// in between serialization.
capnp::message::Builder::new(ScratchSpaceHeapAllocator::new(unsafe {
std::mem::transmute(&mut self.scratch)
}))
}
}
impl<'a> RunnerSerialize for CapnpWriter<'a> {
fn serialize(&mut self, payload: &IexPayload, mut output: &mut Vec<u8>) {
// First, count the messages we actually care about.
let num_msgs = payload
.messages
.iter()
.map(|m| match m {
IexMessage::TradeReport(_) | IexMessage::PriceLevelUpdate(_) => 1,
_ => 0,
})
.fold(0, |sum, i| sum + i);
if num_msgs == 0 {
return;
}
// And actually serialize the IEX payload to CapNProto format
// This is the safe builder used for testing
//let mut builder = capnp::message::Builder::new_default();
//let mut multimsg = builder.init_root::<multi_message::Builder>();
// This is the unsafe (but faster) version
let mut builder = self.builder();
let mut multimsg = builder.init_root::<multi_message::Builder>();
multimsg.set_seq_no(payload.first_seq_no);
let mut messages = multimsg.init_messages(num_msgs as u32);
let mut current_msg_no = 0;
for iex_msg in payload.messages.iter() {
match iex_msg {
IexMessage::TradeReport(tr) => {
let mut message = messages.reborrow().get(current_msg_no);
current_msg_no += 1;
message.set_ts(tr.timestamp);
let sym = crate::parse_symbol(&tr.symbol);
message.reborrow().init_symbol(sym.len() as u32);
message.set_symbol(sym);
let mut msg_tr = message.init_trade();
msg_tr.set_size(tr.size);
msg_tr.set_price(tr.price);
}
IexMessage::PriceLevelUpdate(plu) => {
let mut message = messages.reborrow().get(current_msg_no);
current_msg_no += 1;
message.set_ts(plu.timestamp);
let sym = crate::parse_symbol(&plu.symbol);
message.reborrow().init_symbol(sym.len() as u32);
message.set_symbol(sym);
let mut msg_plu = message.init_quote();
msg_plu.set_price(plu.price);
msg_plu.set_size(plu.size);
msg_plu.set_flags(plu.event_flags);
msg_plu.set_side(if plu.msg_type == 0x38 {
Side::Buy
} else {
Side::Sell
});
}
_ => (),
}
}
let write_fn = if self.packed {
write_message_packed
} else {
write_message
};
write_fn(&mut output, &builder).unwrap();
}
}
pub struct CapnpReader {
read_opts: ReaderOptions,
packed: bool,
}
impl CapnpReader {
pub fn new(packed: bool) -> CapnpReader {
CapnpReader {
read_opts: ReaderOptions::new(),
packed,
}
}
}
impl CapnpReader {
fn deserialize_packed<'a>(
&self,
buf: &'a mut StreamVec,
stats: &mut Summarizer,
) -> Result<(), ()> {
// Because `capnp::serialize_packed::PackedRead` is hidden from us, packed reads
// *have* to both allocate new segments every read, and copy the buffer into
// those same segments, no ability to re-use allocated memory.
let reader = read_message_packed(buf, self.read_opts).map_err(|_| ())?;
let multimsg = reader.get_root::<multi_message::Reader>().unwrap();
for msg in multimsg.get_messages().unwrap().iter() {
match msg.which() {
Ok(message::Trade(tr)) => {
let tr = tr.unwrap();
stats.append_trade_volume(msg.get_symbol().unwrap(), tr.get_size() as u64);
}
Ok(message::Quote(q)) => {
let q = q.unwrap();
let is_bid = match q.get_side().unwrap() {
Side::Buy => true,
_ => false,
};
stats.update_quote_prices(msg.get_symbol().unwrap(), q.get_price(), is_bid);
}
_ => panic!("Unrecognized message type!"),
}
}
Ok(())
}
fn deserialize_unpacked(&self, buf: &mut StreamVec, stats: &mut Summarizer) -> Result<(), ()> {
let mut data = buf.fill_buf().map_err(|_| ())?;
if data.len() == 0 {
return Err(());
}
let orig_data = data;
let reader_opts = ReaderOptions::default();
/*
Read into `OwnedSegments`, which means we copy the entire message into a new Vec. Note that
the `data` pointer is modified underneath us, can figure out the message length by
checking the difference between where we started and what `data` is afterward.
This is a trick you learn only by looking at the fuzzing test cases.
let reader = capnp::serialize::read_message(&mut data, reader_opts)?;
let bytes_consumed = orig_data.len() - data.len();
*/
/*
Read into `SliceSegments`, which allows us to re-use the underlying message storage,
but still forces a Vec allocation for `offsets`. Also requires us to copy code from
Cap'n'Proto because `SliceSegments` has private fields, and `read_segment_table`
is private. And all this because `read_segment_from_words` has a length check
that triggers an error if our buffer is too large. What the hell?
There is no documentation on how to calculate `bytes_consumed` when parsing by hand
that I could find, you just have to guess and check until you figure this one out.
*/
let (num_words, offsets) = read_segment_table(&mut data, reader_opts).map_err(|_| ())?;
let words = unsafe { capnp::Word::bytes_to_words(data) };
let reader = capnp::message::Reader::new(
SliceSegments {
words: words,
segment_slices: offsets,
},
reader_opts,
);
let segment_table_bytes = orig_data.len() - data.len();
let msg_bytes = num_words * size_of::<capnp::Word>();
let bytes_consumed = segment_table_bytes + msg_bytes;
let multimsg = reader.get_root::<multi_message::Reader>().map_err(|_| ())?;
for msg in multimsg.get_messages().map_err(|_| ())?.iter() {
let sym = msg.get_symbol().map_err(|_| ())?;
match msg.which().map_err(|_| ())? {
message::Trade(trade) => {
let trade = trade.unwrap();
stats.append_trade_volume(sym, trade.get_size().into());
}
message::Quote(quote) => {
let quote = quote.unwrap();
let is_buy = match quote.get_side().unwrap() {
Side::Buy => true,
_ => false,
};
stats.update_quote_prices(sym, quote.get_price(), is_buy);
}
}
}
buf.consume(bytes_consumed);
Ok(())
}
}
impl RunnerDeserialize for CapnpReader {
fn deserialize<'a>(&self, buf: &'a mut StreamVec, stats: &mut Summarizer) -> Result<(), ()> {
// While this is an extra branch per call, we're going to assume that the overhead
// is essentially nil in practice
if self.packed {
self.deserialize_packed(buf, stats)
} else {
self.deserialize_unpacked(buf, stats)
}
}
}
pub struct SliceSegments<'a> {
words: &'a [capnp::Word],
segment_slices: Vec<(usize, usize)>,
}
impl<'a> capnp::message::ReaderSegments for SliceSegments<'a> {
fn get_segment<'b>(&'b self, id: u32) -> Option<&'b [capnp::Word]> {
if id < self.segment_slices.len() as u32 {
let (a, b) = self.segment_slices[id as usize];
Some(&self.words[a..b])
} else {
None
}
}
fn len(&self) -> usize {
self.segment_slices.len()
}
}
fn read_segment_table<R>(
read: &mut R,
options: capnp::message::ReaderOptions,
) -> capnp::Result<(usize, Vec<(usize, usize)>)>
where
R: Read,
{
let mut buf: [u8; 8] = [0; 8];
// read the first Word, which contains segment_count and the 1st segment length
read.read_exact(&mut buf)?;
let segment_count = u32::from_le_bytes(buf[0..4].try_into().unwrap()).wrapping_add(1) as usize;
if segment_count >= 512 {
return Err(Error::failed(format!(
"Too many segments: {}",
segment_count
)));
} else if segment_count == 0 {
return Err(Error::failed(format!(
"Too few segments: {}",
segment_count
)));
}
let mut segment_slices = Vec::with_capacity(segment_count);
let mut total_words = u32::from_le_bytes(buf[4..8].try_into().unwrap()) as usize;
segment_slices.push((0, total_words));
if segment_count > 1 {
if segment_count < 4 {
read.read_exact(&mut buf)?;
for idx in 0..(segment_count - 1) {
let segment_len =
u32::from_le_bytes(buf[(idx * 4)..(idx + 1) * 4].try_into().unwrap()) as usize;
segment_slices.push((total_words, total_words + segment_len));
total_words += segment_len;
}
} else {
let mut segment_sizes = vec![0u8; (segment_count & !1) * 4];
read.read_exact(&mut segment_sizes[..])?;
for idx in 0..(segment_count - 1) {
let segment_len =
u32::from_le_bytes(segment_sizes[(idx * 4)..(idx + 1) * 4].try_into().unwrap())
as usize;
segment_slices.push((total_words, total_words + segment_len));
total_words += segment_len;
}
}
}
// Don't accept a message which the receiver couldn't possibly traverse without hitting the
// traversal limit. Without this check, a malicious client could transmit a very large segment
// size to make the receiver allocate excessive space and possibly crash.
if total_words as u64 > options.traversal_limit_in_words {
return Err(Error::failed(format!(
"Message has {} words, which is too large. To increase the limit on the \
receiving end, see capnp::message::ReaderOptions.",
total_words
)));
}
Ok((total_words, segment_slices))
}
You can’t perform that action at this time.