Skip to content
Permalink
Branch: master
Find file Copy path
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
154 lines (134 sloc) 6.25 KB
use std::convert::TryInto;
use std::io::{BufRead, Write};
use std::mem::size_of;
use crate::iex::{IexMessage, IexPayload};
use crate::marketdata_generated::md_shootout;
use crate::{RunnerDeserialize, RunnerSerialize, StreamVec, Summarizer};
pub struct FlatbuffersWriter<'a> {
builder: flatbuffers::FlatBufferBuilder<'a>,
message_buffer: Vec<flatbuffers::WIPOffset<md_shootout::Message<'a>>>,
}
impl<'a> FlatbuffersWriter<'a> {
pub fn new() -> FlatbuffersWriter<'static> {
FlatbuffersWriter {
builder: flatbuffers::FlatBufferBuilder::new(),
message_buffer: Vec::new(),
}
}
}
impl<'a> RunnerSerialize for FlatbuffersWriter<'a> {
fn serialize(&mut self, payload: &IexPayload, output: &mut Vec<u8>) {
// Because FlatBuffers can't handle nested vectors (specifically, we can't track
// both the variable-length vector of messages, and the variable-length strings
// within those messages), we have to cache the messages as they get built
// so they can be added all at once later.
for iex_msg in &payload.messages {
let msg_args = match iex_msg {
IexMessage::TradeReport(tr) => {
// The `Args` objects used are wrappers over an underlying `Builder`.
// We trust release builds to optimize out the wrapper.
let trade = md_shootout::Trade::create(
&mut self.builder,
&md_shootout::TradeArgs {
price: tr.price,
size_: tr.size,
},
);
let sym_str = self.builder.create_string(crate::parse_symbol(&tr.symbol));
Some(md_shootout::MessageArgs {
ts_nanos: tr.timestamp,
symbol: Some(sym_str),
body_type: md_shootout::MessageBody::Trade,
body: Some(trade.as_union_value()),
})
}
IexMessage::PriceLevelUpdate(plu) => {
let level_update = md_shootout::LevelUpdate::create(
&mut self.builder,
&md_shootout::LevelUpdateArgs {
price: plu.price,
size_: plu.size,
flags: plu.event_flags,
side: if plu.msg_type == 0x38 {
md_shootout::Side::Buy
} else {
md_shootout::Side::Sell
},
},
);
let sym_str = self.builder.create_string(crate::parse_symbol(&plu.symbol));
Some(md_shootout::MessageArgs {
ts_nanos: plu.timestamp,
symbol: Some(sym_str),
body_type: md_shootout::MessageBody::LevelUpdate,
body: Some(level_update.as_union_value()),
})
}
_ => None,
};
msg_args.map(|a| {
let msg = md_shootout::Message::create(&mut self.builder, &a);
self.message_buffer.push(msg);
});
}
let messages = self.builder.create_vector(&self.message_buffer[..]);
// Now that we've finished building all the messages, time to set up the final buffer
let mut multimsg_builder = md_shootout::MultiMessageBuilder::new(&mut self.builder);
multimsg_builder.add_seq_no(payload.first_seq_no);
multimsg_builder.add_messages(messages);
let multimsg = multimsg_builder.finish();
// IMPORTANT NOTE: If you just `finish`, Flatbuffers has no idea where
// an object ends in memory. To support streaming reads, you *must*
// use `finish_size_prefixed`. This adds a LE u32 to the front of the payload.
self.builder.finish_size_prefixed(multimsg, None);
output.write(self.builder.finished_data()).unwrap();
self.builder.reset();
self.message_buffer.clear();
}
}
pub struct FlatbuffersReader;
impl FlatbuffersReader {
pub fn new() -> FlatbuffersReader {
FlatbuffersReader {}
}
}
impl RunnerDeserialize for FlatbuffersReader {
fn deserialize<'a>(&self, buf: &'a mut StreamVec, stats: &mut Summarizer) -> Result<(), ()> {
// Flatbuffers has kinda ad-hoc support for streaming: https://github.com/google/flatbuffers/issues/3898
// Essentially, you can write an optional `u32` value to the front of each message
// (`finish_size_prefixed` above) to figure out how long that message actually is.
// Ultimately, end-users are responsible for all buffer management, "reading" is just
// a view over the underlying buffer.
let data = buf.fill_buf().unwrap();
if data.len() == 0 {
return Err(());
}
let msg_len_buf: [u8; 4] = data[..size_of::<u32>()].try_into().unwrap();
let msg_len = u32::from_le_bytes(msg_len_buf) as usize;
let multimsg = flatbuffers::get_size_prefixed_root::<md_shootout::MultiMessage>(data);
let msg_vec = match multimsg.messages() {
Some(m) => m,
None => panic!("Couldn't find messages"),
};
for i in 0..msg_vec.len() {
let msg: md_shootout::Message = msg_vec.get(i);
match msg.body_type() {
md_shootout::MessageBody::Trade => {
let trade = msg.body_as_trade().unwrap();
stats.append_trade_volume(msg.symbol().unwrap(), trade.size_().into());
}
md_shootout::MessageBody::LevelUpdate => {
let lu = msg.body_as_level_update().unwrap();
let is_bid = match lu.side() {
md_shootout::Side::Buy => true,
_ => false,
};
stats.update_quote_prices(msg.symbol().unwrap(), lu.price(), is_bid);
}
md_shootout::MessageBody::NONE => panic!("Unrecognized message type"),
}
}
buf.consume(msg_len + size_of::<u32>());
Ok(())
}
}
You can’t perform that action at this time.