feat(ws): add tape-based book update processor

src/book.rs

@@ -396,6 +396,58 @@ 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> {
+        if asset_id != self.token_id {
+            return Err(PolyfillError::validation("Token ID mismatch"));
+        }
+
+        if timestamp <= self.sequence {
+            return Ok(false);
+        }
+
+        self.sequence = timestamp;
+        self.timestamp =
+            chrono::DateTime::<Utc>::from_timestamp(timestamp as i64, 0).unwrap_or_else(Utc::now);
+
+        Ok(true)
+    }
+
+    /// Apply a single WS `book` level (already converted to internal fixed-point).
+    ///
+    /// Note: Insertions of new price levels may allocate (BTreeMap node growth). 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(
+        &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"));
+            }
+        }
+
+        match side {
+            Side::BUY => self.apply_bid_delta_fast(price_ticks, size_units),
+            Side::SELL => self.apply_ask_delta_fast(price_ticks, size_units),
+        }
+
+        Ok(())
+    }
+
+    /// Finish applying a WS `book` update.
+    pub(crate) fn finish_ws_book_update(&mut self) {
+        self.trim_depth();
+    }
+
     /// Apply a WebSocket `book` update for this token.
     ///
     /// The official Polymarket CLOB WebSocket `book` event contains batches of
@@ -761,6 +813,30 @@ impl OrderBookManager {
         }
     }
 
