rusty_strategy/

engine.rs

use rusty_common::collections::{FxHashMap, SmallSignalVec, SmallStrategyVec};
use std::sync::Arc;

use anyhow::{Result, anyhow};
use flume::{Receiver, Sender};
use parking_lot::RwLock;
use quanta::Clock;
use smartstring::alias::String;
use tokio::sync::mpsc;
use tokio::time::{self, Duration};

use rusty_model::{
    data::{bar::Bar, book_snapshot::OrderBookSnapshot, market_trade::MarketTrade},
    instruments::InstrumentId,
};

use crate::signals::Signal;
use crate::strategy::Strategy;

/// Configuration options for the strategy engine
#[derive(Debug, Clone)]
pub struct StrategyEngineConfig {
    /// Timer interval in milliseconds
    pub timer_interval_ms: u64,

    /// Channel capacity for trade messages
    pub trade_channel_capacity: usize,

    /// Channel capacity for order book depth messages
    pub depth_channel_capacity: usize,

    /// Channel capacity for bar messages
    pub bar_channel_capacity: usize,

    /// Channel capacity for signal messages
    pub signal_channel_capacity: usize,
}

impl Default for StrategyEngineConfig {
    fn default() -> Self {
        Self {
            timer_interval_ms: 100,
            trade_channel_capacity: 10000,
            depth_channel_capacity: 1000,
            bar_channel_capacity: 1000,
            signal_channel_capacity: 1000,
        }
    }
}

/// Strategy engine that coordinates data flow and strategy execution
pub struct StrategyEngine {
    /// Engine configuration
    config: StrategyEngineConfig,

    /// Registered strategies
    strategies: Arc<RwLock<FxHashMap<String, Arc<dyn Strategy>>>>,

    /// Map of instrument IDs to interested strategies (using `SmallStrategyVec` for performance)
    instrument_map: Arc<RwLock<FxHashMap<InstrumentId, SmallStrategyVec<String>>>>,

    /// Clock for precise timing
    clock: Clock,

    /// Sender for trade messages
    trade_sender: mpsc::Sender<MarketTrade>,

    /// Sender for order book depth messages
    depth_sender: mpsc::Sender<OrderBookSnapshot>,

    /// Sender for bar messages
    bar_sender: mpsc::Sender<Bar>,

    /// Sender for signal messages (output)
    signal_sender: Sender<Signal>,

    /// Receiver for signal messages (output)
    signal_receiver: Receiver<Signal>,

    /// Flag indicating if the engine is running
    running: Arc<RwLock<bool>>,
}

impl Default for StrategyEngine {
    fn default() -> Self {
        Self::new()
    }
}

impl StrategyEngine {
    /// Creates a new strategy engine with the default configuration
    #[must_use]
    pub fn new() -> Self {
        Self::with_config(StrategyEngineConfig::default())
    }

    /// Creates a new strategy engine with a custom configuration
    #[must_use]
    pub fn with_config(config: StrategyEngineConfig) -> Self {
        // Create channels for data flow
        let (trade_sender, _trade_receiver) = mpsc::channel(config.trade_channel_capacity);
        let (depth_sender, _depth_receiver) = mpsc::channel(config.depth_channel_capacity);
        let (bar_sender, _bar_receiver) = mpsc::channel(config.bar_channel_capacity);
        let (signal_sender, signal_receiver) = flume::bounded(config.signal_channel_capacity);

        Self {
            config,
            strategies: Arc::new(RwLock::new(FxHashMap::default())),
            instrument_map: Arc::new(RwLock::new(FxHashMap::default())),
            clock: Clock::new(),
            trade_sender,
            depth_sender,
            bar_sender,
            signal_sender,
            signal_receiver,
            running: Arc::new(RwLock::new(false)),
        }
    }

    /// Registers a strategy with the engine
    pub fn register_strategy(&self, strategy: Arc<dyn Strategy>) -> Result<()> {
        let strategy_id: String = String::from(strategy.id());

        // Register strategy
        self.strategies
            .write()
            .insert(strategy_id.clone(), strategy.clone());

        // Update instrument map
        let instruments = strategy.instruments();
        let mut instrument_map = self.instrument_map.write();

        for instrument in instruments {
            let entry = instrument_map.entry(instrument.clone()).or_default();
            entry.push(strategy_id.clone());
        }

        // Initialize the strategy
        tokio::spawn(async move {
            if let Err(e) = strategy.initialize().await {
                log::error!("Failed to initialize strategy {strategy_id}: {e}");
            }
        });

        Ok(())
    }

    /// Unregisters a strategy from the engine
    pub async fn unregister_strategy(&self, strategy_id: &str) -> Result<()> {
        // Get the strategy first to shut it down
        let strategy = match self.strategies.read().get(strategy_id) {
            Some(s) => s.clone(),
            None => return Err(anyhow!("Strategy not found: {}", strategy_id)),
        };

        // Shut down the strategy
        strategy.shutdown().await?;

        // Remove from strategies map
        self.strategies.write().remove(strategy_id);

        // Update instrument map
        let mut instrument_map = self.instrument_map.write();
        for strategies in instrument_map.values_mut() {
            strategies.retain(|id| id != strategy_id);
        }

        // Clean up empty entries
        instrument_map.retain(|_, strategies| !strategies.is_empty());

        Ok(())
    }

