rusty_ems/

execution_engine_lockfree.rs

//! Lock-free execution engine implementation using `crossbeam::queue::SegQueue`
//!
//! This implementation eliminates contention by using lock-free data structures
//! for all critical paths in the execution pipeline.

use std::sync::{
    Arc,
    atomic::{AtomicBool, AtomicUsize, Ordering},
};

use anyhow::{Result, anyhow};
use crossbeam::queue::SegQueue;
use dashmap::DashMap;
use flume::{Receiver, Sender};
use quanta::Clock;
use rust_decimal::Decimal;
use rusty_common::SmartString;

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

use crate::{Exchange, ExecutionEvent, ExecutionReport};

/// Work item for exchange operations with lock-free passing
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,
    },
    ModifyOrder {
        order_id: SmartString,
        new_price: Option<Decimal>,
        new_quantity: Option<Decimal>,
    },
    CancelAllOrders(Option<InstrumentId>),
}

/// Lock-free execution engine using `crossbeam::queue::SegQueue`
///
/// This implementation provides:
/// - Zero contention on the critical path
/// - Lock-free work distribution
/// - Efficient batch processing
/// - Cache-friendly memory layout
pub struct LockFreeExecutionEngine {
    /// Map of venue to exchange adapter
    exchanges: Arc<DashMap<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>>,

    /// Lock-free work queue for exchange operations
    work_queue: Arc<SegQueue<ExchangeWork>>,

    /// Atomic flag to signal shutdown
    shutdown: Arc<AtomicBool>,

    /// Worker thread count
    worker_count: Arc<AtomicUsize>,
}

impl LockFreeExecutionEngine {
    /// Create a new lock-free execution engine
    #[must_use]
    pub fn new(
        event_receiver: Receiver<ExecutionEvent>,
        report_sender: Sender<ExecutionReport>,
    ) -> Self {
        Self {
            exchanges: Arc::new(DashMap::new()),
            event_receiver,
            report_sender,
            clock: Clock::new(),
            order_venues: Arc::new(DashMap::new()),
            work_queue: Arc::new(SegQueue::new()),
            shutdown: Arc::new(AtomicBool::new(false)),
            worker_count: Arc::new(AtomicUsize::new(0)),
        }
    }

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

    /// Start the lock-free 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_queue = self.work_queue.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();
        let shutdown = self.shutdown.clone();
        let worker_count = self.worker_count.clone();

        // Start dedicated worker pool for exchange operations
        let worker_threads = std::env::var("EMS_WORKER_COUNT")
            .ok()
            .and_then(|s| s.parse::<usize>().ok())
            .unwrap_or(4)
            .max(1); // Ensure at least one worker

        for worker_id in 0..worker_threads {
            let work_queue = work_queue.clone();
            let shutdown = shutdown.clone();
            let worker_count = worker_count.clone();

            // Increment worker count
            worker_count.fetch_add(1, Ordering::SeqCst);

            tokio::spawn(async move {
                // Worker main loop - process work items from lock-free queue
                loop {
                    // Check shutdown flag
                    if shutdown.load(Ordering::Relaxed) {
                        break;
                    }

                    // Try to get work from the lock-free queue
                    if let Some(work) = work_queue.pop() {
                        // Process the work item
                        match work.operation {
                            ExchangeOperation::PlaceOrder(order) => {
                                if let Err(e) =
                                    work.exchange.place_order(order, work.report_sender).await
                                {
                                    log::error!("Worker {worker_id}: Failed to place order: {e}");
                                }
                            }
                            ExchangeOperation::CancelOrder { order_id } => {
                                if let Err(e) = work
                                    .exchange
                                    .cancel_order(order_id, work.report_sender)
                                    .await
                                {
                                    log::error!("Worker {worker_id}: Failed to cancel order: {e}");
                                }
                            }
                            ExchangeOperation::ModifyOrder {
                                order_id,
                                new_price,
                                new_quantity,
                            } => {
                                if let Err(e) = work
                                    .exchange
                                    .modify_order(
                                        order_id,
                                        new_price,
                                        new_quantity,
                                        work.report_sender,
                                    )
                                    .await
                                {
                                    log::error!("Worker {worker_id}: Failed to modify order: {e}");
                                }
                            }
                            ExchangeOperation::CancelAllOrders(instrument_id) => {
                                if let Err(e) = work
                                    .exchange
                                    .cancel_all_orders(instrument_id, work.report_sender)
                                    .await
                                {
                                    log::error!(
                                        "Worker {worker_id}: Failed to cancel all orders: {e}"
                                    );
                                }
                            }
                        }
                    } else {
                        // No work available, yield to prevent busy-waiting
                        tokio::time::sleep(tokio::time::Duration::from_micros(10)).await;
                    }
                }

                // Decrement worker count on exit
                worker_count.fetch_sub(1, Ordering::SeqCst);
                log::info!("Worker {worker_id} shutting down");
            });
        }

