mirror of
https://github.com/floor-licker/polyfill-rs.git
synced 2026-08-25 10:28:07 +00:00
fix: deliver stream manager broadcasts
This commit is contained in:
+32
-13
@@ -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]
|
||||||
|
|||||||
Reference in New Issue
Block a user