rusty_ems/

execution_engine.rs

use rusty_common::SmartString;
use rusty_common::collections::FxHashMap;
use rusty_common::zerocopy::{BufferedStringOps, ZeroCopyStringUtils};
use std::sync::Arc;

use anyhow::{Result, anyhow};
use async_trait::async_trait;
use dashmap::DashMap;
use flume::{Receiver, Sender};
use quanta::Clock;
use rust_decimal::Decimal;
use smallvec::SmallVec;

use rusty_model::{
    enums::OrderStatus, instruments::InstrumentId, trading_order::Order, venues::Venue,
};

/// Events related to order execution
#[derive(Debug, Clone)]
pub enum ExecutionEvent {
    /// Order accepted by the exchange
    OrderAccepted(Order),

    /// Request to cancel an existing order
    CancelOrder {
        /// Order ID to cancel
        order_id: SmartString,

        /// Original order details
        original_order: Order,
    },

    /// Request to modify an existing order
    ModifyOrder {
        /// Order ID to modify
        order_id: SmartString,

        /// Original order details
        original_order: Order,

        /// New price (if modifying price)
        new_price: Option<Decimal>,

        /// New quantity (if modifying quantity)
        new_quantity: Option<Decimal>,
    },

    /// Request to cancel all open orders for a specific instrument
    CancelAllOrders(InstrumentId),

    /// Request to cancel all open orders across all instruments
    CancelAllOrdersGlobally,
}

/// Execution report from the exchange
#[derive(Debug, Clone)]
pub struct ExecutionReport {
    /// Unique identifier for this execution report
    pub id: SmartString,

    /// Order ID this report is for
    pub order_id: SmartString,

    /// Exchange timestamp (in nanoseconds)
    pub exchange_timestamp: u64,

    /// System timestamp (in nanoseconds)
    pub system_timestamp: u64,

    /// Instrument ID
    pub instrument_id: InstrumentId,

    /// New status of the order
    pub status: OrderStatus,

    /// Filled quantity in this execution (0 for non-fills)
    pub filled_quantity: Decimal,

    /// Remaining quantity to be filled
    pub remaining_quantity: Decimal,

    /// Price of this execution (for fills)
    pub execution_price: Option<Decimal>,

    /// Reason for rejection (if applicable)
    pub reject_reason: Option<SmartString>,

    /// Exchange-specific execution ID
    pub exchange_execution_id: Option<SmartString>,

    /// Whether this is the final report for the order
    pub is_final: bool,
}

/// Work item for exchange operations
struct ExchangeWork {
    exchange: Arc<dyn Exchange>,
    operation: ExchangeOperation,
    report_sender: Sender<ExecutionReport>,
}

/// Types of exchange operations
#[derive(Debug)]
enum ExchangeOperation {
    PlaceOrder(Order),
    CancelOrder {
        order_id: SmartString,
        instrument_id: InstrumentId,
    },
    ModifyOrder {
        order_id: SmartString,
        instrument_id: InstrumentId,
        new_price: Option<Decimal>,
        new_quantity: Option<Decimal>,
    },
    CancelAllOrders(Option<InstrumentId>),
}

/// An execution management engine that routes orders to appropriate exchange connections
pub struct ExecutionEngine {
    /// Map of venue to exchange adapter
    exchanges: FxHashMap<Venue, Arc<dyn Exchange>>,

    /// Receiver for execution events from OMS
    event_receiver: Receiver<ExecutionEvent>,

    /// Sender for execution reports to OMS
    report_sender: Sender<ExecutionReport>,

    /// High-precision clock
    clock: Clock,

    /// Lock-free map of order ID to venue for routing
    order_venues: Arc<DashMap<SmartString, Venue>>,

    /// Work queue for exchange operations
    work_sender: Sender<ExchangeWork>,
    work_receiver: Receiver<ExchangeWork>,

    /// Number of worker threads for processing exchange operations
    worker_count: usize,
}

