simulator_client/managed/

session.rs

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

use simulator_api::{
    AgentStatsReport, BacktestError, BacktestStatus, ContinueParams, ContinueToParams,
    CreateBacktestSessionRequest, DiscoveryBatchEvent, PausedEvent, SessionSummary,
};
use solana_transaction_status::EncodedConfirmedTransactionWithStatusMeta;
use thiserror::Error;
use tokio::sync::watch;
use tokio_util::sync::CancellationToken;

use super::{
    ConnectionStatus, ControlEvent, ControlHandle, ReconnectCoordinator, SessionInfo,
    SubscriptionHandle, SubscriptionNotification, spawn_account_diff_subscription_manager,
    spawn_control_manager, spawn_transaction_subscription_manager,
};
use crate::subscriptions::AccountDiffNotification;

/// Error returned by the high-level managed session wrapper.
#[derive(Debug, Error)]
pub enum ManagedSessionError {
    #[error("session create failed: {0}")]
    Create(String),

    #[error("control channel closed")]
    ControlClosed,

    #[error("control failed: {0}")]
    ControlFailed(String),

    #[error("subscription failed: {0}")]
    SubscriptionFailed(String),

    #[error("cancelled")]
    Cancelled,

    #[error("control closed while sending continue: {0}")]
    ContinueSend(String),
}

#[derive(Debug)]
pub enum ManagedEvent {
    ReadyForContinue,
    /// Server paused at a `ContinueTo` target. The session is ready for
    /// another `Continue` or `ContinueTo` from this point.
    Paused(PausedEvent),
    /// Server discovered an upcoming batch matching a registered
    /// `DiscoveryFilter`. Send `send_continue_to(slot, batch_index)` to pause
    /// immediately before it executes.
    DiscoveryBatch(DiscoveryBatchEvent),
    Slot(u64),
    Status(BacktestStatus),
    /// `agent_stats` is always `None` for [`super::ParallelSubSession`] —
    /// the multiplexed wire carries only the summary.
    Completed {
        summary: Option<SessionSummary>,
        agent_stats: Option<Vec<AgentStatsReport>>,
    },
    Error(BacktestError),
    Transaction(Box<EncodedConfirmedTransactionWithStatusMeta>),
    AccountDiff(AccountDiffNotification),
}

/// Idle backstop for the completion drain: if no trailing notification arrives
/// for this long after `Completed`, assume the data plane is stalled and stop
/// waiting. Bounds the gap *between* notifications, not total drain time, so a
/// slow-but-flowing stream (large backlog over a slow link) is never truncated.
const DEFAULT_COMPLETION_DRAIN_TIMEOUT: Duration = Duration::from_secs(60);

/// Result of [`ManagedBacktestSession::drain_until_subscriptions_complete`].
pub(super) enum DrainOutcome {
    /// Every subscription delivered its terminal (channels closed), or the
    /// session was cancelled. The drained events are complete up to the stop.
    Complete(Vec<ManagedEvent>),
    /// The idle backstop fired with subscriptions still open — trailing
    /// notifications never arrived, so the run is incomplete. Carries whatever
    /// was drained before the stall.
    Stalled(Vec<ManagedEvent>),
}

/// High-level managed backtest session.
///
/// This wrapper owns the control manager, supported subscription managers,
/// cancellation, status gating, and shutdown order. Callers only need to react
/// to [`ManagedEvent`]s and send [`ContinueParams`] after `ReadyForContinue`.
pub struct ManagedBacktestSession {
    session_info: SessionInfo,
    control: Option<ControlHandle>,
    subscriptions: Vec<SubscriptionHandle>,
    session_cancel: CancellationToken,
    /// Notifications drained on `Completed`, followed by `Completed`; served in
    /// order by `next_event`. `None` until completion.
    post_completion: Option<VecDeque<ManagedEvent>>,
    /// Surfaced after `post_completion` drains: set when the completion drain
    /// stalled (idle backstop fired with subscriptions still open), so a
    /// silently-truncated run fails loudly instead of reporting `Completed`.
    post_completion_error: Option<ManagedSessionError>,
    completion_drain_timeout: Duration,
    /// Shared across a parallel batch so this session's subscription reconnects
    /// step aside for still-streaming siblings. `None` for a standalone session.
    reconnect_coordinator: Option<Arc<ReconnectCoordinator>>,
}