+    /// Execute a closure with mutable access to a managed book.
+    ///
+    /// This is useful for hot-path update ingestion where you want to avoid allocating
+    /// intermediate update structs (e.g., applying WS updates directly).
+    pub fn with_book_mut<R>(
+        &self,
+        token_id: &str,
+        f: impl FnOnce(&mut OrderBook) -> Result<R>,
+    ) -> Result<R> {
+        let mut books = self
+            .books
+            .write()
+            .map_err(|_| PolyfillError::internal_simple("Failed to acquire book lock"))?;
+
+        let book = books.get_mut(token_id).ok_or_else(|| {
+            PolyfillError::market_data(
+                format!("No book found for token: {}", token_id),
+                crate::errors::MarketDataErrorKind::TokenNotFound,
+            )
+        })?;
+
+        f(book)
+    }
+
     /// Update a book with a delta
     /// This is called when we receive real-time updates from the exchange
     pub fn apply_delta(&self, delta: OrderDelta) -> Result<()> {

src/lib.rs

@@ -146,6 +146,7 @@ pub use crate::book::{OrderBook as OrderBookImpl, OrderBookManager};
 pub use crate::decode::Decoder;
 pub use crate::fill::{FillEngine, FillResult};
 pub use crate::stream::{MarketStream, StreamManager, WebSocketStream};
+pub use crate::ws_hot_path::{WsBookApplyStats, WsBookUpdateProcessor};
 
 // Re-export utilities
 pub use crate::utils::{crypto, math, rate_limit, retry, time, url};
@@ -165,6 +166,7 @@ pub mod orders;
 pub mod stream;
 pub mod types;
 pub mod utils;
+pub mod ws_hot_path;
 
 // Benchmarks
 #[cfg(test)]

src/ws_hot_path.rs

@@ -0,0 +1,195 @@
+//! Zero-allocation-ish WebSocket hot-path processing.
+//!
+//! This module is focused on the "decode + apply" path for WS `book` events:
+//! after warmup, processing a message should not perform heap allocations.
+//!
+//! Important: using the current tokio-tungstenite transport, the *network layer*
+//! 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::errors::{PolyfillError, Result};
+use crate::types::{decimal_to_price, decimal_to_qty, Side};
+use rust_decimal::Decimal;
+use simd_json::prelude::*;
+use std::str::FromStr;
+
+/// Summary of what happened while processing a WS payload.
+#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
+pub struct WsBookApplyStats {
+    pub book_messages: usize,
+    pub book_levels_applied: usize,
+}
+
+/// In-place WS `book` message processor built on `simd-json`'s tape API.
+///
+/// This avoids building a DOM (which allocates for arrays/objects) by decoding into a
+/// reusable tape, then traversing it to extract the fields needed for order book updates.
+pub struct WsBookUpdateProcessor {
+    buffers: simd_json::Buffers,
+    tape: Option<simd_json::Tape<'static>>,
+}
+
+impl WsBookUpdateProcessor {
+    /// Create a new processor.
+    ///
+    /// `input_len_hint` should be set to the typical WS message size to reduce warmup reallocs.
+    pub fn new(input_len_hint: usize) -> Self {
+        Self {
+            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()),
+        }
+    }
+
+    /// Process a WS payload in-place (bytes will be mutated by the JSON parser).
+    pub fn process_bytes(
+        &mut self,
+        bytes: &mut [u8],
+        books: &OrderBookManager,
+    ) -> Result<WsBookApplyStats> {
+        let mut tape = self
+            .tape
+            .take()
+            .expect("WsBookUpdateProcessor tape must be present")
+            .reset();
+
+        simd_json::fill_tape(bytes, &mut self.buffers, &mut tape).map_err(|e| {
+            PolyfillError::parse("Failed to parse WebSocket JSON", Some(Box::new(e)))
+        })?;
+
+        let root = tape.as_value();
+        let stats = process_root_value(root, books)?;
+
+        // Reset the tape to detach lifetimes and keep capacity for reuse.
+        self.tape = Some(tape.reset());
+        Ok(stats)
+    }
+
+    /// Convenience: process an owned text message without allocating an additional buffer.
+    pub fn process_text(
+        &mut self,
+        text: String,
+        books: &OrderBookManager,
+    ) -> Result<WsBookApplyStats> {
+        let mut bytes = text.into_bytes();
+        self.process_bytes(bytes.as_mut_slice(), books)
+    }
+}
+
+fn process_root_value<'tape, 'input>(
+    value: simd_json::tape::Value<'tape, 'input>,
+    books: &OrderBookManager,
+) -> Result<WsBookApplyStats> {
+    if let Some(obj) = value.as_object() {
+        return process_stream_object(obj, books);
+    }
+
+    let Some(arr) = value.as_array() else {
+        return Ok(WsBookApplyStats::default());
+    };
+
+    let mut total = WsBookApplyStats::default();
+    for elem in arr.iter() {
+        let Some(obj) = elem.as_object() else {
+            continue;
+        };
+        let stats = process_stream_object(obj, books)?;
+        total.book_messages += stats.book_messages;
+        total.book_levels_applied += stats.book_levels_applied;
+    }
+
+    Ok(total)
+}
+
+fn process_stream_object<'tape, 'input>(
+    obj: simd_json::tape::Object<'tape, 'input>,
+    books: &OrderBookManager,
+) -> Result<WsBookApplyStats> {
+    let Some(event_type) = obj.get("event_type").and_then(|v| v.into_string()) else {
+        return Ok(WsBookApplyStats::default());
+    };
+
+    if event_type != "book" {
+        return Ok(WsBookApplyStats::default());
+    }
+
+    let asset_id = obj
+        .get("asset_id")
+        .and_then(|v| v.into_string())
+        .ok_or_else(|| PolyfillError::parse("Missing asset_id", None))?;
+
+    let timestamp_value = obj
+        .get("timestamp")
+        .ok_or_else(|| PolyfillError::parse("Missing timestamp", None))?;
+    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 levels_applied = books.with_book_mut(asset_id, |book| {
+        if !book.begin_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)?;
+        }
+
+        book.finish_ws_book_update();
+        Ok(applied)
+    })?;
+
+    Ok(WsBookApplyStats {
+        book_messages: 1,
+        book_levels_applied: levels_applied,
+    })
+}
+
+fn parse_u64<'tape, 'input>(value: simd_json::tape::Value<'tape, 'input>) -> Option<u64> {
+    value
+        .as_u64()
+        .or_else(|| value.into_string().and_then(|s| s.parse::<u64>().ok()))
+}
+
+fn apply_levels<'tape, 'input>(
+    book: &mut crate::book::OrderBook,
+    side: Side,
+    levels: simd_json::tape::Array<'tape, 'input>,
+) -> Result<usize> {
+    let mut applied = 0usize;
+    for level in levels.iter() {
+        let Some(obj) = level.as_object() else {
+            continue;
+        };
+
+        let price_str = obj
+            .get("price")
+            .and_then(|v| v.into_string())
+            .ok_or_else(|| PolyfillError::parse("Missing price", None))?;
+        let size_str = obj
+            .get("size")
+            .and_then(|v| v.into_string())
+            .ok_or_else(|| PolyfillError::parse("Missing size", None))?;
+
+        let price_decimal =
+            Decimal::from_str(price_str).map_err(|_| PolyfillError::validation("Invalid price"))?;
+        let size_decimal =
+            Decimal::from_str(size_str).map_err(|_| PolyfillError::validation("Invalid size"))?;
+
+        let price_ticks = decimal_to_price(price_decimal)
+            .map_err(|_| PolyfillError::validation("Invalid price"))?;
+        let size_units =
+            decimal_to_qty(size_decimal).map_err(|_| PolyfillError::validation("Invalid size"))?;
+
+        book.apply_ws_book_level_fast(side, price_ticks, size_units)?;
+        applied += 1;
+    }
+
+    Ok(applied)
+}