        // Start the main event processing loop
        tokio::spawn(async move {
            while let Ok(event) = event_receiver.recv_async().await {
                // Check shutdown flag
                if shutdown.load(Ordering::Relaxed) {
                    break;
                }

                match event {
                    ExecutionEvent::OrderAccepted(order) => {
                        let venue = order.venue;

                        // Route to the appropriate exchange
                        if let Some(exchange) = exchanges.get(&venue) {
                            // Store the venue for this order only after confirming exchange exists
                            order_venues.insert(order.id.to_string().into(), venue);

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

                            // Push to lock-free queue
                            work_queue.push(work);
                            // Note: With lock-free queue, we can't check if push failed
                            // as push() always succeeds (unbounded queue)
                        } else {
                            // Send rejection report if no exchange is available
                            let report = ExecutionReport {
                                id: {
                                    let mut id = SmartString::new();
                                    id.push_str("rej_");
                                    id.push_str(&order.id.to_string());
                                    id
                                },
                                order_id: order.id.to_string().into(),
                                exchange_timestamp: 0,
                                system_timestamp: rusty_common::time::get_epoch_timestamp_ns(),
                                instrument_id: InstrumentId::new(order.symbol.clone(), order.venue),
                                status: OrderStatus::Rejected,
                                filled_quantity: Decimal::ZERO,
                                remaining_quantity: order.quantity,
                                execution_price: None,
                                reject_reason: Some(
                                    format!("No exchange adapter for venue: {venue:?}").into(),
                                ),
                                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 = order_venues
                            .get(&order_id)
                            .map_or(original_order.venue, |v| *v);

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

                            // Push to lock-free queue
                            work_queue.push(work);
                        }
                    }

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

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

                            // Push to lock-free queue
                            work_queue.push(work);
                        }
                    }

                    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(),
                            };

                            // Push to lock-free queue
                            work_queue.push(work);
                        }
                    }

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

                            // Push to lock-free queue
                            work_queue.push(work);
                        }
                    }
                }
            }

            log::info!("Event processor shutting down");
        });

        Ok(())
    }

    /// Shutdown the execution engine gracefully
    pub async fn shutdown(&self) {
        // Signal shutdown
        self.shutdown.store(true, Ordering::SeqCst);

        // Wait for all workers to finish
        let mut retries = 0;
        while self.worker_count.load(Ordering::SeqCst) > 0 && retries < 100 {
            tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
            retries += 1;
        }

        if retries >= 100 {
            log::warn!("Timeout waiting for workers to shutdown");
        }
    }

    /// Get queue depth for monitoring
    #[must_use]
    pub fn queue_depth(&self) -> usize {
        self.work_queue.len()
    }

    /// 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()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_lock_free_engine_creation() {
        let (_event_tx, event_rx) = flume::unbounded();
        let (report_tx, _report_rx) = flume::unbounded();

        let engine = LockFreeExecutionEngine::new(event_rx, report_tx);
        assert_eq!(engine.queue_depth(), 0);
    }
}