impl ManagedBacktestSession {
    /// Start a managed session with an internally owned cancellation token.
    pub async fn start(
        url: String,
        api_key: String,
        create: CreateBacktestSessionRequest,
    ) -> Result<Self, ManagedSessionError> {
        Self::start_with_cancel(url, api_key, create, CancellationToken::new(), None).await
    }

    /// Start a managed session tied to a caller-owned cancellation token.
    ///
    /// Cancelling `parent_cancel` aborts startup and stops manager tasks.
    ///
    /// `reconnect_coordinator` is an optional coordinator shared across sibling
    /// sessions in a parallel batch; a dropped subscription parks on it until
    /// streaming siblings finish, then reconnects into the freed bandwidth. It
    /// is handed to the subscription managers spawned by `subscribe_*`. Pass
    /// `None` for a standalone session.
    pub async fn start_with_cancel(
        url: String,
        api_key: String,
        create: CreateBacktestSessionRequest,
        parent_cancel: CancellationToken,
        reconnect_coordinator: Option<Arc<ReconnectCoordinator>>,
    ) -> Result<Self, ManagedSessionError> {
        let session_cancel = parent_cancel.child_token();
        let mut control = spawn_control_manager(url, api_key, create, session_cancel.clone());

        let session_info = tokio::select! {
            biased;
            _ = parent_cancel.cancelled() => {
                session_cancel.cancel();
                control.join().await;
                return Err(ManagedSessionError::Cancelled);
            }
            result = control.wait_for_session() => {
                result.map_err(ManagedSessionError::Create)?
            }
        };

        Ok(Self {
            session_info,
            control: Some(control),
            subscriptions: Vec::new(),
            session_cancel,
            post_completion: None,
            post_completion_error: None,
            completion_drain_timeout: DEFAULT_COMPLETION_DRAIN_TIMEOUT,
            reconnect_coordinator,
        })
    }

    /// Metadata reported by the server when the session was created.
    pub fn session_info(&self) -> &SessionInfo {
        &self.session_info
    }

    /// Override the completion-drain idle timeout (default
    /// [`DEFAULT_COMPLETION_DRAIN_TIMEOUT`]). The drain gives up only after this
    /// long with no trailing notification; it does not cap total drain time.
    pub fn set_completion_drain_timeout(&mut self, idle_timeout: std::time::Duration) {
        self.completion_drain_timeout = idle_timeout;
    }

    /// Subscribe to transaction notifications for the configured programs.
    pub fn subscribe_transactions(&mut self, program_ids: Vec<String>) {
        self.subscriptions
            .push(spawn_transaction_subscription_manager(
                self.session_info.rpc_endpoint.clone(),
                program_ids,
                self.session_cancel.clone(),
                self.reconnect_coordinator.clone(),
            ));
    }

    /// Subscribe to account-diff notifications for the configured programs.
    pub fn subscribe_account_diffs(&mut self, program_ids: Vec<String>) {
        self.subscriptions
            .push(spawn_account_diff_subscription_manager(
                self.session_info.rpc_endpoint.clone(),
                program_ids,
                self.session_cancel.clone(),
                self.reconnect_coordinator.clone(),
            ));
    }

    /// Drain trailing notifications until every subscription delivers its
    /// end-of-stream terminal (closing its channel), the session is cancelled,
    /// or the data plane is found to have failed. See the `idle_timeout` arm for
    /// how reconnecting/parked subscriptions are distinguished from a stall.
    async fn drain_until_subscriptions_complete(
        &mut self,
        idle_timeout: std::time::Duration,
    ) -> DrainOutcome {
        drain_subscriptions_until_complete(
            &mut self.subscriptions,
            &self.session_cancel,
            idle_timeout,
        )
        .await
    }

