foxhunt/ctrader-openapi/src/dispatch.rs

//! 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<CTraderConnection>,

    /// Pending requests keyed by clientMsgId.
    pending: Arc<Mutex<HashMap<String, oneshot::Sender<ProtoMessage>>>>,

    /// Broadcast channel for execution events (payloadType 2126).
    execution_tx: broadcast::Sender<ProtoMessage>,

    /// Broadcast channel for spot events (payloadType 2131).
    spot_tx: broadcast::Sender<ProtoMessage>,

    /// Broadcast channel for error events (payloadType 2132, 2142, 50).
    error_tx: broadcast::Sender<ProtoMessage>,

    /// Handle to the reader task.
    reader_handle: Option<tokio::task::JoinHandle<()>>,
}

impl MessageDispatcher {
    /// Create a new dispatcher and start the reader task.
    pub fn new(mut conn: CTraderConnection) -> Self {
        let pending: Arc<Mutex<HashMap<String, oneshot::Sender<ProtoMessage>>>> =
            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<ProtoMessage> {
        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<ProtoMessage> {
        self.execution_tx.subscribe()
    }

    /// Subscribe to spot price events.
    pub fn subscribe_spots(&self) -> broadcast::Receiver<ProtoMessage> {
        self.spot_tx.subscribe()
    }

    /// Subscribe to error events.
    pub fn subscribe_errors(&self) -> broadcast::Receiver<ProtoMessage> {
        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<tokio::net::TcpStream>,
                crate::codec::CTraderCodec,
            >,
        >,
        pending: Arc<Mutex<HashMap<String, oneshot::Sender<ProtoMessage>>>>,
        execution_tx: broadcast::Sender<ProtoMessage>,
        spot_tx: broadcast::Sender<ProtoMessage>,
        error_tx: broadcast::Sender<ProtoMessage>,
    ) {
        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() {
        assert!(BROADCAST_CAPACITY >= 64);
    }
}