impl ExecutionEngine {
    /// Create a new execution engine
    #[must_use]
    pub fn new(
        event_receiver: Receiver<ExecutionEvent>,
        report_sender: Sender<ExecutionReport>,
    ) -> Self {
        Self::with_worker_count(event_receiver, report_sender, 4)
    }

    /// Create a new execution engine with a specific worker count
    #[must_use]
    pub fn with_worker_count(
        event_receiver: Receiver<ExecutionEvent>,
        report_sender: Sender<ExecutionReport>,
        worker_count: usize,
    ) -> Self {
        let (work_sender, work_receiver) = flume::unbounded();

        Self {
            exchanges: FxHashMap::default(),
            event_receiver,
            report_sender,
            clock: Clock::new(),
            order_venues: Arc::new(DashMap::new()),
            work_sender,
            work_receiver,
            worker_count: worker_count.max(1), // Ensure at least one worker
        }
    }

    /// Register an exchange adapter with this execution engine
    pub fn register_exchange(&mut self, exchange: Arc<dyn Exchange>) {
        let venue = exchange.venue();
        self.exchanges.insert(venue, exchange);
    }

    /// Start the execution engine
    pub async fn start(&self) -> Result<()> {
        if self.exchanges.is_empty() {
            return Err(anyhow!("No exchanges registered"));
        }

        // Clone values needed for the loops
        let event_receiver = self.event_receiver.clone();
        let work_sender = self.work_sender.clone();
        let exchanges = self.exchanges.clone();
        let order_venues = self.order_venues.clone();
        let clock = self.clock.clone();
        let report_sender = self.report_sender.clone();

        // Start dedicated worker pool for exchange operations
        let work_receiver = self.work_receiver.clone();
        let worker_count = self.worker_count;

        for worker_id in 0..worker_count {
            let work_receiver = work_receiver.clone();
            tokio::spawn(async move {
                while let Ok(work) = work_receiver.recv_async().await {
                    match work.operation {
                        ExchangeOperation::PlaceOrder(order) => {
                            // Create references for rejection report before moving order
                            let order_id_for_rejection = order.id;
                            let symbol_for_rejection = order.symbol.clone();
                            let venue_for_rejection = order.venue;
                            let quantity_for_rejection = order.quantity;

                            if let Err(e) = work
                                .exchange
                                .place_order(order, work.report_sender.clone())
                                .await
                            {
                                log::error!("Worker {worker_id}: Failed to place order: {e}");

                                // Send rejection ExecutionReport using zero-copy utilities
                                let rejection = ExecutionReport {
                                    id: BufferedStringOps::create_rejection_id_buffered(
                                        &order_id_for_rejection,
                                    ),
                                    order_id: ZeroCopyStringUtils::display_to_smart_string(
                                        &order_id_for_rejection,
                                    ),
                                    exchange_timestamp: 0, // No exchange interaction occurred
                                    system_timestamp: 0,   // Will be set by receiving system
                                    instrument_id: InstrumentId::new(
                                        symbol_for_rejection,
                                        venue_for_rejection,
                                    ),
                                    status: OrderStatus::Rejected,
                                    filled_quantity: Decimal::ZERO,
                                    remaining_quantity: quantity_for_rejection,
                                    execution_price: None,
                                    reject_reason: Some(
                                        ZeroCopyStringUtils::error_to_smart_string(&e),
                                    ),
                                    exchange_execution_id: None,
                                    is_final: true,
                                };

                                // Try to send rejection report (ignore send errors to avoid infinite loops)
                                if let Err(send_err) =
                                    work.report_sender.send_async(rejection).await
                                {
                                    log::warn!(
                                        "Worker {worker_id}: Failed to send rejection report: {send_err}"
                                    );
                                }
                            }
                        }
                        ExchangeOperation::CancelOrder {
                            order_id,
                            instrument_id,
                        } => {
                            // Create reference for failure report before moving order_id
                            let order_id_for_failure = order_id.clone();

                            if let Err(e) = work
                                .exchange
                                .cancel_order(order_id, work.report_sender.clone())
                                .await
                            {
                                log::error!("Worker {worker_id}: Failed to cancel order: {e}");

                                // Send cancellation failure report using zero-copy utilities
                                let cancel_failure = ExecutionReport {
                                    id: {
                                        let mut id = SmartString::new();
                                        id.push_str("cancel_fail_");
                                        id.push_str(&order_id_for_failure);
                                        id
                                    },
                                    order_id: order_id_for_failure,
                                    exchange_timestamp: 0,
                                    system_timestamp: 0,
                                    instrument_id,
                                    status: OrderStatus::Rejected,
                                    filled_quantity: Decimal::ZERO,
                                    remaining_quantity: Decimal::ZERO,
                                    execution_price: None,
                                    reject_reason: Some({
                                        let mut reason = SmartString::new();
                                        reason.push_str("Cancel failed: ");
                                        reason.push_str(&e.to_string());
                                        reason
                                    }),
                                    exchange_execution_id: None,
                                    is_final: false, // Order may still be active
                                };

                                if let Err(send_err) =
                                    work.report_sender.send_async(cancel_failure).await
                                {
                                    log::warn!(
                                        "Worker {worker_id}: Failed to send cancel failure report: {send_err}"
                                    );
                                }
                            }
                        }
                        ExchangeOperation::ModifyOrder {
                            order_id,
                            instrument_id,
                            new_price,
                            new_quantity,
                        } => {
                            // Create reference for failure report before moving order_id
                            let order_id_for_failure = order_id.clone();

                            if let Err(e) = work
                                .exchange
                                .modify_order(
                                    order_id,
                                    new_price,
                                    new_quantity,
                                    work.report_sender.clone(),
                                )
                                .await
                            {
                                log::error!("Worker {worker_id}: Failed to modify order: {e}");

                                // Send modification failure report using zero-copy utilities
                                let modify_failure = ExecutionReport {
                                    id: {
                                        let mut id = SmartString::new();
                                        id.push_str("modify_fail_");
                                        id.push_str(&order_id_for_failure);
                                        id
                                    },
                                    order_id: order_id_for_failure,
                                    exchange_timestamp: 0,
                                    system_timestamp: 0,
                                    instrument_id,
                                    status: OrderStatus::Rejected, // Order modification failed
                                    filled_quantity: Decimal::ZERO,
                                    remaining_quantity: Decimal::ZERO,
                                    execution_price: None,
                                    reject_reason: Some({
                                        let mut reason = SmartString::new();
                                        reason.push_str("Modify failed: ");
                                        reason.push_str(&e.to_string());
                                        reason
                                    }),
                                    exchange_execution_id: None,
                                    is_final: false,
                                };

                                if let Err(send_err) =
                                    work.report_sender.send_async(modify_failure).await
                                {
                                    log::warn!(
                                        "Worker {worker_id}: Failed to send modify failure report: {send_err}"
                                    );
                                }
                            }
                        }
                        ExchangeOperation::CancelAllOrders(instrument_id) => {
                            if let Err(e) = work
                                .exchange
                                .cancel_all_orders(
                                    instrument_id.clone(),
                                    work.report_sender.clone(),
                                )
                                .await
                            {
                                log::error!("Worker {worker_id}: Failed to cancel all orders: {e}");

                                // Send cancel all failure report
                                let cancel_all_failure = ExecutionReport {
                                    id: {
                                        let mut id = SmartString::new();
                                        id.push_str("cancel_all_fail_");
                                        id.push_str(
                                            instrument_id
                                                .as_ref()
                                                .map_or("ALL", |i| i.symbol.as_str()),
                                        );
                                        id
                                    },
                                    order_id: "CANCEL_ALL_FAILED".into(),
                                    exchange_timestamp: 0,
                                    system_timestamp: 0,
                                    instrument_id: instrument_id
                                        .unwrap_or_else(|| InstrumentId::new("ALL", Venue::Test)),
                                    status: OrderStatus::Rejected,
                                    filled_quantity: Decimal::ZERO,
                                    remaining_quantity: Decimal::ZERO,
                                    execution_price: None,
                                    reject_reason: Some({
                                        let mut reason = SmartString::new();
                                        reason.push_str("Cancel all failed: ");
                                        reason.push_str(&e.to_string());
                                        reason
                                    }),
                                    exchange_execution_id: None,
                                    is_final: false,
                                };

                                if let Err(send_err) =
                                    work.report_sender.send_async(cancel_all_failure).await
                                {
                                    log::warn!(
                                        "Worker {worker_id}: Failed to send cancel all failure report: {send_err}"
                                    );
                                }
                            }
                        }
                    }
                }
            });
        }

        // Start the main event processing loop
        tokio::spawn(async move {
            while let Ok(event) = event_receiver.recv_async().await {
                match event {
                    ExecutionEvent::OrderAccepted(order) => {
                        let venue = order.venue;

                        // Route to the appropriate exchange
                        if let Some(exchange) = exchanges.get(&venue) {
                            let order_id_str =
                                ZeroCopyStringUtils::display_to_smart_string(&order.id);

                            // Store the venue for this order only after confirming exchange exists
                            order_venues.insert(order_id_str.clone(), venue);

                            let work = ExchangeWork {
                                exchange: exchange.clone(),
                                operation: ExchangeOperation::PlaceOrder(order),
                                report_sender: report_sender.clone(),
                            };

                            if let Err(e) = work_sender.send_async(work).await {
                                log::error!("Failed to queue place order work: {e}");
                                // Remove from venues map if we failed to queue the work
                                order_venues.remove(&order_id_str);
                            }
                        } else {
                            // Send rejection report if no exchange is available using zero-copy utilities
                            let report = ExecutionReport {
                                id: BufferedStringOps::create_rejection_id_buffered(&order.id),
                                order_id: ZeroCopyStringUtils::display_to_smart_string(&order.id),
                                exchange_timestamp: 0,
                                system_timestamp: rusty_common::time::get_epoch_timestamp_ns(),
                                instrument_id: InstrumentId::new(order.symbol, order.venue),
                                status: OrderStatus::Rejected,
                                filled_quantity: Decimal::ZERO,
                                remaining_quantity: order.quantity,
                                execution_price: None,
                                reject_reason: Some({
                                    let mut reason = SmartString::new();
                                    reason.push_str("No exchange adapter for venue: ");
                                    reason.push_str(&format!("{venue:?}"));
                                    reason
                                }),
                                exchange_execution_id: None,
                                is_final: true,
                            };

                            if let Err(e) = report_sender.send_async(report).await {
                                log::error!("Failed to send rejection report: {e}");
                            }
                        }
                    }

                    ExecutionEvent::CancelOrder {
                        order_id,
                        original_order,
                    } => {
                        // Lookup the venue for this order (lock-free)
                        let venue = if let Some(v) = order_venues.get(&order_id) {
                            *v
                        } else {
                            original_order.venue
                        };

                        // Route to the appropriate exchange
                        if let Some(exchange) = exchanges.get(&venue) {
                            let work = ExchangeWork {
                                exchange: exchange.clone(),
                                operation: ExchangeOperation::CancelOrder {
                                    order_id,
                                    instrument_id: InstrumentId::new(
                                        original_order.symbol.clone(),
                                        original_order.venue,
                                    ),
                                },
                                report_sender: report_sender.clone(),
                            };

                            if let Err(e) = work_sender.send_async(work).await {
                                log::error!("Failed to queue cancel order work: {e}");
                            }
                        }
                    }

                    ExecutionEvent::ModifyOrder {
                        order_id,
                        original_order,
                        new_price,
                        new_quantity,
                    } => {
                        // Lookup the venue for this order (lock-free)
                        let venue = if let Some(v) = order_venues.get(&order_id) {
                            *v
                        } else {
                            original_order.venue
                        };

                        // Route to the appropriate exchange
                        if let Some(exchange) = exchanges.get(&venue) {
                            let work = ExchangeWork {
                                exchange: exchange.clone(),
                                operation: ExchangeOperation::ModifyOrder {
                                    order_id,
                                    instrument_id: InstrumentId::new(
                                        original_order.symbol.clone(),
                                        original_order.venue,
                                    ),
                                    new_price,
                                    new_quantity,
                                },
                                report_sender: report_sender.clone(),
                            };

                            if let Err(e) = work_sender.send_async(work).await {
                                log::error!("Failed to queue modify order work: {e}");
                            }
                        }
                    }

                    ExecutionEvent::CancelAllOrders(instrument_id) => {
                        let venue = instrument_id.venue;

                        // Route to the appropriate exchange
                        if let Some(exchange) = exchanges.get(&venue) {
                            let work = ExchangeWork {
                                exchange: exchange.clone(),
                                operation: ExchangeOperation::CancelAllOrders(Some(instrument_id)),
                                report_sender: report_sender.clone(),
                            };

                            if let Err(e) = work_sender.send_async(work).await {
                                log::error!("Failed to queue cancel all orders work: {e}");
                            }
                        }
                    }

                    ExecutionEvent::CancelAllOrdersGlobally => {
                        // Send cancel all orders to every registered exchange
                        for exchange in exchanges.values() {
                            let work = ExchangeWork {
                                exchange: exchange.clone(),
                                operation: ExchangeOperation::CancelAllOrders(None),
                                report_sender: report_sender.clone(),
                            };

                            if let Err(e) = work_sender.send_async(work).await {
                                log::error!("Failed to queue global cancel all orders work: {e}");
                            }
                        }
                    }
                }
            }
        });

        Ok(())
    }

