//! Message dispatcher with request/response correlation. //! //! Reads from the connection's receive stream, routes responses to pending //! requests by `clientMsgId`, and broadcasts server-pushed events. use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; use futures::StreamExt; use tokio::sync::{broadcast, oneshot, Mutex}; use tracing::{debug, error, trace, warn}; use crate::connection::CTraderConnection; use crate::error::{CTraderError, Result}; use crate::proto::{self, ProtoMessage}; /// Capacity for broadcast channels (events can be dropped if receiver falls behind). const BROADCAST_CAPACITY: usize = 256; /// Message dispatcher — correlates request/response pairs and broadcasts events. pub struct MessageDispatcher { /// Connection sender for outgoing messages. conn: Arc, /// Pending requests keyed by clientMsgId. pending: Arc>>>, /// Broadcast channel for execution events (payloadType 2126). execution_tx: broadcast::Sender, /// Broadcast channel for spot events (payloadType 2131). spot_tx: broadcast::Sender, /// Broadcast channel for error events (payloadType 2132, 2142, 50). error_tx: broadcast::Sender, /// Handle to the reader task. reader_handle: Option>, } impl MessageDispatcher { /// Create a new dispatcher and start the reader task. pub fn new(mut conn: CTraderConnection) -> Self { let pending: Arc>>> = Arc::new(Mutex::new(HashMap::new())); let (execution_tx, _) = broadcast::channel(BROADCAST_CAPACITY); let (spot_tx, _) = broadcast::channel(BROADCAST_CAPACITY); let (error_tx, _) = broadcast::channel(BROADCAST_CAPACITY); // Take ownership of the receiver stream // Invariant: receiver is always available on a fresh CTraderConnection #[allow(clippy::expect_used)] let receiver = conn .take_receiver() .expect("receiver already taken from connection"); let pending_clone = Arc::clone(&pending); let exec_tx = execution_tx.clone(); let spot_tx_clone = spot_tx.clone(); let error_tx_clone = error_tx.clone(); let reader_handle = tokio::spawn(async move { Self::reader_loop( receiver, pending_clone, exec_tx, spot_tx_clone, error_tx_clone, ) .await; }); Self { conn: Arc::new(conn), pending, execution_tx, spot_tx, error_tx, reader_handle: Some(reader_handle), } } /// Send a request and wait for the correlated response. /// /// Generates a UUID `clientMsgId`, inserts a oneshot receiver, /// sends the message, and awaits the response or timeout. pub async fn send_request( &self, mut msg: ProtoMessage, timeout: Duration, ) -> Result { let msg_id = uuid::Uuid::new_v4().to_string(); msg.client_msg_id = Some(msg_id.clone()); let (tx, rx) = oneshot::channel(); { let mut pending = self.pending.lock().await; pending.insert(msg_id.clone(), tx); } // Send the message if let Err(e) = self.conn.send(msg).await { // Clean up the pending entry on send failure let mut pending = self.pending.lock().await; pending.remove(&msg_id); return Err(e); } // Wait for response or timeout match tokio::time::timeout(timeout, rx).await { Ok(Ok(response)) => Ok(response), Ok(Err(_recv_err)) => { // Sender dropped — connection likely closed Err(CTraderError::NotConnected) } Err(_elapsed) => { // Timeout — clean up pending let mut pending = self.pending.lock().await; pending.remove(&msg_id); Err(CTraderError::Timeout(timeout)) } } } /// Send a message without waiting for a response (fire-and-forget). pub async fn send(&self, msg: ProtoMessage) -> Result<()> { self.conn.send(msg).await } /// Subscribe to execution events. pub fn subscribe_executions(&self) -> broadcast::Receiver { self.execution_tx.subscribe() } /// Subscribe to spot price events. pub fn subscribe_spots(&self) -> broadcast::Receiver { self.spot_tx.subscribe() } /// Subscribe to error events. pub fn subscribe_errors(&self) -> broadcast::Receiver { self.error_tx.subscribe() } /// Get a reference to the underlying connection. pub fn connection(&self) -> &CTraderConnection { &self.conn } /// Reader loop — dispatches incoming messages. #[allow(clippy::let_underscore_must_use)] async fn reader_loop( mut receiver: futures::stream::SplitStream< tokio_util::codec::Framed< tokio_native_tls::TlsStream, crate::codec::CTraderCodec, >, >, pending: Arc>>>, execution_tx: broadcast::Sender, spot_tx: broadcast::Sender, error_tx: broadcast::Sender, ) { while let Some(result) = receiver.next().await { let msg = match result { Ok(msg) => msg, Err(e) => { error!("reader error: {e}"); break; } }; let pt = msg.payload_type; let pt_name = proto::payload_type_name(pt); trace!(payload_type = pt, name = pt_name, "received message"); // Check for correlated request first if let Some(ref msg_id) = msg.client_msg_id { let mut pending_map = pending.lock().await; if let Some(tx) = pending_map.remove(msg_id) { debug!( client_msg_id = %msg_id, payload_type = pt, "routing correlated response" ); let _ = tx.send(msg); continue; } } // Route by payload type match pt { proto::PT_HEARTBEAT_EVENT => { trace!("heartbeat received"); } proto::PT_EXECUTION_EVENT => { debug!("execution event received"); let _ = execution_tx.send(msg); } proto::PT_SPOT_EVENT => { let _ = spot_tx.send(msg); } proto::PT_ORDER_ERROR_EVENT | proto::PT_OA_ERROR_RES | proto::PT_ERROR_RES => { warn!(payload_type = pt, "error event received"); let _ = error_tx.send(msg); } proto::PT_CLIENT_DISCONNECT_EVENT => { warn!("client disconnect event received"); break; } proto::PT_ACCOUNTS_TOKEN_INVALIDATED => { warn!("token invalidated event received"); let _ = error_tx.send(msg); } proto::PT_ACCOUNT_DISCONNECT_EVENT => { warn!("account disconnect event received"); let _ = error_tx.send(msg); } _ => { // Uncorrelated response without a matching pending request. // This can happen for events we don't specifically handle above. debug!( payload_type = pt, name = pt_name, "unhandled message type" ); } } } // Connection closed — cancel all pending requests let mut pending_map = pending.lock().await; let count = pending_map.len(); if count > 0 { warn!(count, "connection closed, dropping pending requests"); pending_map.clear(); } } } impl Drop for MessageDispatcher { fn drop(&mut self) { if let Some(handle) = self.reader_handle.take() { handle.abort(); } } } #[cfg(test)] mod tests { use super::*; #[test] fn broadcast_capacity_is_reasonable() { // Static assertion — if BROADCAST_CAPACITY drops below 64, // message dispatching may stall under burst load. let cap = BROADCAST_CAPACITY; assert!(cap >= 64, "BROADCAST_CAPACITY too small: {cap}"); } }