tests/no_alloc_hot_paths.rs

@@ -5,7 +5,7 @@ use std::hash::{Hash, Hasher};
 use std::str::FromStr;
 
 use chrono::Utc;
-use polyfill_rs::{book::OrderBookManager, OrderBookImpl, Side};
+use polyfill_rs::{book::OrderBookManager, OrderBookImpl, Side, WsBookUpdateProcessor};
 use rust_decimal::Decimal;
 
 thread_local! {
@@ -215,3 +215,57 @@ fn no_alloc_book_manager_apply_book_update_existing_levels() {
     manager.apply_book_update(&update).unwrap();
     guard.assert_no_allocations();
 }
+
+#[test]
+fn no_alloc_ws_book_update_processor_apply_existing_levels() {
+    let asset_id = "test_asset_id";
+    let manager = OrderBookManager::new(100);
+    manager.get_or_create_book(asset_id).unwrap();
+
+    // Warm up the internal book with initial levels (allocations allowed).
+    manager
+        .apply_delta(polyfill_rs::types::OrderDelta {
+            token_id: asset_id.to_string(),
+            timestamp: chrono::Utc::now(),
+            side: Side::BUY,
+            price: Decimal::from_str("0.75").unwrap(),
+            size: Decimal::from_str("100.0").unwrap(),
+            sequence: 1,
+        })
+        .unwrap();
+    manager
+        .apply_delta(polyfill_rs::types::OrderDelta {
+            token_id: asset_id.to_string(),
+            timestamp: chrono::Utc::now(),
+            side: Side::SELL,
+            price: Decimal::from_str("0.76").unwrap(),
+            size: Decimal::from_str("100.0").unwrap(),
+            sequence: 2,
+        })
+        .unwrap();
+
+    let mut processor = WsBookUpdateProcessor::new(1024);
+
+    // Warm up simd-json buffers/tape outside the guarded section.
+    let mut warmup_msg = format!(
+        "{{\"event_type\":\"book\",\"asset_id\":\"{asset_id}\",\"market\":\"0xabc\",\"timestamp\":10,\"bids\":[{{\"price\":\"0.75\",\"size\":\"200.0\"}}],\"asks\":[{{\"price\":\"0.76\",\"size\":\"50.0\"}}]}}"
+    )
+    .into_bytes();
+    processor
+        .process_bytes(warmup_msg.as_mut_slice(), &manager)
+        .unwrap();
+
+    let mut msg = format!(
+        "{{\"event_type\":\"book\",\"asset_id\":\"{asset_id}\",\"market\":\"0xabc\",\"timestamp\":11,\"bids\":[{{\"price\":\"0.75\",\"size\":\"150.0\"}}],\"asks\":[{{\"price\":\"0.76\",\"size\":\"75.0\"}}]}}"
+    )
+    .into_bytes();
+
+    // Warm up TLS access before measuring (defensive).
+    let _ = allocation_count();
+
+    let guard = NoAllocGuard::new();
+    processor
+        .process_bytes(msg.as_mut_slice(), &manager)
+        .unwrap();
+    guard.assert_no_allocations();
+}