polyfill_rs/
stream.rs

1//! Async streaming functionality for Polymarket client
2//!
3//! This module provides high-performance streaming capabilities for
4//! real-time market data and order updates.
5
6use crate::errors::{PolyfillError, Result};
7use crate::types::*;
8use crate::ws_hot_path::{WsBookApplyStats, WsBookUpdateProcessor};
9use chrono::Utc;
10use futures::{SinkExt, Stream, StreamExt};
11use parking_lot::Mutex;
12use serde_json::Value;
13use std::collections::VecDeque;
14use std::pin::Pin;
15use std::task::{Context, Poll};
16use tokio::sync::mpsc;
17use tracing::{debug, error, info, warn};
18
19/// Trait for market data streams
20pub trait MarketStream: Stream<Item = Result<StreamMessage>> + Send + Sync {
21    /// Subscribe to market data for specific tokens
22    fn subscribe(&mut self, subscription: Subscription) -> Result<()>;
23
24    /// Unsubscribe from market data
25    fn unsubscribe(&mut self, token_ids: &[String]) -> Result<()>;
26
27    /// Check if the stream is connected
28    fn is_connected(&self) -> bool;
29
30    /// Get connection statistics
31    fn get_stats(&self) -> StreamStats;
32}
33
34/// WebSocket-based market stream implementation
35#[derive(Debug)]
36#[allow(dead_code)]
37pub struct WebSocketStream {
38    /// WebSocket connection
39    connection: Option<
40        tokio_tungstenite::WebSocketStream<
41            tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
42        >,
43    >,
44    /// URL for the WebSocket connection
45    url: String,
46    /// Authentication credentials
47    auth: Option<WssAuth>,
48    /// Current subscriptions
49    subscriptions: Vec<WssSubscription>,
50    /// Parsed messages awaiting delivery to the caller.
51    ///
52    /// This replaces an internal unbounded channel to avoid per-message
53    /// allocations in the buffering layer and to enforce a bounded backlog.
54    pending: VecDeque<StreamMessage>,
55    pending_capacity: usize,
56    /// Connection statistics
57    stats: StreamStats,
58    /// Reconnection configuration
59    reconnect_config: ReconnectConfig,
60}
61
62/// Stream statistics
63#[derive(Debug, Clone)]
64pub struct StreamStats {
65    pub messages_received: u64,
66    pub messages_sent: u64,
67    pub errors: u64,
68    pub dropped_messages: u64,
69    pub last_message_time: Option<chrono::DateTime<Utc>>,
70    pub connection_uptime: std::time::Duration,
71    pub reconnect_count: u32,
72}
73
74/// Reconnection configuration
75#[derive(Debug, Clone)]
76pub struct ReconnectConfig {
77    pub max_retries: u32,
78    pub base_delay: std::time::Duration,
79    pub max_delay: std::time::Duration,
80    pub backoff_multiplier: f64,
81}
82
83impl Default for ReconnectConfig {
84    fn default() -> Self {
85        Self {
86            max_retries: 5,
87            base_delay: std::time::Duration::from_secs(1),
88            max_delay: std::time::Duration::from_secs(60),
89            backoff_multiplier: 2.0,
90        }
91    }
92}
93
94impl WebSocketStream {
95    /// Create a new WebSocket stream
96    pub fn new(url: &str) -> Self {
97        let pending_capacity = 1024;
98
99        Self {
100            connection: None,
101            url: url.to_string(),
102            auth: None,
103            subscriptions: Vec::new(),
104            pending: VecDeque::with_capacity(pending_capacity),
105            pending_capacity,
106            stats: StreamStats {
107                messages_received: 0,
108                messages_sent: 0,
109                errors: 0,
110                dropped_messages: 0,
111                last_message_time: None,
112                connection_uptime: std::time::Duration::ZERO,
113                reconnect_count: 0,
114            },
115            reconnect_config: ReconnectConfig::default(),
116        }
117    }
118
119    fn enqueue(&mut self, message: StreamMessage) {
120        if self.pending.len() >= self.pending_capacity {
121            let _ = self.pending.pop_front();
122            self.stats.dropped_messages += 1;
123        }
124        self.pending.push_back(message);
125    }
126
127    /// Set authentication credentials
128    pub fn with_auth(mut self, auth: WssAuth) -> Self {
129        self.auth = Some(auth);
130        self
131    }
132
133    /// Connect to the WebSocket
134    async fn connect(&mut self) -> Result<()> {
135        let (ws_stream, _) = tokio_tungstenite::connect_async(&self.url)
136            .await
137            .map_err(|e| {
138                PolyfillError::stream(
139                    format!("WebSocket connection failed: {}", e),
140                    crate::errors::StreamErrorKind::ConnectionFailed,
141                )
142            })?;
143
144        self.connection = Some(ws_stream);
145        info!("Connected to WebSocket stream at {}", self.url);
146        Ok(())
147    }
148
149    /// Send a message to the WebSocket
150    async fn send_message(&mut self, message: Value) -> Result<()> {
151        if let Some(connection) = &mut self.connection {
152            let text = serde_json::to_string(&message).map_err(|e| {
153                PolyfillError::parse(format!("Failed to serialize message: {}", e), None)
154            })?;
155
156            let ws_message = tokio_tungstenite::tungstenite::Message::Text(text);
157            connection.send(ws_message).await.map_err(|e| {
158                PolyfillError::stream(
159                    format!("Failed to send message: {}", e),
160                    crate::errors::StreamErrorKind::MessageCorrupted,
161                )
162            })?;
163
164            self.stats.messages_sent += 1;
165        }
166
167        Ok(())
168    }
169
170    /// Subscribe to market data using official Polymarket WebSocket API
171    pub async fn subscribe_async(&mut self, subscription: WssSubscription) -> Result<()> {
172        // Ensure connection
173        if self.connection.is_none() {
174            self.connect().await?;
175        }
176
177        // Send subscription message in the format expected by Polymarket
178        // The subscription struct will serialize correctly with proper field names
179        let message = serde_json::to_value(&subscription).map_err(|e| {
180            PolyfillError::parse(format!("Failed to serialize subscription: {}", e), None)
181        })?;
182
183        self.send_message(message).await?;
184        self.subscriptions.push(subscription.clone());
185
186        info!("Subscribed to {} channel", subscription.channel_type);
187        Ok(())
188    }
189
190    /// Subscribe to user channel (orders and trades)
191    pub async fn subscribe_user_channel(&mut self, markets: Vec<String>) -> Result<()> {
192        let auth = self
193            .auth
194            .as_ref()
195            .ok_or_else(|| PolyfillError::auth("No authentication provided for WebSocket"))?
196            .clone();
197
198        let subscription = WssSubscription {
199            channel_type: "user".to_string(),
200            operation: Some("subscribe".to_string()),
201            markets,
202            asset_ids: Vec::new(),
203            initial_dump: Some(true),
204            custom_feature_enabled: None,
205            auth: Some(auth),
206        };
207
208        self.subscribe_async(subscription).await
209    }
210
211    /// Subscribe to market channel (order book and trades)
212    /// Market subscriptions do not require authentication
213    pub async fn subscribe_market_channel(&mut self, asset_ids: Vec<String>) -> Result<()> {
214        let subscription = WssSubscription {
215            channel_type: "market".to_string(),
216            operation: Some("subscribe".to_string()),
217            markets: Vec::new(),
218            asset_ids,
219            initial_dump: Some(true),
220            custom_feature_enabled: None,
221            auth: None,
222        };
223
224        self.subscribe_async(subscription).await
225    }
226
227    /// Subscribe to market channel with custom features enabled
228    /// Custom features include: best_bid_ask, new_market, market_resolved events
229    pub async fn subscribe_market_channel_with_features(
230        &mut self,
231        asset_ids: Vec<String>,
232    ) -> Result<()> {
233        let subscription = WssSubscription {
234            channel_type: "market".to_string(),
235            operation: Some("subscribe".to_string()),
236            markets: Vec::new(),
237            asset_ids,
238            initial_dump: Some(true),
239            custom_feature_enabled: Some(true),
240            auth: None,
241        };
242
243        self.subscribe_async(subscription).await
244    }
245
246    /// Unsubscribe from market channel
247    pub async fn unsubscribe_market_channel(&mut self, asset_ids: Vec<String>) -> Result<()> {
248        let subscription = WssSubscription {
249            channel_type: "market".to_string(),
250            operation: Some("unsubscribe".to_string()),
251            markets: Vec::new(),
252            asset_ids,
253            initial_dump: None,
254            custom_feature_enabled: None,
255            auth: None,
256        };
257
258        self.subscribe_async(subscription).await
259    }
260
261    /// Unsubscribe from user channel
262    pub async fn unsubscribe_user_channel(&mut self, markets: Vec<String>) -> Result<()> {
263        let auth = self
264            .auth
265            .as_ref()
266            .ok_or_else(|| PolyfillError::auth("No authentication provided for WebSocket"))?
267            .clone();
268
269        let subscription = WssSubscription {
270            channel_type: "user".to_string(),
271            operation: Some("unsubscribe".to_string()),
272            markets,
273            asset_ids: Vec::new(),
274            initial_dump: None,
275            custom_feature_enabled: None,
276            auth: Some(auth),
277        };
278
279        self.subscribe_async(subscription).await
280    }
281
282    /// Handle incoming WebSocket messages
283    #[allow(dead_code)]
284    async fn handle_message(
285        &mut self,
286        message: tokio_tungstenite::tungstenite::Message,
287    ) -> Result<()> {
288        match message {
289            tokio_tungstenite::tungstenite::Message::Text(text) => {
290                debug!("Received WebSocket message: {}", text);
291
292                // Parse the message according to Polymarket's `event_type` format
293                let stream_messages = crate::decode::parse_stream_messages(&text)?;
294                for stream_message in stream_messages {
295                    self.enqueue(stream_message);
296                }
297
298                self.stats.messages_received += 1;
299                self.stats.last_message_time = Some(Utc::now());
300            },
301            tokio_tungstenite::tungstenite::Message::Close(_) => {
302                info!("WebSocket connection closed by server");
303                self.connection = None;
304            },
305            tokio_tungstenite::tungstenite::Message::Ping(data) => {
306                // Respond with pong
307                if let Some(connection) = &mut self.connection {
308                    let pong = tokio_tungstenite::tungstenite::Message::Pong(data);
309                    if let Err(e) = connection.send(pong).await {
310                        error!("Failed to send pong: {}", e);
311                    }
312                }
313            },
314            tokio_tungstenite::tungstenite::Message::Pong(_) => {
315                // Handle pong if needed
316                debug!("Received pong");
317            },
318            tokio_tungstenite::tungstenite::Message::Binary(_) => {
319                warn!("Received binary message (not supported)");
320            },
321            tokio_tungstenite::tungstenite::Message::Frame(_) => {
322                warn!("Received raw frame (not supported)");
323            },
324        }
325
326        Ok(())
327    }
328
329    /// Parse Polymarket WebSocket message(s) in `event_type` format.
330    #[allow(dead_code)]
331    fn parse_polymarket_messages(&self, text: &str) -> Result<Vec<StreamMessage>> {
332        crate::decode::parse_stream_messages(text)
333    }
334
335    /// Reconnect with exponential backoff
336    #[allow(dead_code)]
337    async fn reconnect(&mut self) -> Result<()> {
338        let mut delay = self.reconnect_config.base_delay;
339        let mut retries = 0;
340
341        while retries < self.reconnect_config.max_retries {
342            warn!("Attempting to reconnect (attempt {})", retries + 1);
343
344            match self.connect().await {
345                Ok(()) => {
346                    info!("Successfully reconnected");
347                    self.stats.reconnect_count += 1;
348
349                    // Resubscribe to all previous subscriptions
350                    let subscriptions = self.subscriptions.clone();
351                    for subscription in subscriptions {
352                        self.send_message(serde_json::to_value(subscription)?)
353                            .await?;
354                    }
355
356                    return Ok(());
357                },
358                Err(e) => {
359                    error!("Reconnection attempt {} failed: {}", retries + 1, e);
360                    retries += 1;
361
362                    if retries < self.reconnect_config.max_retries {
363                        tokio::time::sleep(delay).await;
364                        delay = std::cmp::min(
365                            delay.mul_f64(self.reconnect_config.backoff_multiplier),
366                            self.reconnect_config.max_delay,
367                        );
368                    }
369                },
370            }
371        }
372
373        Err(PolyfillError::stream(
374            format!(
375                "Failed to reconnect after {} attempts",
376                self.reconnect_config.max_retries
377            ),
378            crate::errors::StreamErrorKind::ConnectionFailed,
379        ))
380    }
381}
382
383/// WebSocket stream wrapper that applies `book` updates directly into an [`crate::book::OrderBookManager`].
384///
385/// This bypasses `StreamMessage` decoding (serde/DOM parsing) for the `book` hot path by using
386/// [`WsBookUpdateProcessor`]. Non-`book` WS payloads are ignored.
387///
388/// Note: the underlying WS transport may still allocate when producing `Message::Text(String)`.
389pub struct WebSocketBookApplier<'a> {
390    stream: WebSocketStream,
391    books: &'a crate::book::OrderBookManager,
392    processor: WsBookUpdateProcessor,
393}
394
395impl WebSocketStream {
396    /// Convert this stream into a book-applier stream.
397    ///
398    /// The caller is expected to "warm up" the [`crate::book::OrderBookManager`] by creating books for all
399    /// subscribed asset IDs ahead of time. Missing books are treated as an error.
400    pub fn into_book_applier<'a>(
401        mut self,
402        books: &'a crate::book::OrderBookManager,
403        processor: WsBookUpdateProcessor,
404    ) -> WebSocketBookApplier<'a> {
405        // Drop any pre-parsed messages to avoid mixing the two streaming modes.
406        self.pending.clear();
407        WebSocketBookApplier {
408            stream: self,
409            books,
410            processor,
411        }
412    }
413}
414
415impl<'a> WebSocketBookApplier<'a> {
416    /// Access the underlying WebSocket stream (e.g., for subscribe/unsubscribe calls).
417    pub fn stream_mut(&mut self) -> &mut WebSocketStream {
418        &mut self.stream
419    }
420
421    /// Current WebSocket connection stats.
422    pub fn stream_stats(&self) -> StreamStats {
423        self.stream.stats.clone()
424    }
425
426    /// Access the hot-path processor (e.g., to reuse it across connections).
427    pub fn processor_mut(&mut self) -> &mut WsBookUpdateProcessor {
428        &mut self.processor
429    }
430
431    /// Apply a single WS text payload (useful for custom transports and for testing).
432    ///
433    /// This convenience method consumes an owned `String`, so it may release that
434    /// buffer after processing. Use [`Self::apply_bytes_message`] for the
435    /// allocation-sensitive path when the caller owns a reusable mutable buffer.
436    pub fn apply_text_message(&mut self, text: String) -> Result<WsBookApplyStats> {
437        let stats = self.processor.process_text(text, self.books)?;
438        self.stream.stats.messages_received += 1;
439        self.stream.stats.last_message_time = Some(Utc::now());
440        Ok(stats)
441    }
442
443    /// Apply a single WS payload from a caller-owned mutable byte buffer.
444    ///
445    /// The buffer is mutated by `simd-json`. After processor warmup, this is the
446    /// allocation-sensitive book-applier entry point because no owned message buffer
447    /// is created or dropped by this method.
448    pub fn apply_bytes_message(&mut self, bytes: &mut [u8]) -> Result<WsBookApplyStats> {
449        let stats = self.processor.process_bytes(bytes, self.books)?;
450        self.stream.stats.messages_received += 1;
451        self.stream.stats.last_message_time = Some(Utc::now());
452        Ok(stats)
453    }
454}
455
456impl<'a> Stream for WebSocketBookApplier<'a> {
457    type Item = Result<WsBookApplyStats>;
458
459    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
460        loop {
461            let Some(connection) = &mut self.stream.connection else {
462                return Poll::Ready(None);
463            };
464
465            match connection.poll_next_unpin(cx) {
466                Poll::Pending => return Poll::Pending,
467                Poll::Ready(Some(Ok(msg))) => match msg {
468                    tokio_tungstenite::tungstenite::Message::Text(text) => {
469                        match self.apply_text_message(text) {
470                            Ok(stats) => {
471                                if stats.book_messages == 0 {
472                                    continue;
473                                }
474                                return Poll::Ready(Some(Ok(stats)));
475                            },
476                            Err(e) => {
477                                self.stream.stats.errors += 1;
478                                return Poll::Ready(Some(Err(e)));
479                            },
480                        }
481                    },
482                    tokio_tungstenite::tungstenite::Message::Close(_) => {
483                        info!("WebSocket connection closed by server");
484                        self.stream.connection = None;
485                        return Poll::Ready(None);
486                    },
487                    tokio_tungstenite::tungstenite::Message::Ping(_) => {
488                        // Best-effort: tokio-tungstenite/tungstenite may handle pings internally.
489                        continue;
490                    },
491                    tokio_tungstenite::tungstenite::Message::Pong(_) => continue,
492                    tokio_tungstenite::tungstenite::Message::Binary(_) => continue,
493                    tokio_tungstenite::tungstenite::Message::Frame(_) => continue,
494                },
495                Poll::Ready(Some(Err(e))) => {
496                    error!("WebSocket error: {}", e);
497                    self.stream.stats.errors += 1;
498                    return Poll::Ready(Some(Err(e.into())));
499                },
500                Poll::Ready(None) => {
501                    info!("WebSocket stream ended");
502                    return Poll::Ready(None);
503                },
504            }
505        }
506    }
507}
508
509impl Stream for WebSocketStream {
510    type Item = Result<StreamMessage>;
511
512    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
513        loop {
514            if let Some(message) = self.pending.pop_front() {
515                return Poll::Ready(Some(Ok(message)));
516            }
517
518            let Some(connection) = &mut self.connection else {
519                return Poll::Ready(None);
520            };
521
522            match connection.poll_next_unpin(cx) {
523                Poll::Pending => return Poll::Pending,
524                Poll::Ready(Some(Ok(ws_message))) => match ws_message {
525                    tokio_tungstenite::tungstenite::Message::Text(text) => {
526                        match crate::decode::parse_stream_messages(&text) {
527                            Ok(messages) => {
528                                let mut iter = messages.into_iter();
529                                let Some(first) = iter.next() else {
530                                    continue;
531                                };
532
533                                for msg in iter {
534                                    self.enqueue(msg);
535                                }
536                                self.stats.messages_received += 1;
537                                self.stats.last_message_time = Some(Utc::now());
538                                return Poll::Ready(Some(Ok(first)));
539                            },
540                            Err(e) => {
541                                self.stats.errors += 1;
542                                return Poll::Ready(Some(Err(e)));
543                            },
544                        }
545                    },
546                    tokio_tungstenite::tungstenite::Message::Close(_) => {
547                        info!("WebSocket connection closed by server");
548                        self.connection = None;
549                        return Poll::Ready(None);
550                    },
551                    tokio_tungstenite::tungstenite::Message::Ping(_) => {
552                        // Best-effort: tokio-tungstenite/tungstenite may handle pings internally.
553                        continue;
554                    },
555                    tokio_tungstenite::tungstenite::Message::Pong(_) => continue,
556                    tokio_tungstenite::tungstenite::Message::Binary(_) => continue,
557                    tokio_tungstenite::tungstenite::Message::Frame(_) => continue,
558                },
559                Poll::Ready(Some(Err(e))) => {
560                    error!("WebSocket error: {}", e);
561                    self.stats.errors += 1;
562                    return Poll::Ready(Some(Err(e.into())));
563                },
564                Poll::Ready(None) => {
565                    info!("WebSocket stream ended");
566                    return Poll::Ready(None);
567                },
568            }
569        }
570    }
571}
572
573impl MarketStream for WebSocketStream {
574    fn subscribe(&mut self, _subscription: Subscription) -> Result<()> {
575        // This is for backward compatibility - use subscribe_async for new code
576        Ok(())
577    }
578
579    fn unsubscribe(&mut self, _token_ids: &[String]) -> Result<()> {
580        // This is for backward compatibility - use unsubscribe_async for new code
581        Ok(())
582    }
583
584    fn is_connected(&self) -> bool {
585        self.connection.is_some()
586    }
587
588    fn get_stats(&self) -> StreamStats {
589        self.stats.clone()
590    }
591}
592
593/// Mock stream for testing
594#[derive(Debug)]
595pub struct MockStream {
596    messages: Vec<Result<StreamMessage>>,
597    index: usize,
598    connected: bool,
599}
600
601impl Default for MockStream {
602    fn default() -> Self {
603        Self::new()
604    }
605}
606
607impl MockStream {
608    pub fn new() -> Self {
609        Self {
610            messages: Vec::new(),
611            index: 0,
612            connected: true,
613        }
614    }
615
616    pub fn add_message(&mut self, message: StreamMessage) {
617        self.messages.push(Ok(message));
618    }
619
620    pub fn add_error(&mut self, error: PolyfillError) {
621        self.messages.push(Err(error));
622    }
623
624    pub fn set_connected(&mut self, connected: bool) {
625        self.connected = connected;
626    }
627}
628
629impl Stream for MockStream {
630    type Item = Result<StreamMessage>;
631
632    fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
633        if self.index >= self.messages.len() {
634            Poll::Ready(None)
635        } else {
636            let message = self.messages[self.index].clone();
637            self.index += 1;
638            Poll::Ready(Some(message))
639        }
640    }
641}
642
643impl MarketStream for MockStream {
644    fn subscribe(&mut self, _subscription: Subscription) -> Result<()> {
645        Ok(())
646    }
647
648    fn unsubscribe(&mut self, _token_ids: &[String]) -> Result<()> {
649        Ok(())
650    }
651
652    fn is_connected(&self) -> bool {
653        self.connected
654    }
655
656    fn get_stats(&self) -> StreamStats {
657        StreamStats {
658            messages_received: self.messages.len() as u64,
659            messages_sent: 0,
660            errors: self.messages.iter().filter(|m| m.is_err()).count() as u64,
661            dropped_messages: 0,
662            last_message_time: None,
663            connection_uptime: std::time::Duration::ZERO,
664            reconnect_count: 0,
665        }
666    }
667}
668
669/// Stream manager for handling multiple streams
670#[allow(dead_code)]
671pub struct StreamManager {
672    streams: Vec<Box<dyn MarketStream>>,
673    message_subscribers: Mutex<Vec<mpsc::UnboundedSender<StreamMessage>>>,
674}
675
676impl Default for StreamManager {
677    fn default() -> Self {
678        Self::new()
679    }
680}
681
682impl StreamManager {
683    pub fn new() -> Self {
684        Self {
685            streams: Vec::new(),
686            message_subscribers: Mutex::new(Vec::new()),
687        }
688    }
689
690    pub fn add_stream(&mut self, stream: Box<dyn MarketStream>) {
691        self.streams.push(stream);
692    }
693
694    pub fn get_message_receiver(&mut self) -> mpsc::UnboundedReceiver<StreamMessage> {
695        let (tx, rx) = mpsc::unbounded_channel();
696        self.message_subscribers.lock().push(tx);
697        rx
698    }
699
700    pub fn broadcast_message(&self, message: StreamMessage) -> Result<()> {
701        let mut subscribers = self.message_subscribers.lock();
702        subscribers.retain(|tx| tx.send(message.clone()).is_ok());
703        Ok(())
704    }
705}
706
707#[cfg(test)]
708mod tests {
709    use super::*;
710    use rust_decimal::Decimal;
711    use std::str::FromStr;
712
713    #[test]
714    fn test_mock_stream() {
715        let mut stream = MockStream::new();
716
717        // Add some test messages
718        stream.add_message(StreamMessage::Book(BookUpdate {
719            asset_id: "1".to_string(),
720            market: "0xabc".to_string(),
721            timestamp: 1_234_567_890,
722            bids: vec![],
723            asks: vec![],
724            hash: None,
725        }));
726        stream.add_message(StreamMessage::PriceChange(PriceChange {
727            market: "0xabc".to_string(),
728            timestamp: 1_234_567_891,
729            price_changes: vec![],
730        }));
731
732        assert!(stream.is_connected());
733        assert_eq!(stream.get_stats().messages_received, 2);
734    }
735
736    #[test]
737    fn test_stream_manager() {
738        let mut manager = StreamManager::new();
739        let mock_stream = Box::new(MockStream::new());
740        manager.add_stream(mock_stream);
741        let mut first_rx = manager.get_message_receiver();
742        let mut second_rx = manager.get_message_receiver();
743
744        // Test message broadcasting
745        let message = StreamMessage::Book(BookUpdate {
746            asset_id: "1".to_string(),
747            market: "0xabc".to_string(),
748            timestamp: 1_234_567_890,
749            bids: vec![],
750            asks: vec![],
751            hash: None,
752        });
753        assert!(manager.broadcast_message(message).is_ok());
754
755        assert!(matches!(
756            first_rx.try_recv().unwrap(),
757            StreamMessage::Book(BookUpdate { asset_id, .. }) if asset_id == "1"
758        ));
759        assert!(matches!(
760            second_rx.try_recv().unwrap(),
761            StreamMessage::Book(BookUpdate { asset_id, .. }) if asset_id == "1"
762        ));
763
764        drop(first_rx);
765        let message = StreamMessage::PriceChange(PriceChange {
766            market: "0xabc".to_string(),
767            timestamp: 1_234_567_891,
768            price_changes: vec![],
769        });
770        assert!(manager.broadcast_message(message).is_ok());
771
772        assert!(matches!(
773            second_rx.try_recv().unwrap(),
774            StreamMessage::PriceChange(PriceChange { market, .. }) if market == "0xabc"
775        ));
776    }
777
778    #[test]
779    fn test_websocket_book_applier_apply_text_message_updates_book() {
780        let books = crate::book::OrderBookManager::new(64);
781        let _ = books.get_or_create_book("12345").unwrap();
782
783        let processor = WsBookUpdateProcessor::new(1024);
784        let stream = WebSocketStream::new("wss://example.com/ws");
785        let mut applier = stream.into_book_applier(&books, processor);
786
787        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();
788        let stats = applier.apply_text_message(msg).unwrap();
789        assert_eq!(stats.book_messages, 1);
790        assert_eq!(stats.book_levels_applied, 2);
791
792        let snapshot = books.get_book("12345").unwrap();
793        assert_eq!(snapshot.bids.len(), 1);
794        assert_eq!(snapshot.asks.len(), 1);
795        assert_eq!(snapshot.bids[0].price, Decimal::from_str("0.75").unwrap());
796        assert_eq!(snapshot.bids[0].size, Decimal::from_str("10").unwrap());
797        assert_eq!(snapshot.asks[0].price, Decimal::from_str("0.76").unwrap());
798        assert_eq!(snapshot.asks[0].size, Decimal::from_str("5").unwrap());
799    }
800
801    #[test]
802    fn test_websocket_book_applier_replaces_snapshot() {
803        let books = crate::book::OrderBookManager::new(64);
804        let _ = books.get_or_create_book("12345").unwrap();
805
806        let processor = WsBookUpdateProcessor::new(1024);
807        let stream = WebSocketStream::new("wss://example.com/ws");
808        let mut applier = stream.into_book_applier(&books, processor);
809
810        let first = r#"{"event_type":"book","asset_id":"12345","timestamp":"1757908892351","bids":[{"price":"0.74","size":"8"},{"price":"0.75","size":"10"}],"asks":[{"price":"0.76","size":"5"},{"price":"0.77","size":"4"}]}"#.to_string();
811        applier.apply_text_message(first).unwrap();
812
813        let second = r#"{"event_type":"book","asset_id":"12345","timestamp":"1757908892352","bids":[{"price":"0.75","size":"11"}],"asks":[{"price":"0.76","size":"6"}]}"#.to_string();
814        let stats = applier.apply_text_message(second).unwrap();
815        assert_eq!(stats.book_messages, 1);
816        assert_eq!(stats.book_levels_applied, 2);
817
818        let snapshot = books.get_book("12345").unwrap();
819        assert_eq!(snapshot.timestamp.timestamp_millis(), 1_757_908_892_352);
820        assert_eq!(snapshot.bids.len(), 1);
821        assert_eq!(snapshot.asks.len(), 1);
822        assert_eq!(snapshot.bids[0].price, Decimal::from_str("0.75").unwrap());
823        assert_eq!(snapshot.bids[0].size, Decimal::from_str("11").unwrap());
824        assert_eq!(snapshot.asks[0].price, Decimal::from_str("0.76").unwrap());
825        assert_eq!(snapshot.asks[0].size, Decimal::from_str("6").unwrap());
826    }
827}
polyfill_rs/stream.rs

polyfill_rs/
stream.rs
