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::{ready, 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
383fn poll_send_pong(
384    connection: &mut tokio_tungstenite::WebSocketStream<
385        tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
386    >,
387    cx: &mut Context<'_>,
388    data: Vec<u8>,
389) -> Poll<Result<()>> {
390    ready!(connection.poll_ready_unpin(cx)).map_err(|e| {
391        PolyfillError::stream(
392            format!("Failed to prepare pong: {}", e),
393            crate::errors::StreamErrorKind::MessageCorrupted,
394        )
395    })?;
396    connection
397        .start_send_unpin(tokio_tungstenite::tungstenite::Message::Pong(data))
398        .map_err(|e| {
399            PolyfillError::stream(
400                format!("Failed to send pong: {}", e),
401                crate::errors::StreamErrorKind::MessageCorrupted,
402            )
403        })?;
404    ready!(connection.poll_flush_unpin(cx)).map_err(|e| {
405        PolyfillError::stream(
406            format!("Failed to flush pong: {}", e),
407            crate::errors::StreamErrorKind::MessageCorrupted,
408        )
409    })?;
410    Poll::Ready(Ok(()))
411}
412
413/// WebSocket stream wrapper that applies `book` updates directly into an [`crate::book::OrderBookManager`].
414///
415/// This bypasses `StreamMessage` decoding (serde/DOM parsing) for the `book` hot path by using
416/// [`WsBookUpdateProcessor`]. Non-`book` WS payloads are ignored.
417///
418/// Note: the underlying WS transport may still allocate when producing `Message::Text(String)`.
419pub struct WebSocketBookApplier<'a> {
420    stream: WebSocketStream,
421    books: &'a crate::book::OrderBookManager,
422    processor: WsBookUpdateProcessor,
423}
424
425impl WebSocketStream {
426    /// Convert this stream into a book-applier stream.
427    ///
428    /// The caller is expected to "warm up" the [`crate::book::OrderBookManager`] by creating books for all
429    /// subscribed asset IDs ahead of time. Missing books are treated as an error.
430    pub fn into_book_applier<'a>(
431        mut self,
432        books: &'a crate::book::OrderBookManager,
433        processor: WsBookUpdateProcessor,
434    ) -> WebSocketBookApplier<'a> {
435        // Drop any pre-parsed messages to avoid mixing the two streaming modes.
436        self.pending.clear();
437        WebSocketBookApplier {
438            stream: self,
439            books,
440            processor,
441        }
442    }
443}
444
445impl<'a> WebSocketBookApplier<'a> {
446    /// Access the underlying WebSocket stream (e.g., for subscribe/unsubscribe calls).
447    pub fn stream_mut(&mut self) -> &mut WebSocketStream {
448        &mut self.stream
449    }
450
451    /// Current WebSocket connection stats.
452    pub fn stream_stats(&self) -> StreamStats {
453        self.stream.stats.clone()
454    }
455
456    /// Access the hot-path processor (e.g., to reuse it across connections).
457    pub fn processor_mut(&mut self) -> &mut WsBookUpdateProcessor {
458        &mut self.processor
459    }
460
461    /// Apply a single WS text payload (useful for custom transports and for testing).
462    ///
463    /// This convenience method consumes an owned `String`, so it may release that
464    /// buffer after processing. Use [`Self::apply_bytes_message`] for the
465    /// allocation-sensitive path when the caller owns a reusable mutable buffer.
466    pub fn apply_text_message(&mut self, text: String) -> Result<WsBookApplyStats> {
467        let stats = self.processor.process_text(text, self.books)?;
468        self.stream.stats.messages_received += 1;
469        self.stream.stats.last_message_time = Some(Utc::now());
470        Ok(stats)
471    }
472
473    /// Apply a single WS payload from a caller-owned mutable byte buffer.
474    ///
475    /// The buffer is mutated by `simd-json`. After processor warmup, this is the
476    /// allocation-sensitive book-applier entry point because no owned message buffer
477    /// is created or dropped by this method.
478    pub fn apply_bytes_message(&mut self, bytes: &mut [u8]) -> Result<WsBookApplyStats> {
479        let stats = self.processor.process_bytes(bytes, self.books)?;
480        self.stream.stats.messages_received += 1;
481        self.stream.stats.last_message_time = Some(Utc::now());
482        Ok(stats)
483    }
484}
485
486impl<'a> Stream for WebSocketBookApplier<'a> {
487    type Item = Result<WsBookApplyStats>;
488
489    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
490        loop {
491            let Some(connection) = &mut self.stream.connection else {
492                return Poll::Ready(None);
493            };
494
495            match connection.poll_next_unpin(cx) {
496                Poll::Pending => return Poll::Pending,
497                Poll::Ready(Some(Ok(msg))) => match msg {
498                    tokio_tungstenite::tungstenite::Message::Text(text) => {
499                        match self.apply_text_message(text) {
500                            Ok(stats) => {
501                                if stats.book_messages == 0 {
502                                    continue;
503                                }
504                                return Poll::Ready(Some(Ok(stats)));
505                            },
506                            Err(e) => {
507                                self.stream.stats.errors += 1;
508                                return Poll::Ready(Some(Err(e)));
509                            },
510                        }
511                    },
512                    tokio_tungstenite::tungstenite::Message::Close(_) => {
513                        info!("WebSocket connection closed by server");
514                        self.stream.connection = None;
515                        return Poll::Ready(None);
516                    },
517                    tokio_tungstenite::tungstenite::Message::Ping(data) => {
518                        match poll_send_pong(connection, cx, data) {
519                            Poll::Ready(Ok(())) => continue,
520                            Poll::Ready(Err(e)) => {
521                                self.stream.stats.errors += 1;
522                                return Poll::Ready(Some(Err(e)));
523                            },
524                            Poll::Pending => return Poll::Pending,
525                        }
526                    },
527                    tokio_tungstenite::tungstenite::Message::Pong(_) => continue,
528                    tokio_tungstenite::tungstenite::Message::Binary(_) => continue,
529                    tokio_tungstenite::tungstenite::Message::Frame(_) => continue,
530                },
531                Poll::Ready(Some(Err(e))) => {
532                    error!("WebSocket error: {}", e);
533                    self.stream.stats.errors += 1;
534                    return Poll::Ready(Some(Err(e.into())));
535                },
536                Poll::Ready(None) => {
537                    info!("WebSocket stream ended");
538                    return Poll::Ready(None);
539                },
540            }
541        }
542    }
543}
544
545impl Stream for WebSocketStream {
546    type Item = Result<StreamMessage>;
547
548    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
549        loop {
550            if let Some(message) = self.pending.pop_front() {
551                return Poll::Ready(Some(Ok(message)));
552            }
553
554            let Some(connection) = &mut self.connection else {
555                return Poll::Ready(None);
556            };
557
558            match connection.poll_next_unpin(cx) {
559                Poll::Pending => return Poll::Pending,
560                Poll::Ready(Some(Ok(ws_message))) => match ws_message {
561                    tokio_tungstenite::tungstenite::Message::Text(text) => {
562                        match crate::decode::parse_stream_messages(&text) {
563                            Ok(messages) => {
564                                let mut iter = messages.into_iter();
565                                let Some(first) = iter.next() else {
566                                    continue;
567                                };
568
569                                for msg in iter {
570                                    self.enqueue(msg);
571                                }
572                                self.stats.messages_received += 1;
573                                self.stats.last_message_time = Some(Utc::now());
574                                return Poll::Ready(Some(Ok(first)));
575                            },
576                            Err(e) => {
577                                self.stats.errors += 1;
578                                return Poll::Ready(Some(Err(e)));
579                            },
580                        }
581                    },
582                    tokio_tungstenite::tungstenite::Message::Close(_) => {
583                        info!("WebSocket connection closed by server");
584                        self.connection = None;
585                        return Poll::Ready(None);
586                    },
587                    tokio_tungstenite::tungstenite::Message::Ping(data) => {
588                        match poll_send_pong(connection, cx, data) {
589                            Poll::Ready(Ok(())) => continue,
590                            Poll::Ready(Err(e)) => {
591                                self.stats.errors += 1;
592                                return Poll::Ready(Some(Err(e)));
593                            },
594                            Poll::Pending => return Poll::Pending,
595                        }
596                    },
597                    tokio_tungstenite::tungstenite::Message::Pong(_) => continue,
598                    tokio_tungstenite::tungstenite::Message::Binary(_) => continue,
599                    tokio_tungstenite::tungstenite::Message::Frame(_) => continue,
600                },
601                Poll::Ready(Some(Err(e))) => {
602                    error!("WebSocket error: {}", e);
603                    self.stats.errors += 1;
604                    return Poll::Ready(Some(Err(e.into())));
605                },
606                Poll::Ready(None) => {
607                    info!("WebSocket stream ended");
608                    return Poll::Ready(None);
609                },
610            }
611        }
612    }
613}
614
615impl MarketStream for WebSocketStream {
616    fn subscribe(&mut self, _subscription: Subscription) -> Result<()> {
617        // This is for backward compatibility - use subscribe_async for new code
618        Ok(())
619    }
620
621    fn unsubscribe(&mut self, _token_ids: &[String]) -> Result<()> {
622        // This is for backward compatibility - use unsubscribe_async for new code
623        Ok(())
624    }
625
626    fn is_connected(&self) -> bool {
627        self.connection.is_some()
628    }
629
630    fn get_stats(&self) -> StreamStats {
631        self.stats.clone()
632    }
633}
634
635/// Mock stream for testing
636#[derive(Debug)]
637pub struct MockStream {
638    messages: Vec<Result<StreamMessage>>,
639    index: usize,
640    connected: bool,
641}
642
643impl Default for MockStream {
644    fn default() -> Self {
645        Self::new()
646    }
647}
648
649impl MockStream {
650    pub fn new() -> Self {
651        Self {
652            messages: Vec::new(),
653            index: 0,
654            connected: true,
655        }
656    }
657
658    pub fn add_message(&mut self, message: StreamMessage) {
659        self.messages.push(Ok(message));
660    }
661
662    pub fn add_error(&mut self, error: PolyfillError) {
663        self.messages.push(Err(error));
664    }
665
666    pub fn set_connected(&mut self, connected: bool) {
667        self.connected = connected;
668    }
669}
670
671impl Stream for MockStream {
672    type Item = Result<StreamMessage>;
673
674    fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
675        if self.index >= self.messages.len() {
676            Poll::Ready(None)
677        } else {
678            let message = self.messages[self.index].clone();
679            self.index += 1;
680            Poll::Ready(Some(message))
681        }
682    }
683}
684
685impl MarketStream for MockStream {
686    fn subscribe(&mut self, _subscription: Subscription) -> Result<()> {
687        Ok(())
688    }
689
690    fn unsubscribe(&mut self, _token_ids: &[String]) -> Result<()> {
691        Ok(())
692    }
693
694    fn is_connected(&self) -> bool {
695        self.connected
696    }
697
698    fn get_stats(&self) -> StreamStats {
699        StreamStats {
700            messages_received: self.messages.len() as u64,
701            messages_sent: 0,
702            errors: self.messages.iter().filter(|m| m.is_err()).count() as u64,
703            dropped_messages: 0,
704            last_message_time: None,
705            connection_uptime: std::time::Duration::ZERO,
706            reconnect_count: 0,
707        }
708    }
709}
710
711/// Stream manager for handling multiple streams
712#[allow(dead_code)]
713pub struct StreamManager {
714    streams: Vec<Box<dyn MarketStream>>,
715    message_subscribers: Mutex<Vec<mpsc::UnboundedSender<StreamMessage>>>,
716}
717
718impl Default for StreamManager {
719    fn default() -> Self {
720        Self::new()
721    }
722}
723
724impl StreamManager {
725    pub fn new() -> Self {
726        Self {
727            streams: Vec::new(),
728            message_subscribers: Mutex::new(Vec::new()),
729        }
730    }
731
732    pub fn add_stream(&mut self, stream: Box<dyn MarketStream>) {
733        self.streams.push(stream);
734    }
735
736    pub fn get_message_receiver(&mut self) -> mpsc::UnboundedReceiver<StreamMessage> {
737        let (tx, rx) = mpsc::unbounded_channel();
738        self.message_subscribers.lock().push(tx);
739        rx
740    }
741
742    pub fn broadcast_message(&self, message: StreamMessage) -> Result<()> {
743        let mut subscribers = self.message_subscribers.lock();
744        subscribers.retain(|tx| tx.send(message.clone()).is_ok());
745        Ok(())
746    }
747}
748
749#[cfg(test)]
750mod tests {
751    use super::*;
752    use rust_decimal::Decimal;
753    use std::str::FromStr;
754
755    #[test]
756    fn test_mock_stream() {
757        let mut stream = MockStream::new();
758
759        // Add some test messages
760        stream.add_message(StreamMessage::Book(BookUpdate {
761            asset_id: "1".to_string(),
762            market: "0xabc".to_string(),
763            timestamp: 1_234_567_890,
764            bids: vec![],
765            asks: vec![],
766            hash: None,
767        }));
768        stream.add_message(StreamMessage::PriceChange(PriceChange {
769            market: "0xabc".to_string(),
770            timestamp: 1_234_567_891,
771            price_changes: vec![],
772        }));
773
774        assert!(stream.is_connected());
775        assert_eq!(stream.get_stats().messages_received, 2);
776    }
777
778    #[test]
779    fn test_stream_manager() {
780        let mut manager = StreamManager::new();
781        let mock_stream = Box::new(MockStream::new());
782        manager.add_stream(mock_stream);
783        let mut first_rx = manager.get_message_receiver();
784        let mut second_rx = manager.get_message_receiver();
785
786        // Test message broadcasting
787        let message = StreamMessage::Book(BookUpdate {
788            asset_id: "1".to_string(),
789            market: "0xabc".to_string(),
790            timestamp: 1_234_567_890,
791            bids: vec![],
792            asks: vec![],
793            hash: None,
794        });
795        assert!(manager.broadcast_message(message).is_ok());
796
797        assert!(matches!(
798            first_rx.try_recv().unwrap(),
799            StreamMessage::Book(BookUpdate { asset_id, .. }) if asset_id == "1"
800        ));
801        assert!(matches!(
802            second_rx.try_recv().unwrap(),
803            StreamMessage::Book(BookUpdate { asset_id, .. }) if asset_id == "1"
804        ));
805
806        drop(first_rx);
807        let message = StreamMessage::PriceChange(PriceChange {
808            market: "0xabc".to_string(),
809            timestamp: 1_234_567_891,
810            price_changes: vec![],
811        });
812        assert!(manager.broadcast_message(message).is_ok());
813
814        assert!(matches!(
815            second_rx.try_recv().unwrap(),
816            StreamMessage::PriceChange(PriceChange { market, .. }) if market == "0xabc"
817        ));
818    }
819
820    #[test]
821    fn test_websocket_book_applier_apply_text_message_updates_book() {
822        let books = crate::book::OrderBookManager::new(64);
823        let _ = books.get_or_create_book("12345").unwrap();
824
825        let processor = WsBookUpdateProcessor::new(1024);
826        let stream = WebSocketStream::new("wss://example.com/ws");
827        let mut applier = stream.into_book_applier(&books, processor);
828
829        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();
830        let stats = applier.apply_text_message(msg).unwrap();
831        assert_eq!(stats.book_messages, 1);
832        assert_eq!(stats.book_levels_applied, 2);
833
834        let snapshot = books.get_book("12345").unwrap();
835        assert_eq!(snapshot.bids.len(), 1);
836        assert_eq!(snapshot.asks.len(), 1);
837        assert_eq!(snapshot.bids[0].price, Decimal::from_str("0.75").unwrap());
838        assert_eq!(snapshot.bids[0].size, Decimal::from_str("10").unwrap());
839        assert_eq!(snapshot.asks[0].price, Decimal::from_str("0.76").unwrap());
840        assert_eq!(snapshot.asks[0].size, Decimal::from_str("5").unwrap());
841    }
842
843    #[test]
844    fn test_websocket_book_applier_replaces_snapshot() {
845        let books = crate::book::OrderBookManager::new(64);
846        let _ = books.get_or_create_book("12345").unwrap();
847
848        let processor = WsBookUpdateProcessor::new(1024);
849        let stream = WebSocketStream::new("wss://example.com/ws");
850        let mut applier = stream.into_book_applier(&books, processor);
851
852        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();
853        applier.apply_text_message(first).unwrap();
854
855        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();
856        let stats = applier.apply_text_message(second).unwrap();
857        assert_eq!(stats.book_messages, 1);
858        assert_eq!(stats.book_levels_applied, 2);
859
860        let snapshot = books.get_book("12345").unwrap();
861        assert_eq!(snapshot.timestamp.timestamp_millis(), 1_757_908_892_352);
862        assert_eq!(snapshot.bids.len(), 1);
863        assert_eq!(snapshot.asks.len(), 1);
864        assert_eq!(snapshot.bids[0].price, Decimal::from_str("0.75").unwrap());
865        assert_eq!(snapshot.bids[0].size, Decimal::from_str("11").unwrap());
866        assert_eq!(snapshot.asks[0].price, Decimal::from_str("0.76").unwrap());
867        assert_eq!(snapshot.asks[0].size, Decimal::from_str("6").unwrap());
868    }
869}
polyfill_rs/stream.rs

polyfill_rs/
stream.rs