    /// Starts the strategy engine
    pub async fn start(&mut self) -> Result<()> {
        // Check if already running
        if *self.running.read() {
            return Err(anyhow!("Strategy engine is already running"));
        }

        // Set running flag
        *self.running.write() = true;

        // Create separate clones for each task to avoid ownership issues
        let running_trade = self.running.clone();
        let strategies_trade = self.strategies.clone();
        let instrument_map_trade = self.instrument_map.clone();
        let signal_sender_trade = self.signal_sender.clone();

        let running_depth = self.running.clone();
        let strategies_depth = self.strategies.clone();
        let instrument_map_depth = self.instrument_map.clone();
        let signal_sender_depth = self.signal_sender.clone();

        let running_bar = self.running.clone();
        let strategies_bar = self.strategies.clone();
        let instrument_map_bar = self.instrument_map.clone();
        let signal_sender_bar = self.signal_sender.clone();

        let timer_interval = self.config.timer_interval_ms;

        // Create new channel pairs since we need to move the receivers
        let (new_trade_sender, mut trade_receiver) =
            mpsc::channel::<MarketTrade>(self.config.trade_channel_capacity);
        let (new_depth_sender, mut depth_receiver) =
            mpsc::channel::<OrderBookSnapshot>(self.config.depth_channel_capacity);
        let (new_bar_sender, mut bar_receiver) =
            mpsc::channel::<Bar>(self.config.bar_channel_capacity);

        // Spawn trade processing task
        tokio::spawn(async move {
            while *running_trade.read() {
                if let Some(trade) = trade_receiver.recv().await {
                    // Find strategies interested in this instrument
                    let interested_strategies = {
                        let map = instrument_map_trade.read();
                        match map.get(&trade.instrument_id) {
                            Some(strategies) => strategies.clone(),
                            None => continue,
                        }
                    };

                    // Process trade with each interested strategy
                    for strategy_id in interested_strategies {
                        let strategies_guard = strategies_trade.read();
                        let strategy = match strategies_guard.get(&strategy_id) {
                            Some(s) => s.clone(),
                            None => continue,
                        };

                        // Clone values for async task
                        let trade_clone = trade.clone();
                        let signal_sender_clone = signal_sender_trade.clone();

                        // Process trade asynchronously
                        tokio::spawn(Self::process_trade_with_signals(
                            strategy,
                            trade_clone,
                            signal_sender_clone,
                            strategy_id,
                        ));
                    }
                }
            }
        });

        // Spawn depth processing task
        tokio::spawn(async move {
            while *running_depth.read() {
                if let Some(depth) = depth_receiver.recv().await {
                    // Find strategies interested in this instrument
                    let interested_strategies = {
                        let map = instrument_map_depth.read();
                        match map.get(&depth.instrument_id) {
                            Some(strategies) => strategies.clone(),
                            None => continue,
                        }
                    };

                    // Process depth with each interested strategy
                    for strategy_id in interested_strategies {
                        let strategies_guard = strategies_depth.read();
                        let strategy = match strategies_guard.get(&strategy_id) {
                            Some(s) => s.clone(),
                            None => continue,
                        };

                        // Clone values for async task
                        let depth_clone = depth.clone();
                        let signal_sender_clone = signal_sender_depth.clone();

                        // Process depth asynchronously
                        tokio::spawn(Self::process_depth_with_signals(
                            strategy,
                            depth_clone,
                            signal_sender_clone,
                            strategy_id,
                        ));
                    }
                }
            }
        });

        // Spawn bar processing task
        tokio::spawn(async move {
            while *running_bar.read() {
                if let Some(bar) = bar_receiver.recv().await {
                    // Find strategies interested in this instrument
                    // For now, we'll check all strategies since Bar doesn't have venue info
                    // TODO: Consider adding venue to Bar or using a different mapping strategy
                    let interested_strategies = {
                        let map = instrument_map_bar.read();
                        map.iter()
                            .filter(|(instrument_id, _)| {
                                instrument_id.symbol == bar.bar_type.symbol
                            })
                            .flat_map(|(_, strategy_list)| strategy_list.clone())
                            .collect::<SmallStrategyVec<_>>()
                    };

                    // Process bar with each interested strategy
                    for strategy_id in interested_strategies {
                        let strategies_guard = strategies_bar.read();
                        let strategy = match strategies_guard.get(&strategy_id) {
                            Some(s) => s.clone(),
                            None => continue,
                        };

                        // Clone values for async task
                        let bar_clone = bar.clone();
                        let signal_sender_clone = signal_sender_bar.clone();

                        // Process bar asynchronously
                        tokio::spawn(Self::process_bar_with_signals(
                            strategy,
                            bar_clone,
                            signal_sender_clone,
                            strategy_id,
                        ));
                    }
                }
            }
        });

        // Clone for timer task
        let running_timer = self.running.clone();
        let strategies_timer = self.strategies.clone();
        let signal_sender_timer = self.signal_sender.clone();
        let clock = self.clock.clone();

        // Spawn timer task
        tokio::spawn(async move {
            let mut interval = time::interval(Duration::from_millis(timer_interval));

            while *running_timer.read() {
                interval.tick().await;

                // Get current timestamp in nanoseconds
                let timestamp_ns = clock.raw();

                // Process timer event with each strategy
                let strategy_ids: SmallStrategyVec<String> =
                    strategies_timer.read().keys().cloned().collect();

                for strategy_id in strategy_ids {
                    let strategies_guard = strategies_timer.read();
                    let strategy = match strategies_guard.get(&strategy_id) {
                        Some(s) => s.clone(),
                        None => continue,
                    };

                    // Clone values for async task
                    let signal_sender_clone = signal_sender_timer.clone();

                    // Process timer event asynchronously
                    tokio::spawn(Self::process_timer_signals(
                        strategy,
                        timestamp_ns,
                        signal_sender_clone,
                        strategy_id,
                    ));
                }
            }
        });

        // Save the new channel senders
        self.trade_sender = new_trade_sender;
        self.depth_sender = new_depth_sender;
        self.bar_sender = new_bar_sender;

        Ok(())
    }

