use std::collections::BTreeMap;
use rust_decimal::Decimal;
use crate::{
http::{Orderbook, OrderbookLevel},
ws::{OrderbookDataMsg, OrderbookMsg},
};
type Side = BTreeMap<Decimal, Decimal>;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ApplyOutcome {
Buffered,
Synced,
Applied,
Ignored,
ResyncRequired,
}
#[derive(Debug, Default)]
pub struct OrderBookState {
inner: Inner,
}
#[derive(Debug)]
enum Inner {
Empty {
buffered: BTreeMap<i64, OrderbookDataMsg>,
},
Synced {
last_id: i64,
bids: Side,
asks: Side,
},
}
impl Default for Inner {
fn default() -> Self {
Self::Empty {
buffered: BTreeMap::new(),
}
}
}
impl OrderBookState {
pub fn new() -> Self {
Self::default()
}
pub fn apply_snapshot(&mut self, snapshot: Orderbook) -> ApplyOutcome {
self.set_snapshot(snapshot.update_id, snapshot.bids, snapshot.asks)
}
pub fn apply_diff(&mut self, diff: OrderbookDataMsg) -> ApplyOutcome {
match &mut self.inner {
Inner::Empty { buffered } => {
buffered.insert(diff.update_id, diff);
ApplyOutcome::Buffered
}
Inner::Synced {
last_id,
bids,
asks,
} => {
if diff.update_id == 1 {
self.inner = Inner::default();
return ApplyOutcome::ResyncRequired;
}
if diff.update_id <= *last_id {
return ApplyOutcome::Ignored;
}
if diff.update_id != *last_id + 1 {
self.inner = Inner::Empty {
buffered: BTreeMap::from([(diff.update_id, diff)]),
};
return ApplyOutcome::ResyncRequired;
}
apply_levels(bids, &diff.bids);
apply_levels(asks, &diff.asks);
*last_id = diff.update_id;
ApplyOutcome::Applied
}
}
}
pub fn apply(&mut self, msg: OrderbookMsg) -> ApplyOutcome {
match msg {
OrderbookMsg::Snapshot { data, .. } => {
self.set_snapshot(data.update_id, data.bids, data.asks)
}
OrderbookMsg::Delta { data, .. } => self.apply_diff(data),
}
}
pub fn is_synced(&self) -> bool {
matches!(self.inner, Inner::Synced { .. })
}
pub fn last_update_id(&self) -> Option<i64> {
match &self.inner {
Inner::Synced { last_id, .. } => Some(*last_id),
Inner::Empty { .. } => None,
}
}
pub fn bids(&self) -> Option<&Side> {
match &self.inner {
Inner::Synced { bids, .. } => Some(bids),
Inner::Empty { .. } => None,
}
}
pub fn asks(&self) -> Option<&Side> {
match &self.inner {
Inner::Synced { asks, .. } => Some(asks),
Inner::Empty { .. } => None,
}
}
pub fn best_bid(&self) -> Option<(Decimal, Decimal)> {
self.bids()
.and_then(|bids| bids.iter().next_back())
.map(|(&price, &size)| (price, size))
}
pub fn best_ask(&self) -> Option<(Decimal, Decimal)> {
self.asks()
.and_then(|asks| asks.iter().next())
.map(|(&price, &size)| (price, size))
}
fn set_snapshot(
&mut self,
new_id: i64,
bids: Vec<OrderbookLevel>,
asks: Vec<OrderbookLevel>,
) -> ApplyOutcome {
if let Inner::Synced { last_id, .. } = &self.inner {
if new_id != 1 && new_id <= *last_id {
return ApplyOutcome::Ignored;
}
}
let mut bids: Side = bids.into_iter().map(|l| (l.price, l.size)).collect();
let mut asks: Side = asks.into_iter().map(|l| (l.price, l.size)).collect();
let buffered = match std::mem::take(&mut self.inner) {
Inner::Empty { buffered } => buffered,
Inner::Synced { .. } => BTreeMap::new(),
};
let mut last_id = new_id;
for (&id, diff) in buffered.range((new_id + 1)..) {
if id != last_id + 1 {
break;
}
apply_levels(&mut bids, &diff.bids);
apply_levels(&mut asks, &diff.asks);
last_id = id;
}
self.inner = Inner::Synced {
last_id,
bids,
asks,
};
ApplyOutcome::Synced
}
}
fn apply_levels(side: &mut Side, levels: &[OrderbookLevel]) {
for level in levels {
if level.size.is_zero() {
side.remove(&level.price);
} else {
side.insert(level.price, level.size);
}
}
}
#[cfg(test)]
mod tests {
use rust_decimal::dec;
use crate::{DepthLevel, Topic};
use super::*;
fn level(price: Decimal, size: Decimal) -> OrderbookLevel {
OrderbookLevel { price, size }
}
fn snapshot(update_id: i64, bids: Vec<OrderbookLevel>, asks: Vec<OrderbookLevel>) -> Orderbook {
Orderbook {
symbol: "BTCUSDT".into(),
bids,
asks,
ts: 0,
update_id,
seq: update_id,
cts: None,
}
}
fn diff(
update_id: i64,
bids: Vec<OrderbookLevel>,
asks: Vec<OrderbookLevel>,
) -> OrderbookDataMsg {
OrderbookDataMsg {
symbol: "BTCUSDT".into(),
bids,
asks,
update_id,
seq: update_id,
}
}
#[test]
fn snapshot_then_contiguous_diff_is_applied() {
let mut book = OrderBookState::new();
assert_eq!(
book.apply_snapshot(snapshot(
1,
vec![level(dec!(100), dec!(1))],
vec![level(dec!(101), dec!(2))]
)),
ApplyOutcome::Synced
);
assert!(book.is_synced());
assert_eq!(
book.apply_diff(diff(2, vec![level(dec!(100), dec!(1.5))], vec![])),
ApplyOutcome::Applied
);
assert_eq!(book.last_update_id(), Some(2));
assert_eq!(book.best_bid(), Some((dec!(100), dec!(1.5))));
assert_eq!(book.best_ask(), Some((dec!(101), dec!(2))));
}
#[test]
fn diff_buffered_before_snapshot_bridges_on_arrival() {
let mut book = OrderBookState::new();
assert_eq!(
book.apply_diff(diff(2, vec![level(dec!(100), dec!(1.5))], vec![])),
ApplyOutcome::Buffered
);
assert!(!book.is_synced());
assert_eq!(
book.apply_snapshot(snapshot(
1,
vec![level(dec!(100), dec!(1))],
vec![level(dec!(101), dec!(2))]
)),
ApplyOutcome::Synced
);
assert_eq!(book.last_update_id(), Some(2));
assert_eq!(book.best_bid(), Some((dec!(100), dec!(1.5))));
}
#[test]
fn stale_snapshot_is_ignored() {
let mut book = OrderBookState::new();
book.apply_snapshot(snapshot(5, vec![], vec![]));
assert_eq!(
book.apply_snapshot(snapshot(3, vec![], vec![])),
ApplyOutcome::Ignored
);
assert_eq!(book.last_update_id(), Some(5));
}
#[test]
fn stale_diff_is_ignored() {
let mut book = OrderBookState::new();
book.apply_snapshot(snapshot(5, vec![], vec![]));
assert_eq!(
book.apply_diff(diff(4, vec![], vec![])),
ApplyOutcome::Ignored
);
assert_eq!(book.last_update_id(), Some(5));
}
#[test]
fn gap_in_diff_stream_requires_resync_and_a_fresh_snapshot_bridges_it() {
let mut book = OrderBookState::new();
book.apply_snapshot(snapshot(1, vec![], vec![]));
assert_eq!(
book.apply_diff(diff(5, vec![level(dec!(100), dec!(1))], vec![])),
ApplyOutcome::ResyncRequired
);
assert!(!book.is_synced());
assert_eq!(
book.apply_snapshot(snapshot(4, vec![], vec![])),
ApplyOutcome::Synced
);
assert_eq!(book.last_update_id(), Some(5));
assert_eq!(book.best_bid(), Some((dec!(100), dec!(1))));
}
#[test]
fn restart_signal_resets_the_book_even_if_id_is_smaller() {
let mut book = OrderBookState::new();
book.apply_snapshot(snapshot(50, vec![level(dec!(100), dec!(1))], vec![]));
assert!(book.is_synced());
assert_eq!(
book.apply_snapshot(snapshot(1, vec![level(dec!(200), dec!(9))], vec![])),
ApplyOutcome::Synced
);
assert_eq!(book.last_update_id(), Some(1));
assert_eq!(book.best_bid(), Some((dec!(200), dec!(9))));
}
#[test]
fn zero_size_level_removes_the_price_from_the_book() {
let mut book = OrderBookState::new();
book.apply_snapshot(snapshot(
1,
vec![level(dec!(100), dec!(1)), level(dec!(99), dec!(2))],
vec![],
));
assert_eq!(
book.apply_diff(diff(2, vec![level(dec!(100), dec!(0))], vec![])),
ApplyOutcome::Applied
);
assert_eq!(book.best_bid(), Some((dec!(99), dec!(2))));
}
#[test]
fn apply_dispatches_snapshot_and_delta_messages() {
let mut book = OrderBookState::new();
let topic = Topic::Orderbook {
symbol: "BTCUSDT".into(),
depth: DepthLevel::Level50,
};
let snapshot_msg = OrderbookMsg::Snapshot {
topic: topic.clone(),
ts: 0,
data: diff(
1,
vec![level(dec!(100), dec!(1))],
vec![level(dec!(101), dec!(2))],
),
cts: 0,
};
assert_eq!(book.apply(snapshot_msg), ApplyOutcome::Synced);
let delta_msg = OrderbookMsg::Delta {
topic,
ts: 0,
data: diff(2, vec![level(dec!(100), dec!(1.5))], vec![]),
cts: 0,
};
assert_eq!(book.apply(delta_msg), ApplyOutcome::Applied);
assert_eq!(book.best_bid(), Some((dec!(100), dec!(1.5))));
}
}