    /// Receive the next control or subscription event.
    ///
    /// On `Completed`, trailing subscription notifications are drained and
    /// delivered before the `Completed` event.
    pub async fn next_event(&mut self) -> Result<ManagedEvent, ManagedSessionError> {
        // Serve buffered post-completion events (trailing notifications, then
        // `Completed`); the control stream is gone once they're exhausted.
        if let Some(buffered) = self.post_completion.as_mut() {
            if let Some(event) = buffered.pop_front() {
                return Ok(event);
            }
            // Buffer drained: surface a pending stall error so an incomplete
            // run fails loudly, else the control stream is simply done.
            return Err(self
                .post_completion_error
                .take()
                .unwrap_or(ManagedSessionError::ControlClosed));
        }

        if let Some(event) = try_next_subscription_event(&mut self.subscriptions) {
            return Ok(event);
        }

        // Scope the borrows to the `select!` so the completion drain below can
        // re-borrow `self`.
        let event = {
            let cancel = self.session_cancel.clone();
            let control = self
                .control
                .as_mut()
                .ok_or(ManagedSessionError::ControlClosed)?;
            let subscriptions = &mut self.subscriptions;
            tokio::select! {
                biased;
                _ = cancel.cancelled() => return Err(ManagedSessionError::Cancelled),
                event = control.events.recv() => {
                    event.map(ManagedEvent::from).ok_or(ManagedSessionError::ControlClosed)?
                }
                event = wait_any_subscription_event(subscriptions) => event,
            }
        };

        // Bind the payload so the re-emitted `Completed` below carries it.
        let ManagedEvent::Completed {
            summary,
            agent_stats,
        } = event
        else {
            return Ok(event);
        };

        // Flush trailing notifications up to each subscription's terminal,
        // delivering them before `Completed` so none are dropped.
        let (mut buffered, terminal): (VecDeque<ManagedEvent>, _) = match self
            .drain_until_subscriptions_complete(self.completion_drain_timeout)
            .await
        {
            DrainOutcome::Complete(events) => (
                events.into(),
                Ok(ManagedEvent::Completed {
                    summary,
                    agent_stats,
                }),
            ),
            // The data plane stalled before every subscription finished:
            // trailing notifications are missing, so report failure rather
            // than a silently-truncated `Completed`. Deliver whatever was
            // drained first, then surface the error once the buffer empties.
            DrainOutcome::Stalled(events) => (
                events.into(),
                Err(ManagedSessionError::SubscriptionFailed(
                    "completion drain stalled: subscriptions did not deliver their \
                     end-of-stream terminals; the captured stream is incomplete"
                        .to_string(),
                )),
            ),
        };
        match terminal {
            Ok(completed) => buffered.push_back(completed),
            Err(err) => self.post_completion_error = Some(err),
        }
        let first = buffered.pop_front();
        self.post_completion = Some(buffered);
        match first {
            Some(event) => Ok(event),
            // Nothing buffered and the run failed: surface the error now.
            None => Err(self
                .post_completion_error
                .take()
                .unwrap_or(ManagedSessionError::ControlClosed)),
        }
    }

    /// Wait until the control connection and all subscription connections are
    /// up, then send a `Continue` request.
    ///
    /// Call this after receiving [`ManagedEvent::ReadyForContinue`] or
    /// [`ManagedEvent::Paused`]. If there are no subscriptions, only the
    /// control connection is gated.
    pub async fn send_continue(
        &mut self,
        params: ContinueParams,
    ) -> Result<(), ManagedSessionError> {
        self.wait_all_up().await?;
        self.control_mut()?
            .send_continue(params)
            .await
            .map_err(|e| ManagedSessionError::ContinueSend(e.to_string()))
    }

    /// Wait until the control connection and all subscription connections are
    /// up, then send a `ContinueTo` request to step to a specific slot/batch
    /// boundary.
    ///
    /// Pair with [`ManagedEvent::DiscoveryBatch`] to drive a discovery-paced
    /// loop: receive a discovery event, send `ContinueTo(slot, batch_index)`,
    /// and wait for [`ManagedEvent::Paused`] before inspecting state.
    pub async fn send_continue_to(
        &mut self,
        params: ContinueToParams,
    ) -> Result<(), ManagedSessionError> {
        self.wait_all_up().await?;
        self.control_mut()?
            .send_continue_to(params)
            .await
            .map_err(|e| ManagedSessionError::ContinueSend(e.to_string()))
    }

