rustrade-execution 0.15.0

use crate::{
    AccountEventKind, UnindexedAccountEvent, UnindexedAccountSnapshot,
    balance::AssetBalance,
    client::{ExecutionClient, OrderStatusClient},
    error::{
        ConnectivityError, OrderError, StreamTerminationReason, UnindexedClientError,
        UnindexedOrderError,
    },
    exchange::mock::request::MockExchangeRequest,
    fee::{FeeAssetConfig, FeeModelConfig, SpotFeeAsset},
    fill::{MarketOrderTiming, SimFillConfig, StaleMarketPolicy},
    margin::SpotMarginConfig,
    order::{
        Order, OrderKey, UnindexedInactiveOrder, UnindexedOrderKey,
        request::{OrderRequestCancel, OrderRequestOpen, UnindexedOrderResponseCancel},
        state::{Open, OrderState, UnindexedOrderState},
    },
    slippage::SlippageConfig,
    trade::TradesRead,
};
use chrono::{DateTime, Utc};
use derive_more::Constructor;
use futures::{StreamExt, stream::BoxStream};
use rustrade_instrument::{
    asset::name::AssetNameExchange,
    exchange::ExchangeId,
    instrument::{kind::InstrumentKindDiscriminant, name::InstrumentNameExchange},
};
use serde::{Deserialize, Serialize};
use tokio::sync::{broadcast, mpsc, oneshot};
use tokio_stream::wrappers::BroadcastStream;
use tracing::error;

#[derive(
    Debug, Clone, Eq, PartialEq, Ord, PartialOrd, Hash, Deserialize, Serialize, Constructor,
)]
pub struct MockExecutionConfig {
    pub mocked_exchange: ExchangeId,
    pub initial_state: UnindexedAccountSnapshot,
    pub latency_ms: u64,
    /// Fee model used by the mock exchange to compute trading fees.
    ///
    /// Defaults to [`FeeModelConfig::Zero`]. Use [`FeeModelConfig::Percentage`]
    /// for spot/futures simulation (e.g. 0.1% taker fee).
    #[serde(default)]
    pub fee_model: FeeModelConfig,
    /// Fill model used by the mock exchange to compute execution prices.
    ///
    /// Defaults to [`SimFillConfig::LastPrice`], which fills at the snapshot's `last_price`,
    /// falling back to the side of the book a taker would cross to. Switch to
    /// [`SimFillConfig::BidAsk`] or [`SimFillConfig::Midpoint`] for more realistic spread-cost
    /// simulation when market prices are injected alongside orders.
    #[serde(default)]
    pub fill_model: SimFillConfig,
    /// Which asset the mock exchange charges a spot fill's fee in.
    ///
    /// Defaults to [`SpotFeeAsset::Quote`], on both sides. [`SpotFeeAsset::Received`] takes a
    /// buy's fee out of the base it receives, as Binance spot does by default.
    #[serde(default)]
    pub spot_fee_asset: SpotFeeAsset,
    /// How far the mock exchange moves each taker fill against the order, after the
    /// [`fill_model`](Self::fill_model) has priced it.
    ///
    /// Defaults to [`SlippageConfig::none`]. A seeded [`SlippageConfig::uniform`] draws a fresh
    /// rate for each fill, reproducibly from its seed. See [`SlippageConfig`] for what it applies
    /// to and how.
    #[serde(default)]
    pub slippage: SlippageConfig,
    /// When the mock exchange fills a market order.
    ///
    /// Defaults to [`MarketOrderTiming::OnArrival`]: against the market the venue holds when the
    /// order arrives, which on bar data is the close of the bar the sender has just seen.
    /// [`MarketOrderTiming::NextBar`] holds it for the instrument's next bar instead, and needs a
    /// driver that feeds the venue bars, as `SimRunner` does on a candle feed.
    #[serde(default)]
    pub market_order_timing: MarketOrderTiming,
    /// What the mock exchange does with an order when the instrument's market is stale.
    ///
    /// Defaults to [`StaleMarketPolicy::Off`]: it fills against the market it holds, however old.
    /// [`StaleMarketPolicy::Wait`] makes an order on an instrument with no fresh observation within
    /// an age wait for the next one, and needs a driver that feeds the venue a market, as
    /// `SimRunner` does.
    #[serde(default)]
    pub stale_market: StaleMarketPolicy,
    /// Whether the mock exchange lends a spot order what the account lacks, as a margin account
    /// does.
    ///
    /// Defaults to `None`: a cash account, whose spot orders are paid for with what it holds, and
    /// whose balances report no loans. A [`SpotMarginConfig`] makes it a cross margin account that
    /// borrows under each order's [`MarginEffect`](crate::margin::MarginEffect), up to a
    /// leverage limit.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub spot_margin: Option<SpotMarginConfig>,
    /// A third asset the mock exchange charges spot fees in, at a discount, while the account
    /// holds enough of it, as Binance's "pay fees with BNB" does.
    ///
    /// Defaults to `None`: every fee is charged under [`spot_fee_asset`](Self::spot_fee_asset).
    /// See [`FeeAssetConfig`] for what a fill pays and when it falls back to that rule.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub fee_asset: Option<FeeAssetConfig>,
}

