diff --git a/src/book.rs b/src/book.rs index a866f71..81ed0a4 100644 --- a/src/book.rs +++ b/src/book.rs @@ -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 { + /// Return whether a WebSocket `book` snapshot should be applied. + pub(crate) fn should_apply_ws_book_update( + &self, + asset_id: &str, + timestamp: u64, + ) -> Result { 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::::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 { + 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::::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 { + 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::::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); diff --git a/src/ws_hot_path.rs b/src/ws_hot_path.rs index ecaee94..494feb4 100644 --- a/src/ws_hot_path.rs +++ b/src/ws_hot_path.rs @@ -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>, + parsed_levels: Vec, } 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, ) -> Result { 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, ) -> Result { 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::().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, ) -> Result { 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 { #[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)); + } }