fix: deliver stream manager broadcasts

This commit is contained in:
floor-licker
2026-06-22 19:14:51 -04:00
parent cb8f19c246
commit 02fa74d2f5
+32 -13
View File
@@ -8,6 +8,7 @@ use crate::types::*;
use crate::ws_hot_path::{WsBookApplyStats, WsBookUpdateProcessor}; use crate::ws_hot_path::{WsBookApplyStats, WsBookUpdateProcessor};
use chrono::Utc; use chrono::Utc;
use futures::{SinkExt, Stream, StreamExt}; use futures::{SinkExt, Stream, StreamExt};
use parking_lot::Mutex;
use serde_json::Value; use serde_json::Value;
use std::collections::VecDeque; use std::collections::VecDeque;
use std::pin::Pin; use std::pin::Pin;
@@ -653,8 +654,7 @@ impl MarketStream for MockStream {
#[allow(dead_code)] #[allow(dead_code)]
pub struct StreamManager { pub struct StreamManager {
streams: Vec<Box<dyn MarketStream>>, streams: Vec<Box<dyn MarketStream>>,
message_tx: mpsc::UnboundedSender<StreamMessage>, message_subscribers: Mutex<Vec<mpsc::UnboundedSender<StreamMessage>>>,
message_rx: mpsc::UnboundedReceiver<StreamMessage>,
} }
impl Default for StreamManager { impl Default for StreamManager {
@@ -665,12 +665,9 @@ impl Default for StreamManager {
impl StreamManager { impl StreamManager {
pub fn new() -> Self { pub fn new() -> Self {
let (message_tx, message_rx) = mpsc::unbounded_channel();
Self { Self {
streams: Vec::new(), streams: Vec::new(),
message_tx, message_subscribers: Mutex::new(Vec::new()),
message_rx,
} }
} }
@@ -679,17 +676,15 @@ impl StreamManager {
} }
pub fn get_message_receiver(&mut self) -> mpsc::UnboundedReceiver<StreamMessage> { pub fn get_message_receiver(&mut self) -> mpsc::UnboundedReceiver<StreamMessage> {
// Note: UnboundedReceiver doesn't implement Clone let (tx, rx) = mpsc::unbounded_channel();
// In a real implementation, you'd want to use a different approach self.message_subscribers.lock().push(tx);
// For now, we'll return a dummy receiver
let (_, rx) = mpsc::unbounded_channel();
rx rx
} }
pub fn broadcast_message(&self, message: StreamMessage) -> Result<()> { pub fn broadcast_message(&self, message: StreamMessage) -> Result<()> {
self.message_tx let mut subscribers = self.message_subscribers.lock();
.send(message) subscribers.retain(|tx| tx.send(message.clone()).is_ok());
.map_err(|e| PolyfillError::internal("Failed to broadcast message", e)) Ok(())
} }
} }
@@ -727,6 +722,8 @@ mod tests {
let mut manager = StreamManager::new(); let mut manager = StreamManager::new();
let mock_stream = Box::new(MockStream::new()); let mock_stream = Box::new(MockStream::new());
manager.add_stream(mock_stream); manager.add_stream(mock_stream);
let mut first_rx = manager.get_message_receiver();
let mut second_rx = manager.get_message_receiver();
// Test message broadcasting // Test message broadcasting
let message = StreamMessage::Book(BookUpdate { let message = StreamMessage::Book(BookUpdate {
@@ -738,6 +735,28 @@ mod tests {
hash: None, hash: None,
}); });
assert!(manager.broadcast_message(message).is_ok()); 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] #[test]