#[derive(Debug, Constructor)]
pub struct MockExecutionClientConfig<FnTime> {
    pub mocked_exchange: ExchangeId,
    pub clock: FnTime,
    pub request_tx: mpsc::UnboundedSender<MockExchangeRequest>,
    pub event_rx: broadcast::Receiver<UnindexedAccountEvent>,
}

impl<FnTime> Clone for MockExecutionClientConfig<FnTime>
where
    FnTime: Clone,
{
    fn clone(&self) -> Self {
        Self {
            mocked_exchange: self.mocked_exchange,
            clock: self.clock.clone(),
            request_tx: self.request_tx.clone(),
            event_rx: self.event_rx.resubscribe(),
        }
    }
}

#[derive(Debug, Constructor)]
pub struct MockExecution<FnTime> {
    pub mocked_exchange: ExchangeId,
    pub clock: FnTime,
    pub request_tx: mpsc::UnboundedSender<MockExchangeRequest>,
    pub event_rx: broadcast::Receiver<UnindexedAccountEvent>,
}

impl<FnTime> Clone for MockExecution<FnTime>
where
    FnTime: Clone,
{
    fn clone(&self) -> Self {
        Self {
            mocked_exchange: self.mocked_exchange,
            clock: self.clock.clone(),
            request_tx: self.request_tx.clone(),
            event_rx: self.event_rx.resubscribe(),
        }
    }
}

impl<FnTime> MockExecution<FnTime>
where
    FnTime: Fn() -> DateTime<Utc>,
{
    pub fn time_request(&self) -> DateTime<Utc> {
        (self.clock)()
    }
}

