diff --git a/src/book.rs b/src/book.rs index 3de5c6c..d04a568 100644 --- a/src/book.rs +++ b/src/book.rs @@ -781,6 +781,30 @@ impl OrderBookManager { book.apply_delta(delta) } + /// Apply a WebSocket `book` update to a managed book. + /// + /// This is the preferred way to ingest `StreamMessage::Book` updates into + /// the in-memory order books (avoids rebuilding snapshots via per-level deltas). + pub fn apply_book_update(&self, update: &BookUpdate) -> Result<()> { + let mut books = self + .books + .write() + .map_err(|_| PolyfillError::internal_simple("Failed to acquire book lock"))?; + + if let Some(book) = books.get_mut(update.asset_id.as_str()) { + return book.apply_book_update(update); + } + + // First time we've seen this token; allocating the key and book is part of warmup. + let token_id = update.asset_id.clone(); + books.insert(token_id.clone(), OrderBook::new(token_id, self.max_depth)); + + books + .get_mut(update.asset_id.as_str()) + .ok_or_else(|| PolyfillError::internal_simple("Failed to insert order book"))? + .apply_book_update(update) + } + /// Get a book snapshot /// Returns a copy of the current book state that won't change pub fn get_book(&self, token_id: &str) -> Result { diff --git a/tests/no_alloc_hot_paths.rs b/tests/no_alloc_hot_paths.rs index 6e7465e..51b3bc9 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::{OrderBookImpl, Side}; +use polyfill_rs::{book::OrderBookManager, OrderBookImpl, Side}; use rust_decimal::Decimal; thread_local! { @@ -164,3 +164,54 @@ fn no_alloc_apply_book_update_existing_levels() { book.apply_book_update(&update).unwrap(); guard.assert_no_allocations(); } + +#[test] +fn no_alloc_book_manager_apply_book_update_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 update = polyfill_rs::types::BookUpdate { + asset_id: asset_id.to_string(), + market: "0xabc".to_string(), + timestamp: 10, + bids: vec![polyfill_rs::types::OrderSummary { + price: Decimal::from_str("0.75").unwrap(), + size: Decimal::from_str("200.0").unwrap(), + }], + asks: vec![polyfill_rs::types::OrderSummary { + price: Decimal::from_str("0.76").unwrap(), + size: Decimal::from_str("50.0").unwrap(), + }], + hash: None, + }; + + // Warm up TLS access before measuring (defensive). + let _ = allocation_count(); + + let guard = NoAllocGuard::new(); + manager.apply_book_update(&update).unwrap(); + guard.assert_no_allocations(); +}