    /// Cancel the session and join all manager tasks.
    pub async fn shutdown(mut self) {
        self.session_cancel.cancel();
        if let Some(control) = self.control.take() {
            control.join().await;
        }
        for sub in std::mem::take(&mut self.subscriptions) {
            let _ = sub.join.await;
        }
    }

    fn control_mut(&mut self) -> Result<&mut ControlHandle, ManagedSessionError> {
        self.control
            .as_mut()
            .ok_or(ManagedSessionError::ControlClosed)
    }

    async fn wait_all_up(&self) -> Result<(), ManagedSessionError> {
        let control = self
            .control
            .as_ref()
            .ok_or(ManagedSessionError::ControlClosed)?
            .status
            .clone();
        let subscriptions = self
            .subscriptions
            .iter()
            .map(|s| s.status.clone())
            .collect();
        wait_connections_up(control, subscriptions, &self.session_cancel).await
    }
}

/// Block until the control connection and every subscription connection report
/// `Up`, returning an error if any reports `Failed` or the token is cancelled.
/// Shared by the single-session and parallel sub-session drivers.
pub(super) async fn wait_connections_up(
    mut control: watch::Receiver<ConnectionStatus>,
    mut subscriptions: Vec<watch::Receiver<ConnectionStatus>>,
    cancel: &CancellationToken,
) -> Result<(), ManagedSessionError> {
    loop {
        let control_status = control.borrow().clone();
        if let ConnectionStatus::Failed(why) = &control_status {
            return Err(ManagedSessionError::ControlFailed(why.clone()));
        }

        let mut all_subscriptions_up = true;
        for subscription in &subscriptions {
            match &*subscription.borrow() {
                ConnectionStatus::Failed(why) => {
                    return Err(ManagedSessionError::SubscriptionFailed(why.clone()));
                }
                ConnectionStatus::Up => {}
                _ => all_subscriptions_up = false,
            }
        }

        if control_status == ConnectionStatus::Up && all_subscriptions_up {
            return Ok(());
        }

        tokio::select! {
            _ = cancel.cancelled() => return Err(ManagedSessionError::Cancelled),
            _ = control.changed() => {}
            _ = wait_any_subscription_change(&mut subscriptions) => {}
        }
    }
}

/// Drain trailing notifications until every subscription delivers its
/// end-of-stream terminal (closing its channel), `cancel` fires, or the data
/// plane is found to have stalled. Shared by the single-session and parallel
/// sub-session drivers. See
/// [`ManagedBacktestSession::drain_until_subscriptions_complete`].
///
/// `idle_timeout` bounds the gap *between* notifications, not total drain time,
/// so a slow-but-flowing stream (large backlog over a slow link) is never
/// truncated. A subscription that's `Down` — reconnecting, or parked behind the
/// reconnect coordinator waiting for siblings to free bandwidth — is *not* a
/// stall: it resumes via `replayFromSlot` or exhausts its budget and reports
/// `Failed`. So the idle gap only declares a `Stalled` outcome when an open
/// subscription is still `Up` (connected but not delivering — a genuinely hung
/// data plane); a subscription that closed in `Failed` left a truncated stream.
pub(super) async fn drain_subscriptions_until_complete(
    subscriptions: &mut [SubscriptionHandle],
    cancel: &CancellationToken,
    idle_timeout: std::time::Duration,
) -> DrainOutcome {
    let mut events = Vec::new();
    if subscriptions.is_empty() {
        return DrainOutcome::Complete(events);
    }
    loop {
        while let Some(event) = try_next_subscription_event(subscriptions) {
            events.push(event);
        }
        if subscriptions.iter().all(|s| s.notifications.is_closed()) {
            let any_failed = subscriptions
                .iter()
                .any(|s| matches!(*s.status.borrow(), ConnectionStatus::Failed(_)));
            return if any_failed {
                DrainOutcome::Stalled(events)
            } else {
                DrainOutcome::Complete(events)
            };
        }
        tokio::select! {
            biased;
            // A cancel is an intentional stop, not a data-plane failure.
            _ = cancel.cancelled() => return DrainOutcome::Complete(events),
            // Idle gap elapsed. A still-`Up` subscription should be delivering
            // but isn't — a hung data plane (stall). If the open subscriptions
            // are all reconnecting/parked (`Down`), keep waiting: the
            // coordinator is deferring them on purpose, not failing.
            _ = tokio::time::sleep(idle_timeout) => {
                let any_up = subscriptions.iter().any(|s| {
                    !s.notifications.is_closed()
                        && matches!(*s.status.borrow(), ConnectionStatus::Up)
                });
                if any_up {
                    return DrainOutcome::Stalled(events);
                }
            }
            received = recv_any_open_subscription(subscriptions) => {
                // `None` means a channel closed; the loop re-checks all-closed.
                if let Some(event) = received {
                    events.push(event);
                }
            }
        }
    }
}