impl<FnTime> ExecutionClient for MockExecution<FnTime>
where
    FnTime: Fn() -> DateTime<Utc> + Clone + Send + Sync,
{
    const EXCHANGE: ExchangeId = ExchangeId::Mock;

    // `MockExchange` models no expiry, funding schedule or contract chain, so `Perpetual`, `Future`
    // and `Option` have no faithful projection onto it. `Spot` and `Cfd` need none: both are
    // positions in an instrument that the mock can open, close and mark without simulating a
    // lifecycle. Kept in step with the projection in `generate_mock_exchange_instruments`.
    const SUPPORTED_KINDS: &'static [InstrumentKindDiscriminant] = &[
        InstrumentKindDiscriminant::Spot,
        InstrumentKindDiscriminant::Cfd,
    ];
    type Config = MockExecutionClientConfig<FnTime>;
    type AccountStream = BoxStream<'static, UnindexedAccountEvent>;

    fn new(config: Self::Config) -> Self {
        Self {
            mocked_exchange: config.mocked_exchange,
            clock: config.clock,
            request_tx: config.request_tx,
            event_rx: config.event_rx,
        }
    }

    async fn account_snapshot(
        &self,
        _: &[AssetNameExchange],
        _: &[InstrumentNameExchange],
    ) -> Result<UnindexedAccountSnapshot, UnindexedClientError> {
        let (response_tx, response_rx) = oneshot::channel();

        self.request_tx
            .send(MockExchangeRequest::fetch_account_snapshot(
                self.time_request(),
                response_tx,
            ))
            .map_err(|_| {
                UnindexedClientError::Connectivity(ConnectivityError::ExchangeOffline(
                    self.mocked_exchange,
                ))
            })?;

        response_rx.await.map_err(|_| {
            UnindexedClientError::Connectivity(ConnectivityError::ExchangeOffline(
                self.mocked_exchange,
            ))
        })
    }

    async fn account_stream(
        &self,
        _: &[AssetNameExchange],
        _: &[InstrumentNameExchange],
    ) -> Result<Self::AccountStream, UnindexedClientError> {
        // A lag is delivered in-band as one StreamTerminated, where `map_while` would end the
        // stream with a silent EOF. Whatever its source, a StreamTerminated is terminal by
        // contract, so the stream ends on the poll after it: a lag, or one the venue sends itself,
        // as a market-driven venue does when its market feed lags. It ends there rather than on the
        // next broadcast item, which a quiet venue may never send.
        let exchange = self.mocked_exchange;
        let events = BroadcastStream::new(self.event_rx.resubscribe()).map(move |result| {
            result.unwrap_or_else(|error| {
                error!(
                    ?error,
                    "MockExchange Broadcast AccountStream lagged - terminating"
                );
                UnindexedAccountEvent::stream_terminated(
                    exchange,
                    StreamTerminationReason::Error("mock broadcast stream lagged".to_string()),
                )
            })
        });
        Ok(
            futures::stream::unfold((events, false), |(mut events, ended)| async move {
                if ended {
                    return None;
                }
                let event = events.next().await?;
                let ended = matches!(event.kind, AccountEventKind::StreamTerminated(_));
                Some((event, (events, ended)))
            })
            // `unfold` panics when polled again after it ends, which a combinator downstream may do.
            .fuse()
            .boxed(),
        )
    }

    async fn cancel_order(
        &self,
        request: OrderRequestCancel<ExchangeId, &InstrumentNameExchange>,
    ) -> UnindexedOrderResponseCancel {
        let (response_tx, response_rx) = oneshot::channel();

        let key = OrderKey {
            exchange: request.key.exchange,
            instrument: request.key.instrument.clone(),
            strategy: request.key.strategy.clone(),
            cid: request.key.cid.clone(),
        };

        if self
            .request_tx
            .send(MockExchangeRequest::cancel_order(
                self.time_request(),
                response_tx,
                request.into_owned_instrument(),
            ))
            .is_err()
        {
            return UnindexedOrderResponseCancel {
                key,
                state: Err(UnindexedOrderError::Connectivity(
                    ConnectivityError::ExchangeOffline(self.mocked_exchange),
                )),
            };
        }

        match response_rx.await {
            Ok(response) => response,
            Err(_) => UnindexedOrderResponseCancel {
                key,
                state: Err(UnindexedOrderError::Connectivity(
                    ConnectivityError::ExchangeOffline(self.mocked_exchange),
                )),
            },
        }
    }

    async fn open_order(
        &self,
        request: OrderRequestOpen<ExchangeId, &InstrumentNameExchange>,
    ) -> Order<ExchangeId, InstrumentNameExchange, UnindexedOrderState> {
        let (response_tx, response_rx) = oneshot::channel();

        let request = request.into_owned_instrument();

        if self
            .request_tx
            .send(MockExchangeRequest::open_order(
                self.time_request(),
                response_tx,
                // Carries the snapshot the engine sampled when it decided to send this order on
                // `RequestOpen::market`. It is the venue's only price source for a Market order,
                // which carries no limit price of its own. `None` there means nothing sampled
                // one, and the venue rejects rather than guessing.
                request.clone(),
            ))
            .is_err()
        {
            return Order {
                key: request.key,
                side: request.state.side,
                price: request.state.price,
                quantity: request.state.quantity,
                kind: request.state.kind,
                time_in_force: request.state.time_in_force,
                state: OrderState::inactive(OrderError::Connectivity(
                    ConnectivityError::ExchangeOffline(self.mocked_exchange),
                )),
            };
        }

        match response_rx.await {
            Ok(response) => response,
            Err(_) => Order {
                key: request.key,
                side: request.state.side,
                price: request.state.price,
                quantity: request.state.quantity,
                kind: request.state.kind,
                time_in_force: request.state.time_in_force,
                state: OrderState::inactive(OrderError::Connectivity(
                    ConnectivityError::ExchangeOffline(self.mocked_exchange),
                )),
            },
        }
    }

    async fn fetch_balances(
        &self,
        assets: &[AssetNameExchange],
    ) -> Result<Vec<AssetBalance<AssetNameExchange>>, UnindexedClientError> {
        let (response_tx, response_rx) = oneshot::channel();

        self.request_tx
            .send(MockExchangeRequest::fetch_balances(
                self.time_request(),
                assets.to_vec(),
                response_tx,
            ))
            .map_err(|_| {
                UnindexedClientError::Connectivity(ConnectivityError::ExchangeOffline(
                    self.mocked_exchange,
                ))
            })?;

        response_rx.await.map_err(|_| {
            UnindexedClientError::Connectivity(ConnectivityError::ExchangeOffline(
                self.mocked_exchange,
            ))
        })
    }

    async fn fetch_open_orders(
        &self,
        instruments: &[InstrumentNameExchange],
    ) -> Result<Vec<Order<ExchangeId, InstrumentNameExchange, Open>>, UnindexedClientError> {
        let (response_tx, response_rx) = oneshot::channel();

        self.request_tx
            .send(MockExchangeRequest::fetch_orders_open(
                self.time_request(),
                instruments.to_vec(),
                response_tx,
            ))
            .map_err(|_| {
                UnindexedClientError::Connectivity(ConnectivityError::ExchangeOffline(
                    self.mocked_exchange,
                ))
            })?;

        response_rx.await.map_err(|_| {
            UnindexedClientError::Connectivity(ConnectivityError::ExchangeOffline(
                self.mocked_exchange,
            ))
        })
    }

    /// The mock venue holds every trade and answers a span whole, so the read is always
    /// complete.
    async fn fetch_trades(
        &self,
        start: DateTime<Utc>,
        end: DateTime<Utc>,
        instruments: &[InstrumentNameExchange],
    ) -> Result<TradesRead<AssetNameExchange, InstrumentNameExchange>, UnindexedClientError> {
        let (response_tx, response_rx) = oneshot::channel();

        self.request_tx
            .send(MockExchangeRequest::fetch_trades(
                self.time_request(),
                response_tx,
                start,
                end,
            ))
            .map_err(|_| {
                UnindexedClientError::Connectivity(ConnectivityError::ExchangeOffline(
                    self.mocked_exchange,
                ))
            })?;

        let mut trades = response_rx.await.map_err(|_| {
            UnindexedClientError::Connectivity(ConnectivityError::ExchangeOffline(
                self.mocked_exchange,
            ))
        })?;
        if !instruments.is_empty() {
            trades.retain(|trade| instruments.contains(&trade.instrument));
        }
        Ok(TradesRead::complete(trades))
    }
}

