From 02fa74d2f5b561634b83ed5859090e8f49c699d6 Mon Sep 17 00:00:00 2001 From: floor-licker Date: Mon, 22 Jun 2026 19:14:51 -0400 Subject: [PATCH] fix: deliver stream manager broadcasts --- src/stream.rs | 45 ++++++++++++++++++++++++++++++++------------- 1 file changed, 32 insertions(+), 13 deletions(-) diff --git a/src/stream.rs b/src/stream.rs index 5452b1e..6667535 100644 --- a/src/stream.rs +++ b/src/stream.rs @@ -8,6 +8,7 @@ use crate::types::*; use crate::ws_hot_path::{WsBookApplyStats, WsBookUpdateProcessor}; use chrono::Utc; use futures::{SinkExt, Stream, StreamExt}; +use parking_lot::Mutex; use serde_json::Value; use std::collections::VecDeque; use std::pin::Pin; @@ -653,8 +654,7 @@ impl MarketStream for MockStream { #[allow(dead_code)] pub struct StreamManager { streams: Vec>, - message_tx: mpsc::UnboundedSender, - message_rx: mpsc::UnboundedReceiver, + message_subscribers: Mutex>>, } impl Default for StreamManager { @@ -665,12 +665,9 @@ impl Default for StreamManager { impl StreamManager { pub fn new() -> Self { - let (message_tx, message_rx) = mpsc::unbounded_channel(); - Self { streams: Vec::new(), - message_tx, - message_rx, + message_subscribers: Mutex::new(Vec::new()), } } @@ -679,17 +676,15 @@ impl StreamManager { } pub fn get_message_receiver(&mut self) -> mpsc::UnboundedReceiver { - // Note: UnboundedReceiver doesn't implement Clone - // In a real implementation, you'd want to use a different approach - // For now, we'll return a dummy receiver - let (_, rx) = mpsc::unbounded_channel(); + let (tx, rx) = mpsc::unbounded_channel(); + self.message_subscribers.lock().push(tx); rx } pub fn broadcast_message(&self, message: StreamMessage) -> Result<()> { - self.message_tx - .send(message) - .map_err(|e| PolyfillError::internal("Failed to broadcast message", e)) + let mut subscribers = self.message_subscribers.lock(); + subscribers.retain(|tx| tx.send(message.clone()).is_ok()); + Ok(()) } } @@ -727,6 +722,8 @@ mod tests { let mut manager = StreamManager::new(); let mock_stream = Box::new(MockStream::new()); manager.add_stream(mock_stream); + let mut first_rx = manager.get_message_receiver(); + let mut second_rx = manager.get_message_receiver(); // Test message broadcasting let message = StreamMessage::Book(BookUpdate { @@ -738,6 +735,28 @@ mod tests { hash: None, }); assert!(manager.broadcast_message(message).is_ok()); + + assert!(matches!( + first_rx.try_recv().unwrap(), + StreamMessage::Book(BookUpdate { asset_id, .. }) if asset_id == "1" + )); + assert!(matches!( + second_rx.try_recv().unwrap(), + StreamMessage::Book(BookUpdate { asset_id, .. }) if asset_id == "1" + )); + + drop(first_rx); + let message = StreamMessage::PriceChange(PriceChange { + market: "0xabc".to_string(), + timestamp: 1_234_567_891, + price_changes: vec![], + }); + assert!(manager.broadcast_message(message).is_ok()); + + assert!(matches!( + second_rx.try_recv().unwrap(), + StreamMessage::PriceChange(PriceChange { market, .. }) if market == "0xabc" + )); } #[test]