diff --git a/src/lib.rs b/src/lib.rs index d7c7a8c..070ee7e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -145,7 +145,7 @@ pub use crate::errors::{PolyfillError, Result}; 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::stream::{MarketStream, StreamManager, WebSocketBookApplier, WebSocketStream}; pub use crate::ws_hot_path::{WsBookApplyStats, WsBookUpdateProcessor}; // Re-export utilities diff --git a/src/stream.rs b/src/stream.rs index 0ba3b9a..a3e2613 100644 --- a/src/stream.rs +++ b/src/stream.rs @@ -5,6 +5,7 @@ use crate::errors::{PolyfillError, Result}; use crate::types::*; +use crate::ws_hot_path::{WsBookApplyStats, WsBookUpdateProcessor}; use chrono::Utc; use futures::{SinkExt, Stream, StreamExt}; use serde_json::Value; @@ -378,6 +379,116 @@ impl WebSocketStream { } } +/// WebSocket stream wrapper that applies `book` updates directly into an [`crate::book::OrderBookManager`]. +/// +/// This bypasses `StreamMessage` decoding (serde/DOM parsing) for the `book` hot path by using +/// [`WsBookUpdateProcessor`]. Non-`book` WS payloads are ignored. +/// +/// Note: the underlying WS transport may still allocate when producing `Message::Text(String)`. +pub struct WebSocketBookApplier<'a> { + stream: WebSocketStream, + books: &'a crate::book::OrderBookManager, + processor: WsBookUpdateProcessor, +} + +impl WebSocketStream { + /// Convert this stream into a book-applier stream. + /// + /// The caller is expected to "warm up" the [`crate::book::OrderBookManager`] by creating books for all + /// subscribed asset IDs ahead of time. Missing books are treated as an error. + pub fn into_book_applier<'a>( + mut self, + books: &'a crate::book::OrderBookManager, + processor: WsBookUpdateProcessor, + ) -> WebSocketBookApplier<'a> { + // Drop any pre-parsed messages to avoid mixing the two streaming modes. + self.pending.clear(); + WebSocketBookApplier { + stream: self, + books, + processor, + } + } +} + +impl<'a> WebSocketBookApplier<'a> { + /// Access the underlying WebSocket stream (e.g., for subscribe/unsubscribe calls). + pub fn stream_mut(&mut self) -> &mut WebSocketStream { + &mut self.stream + } + + /// Current WebSocket connection stats. + pub fn stream_stats(&self) -> StreamStats { + self.stream.stats.clone() + } + + /// Access the hot-path processor (e.g., to reuse it across connections). + pub fn processor_mut(&mut self) -> &mut WsBookUpdateProcessor { + &mut self.processor + } + + /// Apply a single WS text payload (useful for custom transports and for testing). + pub fn apply_text_message(&mut self, text: String) -> Result { + let stats = self.processor.process_text(text, self.books)?; + self.stream.stats.messages_received += 1; + self.stream.stats.last_message_time = Some(Utc::now()); + Ok(stats) + } +} + +impl<'a> Stream for WebSocketBookApplier<'a> { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + loop { + let Some(connection) = &mut self.stream.connection else { + return Poll::Ready(None); + }; + + match connection.poll_next_unpin(cx) { + Poll::Pending => return Poll::Pending, + Poll::Ready(Some(Ok(msg))) => match msg { + tokio_tungstenite::tungstenite::Message::Text(text) => { + match self.apply_text_message(text) { + Ok(stats) => { + if stats.book_messages == 0 { + continue; + } + return Poll::Ready(Some(Ok(stats))); + }, + Err(e) => { + self.stream.stats.errors += 1; + return Poll::Ready(Some(Err(e))); + }, + } + }, + tokio_tungstenite::tungstenite::Message::Close(_) => { + info!("WebSocket connection closed by server"); + self.stream.connection = None; + return Poll::Ready(None); + }, + tokio_tungstenite::tungstenite::Message::Ping(_) => { + // Best-effort: tokio-tungstenite/tungstenite may handle pings internally. + continue; + }, + tokio_tungstenite::tungstenite::Message::Pong(_) => continue, + tokio_tungstenite::tungstenite::Message::Binary(_) => continue, + tokio_tungstenite::tungstenite::Message::Frame(_) => continue, + }, + Poll::Ready(Some(Err(e))) => { + error!("WebSocket error: {}", e); + self.stream.stats.errors += 1; + return Poll::Ready(Some(Err(e.into()))); + }, + Poll::Ready(None) => { + info!("WebSocket stream ended"); + return Poll::Ready(None); + }, + } + } + } +} + impl Stream for WebSocketStream { type Item = Result; @@ -585,6 +696,8 @@ impl StreamManager { #[cfg(test)] mod tests { use super::*; + use rust_decimal::Decimal; + use std::str::FromStr; #[test] fn test_mock_stream() { @@ -626,4 +739,27 @@ mod tests { }); assert!(manager.broadcast_message(message).is_ok()); } + + #[test] + fn test_websocket_book_applier_apply_text_message_updates_book() { + let books = crate::book::OrderBookManager::new(64); + let _ = books.get_or_create_book("12345").unwrap(); + + let processor = WsBookUpdateProcessor::new(1024); + let stream = WebSocketStream::new("wss://example.com/ws"); + let mut applier = stream.into_book_applier(&books, processor); + + let msg = r#"{"event_type":"book","asset_id":"12345","timestamp":1,"bids":[{"price":"0.75","size":"10"}],"asks":[{"price":"0.76","size":"5"}]}"#.to_string(); + let stats = applier.apply_text_message(msg).unwrap(); + assert_eq!(stats.book_messages, 1); + assert_eq!(stats.book_levels_applied, 2); + + let snapshot = books.get_book("12345").unwrap(); + assert_eq!(snapshot.bids.len(), 1); + assert_eq!(snapshot.asks.len(), 1); + assert_eq!(snapshot.bids[0].price, Decimal::from_str("0.75").unwrap()); + assert_eq!(snapshot.bids[0].size, Decimal::from_str("10").unwrap()); + assert_eq!(snapshot.asks[0].price, Decimal::from_str("0.76").unwrap()); + assert_eq!(snapshot.asks[0].size, Decimal::from_str("5").unwrap()); + } } diff --git a/tests/no_alloc_hot_paths.rs b/tests/no_alloc_hot_paths.rs index c56a310..0e47168 100644 --- a/tests/no_alloc_hot_paths.rs +++ b/tests/no_alloc_hot_paths.rs @@ -5,7 +5,9 @@ use std::hash::{Hash, Hasher}; use std::str::FromStr; use chrono::Utc; -use polyfill_rs::{book::OrderBookManager, OrderBookImpl, Side, WsBookUpdateProcessor}; +use polyfill_rs::{ + book::OrderBookManager, OrderBookImpl, Side, WebSocketStream, WsBookUpdateProcessor, +}; use rust_decimal::Decimal; thread_local! { @@ -269,3 +271,53 @@ fn no_alloc_ws_book_update_processor_apply_existing_levels() { .unwrap(); guard.assert_no_allocations(); } + +#[test] +fn no_alloc_websocket_book_applier_apply_text_message_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 processor = WsBookUpdateProcessor::new(1024); + let stream = WebSocketStream::new("wss://example.com/ws"); + let mut applier = stream.into_book_applier(&manager, processor); + + // Warm up simd-json buffers/tape outside the guarded section. + let 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\"}}]}}" + ); + applier.apply_text_message(warmup_msg).unwrap(); + + let 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\"}}]}}" + ); + + // Warm up TLS access before measuring (defensive). + let _ = allocation_count(); + + let guard = NoAllocGuard::new(); + applier.apply_text_message(msg).unwrap(); + guard.assert_no_allocations(); +}