navi_vm/

run_handle.rs

use std::{collections::VecDeque, sync::Arc};

use futures_util::{StreamExt as _, stream::LocalBoxStream};

use crate::{
    Candlestick, CandlestickItem, Error, Event, TimeFrame,
    data_provider::InternalProviderError,
    instance::{ExecuteMode, Instance},
};

enum Phase {
    SendScriptInfo,
    Running,
    FinalConfirm,
    Done,
}

/// A handle to an active execution run returned by [`Instance::run`].
///
/// Call [`next_event`](Self::next_event) to iterate events one at a time.
/// The handle implements [`Deref`](std::ops::Deref) and
/// [`DerefMut`](std::ops::DerefMut) targeting [`Instance`], so all instance
/// methods — including
/// [`save_state`](crate::Instance::save_state)
/// — are available directly through the handle while the run is in progress.
///
/// # Example
///
/// ```rust,ignore
/// let mut handle = instance.run("NASDAQ:AAPL", TimeFrame::days(1));
/// while let Some(result) = handle.next_event().await {
///     let event = result?;
///     if matches!(event, Event::BarEnd) && handle.current_bar_confirmed() {
///         let snapshot = handle.save_state()?;
///         persist_snapshot(snapshot);
///     }
/// }
/// let instance = handle.into_instance();
/// ```
pub struct RunHandle {
    instance: Instance,
    stream: Option<LocalBoxStream<'static, Result<CandlestickItem, InternalProviderError>>>,
    pending: VecDeque<Event>,
    pending_realtime: Option<Candlestick>,
    last_confirmed: bool,
    /// `true` after a [`CandlestickItem::HistoryEnd`] has been received.
    /// Determines whether subsequent [`CandlestickItem::Bar`] items are routed
    /// as historical (confirmed) or realtime bars.
    realtime_mode: bool,
    phase: Phase,
    /// An error encountered while creating the main stream in
    /// [`RunHandle::new`].
    ///
    /// Returned as the first event from [`next_event`] so the caller gets a
    /// clean `Err` rather than silently receiving no data.
    initial_error: Option<Error>,
}

impl RunHandle {
    pub(crate) fn new(
        instance: Instance,
        symbol: String,
        from_time: i64,
        timeframe: TimeFrame,
    ) -> Self {
        let (stream, initial_error) = match instance
            .data_provider
            .as_ref()
            .map(|p| p.candlesticks(&symbol, timeframe, from_time, None))
        {
            None => (None, None),
            Some(Ok(s)) => (Some(s), None),
            Some(Err(e)) => (
                None,
                Some(Error::DataProvider {
                    message: e.to_string(),
                }),
            ),
        };
        let phase = if stream.is_some() {
            Phase::SendScriptInfo
        } else {
            Phase::Done
        };
        Self {
            instance,
            stream,
            pending: VecDeque::new(),
            pending_realtime: None,
            last_confirmed: true,
            realtime_mode: false,
            phase,
            initial_error,
        }
    }

    /// Consumes the handle and returns the instance.
    #[inline]
    pub fn into_instance(self) -> Instance {
        self.instance
    }

    fn drain_instance_events(&mut self) {
        self.pending.extend(self.instance.events.drain(..));
    }

    fn should_split_realtime(&self, candlestick: &Candlestick) -> bool {
        matches!(
            self.instance.candlesticks.last(),
            Some(last_candlestick)
                if last_candlestick.time != candlestick.time && !self.instance.last_bar_confirmed
        )
    }

    /// Returns the next event from the run, or `None` when execution is
    /// complete.
    pub async fn next_event(&mut self) -> Option<Result<Event, Error>> {
        if let Some(e) = self.initial_error.take() {
            return Some(Err(e));
        }
        loop {
            if let Some(event) = self.pending.pop_front() {
                return Some(Ok(event));
            }
            if let Some(candlestick) = self.pending_realtime.take() {
                if !self.instance.mtf_wait.is_zero()
                    && let Err(e) = self.instance.pre_drain_sub_streams().await
                {
                    return Some(Err(Error::DataProvider {
                        message: e.to_string(),
                    }));
                }
                self.last_confirmed = false;
                if let Err(e) = self
                    .instance
                    .execute(ExecuteMode::Realtime(candlestick))
                    .await
                {
                    return Some(Err(e));
                }
                self.drain_instance_events();
                continue;
            }
            match self.phase {
                Phase::Done => return None,

                Phase::SendScriptInfo => {
                    let event = Event::SessionInfo {
                        script_info: Arc::clone(&self.instance.script_info),
                        symbol_info: Box::new(self.instance.symbol_info.clone()),
                    };
                    let warnings = self.instance.warnings();
                    self.phase = Phase::Running;
                    if !warnings.is_empty() {
                        self.pending.push_back(Event::Warning(warnings));
                    }
                    return Some(Ok(event));
                }

                Phase::Running => {
                    // Poll the candlestick stream. The borrow of `self.stream` ends
                    // once `raw` is assigned (an owned value), before we touch
                    // `self.instance` below.
                    let raw = match self.stream.as_mut() {
                        Some(s) => s.next().await,
                        None => None,
                    };
                    let Some(result) = raw else {
                        if !self.last_confirmed {
                            self.phase = Phase::FinalConfirm;
                            continue;
                        }
                        self.phase = Phase::Done;
                        return None;
                    };
                    let item = match result {
                        Ok(item) => item,
                        Err(e) => {
                            return Some(Err(Error::DataProvider {
                                message: e.to_string(),
                            }));
                        }
                    };
                    match item {
                        CandlestickItem::HistoryEnd => {
                            self.realtime_mode = true;
                            self.pending.push_back(Event::HistoryEnd);
                        }
                        CandlestickItem::Bar(c) if !self.realtime_mode => {
                            self.last_confirmed = true;
                            if let Err(e) = self.instance.execute(ExecuteMode::Confirmed(c)).await {
                                return Some(Err(e));
                            }
                        }
                        CandlestickItem::Bar(c) => {
                            if self.should_split_realtime(&c) {
                                if let Err(e) = self.instance.execute(ExecuteMode::Confirm).await {
                                    return Some(Err(e));
                                }
                                self.last_confirmed = true;
                                self.pending_realtime = Some(c);
                            } else {
                                if !self.instance.mtf_wait.is_zero()
                                    && let Err(e) = self.instance.pre_drain_sub_streams().await
                                {
                                    return Some(Err(Error::DataProvider {
                                        message: e.to_string(),
                                    }));
                                }
                                self.last_confirmed = false;
                                if let Err(e) =
                                    self.instance.execute(ExecuteMode::Realtime(c)).await
                                {
                                    return Some(Err(e));
                                }
                            }
                        }
                    }
                    self.drain_instance_events();
                }

                Phase::FinalConfirm => {
                    if let Err(e) = self.instance.execute(ExecuteMode::Confirm).await {
                        return Some(Err(e));
                    }
                    self.drain_instance_events();
                    self.phase = Phase::Done;
                }
            }
        }
    }
}

impl std::ops::Deref for RunHandle {
    type Target = Instance;

    fn deref(&self) -> &Instance {
        &self.instance
    }
}

impl std::ops::DerefMut for RunHandle {
    fn deref_mut(&mut self) -> &mut Instance {
        &mut self.instance
    }
}