use crate::application::client::StreamerClient;
use crate::error::AppError;
use crate::model::streaming::StreamingMarketField;
use crate::presentation::price::PriceData;
use std::collections::HashSet;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use tokio::sync::{Notify, RwLock, mpsc};
use tracing::{debug, info, warn};
pub struct DynamicMarketStreamer {
epics: Arc<RwLock<HashSet<String>>>,
fields: HashSet<StreamingMarketField>,
price_tx: Arc<RwLock<Option<mpsc::UnboundedSender<PriceData>>>>,
price_rx: Arc<RwLock<Option<mpsc::UnboundedReceiver<PriceData>>>>,
is_connected: Arc<RwLock<bool>>,
shutdown_signal: Arc<RwLock<Option<Arc<Notify>>>>,
generation: Arc<AtomicU64>,
}
impl DynamicMarketStreamer {
#[must_use]
pub fn new(fields: HashSet<StreamingMarketField>) -> Self {
let (price_tx, price_rx) = mpsc::unbounded_channel();
Self {
epics: Arc::new(RwLock::new(HashSet::new())),
fields,
price_tx: Arc::new(RwLock::new(Some(price_tx))),
price_rx: Arc::new(RwLock::new(Some(price_rx))),
is_connected: Arc::new(RwLock::new(false)),
shutdown_signal: Arc::new(RwLock::new(None)),
generation: Arc::new(AtomicU64::new(0)),
}
}
pub async fn add(&self, epic: String) -> Result<(), AppError> {
let mut epics = self.epics.write().await;
if epics.contains(&epic) {
debug!("EPIC {} already subscribed", epic);
return Ok(());
}
epics.insert(epic.clone());
info!("Added EPIC {} to subscription list", epic);
drop(epics);
let is_connected = *self.is_connected.read().await;
if is_connected {
self.reconnect().await?;
}
Ok(())
}
pub async fn remove(&self, epic: String) -> Result<(), AppError> {
let mut epics = self.epics.write().await;
let was_removed = epics.remove(&epic);
if was_removed {
info!("Removed EPIC {} from subscription list", epic);
} else {
debug!("EPIC {} was not in subscription list", epic);
}
drop(epics);
if was_removed {
let is_connected = *self.is_connected.read().await;
if is_connected {
self.reconnect().await?;
}
}
Ok(())
}
pub async fn clear(&self) -> Result<(), AppError> {
let mut epics = self.epics.write().await;
let count = epics.len();
epics.clear();
info!("Cleared {} EPICs from subscription list", count);
drop(epics);
let is_connected = *self.is_connected.read().await;
if is_connected {
self.reconnect().await?;
*self.is_connected.write().await = false;
}
Ok(())
}
pub async fn get_epics(&self) -> Vec<String> {
let epics = self.epics.read().await;
epics.iter().cloned().collect()
}
pub async fn get_receiver(&self) -> Result<mpsc::UnboundedReceiver<PriceData>, AppError> {
let mut rx_lock = self.price_rx.write().await;
rx_lock
.take()
.ok_or_else(|| AppError::InvalidInput("Receiver already taken".to_string()))
}
async fn reconnect(&self) -> Result<(), AppError> {
info!("Reconnecting with updated EPIC list...");
{
let shutdown_lock = self.shutdown_signal.read().await;
if let Some(signal) = shutdown_lock.as_ref() {
signal.notify_one();
}
}
tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
let epics = self.get_epics().await;
if !epics.is_empty() {
self.start_internal().await?;
}
Ok(())
}
async fn start_internal(&self) -> Result<(), AppError> {
let epics = self.get_epics().await;
if epics.is_empty() {
warn!("No EPICs to subscribe to");
return Ok(());
}
info!("Starting connection with {} EPICs", epics.len());
let my_generation = self.generation.fetch_add(1, Ordering::SeqCst) + 1;
let mut new_client = StreamerClient::new().await?;
let fields = self.fields.clone();
let mut receiver = new_client.market_subscribe(epics.clone(), fields).await?;
let price_tx = self.price_tx.read().await;
if let Some(tx) = price_tx.as_ref() {
let tx = tx.clone();
tokio::spawn(async move {
while let Some(price_data) = receiver.recv().await {
if tx.send(price_data).is_err() {
warn!("Failed to send price update: receiver dropped");
break;
}
}
debug!("Subscription forwarding task ended");
});
}
let signal = Arc::new(Notify::new());
*self.shutdown_signal.write().await = Some(Arc::clone(&signal));
*self.is_connected.write().await = true;
let is_connected = Arc::clone(&self.is_connected);
let generation = Arc::clone(&self.generation);
tokio::spawn(async move {
let mut c = new_client;
let result = {
let r = c.connect(Some(signal)).await;
if let Err(e) = c.disconnect().await {
tracing::error!("Error closing streamer session: {}", e);
}
r
};
if generation.load(Ordering::SeqCst) == my_generation {
*is_connected.write().await = false;
}
match result {
Ok(_) => info!("Connection task completed successfully"),
Err(e) => tracing::error!("Connection task failed: {:?}", e),
}
});
info!("Connection task started in background");
Ok(())
}
pub async fn start(&mut self) -> Result<(), AppError> {
self.start_internal().await
}
pub async fn connect(&mut self) -> Result<(), AppError> {
self.start().await?;
use lightstreamer_rs::utils::setup_signal_hook;
let signal = Arc::new(Notify::new());
setup_signal_hook(Arc::clone(&signal)).await;
signal.notified().await;
self.disconnect().await?;
Ok(())
}
pub async fn disconnect(&mut self) -> Result<(), AppError> {
{
let shutdown_lock = self.shutdown_signal.read().await;
if let Some(signal) = shutdown_lock.as_ref() {
signal.notify_one();
}
}
*self.is_connected.write().await = false;
info!("Disconnected from Lightstreamer server");
Ok(())
}
}
impl Clone for DynamicMarketStreamer {
fn clone(&self) -> Self {
Self {
epics: Arc::clone(&self.epics),
fields: self.fields.clone(),
price_tx: Arc::clone(&self.price_tx),
price_rx: Arc::clone(&self.price_rx),
is_connected: Arc::clone(&self.is_connected),
shutdown_signal: Arc::clone(&self.shutdown_signal),
generation: Arc::clone(&self.generation),
}
}
}
#[cfg(test)]
mod tests {
use super::DynamicMarketStreamer;
use std::collections::HashSet;
use std::sync::Arc;
use std::sync::atomic::Ordering;
use std::time::Duration;
use tokio::sync::Notify;
const TEST_EPIC: &str = "IX.D.DAX.DAILY.IP";
const OTHER_EPIC: &str = "IX.D.FTSE.DAILY.IP";
#[tokio::test]
async fn test_add_inserts_epic_when_not_connected() {
let streamer = DynamicMarketStreamer::new(HashSet::new());
let result = streamer.add(TEST_EPIC.to_string()).await;
assert!(result.is_ok(), "add should succeed: {result:?}");
assert_eq!(
streamer.get_epics().await,
vec![TEST_EPIC.to_string()],
"add must insert the EPIC into the subscription set"
);
}
#[tokio::test]
async fn test_add_is_idempotent_for_duplicate_epic() {
let streamer = DynamicMarketStreamer::new(HashSet::new());
for _ in 0..3 {
let result = streamer.add(TEST_EPIC.to_string()).await;
assert!(result.is_ok(), "repeated add should succeed: {result:?}");
}
assert_eq!(
streamer.get_epics().await.len(),
1,
"adding the same EPIC repeatedly must not create duplicates"
);
}
#[tokio::test]
async fn test_remove_absent_epic_is_noop() {
let streamer = DynamicMarketStreamer::new(HashSet::new());
streamer.epics.write().await.insert(TEST_EPIC.to_string());
let result = streamer.remove(OTHER_EPIC.to_string()).await;
assert!(
result.is_ok(),
"removing an absent EPIC should succeed: {result:?}"
);
assert_eq!(
streamer.get_epics().await,
vec![TEST_EPIC.to_string()],
"removing an absent EPIC must leave the set unchanged"
);
}
#[tokio::test]
async fn test_remove_existing_epic_when_not_connected_empties_set() {
let streamer = DynamicMarketStreamer::new(HashSet::new());
streamer.epics.write().await.insert(TEST_EPIC.to_string());
let result = streamer.remove(TEST_EPIC.to_string()).await;
assert!(result.is_ok(), "remove should succeed: {result:?}");
assert!(
streamer.get_epics().await.is_empty(),
"removing the only EPIC must empty the set"
);
}
#[tokio::test]
async fn test_get_receiver_can_only_be_taken_once() {
let streamer = DynamicMarketStreamer::new(HashSet::new());
let first = streamer.get_receiver().await;
assert!(
first.is_ok(),
"first get_receiver should hand out the receiver: {first:?}"
);
let second = streamer.get_receiver().await;
assert!(
second.is_err(),
"second get_receiver must fail once the receiver has been taken"
);
}
#[tokio::test]
async fn test_is_connected_transitions_on_disconnect() {
let mut streamer = DynamicMarketStreamer::new(HashSet::new());
assert!(
!*streamer.is_connected.read().await,
"a new streamer must start disconnected"
);
*streamer.is_connected.write().await = true;
let signal = Arc::new(Notify::new());
*streamer.shutdown_signal.write().await = Some(Arc::clone(&signal));
assert!(
*streamer.is_connected.read().await,
"streamer should report connected once a connection is live"
);
let result = streamer.disconnect().await;
assert!(result.is_ok(), "disconnect should succeed: {result:?}");
assert!(
!*streamer.is_connected.read().await,
"disconnect must transition the streamer back to disconnected"
);
}
#[tokio::test]
async fn test_disconnect_signals_shutdown_and_marks_disconnected() {
let mut streamer = DynamicMarketStreamer::new(HashSet::new());
*streamer.is_connected.write().await = true;
let signal = Arc::new(Notify::new());
*streamer.shutdown_signal.write().await = Some(Arc::clone(&signal));
let waiter = tokio::spawn(async move { signal.notified().await });
let result = streamer.disconnect().await;
assert!(result.is_ok(), "disconnect should succeed: {result:?}");
assert!(
!*streamer.is_connected.read().await,
"disconnect must mark the streamer disconnected"
);
assert!(
tokio::time::timeout(Duration::from_secs(1), waiter)
.await
.is_ok(),
"the parked connection did not observe the shutdown signal from disconnect()"
);
}
#[tokio::test]
async fn test_remove_last_epic_while_connected_signals_reconnect() {
let streamer = DynamicMarketStreamer::new(HashSet::new());
streamer.epics.write().await.insert(TEST_EPIC.to_string());
*streamer.is_connected.write().await = true;
let signal = Arc::new(Notify::new());
*streamer.shutdown_signal.write().await = Some(Arc::clone(&signal));
let waiter = tokio::spawn(async move { signal.notified().await });
let result = streamer.remove(TEST_EPIC.to_string()).await;
assert!(result.is_ok(), "remove should succeed: {result:?}");
assert!(
streamer.get_epics().await.is_empty(),
"removing the last EPIC must empty the subscription set"
);
assert!(
tokio::time::timeout(Duration::from_secs(1), waiter)
.await
.is_ok(),
"the live connection did not observe the reconnect signal from remove()"
);
}
#[tokio::test]
async fn test_clear_when_not_connected_empties_epics() {
let streamer = DynamicMarketStreamer::new(HashSet::new());
streamer.epics.write().await.insert(TEST_EPIC.to_string());
let result = streamer.clear().await;
assert!(result.is_ok(), "clear should succeed: {result:?}");
assert!(
streamer.get_epics().await.is_empty(),
"EPIC set should be empty after clear"
);
assert!(
!*streamer.is_connected.read().await,
"should not report connected when it never was"
);
}
#[tokio::test]
async fn test_clear_when_connected_signals_shutdown_and_reports_stopped() {
let streamer = DynamicMarketStreamer::new(HashSet::new());
streamer.epics.write().await.insert(TEST_EPIC.to_string());
*streamer.is_connected.write().await = true;
let signal = Arc::new(Notify::new());
*streamer.shutdown_signal.write().await = Some(Arc::clone(&signal));
let waiter = tokio::spawn(async move { signal.notified().await });
let result = streamer.clear().await;
assert!(result.is_ok(), "clear should succeed: {result:?}");
assert!(
streamer.get_epics().await.is_empty(),
"EPIC set should be empty after clear"
);
assert!(
!*streamer.is_connected.read().await,
"clear must mark the streamer disconnected so data flow stops"
);
assert!(
tokio::time::timeout(Duration::from_secs(1), waiter)
.await
.is_ok(),
"live connection did not observe the shutdown signal from clear()"
);
}
#[tokio::test]
async fn test_superseded_generation_does_not_clear_is_connected() {
let streamer = DynamicMarketStreamer::new(HashSet::new());
let newer = streamer.generation.fetch_add(1, Ordering::SeqCst) + 1;
*streamer.is_connected.write().await = true;
let older = newer - 1;
if streamer.generation.load(Ordering::SeqCst) == older {
*streamer.is_connected.write().await = false;
}
assert!(
*streamer.is_connected.read().await,
"a superseded generation must not clear is_connected on the newer one"
);
if streamer.generation.load(Ordering::SeqCst) == newer {
*streamer.is_connected.write().await = false;
}
assert!(
!*streamer.is_connected.read().await,
"the current generation's teardown must clear is_connected"
);
}
}