Merge pull request #60 from floor-licker/fix/atomic-book-snapshots

fix: apply book snapshots atomically

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<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/decode.rs

@@ -705,4 +705,17 @@ mod tests {
         let results: Vec<serde_json::Value> = decoder.parse_json_stream(data).unwrap();
         assert_eq!(results.len(), 2);
     }
+
+    #[test]
+    fn stream_book_message_requires_bids_and_asks() {
+        let missing_asks = br#"{"event_type":"book","asset_id":"test_asset_id","market":"0xabc","timestamp":1000,"bids":[]}"#;
+        assert!(parse_stream_messages_bytes(missing_asks).is_err());
+
+        let missing_bids = br#"{"event_type":"book","asset_id":"test_asset_id","market":"0xabc","timestamp":1000,"asks":[]}"#;
+        assert!(parse_stream_messages_bytes(missing_bids).is_err());
+
+        let empty_sides = br#"{"event_type":"book","asset_id":"test_asset_id","market":"0xabc","timestamp":1000,"bids":[],"asks":[]}"#;
+        let messages = parse_stream_messages_bytes(empty_sides).unwrap();
+        assert_eq!(messages.len(), 1);
+    }
 }

src/types.rs

@@ -909,15 +909,9 @@ pub struct BookUpdate {
     pub market: String,
     #[serde(deserialize_with = "crate::decode::deserializers::number_from_string")]
     pub timestamp: u64,
-    #[serde(
-        default,
-        deserialize_with = "crate::decode::deserializers::vec_from_null"
-    )]
+    #[serde(deserialize_with = "crate::decode::deserializers::vec_from_null")]
     pub bids: Vec<OrderSummary>,
-    #[serde(
-        default,
-        deserialize_with = "crate::decode::deserializers::vec_from_null"
-    )]
+    #[serde(deserialize_with = "crate::decode::deserializers::vec_from_null")]
     pub asks: Vec<OrderSummary>,
     #[serde(default)]
     pub hash: Option<String>,

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<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());
@@ -127,25 +131,36 @@ fn process_stream_object<'tape, 'input>(
     let timestamp = parse_u64(timestamp_value)
         .ok_or_else(|| PolyfillError::parse("Invalid timestamp", None))?;
 
-    let bids = obj.get("bids").and_then(|v| v.as_array());
-    let asks = obj.get("asks").and_then(|v| v.as_array());
+    let bids = obj
+        .get("bids")
+        .ok_or_else(|| PolyfillError::parse("Missing bids", None))?
+        .as_array()
+        .ok_or_else(|| PolyfillError::parse("Invalid bids", None))?;
+    let asks = obj
+        .get("asks")
+        .ok_or_else(|| PolyfillError::parse("Missing asks", None))?
+        .as_array()
+        .ok_or_else(|| PolyfillError::parse("Invalid asks", None))?;
 
-    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)?;
-        }
-        if let Some(asks) = asks {
-            applied += apply_levels(book, Side::SELL, asks)?;
-        }
+        collect_levels(Side::BUY, bids, parsed_levels)?;
+        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 +174,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 +197,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 +292,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 +339,76 @@ 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));
+    }
+
+    #[test]
+    fn processor_missing_side_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 missing_asks = br#"{"event_type":"book","asset_id":"test_asset_id","market":"0xabc","timestamp":1001,"bids":[{"price":"0.5100","size":"11.0000"}]}"#.to_vec();
+        assert!(processor
+            .process_bytes(missing_asks.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.asks[0].price, dec!(0.60));
+    }
 }