    /// Stops the strategy engine
    pub async fn stop(&self) -> Result<()> {
        // Check if already stopped
        if !*self.running.read() {
            return Err(anyhow!("Strategy engine is not running"));
        }

        // Set running flag to false
        *self.running.write() = false;

        // Collect all strategies first before awaiting to avoid holding the lock across await points
        let strategy_pairs: SmallStrategyVec<(String, Arc<dyn Strategy>)> = {
            let strategies = self.strategies.read();
            strategies
                .iter()
                .map(|(id, strategy)| (id.clone(), strategy.clone()))
                .collect()
        };

        // Shut down all strategies
        for (id, strategy) in strategy_pairs {
            if let Err(e) = strategy.shutdown().await {
                log::error!("Error shutting down strategy {id}: {e}");
            }
        }

        Ok(())
    }

    /// Gets the trade sender channel
    #[must_use]
    pub fn trade_sender(&self) -> mpsc::Sender<MarketTrade> {
        self.trade_sender.clone()
    }

    /// Gets the depth sender channel
    #[must_use]
    pub fn depth_sender(&self) -> mpsc::Sender<OrderBookSnapshot> {
        self.depth_sender.clone()
    }

    /// Gets the bar sender channel
    #[must_use]
    pub fn bar_sender(&self) -> mpsc::Sender<Bar> {
        self.bar_sender.clone()
    }

    /// Gets the signal receiver channel
    #[must_use]
    pub fn signal_receiver(&self) -> Receiver<Signal> {
        self.signal_receiver.clone()
    }

    // Helper methods to reduce nesting

    async fn process_trade_with_signals(
        strategy: Arc<dyn Strategy>,
        trade: MarketTrade,
        signal_sender: Sender<Signal>,
        strategy_id: String,
    ) {
        match strategy.process_trade(trade).await {
            Ok(signals) => Self::send_signals(signals, signal_sender).await,
            Err(e) => log::error!("Error processing trade in strategy {strategy_id}: {e}"),
        }
    }

    async fn process_depth_with_signals(
        strategy: Arc<dyn Strategy>,
        depth: OrderBookSnapshot,
        signal_sender: Sender<Signal>,
        strategy_id: String,
    ) {
        match strategy.process_depth(depth).await {
            Ok(signals) => Self::send_signals(signals, signal_sender).await,
            Err(e) => log::error!("Error processing depth in strategy {strategy_id}: {e}"),
        }
    }

    async fn process_bar_with_signals(
        strategy: Arc<dyn Strategy>,
        bar: Bar,
        signal_sender: Sender<Signal>,
        strategy_id: String,
    ) {
        match strategy.process_bar(bar).await {
            Ok(signals) => Self::send_signals(signals, signal_sender).await,
            Err(e) => log::error!("Error processing bar in strategy {strategy_id}: {e}"),
        }
    }

    async fn process_timer_signals(
        strategy: Arc<dyn Strategy>,
        timestamp_ns: u64,
        signal_sender: Sender<Signal>,
        strategy_id: String,
    ) {
        match strategy.on_timer(timestamp_ns).await {
            Ok(signals) => Self::send_signals(signals, signal_sender).await,
            Err(e) => log::error!("Error processing timer in strategy {strategy_id}: {e}"),
        }
    }

    async fn send_signals(signals: SmallSignalVec<Signal>, signal_sender: Sender<Signal>) {
        for signal in signals {
            if let Err(e) = signal_sender.send_async(signal).await {
                log::error!("Failed to send signal: {e}");
            }
        }
    }
}