    /// Get a execution report channel that sends to this engine
    #[must_use]
    pub fn report_sender(&self) -> Sender<ExecutionReport> {
        self.report_sender.clone()
    }

    /// Get an execution event channel that this engine listens to
    #[must_use]
    pub fn event_receiver(&self) -> Receiver<ExecutionEvent> {
        self.event_receiver.clone()
    }
}

/// Interface for connecting to trading venues
#[async_trait]
pub trait Exchange: Send + Sync {
    /// Returns the venue this exchange connects to
    fn venue(&self) -> Venue;

    /// Places a new order on the exchange
    async fn place_order(&self, order: Order, report_sender: Sender<ExecutionReport>)
    -> Result<()>;

    /// Cancels an existing order
    async fn cancel_order(
        &self,
        order_id: SmartString,
        report_sender: Sender<ExecutionReport>,
    ) -> Result<()>;

    /// Modifies an existing order
    async fn modify_order(
        &self,
        order_id: SmartString,
        new_price: Option<Decimal>,
        new_quantity: Option<Decimal>,
        report_sender: Sender<ExecutionReport>,
    ) -> Result<()>;

    /// Cancels all open orders, optionally filtered by instrument
    async fn cancel_all_orders(
        &self,
        instrument_id: Option<InstrumentId>,
        report_sender: Sender<ExecutionReport>,
    ) -> Result<()>;

    /// Retrieves the status of an order
    async fn get_order_status(&self, order_id: &str) -> Result<OrderStatus>;

    /// Connects to the exchange
    async fn connect(&self, report_sender: Sender<ExecutionReport>) -> Result<()>;

    /// Disconnects from the exchange
    async fn disconnect(&self) -> Result<()>;

    /// Checks if connected to the exchange
    async fn is_connected(&self) -> bool;

    /// Gets a list of instruments supported by this exchange
    async fn get_instruments(&self) -> Result<SmallVec<[InstrumentId; 32]>>;

    /// Sends a FIX message to the exchange
    async fn send_fix_message(&self, message: Vec<u8>) -> Result<()>;

    /// Receives a FIX message from the exchange
    async fn receive_fix_message(&self) -> Result<Vec<u8>>;
}