impl Drop for ManagedBacktestSession {
    fn drop(&mut self) {
        self.session_cancel.cancel();
    }
}

pub(super) async fn wait_any_subscription_change(
    subscriptions: &mut [watch::Receiver<ConnectionStatus>],
) {
    if subscriptions.is_empty() {
        std::future::pending::<()>().await;
        return;
    }
    let _ =
        futures::future::select_all(subscriptions.iter_mut().map(|s| Box::pin(s.changed()))).await;
}

pub(super) async fn wait_any_subscription_event(
    subscriptions: &mut [SubscriptionHandle],
) -> ManagedEvent {
    loop {
        if let Some(event) = try_next_subscription_event(subscriptions) {
            return event;
        }

        let futures: Vec<_> = subscriptions
            .iter_mut()
            .filter(|s| !s.notifications.is_closed())
            .map(|s| Box::pin(s.notifications.recv()))
            .collect();

        if futures.is_empty() {
            std::future::pending::<()>().await;
        }

        let (notification, _, _) = futures::future::select_all(futures).await;
        if let Some(notification) = notification {
            return notification.into();
        }
    }
}

/// Await the next notification from any still-open subscription channel,
/// returning `None` when one closes. Unlike [`wait_any_subscription_event`],
/// which never resolves on closure, this lets the completion drain observe
/// per-channel end-of-stream.
pub(super) async fn recv_any_open_subscription(
    subscriptions: &mut [SubscriptionHandle],
) -> Option<ManagedEvent> {
    let futures: Vec<_> = subscriptions
        .iter_mut()
        .filter(|s| !s.notifications.is_closed())
        .map(|s| Box::pin(s.notifications.recv()))
        .collect();

    if futures.is_empty() {
        return None;
    }

    let (notification, _, _) = futures::future::select_all(futures).await;
    notification.map(Into::into)
}

pub(super) fn try_next_subscription_event(
    subscriptions: &mut [SubscriptionHandle],
) -> Option<ManagedEvent> {
    for subscription in subscriptions {
        if let Ok(notification) = subscription.notifications.try_recv() {
            return Some(notification.into());
        }
    }
    None
}

impl From<ControlEvent> for ManagedEvent {
    fn from(event: ControlEvent) -> Self {
        match event {
            ControlEvent::ReadyForContinue => Self::ReadyForContinue,
            ControlEvent::Paused(event) => Self::Paused(event),
            ControlEvent::DiscoveryBatch(event) => Self::DiscoveryBatch(event),
            ControlEvent::Slot(slot) => Self::Slot(slot),
            ControlEvent::Status(status) => Self::Status(status),
            ControlEvent::Completed {
                summary,
                agent_stats,
            } => Self::Completed {
                summary,
                agent_stats,
            },
            ControlEvent::Error(error) => Self::Error(error),
        }
    }
}

impl From<SubscriptionNotification> for ManagedEvent {
    fn from(notification: SubscriptionNotification) -> Self {
        match notification {
            SubscriptionNotification::Transaction(transaction) => Self::Transaction(transaction),
            SubscriptionNotification::AccountDiff(diff) => Self::AccountDiff(diff),
        }
    }
}