gavindiaz/polyfill-rs
1
0
Fork 0
mirror of https://github.com/floor-licker/polyfill-rs.git synced 2026-07-27 20:47:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: apply book snapshots atomically

Browse Source
This commit is contained in:
floor-licker
2026-06-22 19:34:37 -04:00
parent c821262af5
commit 43b71fbd1f
2 changed files with 221 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/book.rs
+152 -58
View File
@@ -16,6 +16,13 @@ struct StoredLevel {
generation: u64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct ParsedBookLevel {
pub side: Side,
pub price_ticks: Price,
pub size_units: Qty,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum BookSideKind {
Bid,
@@ -519,14 +526,12 @@ impl OrderBook {
Ok(())
}
/// Begin applying a WebSocket `book` update (hot-path oriented).
///
/// This is intended for in-place WS processing where we *stream* levels out of a decoded
/// message, without constructing intermediate `BookUpdate` structs.
///
/// Returns `Ok(true)` if the update should be applied, or `Ok(false)` if the update is stale
/// and should be skipped.
pub(crate) fn begin_ws_book_update(&mut self, asset_id: &str, timestamp: u64) -> Result<bool> {
/// Return whether a WebSocket `book` snapshot should be applied.
pub(crate) fn should_apply_ws_book_update(
&self,
asset_id: &str,
timestamp: u64,
) -> Result<bool> {
if asset_id != self.token_id {
return Err(PolyfillError::validation("Token ID mismatch"));
}
@@ -535,39 +540,27 @@ impl OrderBook {
return Ok(false);
}
self.sequence = timestamp;
self.timestamp = chrono::DateTime::<Utc>::from_timestamp_millis(timestamp as i64)
.unwrap_or_else(Utc::now);
self.begin_snapshot();
Ok(true)
}
/// Apply a single WS `book` level (already converted to internal fixed-point).
/// Atomically apply a WebSocket `book` snapshot.
///
/// Note: Insertions of new price levels may allocate (Vec growth/shifting). In a strict
/// zero-alloc hot path, all expected levels must be warmed up ahead of time.
pub(crate) fn apply_ws_book_level_fast(
/// The caller should parse levels into `ParsedBookLevel`s before calling this method. This
/// method validates tick alignment before mutating sequence/generation or book levels, so a
/// malformed snapshot leaves the existing book unchanged.
pub(crate) fn apply_ws_book_snapshot_fast(
&mut self,
side: Side,
price_ticks: Price,
size_units: Qty,
) -> Result<()> {
if let Some(tick_size_ticks) = self.tick_size_ticks {
if tick_size_ticks > 0 && !price_ticks.is_multiple_of(tick_size_ticks) {
return Err(PolyfillError::validation("Price not aligned to tick size"));
}
asset_id: &str,
timestamp: u64,
levels: &[ParsedBookLevel],
) -> Result<bool> {
if !self.should_apply_ws_book_update(asset_id, timestamp)? {
return Ok(false);
}
self.apply_snapshot_level(side, price_ticks, size_units);
self.apply_validated_snapshot(timestamp, levels)?;
Ok(())
}
/// Finish applying a WS `book` snapshot.
pub(crate) fn finish_ws_book_update(&mut self) {
self.finish_snapshot();
self.trim_depth();
Ok(true)
}
/// Apply a WebSocket `book` update for this token.
@@ -592,40 +585,38 @@ impl OrderBook {
return Ok(());
}
// Validate the whole snapshot before mutating this book. A malformed level must not leave
// behind a partial generation or an advanced sequence number.
for level in &update.bids {
let parsed = self.parse_snapshot_summary(Side::BUY, level)?;
self.validate_snapshot_level(parsed)?;
}
for level in &update.asks {
let parsed = self.parse_snapshot_summary(Side::SELL, level)?;
self.validate_snapshot_level(parsed)?;
}
self.sequence = update.timestamp;
self.timestamp = chrono::DateTime::<Utc>::from_timestamp_millis(update.timestamp as i64)
.unwrap_or_else(Utc::now);
self.begin_snapshot();
// Apply bids (BUY) and asks (SELL) as level upserts.
// Re-parse after validation to preserve the existing no-allocation behavior for
// `BookUpdate` snapshots. Decimal conversion is deterministic, so these conversions cannot
// fail after the validation pass above.
for level in &update.bids {
let price_ticks = decimal_to_price(level.price)
.map_err(|_| PolyfillError::validation("Invalid price"))?;
let size_units = decimal_to_qty(level.size)
.map_err(|_| PolyfillError::validation("Invalid size"))?;
if let Some(tick_size_ticks) = self.tick_size_ticks {
if tick_size_ticks > 0 && !price_ticks.is_multiple_of(tick_size_ticks) {
return Err(PolyfillError::validation("Price not aligned to tick size"));
}
}
self.apply_snapshot_level(Side::BUY, price_ticks, size_units);
let parsed = self
.parse_snapshot_summary(Side::BUY, level)
.expect("book update bid level was validated before mutation");
self.apply_snapshot_level(parsed.side, parsed.price_ticks, parsed.size_units);
}
for level in &update.asks {
let price_ticks = decimal_to_price(level.price)
.map_err(|_| PolyfillError::validation("Invalid price"))?;
let size_units = decimal_to_qty(level.size)
.map_err(|_| PolyfillError::validation("Invalid size"))?;
if let Some(tick_size_ticks) = self.tick_size_ticks {
if tick_size_ticks > 0 && !price_ticks.is_multiple_of(tick_size_ticks) {
return Err(PolyfillError::validation("Price not aligned to tick size"));
}
}
self.apply_snapshot_level(Side::SELL, price_ticks, size_units);
let parsed = self
.parse_snapshot_summary(Side::SELL, level)
.expect("book update ask level was validated before mutation");
self.apply_snapshot_level(parsed.side, parsed.price_ticks, parsed.size_units);
}
self.finish_snapshot();
@@ -707,6 +698,55 @@ impl OrderBook {
self.snapshot_generation = self.snapshot_generation.wrapping_add(1);
}
#[inline]
fn parse_snapshot_summary(&self, side: Side, level: &OrderSummary) -> Result<ParsedBookLevel> {
let price_ticks = decimal_to_price(level.price)
.map_err(|_| PolyfillError::validation("Invalid price"))?;
let size_units =
decimal_to_qty(level.size).map_err(|_| PolyfillError::validation("Invalid size"))?;
Ok(ParsedBookLevel {
side,
price_ticks,
size_units,
})
}
#[inline]
fn validate_snapshot_level(&self, level: ParsedBookLevel) -> Result<()> {
if let Some(tick_size_ticks) = self.tick_size_ticks {
if tick_size_ticks > 0 && !level.price_ticks.is_multiple_of(tick_size_ticks) {
return Err(PolyfillError::validation("Price not aligned to tick size"));
}
}
Ok(())
}
fn apply_validated_snapshot(
&mut self,
timestamp: u64,
levels: &[ParsedBookLevel],
) -> Result<()> {
for &level in levels {
self.validate_snapshot_level(level)?;
}
self.sequence = timestamp;
self.timestamp = chrono::DateTime::<Utc>::from_timestamp_millis(timestamp as i64)
.unwrap_or_else(Utc::now);
self.begin_snapshot();
for &level in levels {
self.apply_snapshot_level(level.side, level.price_ticks, level.size_units);
}
self.finish_snapshot();
self.trim_depth();
Ok(())
}
#[inline]
fn apply_snapshot_level(&mut self, side: Side, price_ticks: Price, size_units: Qty) {
let generation = self.snapshot_generation;
@@ -1346,6 +1386,60 @@ mod tests {
assert_eq!(book.best_ask().unwrap().size, dec!(45));
}
#[test]
fn test_book_update_error_keeps_existing_snapshot() {
let mut book = OrderBook::new("test_token".to_string(), 10);
book.set_tick_size_ticks(100);
book.apply_book_update(&BookUpdate {
asset_id: "test_token".to_string(),
market: "0xabc".to_string(),
timestamp: 100,
bids: vec![OrderSummary {
price: dec!(0.50),
size: dec!(10),
}],
asks: vec![OrderSummary {
price: dec!(0.60),
size: dec!(20),
}],
hash: None,
})
.unwrap();
let err = book
.apply_book_update(&BookUpdate {
asset_id: "test_token".to_string(),
market: "0xabc".to_string(),
timestamp: 101,
bids: vec![
OrderSummary {
price: dec!(0.51),
size: dec!(11),
},
OrderSummary {
price: dec!(0.515),
size: dec!(12),
},
],
asks: vec![OrderSummary {
price: dec!(0.61),
size: dec!(21),
}],
hash: None,
})
.unwrap_err();
assert!(err.to_string().contains("Price not aligned"));
assert_eq!(book.sequence, 100);
assert_eq!(book.bids(None).len(), 1);
assert_eq!(book.asks(None).len(), 1);
assert_eq!(book.best_bid().unwrap().price, dec!(0.50));
assert_eq!(book.best_bid().unwrap().size, dec!(10));
assert_eq!(book.best_ask().unwrap().price, dec!(0.60));
assert_eq!(book.best_ask().unwrap().size, dec!(20));
}
#[test]
fn test_book_update_depth_keeps_best_prices_independent_of_payload_order() {
let mut book = OrderBook::new("test_token".to_string(), 2);
src/ws_hot_path.rs
+69 -15
View File
@@ -7,7 +7,7 @@
//! may still allocate when producing `Message::Text(String)`. This module aims to
//! make the *processing* layer allocation-free so we can enforce it with tests.
use crate::book::OrderBookManager;
use crate::book::{OrderBookManager, ParsedBookLevel};
use crate::errors::{PolyfillError, Result};
use crate::types::{Price, Qty, Side, MAX_PRICE_TICKS, MAX_QTY, MIN_PRICE_TICKS, SCALE_FACTOR};
use simd_json::prelude::*;
@@ -26,6 +26,7 @@ pub struct WsBookApplyStats {
pub struct WsBookUpdateProcessor {
buffers: simd_json::Buffers,
tape: Option<simd_json::Tape<'static>>,
parsed_levels: Vec<ParsedBookLevel>,
}
impl WsBookUpdateProcessor {
@@ -37,6 +38,7 @@ impl WsBookUpdateProcessor {
buffers: simd_json::Buffers::new(input_len_hint),
// Store an empty tape with a `'static` lifetime so we can reuse its allocation.
tape: Some(simd_json::Tape::null().reset()),
parsed_levels: Vec::with_capacity((input_len_hint / 32).max(8)),
}
}
@@ -55,7 +57,7 @@ impl WsBookUpdateProcessor {
let result = match simd_json::fill_tape(bytes, &mut self.buffers, &mut tape) {
Ok(()) => {
let root = tape.as_value();
process_root_value(root, books)
process_root_value(root, books, &mut self.parsed_levels)
},
Err(e) => Err(PolyfillError::parse(
"Failed to parse WebSocket JSON",
@@ -82,9 +84,10 @@ impl WsBookUpdateProcessor {
fn process_root_value<'tape, 'input>(
value: simd_json::tape::Value<'tape, 'input>,
books: &OrderBookManager,
parsed_levels: &mut Vec<ParsedBookLevel>,
) -> Result<WsBookApplyStats> {
if let Some(obj) = value.as_object() {
return process_stream_object(obj, books);
return process_stream_object(obj, books, parsed_levels);
}
let Some(arr) = value.as_array() else {
@@ -96,7 +99,7 @@ fn process_root_value<'tape, 'input>(
let Some(obj) = elem.as_object() else {
continue;
};
let stats = process_stream_object(obj, books)?;
let stats = process_stream_object(obj, books, parsed_levels)?;
total.book_messages += stats.book_messages;
total.book_levels_applied += stats.book_levels_applied;
}
@@ -107,6 +110,7 @@ fn process_root_value<'tape, 'input>(
fn process_stream_object<'tape, 'input>(
obj: simd_json::tape::Object<'tape, 'input>,
books: &OrderBookManager,
parsed_levels: &mut Vec<ParsedBookLevel>,
) -> Result<WsBookApplyStats> {
let Some(event_type) = obj.get("event_type").and_then(|v| v.into_string()) else {
return Ok(WsBookApplyStats::default());
@@ -130,22 +134,29 @@ fn process_stream_object<'tape, 'input>(
let bids = obj.get("bids").and_then(|v| v.as_array());
let asks = obj.get("asks").and_then(|v| v.as_array());
let levels_applied = books.with_book_mut(asset_id, |book| {
if !book.begin_ws_book_update(asset_id, timestamp)? {
let result = books.with_book_mut(asset_id, |book| {
parsed_levels.clear();
if !book.should_apply_ws_book_update(asset_id, timestamp)? {
return Ok(0);
}
let mut applied = 0usize;
if let Some(bids) = bids {
applied += apply_levels(book, Side::BUY, bids)?;
collect_levels(Side::BUY, bids, parsed_levels)?;
}
if let Some(asks) = asks {
applied += apply_levels(book, Side::SELL, asks)?;
collect_levels(Side::SELL, asks, parsed_levels)?;
}
book.finish_ws_book_update();
Ok(applied)
})?;
let parsed_count = parsed_levels.len();
if book.apply_ws_book_snapshot_fast(asset_id, timestamp, parsed_levels)? {
Ok(parsed_count)
} else {
Ok(0)
}
});
parsed_levels.clear();
let levels_applied = result?;
Ok(WsBookApplyStats {
book_messages: 1,
@@ -159,10 +170,10 @@ fn parse_u64<'tape, 'input>(value: simd_json::tape::Value<'tape, 'input>) -> Opt
.or_else(|| value.into_string().and_then(|s| s.parse::<u64>().ok()))
}
fn apply_levels<'tape, 'input>(
book: &mut crate::book::OrderBook,
fn collect_levels<'tape, 'input>(
side: Side,
levels: simd_json::tape::Array<'tape, 'input>,
parsed_levels: &mut Vec<ParsedBookLevel>,
) -> Result<usize> {
let mut applied = 0usize;
for level in levels.iter() {
@@ -182,7 +193,11 @@ fn apply_levels<'tape, 'input>(
let price_ticks = parse_price_ticks_4dp(price_str)?;
let size_units = parse_qty_scaled_4dp(size_str)?;
book.apply_ws_book_level_fast(side, price_ticks, size_units)?;
parsed_levels.push(ParsedBookLevel {
side,
price_ticks,
size_units,
});
applied += 1;
}
@@ -273,6 +288,8 @@ fn parse_scaled_4_u64(value: &str) -> Result<u64> {
#[cfg(test)]
mod tests {
use super::*;
use crate::types::{BookUpdate, OrderSummary};
use rust_decimal_macros::dec;
#[test]
fn fixed_point_parser_matches_expected_price_ticks() {
@@ -318,4 +335,41 @@ mod tests {
assert_eq!(stats.book_messages, 1);
assert_eq!(stats.book_levels_applied, 1);
}
#[test]
fn processor_error_keeps_existing_snapshot() {
let books = OrderBookManager::new(10);
books.get_or_create_book("test_asset_id").unwrap();
books
.apply_book_update(&BookUpdate {
asset_id: "test_asset_id".to_string(),
market: "0xabc".to_string(),
timestamp: 1000,
bids: vec![OrderSummary {
price: dec!(0.50),
size: dec!(10),
}],
asks: vec![OrderSummary {
price: dec!(0.60),
size: dec!(20),
}],
hash: None,
})
.unwrap();
let mut processor = WsBookUpdateProcessor::new(1024);
let mut invalid_snapshot = br#"{"event_type":"book","asset_id":"test_asset_id","market":"0xabc","timestamp":1001,"bids":[{"price":"0.5100","size":"11.0000"},{"price":"0.51001","size":"12.0000"}],"asks":[{"price":"0.6100","size":"21.0000"}]}"#.to_vec();
assert!(processor
.process_bytes(invalid_snapshot.as_mut_slice(), &books)
.is_err());
let snapshot = books.get_book("test_asset_id").unwrap();
assert_eq!(snapshot.sequence, 1000);
assert_eq!(snapshot.bids.len(), 1);
assert_eq!(snapshot.asks.len(), 1);
assert_eq!(snapshot.bids[0].price, dec!(0.50));
assert_eq!(snapshot.bids[0].size, dec!(10));
assert_eq!(snapshot.asks[0].price, dec!(0.60));
assert_eq!(snapshot.asks[0].size, dec!(20));
}
}