/// Answered from the venue's own ledger: an order it filled, cancelled or retired at its deadline
/// is reported, and one still open or never held is omitted. The venue rejects an order only on
/// arrival, in the response to its open request, so it never reports `OpenFailed`.
impl<FnTime> OrderStatusClient for MockExecution<FnTime>
where
    FnTime: Fn() -> DateTime<Utc> + Clone + Send + Sync,
{
    async fn fetch_ended_orders(
        &self,
        orders: &[UnindexedOrderKey],
    ) -> Result<Vec<UnindexedInactiveOrder>, UnindexedClientError> {
        let (response_tx, response_rx) = oneshot::channel();

        self.request_tx
            .send(MockExchangeRequest::fetch_orders_ended(
                self.time_request(),
                orders.to_vec(),
                response_tx,
            ))
            .map_err(|_| {
                UnindexedClientError::Connectivity(ConnectivityError::ExchangeOffline(
                    self.mocked_exchange,
                ))
            })?;

        response_rx.await.map_err(|_| {
            UnindexedClientError::Connectivity(ConnectivityError::ExchangeOffline(
                self.mocked_exchange,
            ))
        })
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)] // Test code: panics on bad input are acceptable
mod tests {
    use super::*;
    use crate::{AccountEventKind, balance::Balance};
    use rust_decimal::Decimal;
    use rustrade_integration::collection::snapshot::Snapshot;

    fn filler_event() -> UnindexedAccountEvent {
        // Content is irrelevant — these events are overwritten by the lag before they can be
        // delivered; they exist only to overflow the broadcast buffer.
        UnindexedAccountEvent::new(
            ExchangeId::Mock,
            AccountEventKind::BalanceSnapshot(Snapshot::new(AssetBalance {
                asset: AssetNameExchange::new("btc"),
                balance: Balance::new(Decimal::ONE, Decimal::ONE),
                time_exchange: Utc::now(),
            })),
        )
    }

