1use std::collections::BTreeMap;
20
21use rust_decimal::Decimal;
22
23use crate::{
24 http::{Orderbook, OrderbookLevel},
25 ws::{OrderbookDataMsg, OrderbookMsg},
26};
27
28type Side = BTreeMap<Decimal, Decimal>;
30
31#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub enum ApplyOutcome {
35 Buffered,
38 Synced,
42 Applied,
44 Ignored,
47 ResyncRequired,
52}
53
54#[derive(Debug, Default)]
64pub struct OrderBookState {
65 inner: Inner,
66}
67
68#[derive(Debug)]
69enum Inner {
70 Empty {
74 buffered: BTreeMap<i64, OrderbookDataMsg>,
75 },
76 Synced {
78 last_id: i64,
79 bids: Side,
80 asks: Side,
81 },
82}
83
84impl Default for Inner {
85 fn default() -> Self {
86 Self::Empty {
87 buffered: BTreeMap::new(),
88 }
89 }
90}
91
92impl OrderBookState {
93 pub fn new() -> Self {
95 Self::default()
96 }
97
98 pub fn apply_snapshot(&mut self, snapshot: Orderbook) -> ApplyOutcome {
100 self.set_snapshot(snapshot.update_id, snapshot.bids, snapshot.asks)
101 }
102
103 pub fn apply_diff(&mut self, diff: OrderbookDataMsg) -> ApplyOutcome {
110 match &mut self.inner {
111 Inner::Empty { buffered } => {
112 buffered.insert(diff.update_id, diff);
113 ApplyOutcome::Buffered
114 }
115 Inner::Synced {
116 last_id,
117 bids,
118 asks,
119 } => {
120 if diff.update_id == 1 {
124 self.inner = Inner::default();
125 return ApplyOutcome::ResyncRequired;
126 }
127 if diff.update_id <= *last_id {
128 return ApplyOutcome::Ignored;
129 }
130 if diff.update_id != *last_id + 1 {
131 self.inner = Inner::Empty {
135 buffered: BTreeMap::from([(diff.update_id, diff)]),
136 };
137 return ApplyOutcome::ResyncRequired;
138 }
139 apply_levels(bids, &diff.bids);
140 apply_levels(asks, &diff.asks);
141 *last_id = diff.update_id;
142 ApplyOutcome::Applied
143 }
144 }
145 }
146
147 pub fn apply(&mut self, msg: OrderbookMsg) -> ApplyOutcome {
151 match msg {
152 OrderbookMsg::Snapshot { data, .. } => {
153 self.set_snapshot(data.update_id, data.bids, data.asks)
154 }
155 OrderbookMsg::Delta { data, .. } => self.apply_diff(data),
156 }
157 }
158
159 pub fn is_synced(&self) -> bool {
161 matches!(self.inner, Inner::Synced { .. })
162 }
163
164 pub fn last_update_id(&self) -> Option<i64> {
166 match &self.inner {
167 Inner::Synced { last_id, .. } => Some(*last_id),
168 Inner::Empty { .. } => None,
169 }
170 }
171
172 pub fn bids(&self) -> Option<&Side> {
174 match &self.inner {
175 Inner::Synced { bids, .. } => Some(bids),
176 Inner::Empty { .. } => None,
177 }
178 }
179
180 pub fn asks(&self) -> Option<&Side> {
182 match &self.inner {
183 Inner::Synced { asks, .. } => Some(asks),
184 Inner::Empty { .. } => None,
185 }
186 }
187
188 pub fn best_bid(&self) -> Option<(Decimal, Decimal)> {
190 self.bids()
191 .and_then(|bids| bids.iter().next_back())
192 .map(|(&price, &size)| (price, size))
193 }
194
195 pub fn best_ask(&self) -> Option<(Decimal, Decimal)> {
197 self.asks()
198 .and_then(|asks| asks.iter().next())
199 .map(|(&price, &size)| (price, size))
200 }
201
202 fn set_snapshot(
206 &mut self,
207 new_id: i64,
208 bids: Vec<OrderbookLevel>,
209 asks: Vec<OrderbookLevel>,
210 ) -> ApplyOutcome {
211 if let Inner::Synced { last_id, .. } = &self.inner {
212 if new_id != 1 && new_id <= *last_id {
216 return ApplyOutcome::Ignored;
217 }
218 }
219
220 let mut bids: Side = bids.into_iter().map(|l| (l.price, l.size)).collect();
221 let mut asks: Side = asks.into_iter().map(|l| (l.price, l.size)).collect();
222
223 let buffered = match std::mem::take(&mut self.inner) {
224 Inner::Empty { buffered } => buffered,
225 Inner::Synced { .. } => BTreeMap::new(),
226 };
227
228 let mut last_id = new_id;
233 for (&id, diff) in buffered.range((new_id + 1)..) {
234 if id != last_id + 1 {
235 break;
236 }
237 apply_levels(&mut bids, &diff.bids);
238 apply_levels(&mut asks, &diff.asks);
239 last_id = id;
240 }
241
242 self.inner = Inner::Synced {
243 last_id,
244 bids,
245 asks,
246 };
247 ApplyOutcome::Synced
248 }
249}
250
251fn apply_levels(side: &mut Side, levels: &[OrderbookLevel]) {
254 for level in levels {
255 if level.size.is_zero() {
256 side.remove(&level.price);
257 } else {
258 side.insert(level.price, level.size);
259 }
260 }
261}
262
263#[cfg(test)]
264mod tests {
265 use rust_decimal::dec;
266
267 use crate::{DepthLevel, Topic};
268
269 use super::*;
270
271 fn level(price: Decimal, size: Decimal) -> OrderbookLevel {
272 OrderbookLevel { price, size }
273 }
274
275 fn snapshot(update_id: i64, bids: Vec<OrderbookLevel>, asks: Vec<OrderbookLevel>) -> Orderbook {
276 Orderbook {
277 symbol: "BTCUSDT".into(),
278 bids,
279 asks,
280 ts: 0,
281 update_id,
282 seq: update_id,
283 cts: None,
284 }
285 }
286
287 fn diff(
288 update_id: i64,
289 bids: Vec<OrderbookLevel>,
290 asks: Vec<OrderbookLevel>,
291 ) -> OrderbookDataMsg {
292 OrderbookDataMsg {
293 symbol: "BTCUSDT".into(),
294 bids,
295 asks,
296 update_id,
297 seq: update_id,
298 }
299 }
300
301 #[test]
302 fn snapshot_then_contiguous_diff_is_applied() {
303 let mut book = OrderBookState::new();
304 assert_eq!(
305 book.apply_snapshot(snapshot(
306 1,
307 vec![level(dec!(100), dec!(1))],
308 vec![level(dec!(101), dec!(2))]
309 )),
310 ApplyOutcome::Synced
311 );
312 assert!(book.is_synced());
313
314 assert_eq!(
315 book.apply_diff(diff(2, vec![level(dec!(100), dec!(1.5))], vec![])),
316 ApplyOutcome::Applied
317 );
318 assert_eq!(book.last_update_id(), Some(2));
319 assert_eq!(book.best_bid(), Some((dec!(100), dec!(1.5))));
320 assert_eq!(book.best_ask(), Some((dec!(101), dec!(2))));
321 }
322
323 #[test]
324 fn diff_buffered_before_snapshot_bridges_on_arrival() {
325 let mut book = OrderBookState::new();
326 assert_eq!(
327 book.apply_diff(diff(2, vec![level(dec!(100), dec!(1.5))], vec![])),
328 ApplyOutcome::Buffered
329 );
330 assert!(!book.is_synced());
331
332 assert_eq!(
333 book.apply_snapshot(snapshot(
334 1,
335 vec![level(dec!(100), dec!(1))],
336 vec![level(dec!(101), dec!(2))]
337 )),
338 ApplyOutcome::Synced
339 );
340 assert_eq!(book.last_update_id(), Some(2));
341 assert_eq!(book.best_bid(), Some((dec!(100), dec!(1.5))));
342 }
343
344 #[test]
345 fn stale_snapshot_is_ignored() {
346 let mut book = OrderBookState::new();
347 book.apply_snapshot(snapshot(5, vec![], vec![]));
348 assert_eq!(
349 book.apply_snapshot(snapshot(3, vec![], vec![])),
350 ApplyOutcome::Ignored
351 );
352 assert_eq!(book.last_update_id(), Some(5));
353 }
354
355 #[test]
356 fn stale_diff_is_ignored() {
357 let mut book = OrderBookState::new();
358 book.apply_snapshot(snapshot(5, vec![], vec![]));
359 assert_eq!(
360 book.apply_diff(diff(4, vec![], vec![])),
361 ApplyOutcome::Ignored
362 );
363 assert_eq!(book.last_update_id(), Some(5));
364 }
365
366 #[test]
367 fn gap_in_diff_stream_requires_resync_and_a_fresh_snapshot_bridges_it() {
368 let mut book = OrderBookState::new();
369 book.apply_snapshot(snapshot(1, vec![], vec![]));
370
371 assert_eq!(
372 book.apply_diff(diff(5, vec![level(dec!(100), dec!(1))], vec![])),
373 ApplyOutcome::ResyncRequired
374 );
375 assert!(!book.is_synced());
376
377 assert_eq!(
379 book.apply_snapshot(snapshot(4, vec![], vec![])),
380 ApplyOutcome::Synced
381 );
382 assert_eq!(book.last_update_id(), Some(5));
383 assert_eq!(book.best_bid(), Some((dec!(100), dec!(1))));
384 }
385
386 #[test]
387 fn restart_signal_resets_the_book_even_if_id_is_smaller() {
388 let mut book = OrderBookState::new();
389 book.apply_snapshot(snapshot(50, vec![level(dec!(100), dec!(1))], vec![]));
390 assert!(book.is_synced());
391
392 assert_eq!(
393 book.apply_snapshot(snapshot(1, vec![level(dec!(200), dec!(9))], vec![])),
394 ApplyOutcome::Synced
395 );
396 assert_eq!(book.last_update_id(), Some(1));
397 assert_eq!(book.best_bid(), Some((dec!(200), dec!(9))));
398 }
399
400 #[test]
401 fn zero_size_level_removes_the_price_from_the_book() {
402 let mut book = OrderBookState::new();
403 book.apply_snapshot(snapshot(
404 1,
405 vec![level(dec!(100), dec!(1)), level(dec!(99), dec!(2))],
406 vec![],
407 ));
408
409 assert_eq!(
410 book.apply_diff(diff(2, vec![level(dec!(100), dec!(0))], vec![])),
411 ApplyOutcome::Applied
412 );
413 assert_eq!(book.best_bid(), Some((dec!(99), dec!(2))));
414 }
415
416 #[test]
417 fn apply_dispatches_snapshot_and_delta_messages() {
418 let mut book = OrderBookState::new();
419 let topic = Topic::Orderbook {
420 symbol: "BTCUSDT".into(),
421 depth: DepthLevel::Level50,
422 };
423
424 let snapshot_msg = OrderbookMsg::Snapshot {
425 topic: topic.clone(),
426 ts: 0,
427 data: diff(
428 1,
429 vec![level(dec!(100), dec!(1))],
430 vec![level(dec!(101), dec!(2))],
431 ),
432 cts: 0,
433 };
434 assert_eq!(book.apply(snapshot_msg), ApplyOutcome::Synced);
435
436 let delta_msg = OrderbookMsg::Delta {
437 topic,
438 ts: 0,
439 data: diff(2, vec![level(dec!(100), dec!(1.5))], vec![]),
440 cts: 0,
441 };
442 assert_eq!(book.apply(delta_msg), ApplyOutcome::Applied);
443 assert_eq!(book.best_bid(), Some((dec!(100), dec!(1.5))));
444 }
445}