From e0b97c53cb32ea47d4a054b2a0ac73003bbef875 Mon Sep 17 00:00:00 2001 From: floor-licker Date: Fri, 30 Jan 2026 20:22:45 -0500 Subject: [PATCH] feat(ws): add tape-based book update processor --- src/book.rs | 76 ++++++++++++++ src/lib.rs | 2 + src/ws_hot_path.rs | 195 ++++++++++++++++++++++++++++++++++++ tests/no_alloc_hot_paths.rs | 56 ++++++++++- 4 files changed, 328 insertions(+), 1 deletion(-) create mode 100644 src/ws_hot_path.rs diff --git a/src/book.rs b/src/book.rs index d04a568..37bdf35 100644 --- a/src/book.rs +++ b/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 { + 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::::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( + &self, + token_id: &str, + f: impl FnOnce(&mut OrderBook) -> Result, + ) -> Result { + 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<()> { diff --git a/src/lib.rs b/src/lib.rs index dcabb27..d7c7a8c 100644 --- a/src/lib.rs +++ b/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)] diff --git a/src/ws_hot_path.rs b/src/ws_hot_path.rs new file mode 100644 index 0000000..7eaa98c --- /dev/null +++ b/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>, +} + +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 { + 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 { + 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 { + 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 { + 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 { + value + .as_u64() + .or_else(|| value.into_string().and_then(|s| s.parse::().ok())) +} + +fn apply_levels<'tape, 'input>( + book: &mut crate::book::OrderBook, + side: Side, + levels: simd_json::tape::Array<'tape, 'input>, +) -> Result { + 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) +} diff --git a/tests/no_alloc_hot_paths.rs b/tests/no_alloc_hot_paths.rs index 51b3bc9..c56a310 100644 --- a/tests/no_alloc_hot_paths.rs +++ b/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(); +}