    /// `fetch_trades` asks the mock exchange for its span and keeps the requested instruments'
    /// trades, reading the span whole.
    #[tokio::test]
    async fn fetch_trades_reads_the_span_for_the_requested_instruments() {
        use crate::{
            order::id::{OrderId, StrategyId},
            trade::{AssetFees, Trade, TradeId},
        };
        use rustrade_instrument::Side;

        let (_event_tx, event_rx) = broadcast::channel::<UnindexedAccountEvent>(1);
        let (request_tx, mut request_rx) = mpsc::unbounded_channel();
        let client = MockExecution::new(ExchangeId::Mock, Utc::now, request_tx, event_rx);
        let start = Utc::now() - chrono::TimeDelta::hours(1);
        let end = Utc::now();

        let trade = move |instrument: &str| {
            Trade::new(
                TradeId::new(instrument),
                OrderId::new("order"),
                InstrumentNameExchange::new(instrument),
                StrategyId::unknown(),
                start,
                Side::Buy,
                Decimal::ONE,
                Decimal::ONE,
                None,
                AssetFees::new(AssetNameExchange::new("usdt"), Decimal::ZERO, None),
            )
        };
        let exchange = tokio::spawn(async move {
            let request = request_rx.recv().await.expect("a request");
            let crate::exchange::mock::request::MockExchangeRequestKind::FetchTrades {
                response_tx,
                start: asked_start,
                end: asked_end,
            } = request.kind
            else {
                panic!("expected FetchTrades, got {:?}", request.kind);
            };
            assert_eq!((asked_start, asked_end), (start, end));
            response_tx
                .send(vec![trade("btc_usdt"), trade("eth_usdt")])
                .expect("the client awaits the answer");
        });

        let read = client
            .fetch_trades(start, end, &[InstrumentNameExchange::new("eth_usdt")])
            .await
            .unwrap();
        exchange.await.unwrap();
        assert_eq!(read, TradesRead::complete(vec![trade("eth_usdt")]));
    }

    /// Broadcast lag is a terminal stream death: the account stream must deliver an in-band
    /// `StreamTerminated(Error)` (not a silent EOF) and then end.
    #[tokio::test]
    async fn account_stream_emits_stream_terminated_on_broadcast_lag() {
        // Capacity 1 so a burst of sends with no reader immediately overflows → lag.
        let (event_tx, event_rx) = broadcast::channel::<UnindexedAccountEvent>(1);
        let (request_tx, _request_rx) = mpsc::unbounded_channel();
        // `MockExecution` derives `Constructor`, so its inherent `new` takes the four fields
        // directly (the `ExecutionClient::new(config)` trait method is shadowed here).
        let client = MockExecution::new(ExchangeId::Mock, Utc::now, request_tx, event_rx);

        let mut stream = client.account_stream(&[], &[]).await.unwrap();

        // Overflow the buffer before the stream's receiver reads anything → force a lag.
        for _ in 0..3 {
            event_tx.send(filler_event()).unwrap();
        }

        // The lag is surfaced in-band as the terminal event...
        let event = stream.next().await.expect("expected a terminal event");
        assert!(
            matches!(
                event.kind,
                AccountEventKind::StreamTerminated(StreamTerminationReason::Error(_))
            ),
            "expected StreamTerminated(Error) on broadcast lag, got {:?}",
            event.kind
        );
        assert_eq!(event.exchange, ExchangeId::Mock);

        // ...after which the stream ends.
        assert!(
            stream.next().await.is_none(),
            "stream must end after the terminal event"
        );
    }

    /// A terminal event the venue sends itself, as a market-driven one does when its market feed
    /// lags, ends the stream on the next poll, without waiting for an event a quiet venue may
    /// never send.
    #[tokio::test]
    async fn account_stream_ends_right_after_a_terminal_event_from_the_venue() {
        let (event_tx, event_rx) = broadcast::channel::<UnindexedAccountEvent>(16);
        let (request_tx, _request_rx) = mpsc::unbounded_channel();
        let client = MockExecution::new(ExchangeId::Mock, Utc::now, request_tx, event_rx);
        let mut stream = client.account_stream(&[], &[]).await.unwrap();

        event_tx.send(filler_event()).unwrap();
        event_tx
            .send(UnindexedAccountEvent::stream_terminated(
                ExchangeId::Mock,
                StreamTerminationReason::Error("market feed lagged".to_string()),
            ))
            .unwrap();

        assert!(matches!(
            stream.next().await.map(|event| event.kind),
            Some(AccountEventKind::BalanceSnapshot(_))
        ));
        assert!(matches!(
            stream.next().await.map(|event| event.kind),
            Some(AccountEventKind::StreamTerminated(_))
        ));
        assert_eq!(
            futures::FutureExt::now_or_never(stream.next()),
            Some(None),
            "the stream must end on the poll after its terminal event, with the sender still open"
        );
        assert_eq!(
            futures::FutureExt::now_or_never(stream.next()),
            Some(None),
            "and stay ended however often it is polled"
        );
        drop(event_tx);
    }
}