simulator_client/managed/

subscription.rs

//! SubscriptionManager — owns subscription WebSockets and keeps them alive.
//!
//! Two variants:
//! - account-diff subscription (`accountDiffSubscribe`) — account state capture
//! - transaction subscription (`transactionSubscribe`) — full transaction
//!   capture, delivers what `getTransaction` would return in one push so the
//!   client can skip the per-tx fetch entirely
//!
//! Both follow the same reconnect + keepalive pattern. On reconnect, all
//! configured subscriptions are re-established from scratch; we do not attempt
//! to replay notifications missed during the gap.

use std::{collections::HashSet, marker::PhantomData, time::Instant};

use futures::{SinkExt, StreamExt};
use serde::{Deserialize, de::DeserializeOwned};
use solana_transaction_status::EncodedConfirmedTransactionWithStatusMeta;
use tokio::{
    net::TcpStream,
    sync::{mpsc, watch},
    task::JoinHandle,
};
use tokio_tungstenite::{
    MaybeTlsStream, WebSocketStream, connect_async,
    tungstenite::{Message, client::IntoClientRequest},
};
use tokio_util::sync::CancellationToken;
use tracing::{debug, warn};

use super::{
    CONNECT_TIMEOUT, ConnectionStatus, HANDSHAKE_RESPONSE_TIMEOUT, KEEPALIVE_INTERVAL,
    KEEPALIVE_MISS_DEADLINE, RECONNECT_UPTIME_RESET, ReconnectBudget, cancellable_sleep,
};
use crate::{error::err_chain, subscriptions::AccountDiffNotification, urls::http_to_ws_url};

/// Handle to a running subscription manager task.
pub struct SubscriptionHandle {
    pub status: watch::Receiver<ConnectionStatus>,
    pub notifications: mpsc::Receiver<SubscriptionNotification>,
    pub join: JoinHandle<()>,
}

#[derive(Debug)]
pub enum SubscriptionNotification {
    Transaction(Box<EncodedConfirmedTransactionWithStatusMeta>),
    AccountDiff(AccountDiffNotification),
}

/// Per-flavor differences between `accountDiffSubscribe` and `transactionSubscribe`.
trait SubKind: Send + Sync + 'static {
    type Notification: DeserializeOwned + Send + 'static;
    const LABEL: &'static str;
    const SUBSCRIBE_METHOD: &'static str;
    const NOTIFICATION_METHOD: &'static str;
    fn subscribe_params(program_id: &str) -> serde_json::Value;
    fn into_notification(notification: Self::Notification) -> SubscriptionNotification;
}

struct AccountDiff;
impl SubKind for AccountDiff {
    type Notification = AccountDiffNotification;
    const LABEL: &'static str = "account-diff";
    const SUBSCRIBE_METHOD: &'static str = "accountDiffSubscribe";
    const NOTIFICATION_METHOD: &'static str = "accountDiffNotification";
    fn subscribe_params(program_id: &str) -> serde_json::Value {
        serde_json::json!([program_id, {"address_type": "program"}])
    }
    fn into_notification(notification: Self::Notification) -> SubscriptionNotification {
        SubscriptionNotification::AccountDiff(notification)
    }
}

struct Transaction;
impl SubKind for Transaction {
    /// Wire shape is identical to the `getTransaction` RPC response, so we can
    /// reuse `transaction_from_encoded` to build the output record directly
    /// from the push notification — no follow-up fetch required.
    type Notification = EncodedConfirmedTransactionWithStatusMeta;
    const LABEL: &'static str = "transaction";
    const SUBSCRIBE_METHOD: &'static str = "transactionSubscribe";
    const NOTIFICATION_METHOD: &'static str = "transactionNotification";
    fn subscribe_params(program_id: &str) -> serde_json::Value {
        serde_json::json!([{"mentions": [program_id]}, {"commitment": "confirmed"}])
    }
    fn into_notification(notification: Self::Notification) -> SubscriptionNotification {
        SubscriptionNotification::Transaction(Box::new(notification))
    }
}

pub fn spawn_transaction_subscription_manager(
    rpc_endpoint: String,
    program_ids: Vec<String>,
    cancel: CancellationToken,
) -> SubscriptionHandle {
    spawn_subscription_manager::<Transaction>(rpc_endpoint, program_ids, cancel)
}

pub fn spawn_account_diff_subscription_manager(
    rpc_endpoint: String,
    program_ids: Vec<String>,
    cancel: CancellationToken,
) -> SubscriptionHandle {
    spawn_subscription_manager::<AccountDiff>(rpc_endpoint, program_ids, cancel)
}

fn spawn_subscription_manager<K>(
    rpc_endpoint: String,
    program_ids: Vec<String>,
    cancel: CancellationToken,
) -> SubscriptionHandle
where
    K: SubKind,
{
    let (notifications_tx, notifications_rx) = mpsc::channel(1024);
    let (status_tx, status_rx) = watch::channel(ConnectionStatus::Down);
    let task = Task::<K> {
        rpc_endpoint,
        program_ids,
        notifications_tx,
        status_tx,
        cancel,
        _marker: PhantomData,
    };
    let join = tokio::spawn(task.run());
    SubscriptionHandle {
        status: status_rx,
        notifications: notifications_rx,
        join,
    }
}

type Ws = WebSocketStream<MaybeTlsStream<TcpStream>>;
type Subs = HashSet<u64>;

struct Task<K: SubKind> {
    rpc_endpoint: String,
    program_ids: Vec<String>,
    notifications_tx: mpsc::Sender<SubscriptionNotification>,
    status_tx: watch::Sender<ConnectionStatus>,
    /// Session-scoped cancel; fires both on user Ctrl-C *and* on normal
    /// session completion. Stops the connect/message loop either way.
    cancel: CancellationToken,
    _marker: PhantomData<fn() -> K>,
}

impl<K: SubKind> Task<K> {
    async fn run(self) {
        let mut budget = ReconnectBudget::new();

        loop {
            if self.cancel.is_cancelled() {
                break;
            }
            publish(&self.status_tx, ConnectionStatus::Down);

            let connect_result = async {
                let ws = connect_ws(&self.rpc_endpoint).await?;
                subscribe::<K>(ws, &self.program_ids).await
            }
            .await;

            let (ws, subs) = match connect_result {
                Ok(v) => v,
                Err(why) => {
                    if retry_or_fail::<K>(
                        "connect",
                        why,
                        &mut budget,
                        &self.cancel,
                        &self.status_tx,
                    )
                    .await
                    {
                        continue;
                    }
                    break;
                }
            };

            publish(&self.status_tx, ConnectionStatus::Up);
            let connected_at = Instant::now();

            let exit = message_loop::<K>(ws, subs, &self.notifications_tx, &self.cancel).await;

            match exit {
                MessageLoopExit::Cancelled | MessageLoopExit::Completed => break,
                MessageLoopExit::ConnectionLost(why) => {
                    if connected_at.elapsed() >= RECONNECT_UPTIME_RESET {
                        budget.reset();
                    }
                    if retry_or_fail::<K>(
                        "connection lost",
                        why,
                        &mut budget,
                        &self.cancel,
                        &self.status_tx,
                    )
                    .await
                    {
                        continue;
                    }
                    break;
                }
            }
        }
    }
}

enum MessageLoopExit {
    Cancelled,
    ConnectionLost(String),
    /// Every subscription on this connection delivered its end-of-stream
    /// terminal. Stop cleanly without reconnecting.
    Completed,
}

async fn message_loop<K: SubKind>(
    mut ws: Ws,
    subs: Subs,
    notifications_tx: &mpsc::Sender<SubscriptionNotification>,
    cancel: &CancellationToken,
) -> MessageLoopExit {
    let mut ping_timer = tokio::time::interval(KEEPALIVE_INTERVAL);
    ping_timer.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
    let mut last_inbound = Instant::now();
    // Subscriptions whose terminal `subscriptionComplete` marker has arrived.
    // Once this covers every entry in `subs`, the stream is fully drained.
    let mut completed: HashSet<u64> = HashSet::new();

    loop {
        tokio::select! {
            biased;
            _ = cancel.cancelled() => return MessageLoopExit::Cancelled,

            _ = ping_timer.tick() => {
                if last_inbound.elapsed() > KEEPALIVE_MISS_DEADLINE {
                    return MessageLoopExit::ConnectionLost(format!(
                        "no traffic for {:?}", last_inbound.elapsed()
                    ));
                }
                if let Err(e) = ws.send(Message::Ping(vec![])).await {
                    return MessageLoopExit::ConnectionLost(format!("ping send: {}", err_chain(&e)));
                }
            }

            msg = ws.next() => {
                last_inbound = Instant::now();
                match msg {
                    Some(Ok(Message::Text(t))) => {
                        match handle_text::<K>(&t, &subs, notifications_tx, &mut completed).await {
                            TextOutcome::Continue => {}
                            TextOutcome::AllComplete => return MessageLoopExit::Completed,
                            TextOutcome::ChannelClosed => return MessageLoopExit::Cancelled,
                        }
                    }
                    Some(Ok(Message::Binary(b))) => {
                        if let Ok(t) = std::str::from_utf8(&b) {
                            match handle_text::<K>(t, &subs, notifications_tx, &mut completed).await {
                                TextOutcome::Continue => {}
                                TextOutcome::AllComplete => return MessageLoopExit::Completed,
                                TextOutcome::ChannelClosed => return MessageLoopExit::Cancelled,
                            }
                        }
                    }
                    Some(Ok(Message::Pong(_))) | Some(Ok(Message::Ping(_))) => {}
                    Some(Ok(Message::Close(frame))) => {
                        return MessageLoopExit::ConnectionLost(format!("remote close: {frame:?}"));
                    }
                    Some(Ok(Message::Frame(_))) => {}
                    Some(Err(e)) => return MessageLoopExit::ConnectionLost(format!("ws read: {}", err_chain(&e))),
                    None => return MessageLoopExit::ConnectionLost("ws stream ended".into()),
                }
            }
        }
    }
}

/// Sleep for the next backoff interval, or publish `Failed` and return false if
/// the retry budget is exhausted. Returns true if the caller should retry.
async fn retry_or_fail<K: SubKind>(
    phase: &'static str,
    reason: String,
    budget: &mut ReconnectBudget,
    cancel: &CancellationToken,
    status_tx: &watch::Sender<ConnectionStatus>,
) -> bool {
    if let Some(delay) = budget.next_backoff() {
        warn!(
            kind = K::LABEL,
            attempt = budget.attempt(),
            reason = %reason,
            ?delay,
            "subscription {phase}, retrying",
        );
        cancellable_sleep(delay, cancel).await
    } else {
        publish(
            status_tx,
            ConnectionStatus::Failed(format!("{phase}: {reason}")),
        );
        false
    }
}

fn publish(tx: &watch::Sender<ConnectionStatus>, status: ConnectionStatus) {
    tx.send_if_modified(|current| {
        if *current == status {
            false
        } else {
            *current = status;
            true
        }
    });
}

async fn connect_ws(rpc_endpoint: &str) -> Result<Ws, String> {
    let ws_url = http_to_ws_url(rpc_endpoint).map_err(|e| err_chain(&e))?;
    let request = ws_url
        .into_client_request()
        .map_err(|e| format!("build request: {}", err_chain(&e)))?;

    let connect = tokio::time::timeout(CONNECT_TIMEOUT, connect_async(request))
        .await
        .map_err(|_| format!("connect timeout after {CONNECT_TIMEOUT:?}"))?
        .map_err(|e| format!("connect: {}", err_chain(&e)))?;
    Ok(connect.0)
}

async fn subscribe<K: SubKind>(mut ws: Ws, program_ids: &[String]) -> Result<(Ws, Subs), String> {
    let mut subs = Subs::new();
    for (i, program_id) in program_ids.iter().enumerate() {
        let id = (i + 1) as u64;
        let req = serde_json::json!({
            "jsonrpc": "2.0",
            "id": id,
            "method": K::SUBSCRIBE_METHOD,
            "params": K::subscribe_params(program_id),
        });
        ws.send(Message::Text(req.to_string()))
            .await
            .map_err(|e| format!("subscribe send: {}", err_chain(&e)))?;
        subs.insert(read_sub_ack(&mut ws, id).await?);
    }
    debug!(
        kind = K::LABEL,
        count = subs.len(),
        "subscriptions established"
    );
    Ok((ws, subs))
}

#[derive(Deserialize)]
struct SubAck {
    id: u64,
    result: Option<u64>,
    #[serde(default)]
    error: Option<serde_json::Value>,
}

async fn read_sub_ack(ws: &mut Ws, expected_id: u64) -> Result<u64, String> {
    let deadline = tokio::time::Instant::now() + HANDSHAKE_RESPONSE_TIMEOUT;
    loop {
        let msg = tokio::time::timeout_at(deadline, ws.next())
            .await
            .map_err(|_| format!("subscribe ack timeout after {HANDSHAKE_RESPONSE_TIMEOUT:?}"))?;

        let Some(msg) = msg else {
            return Err("ws ended during subscribe".into());
        };
        let msg = msg.map_err(|e| format!("ws read: {}", err_chain(&e)))?;

        if let Message::Text(t) = msg
            && let Ok(ack) = serde_json::from_str::<SubAck>(&t)
        {
            if ack.id != expected_id {
                continue;
            }
            if let Some(err) = ack.error {
                return Err(format!("subscribe rejected: {err}"));
            }
            if let Some(sub_id) = ack.result {
                return Ok(sub_id);
            }
            return Err("subscribe ack missing result".into());
        }
    }
}

/// Result of feeding one inbound text frame to the message loop.
enum TextOutcome {
    /// Keep reading (notification handled, or frame ignored).
    Continue,
    /// Every subscription on this connection has delivered its terminal marker.
    AllComplete,
    /// The downstream notifications channel closed — caller is gone.
    ChannelClosed,
}

/// Handle one inbound text frame: forward a matching notification, or record a
/// terminal `subscriptionComplete` marker. Returns [`TextOutcome::AllComplete`]
/// once a terminal has been seen for every subscription in `subs`.
async fn handle_text<K: SubKind>(
    text: &str,
    subs: &Subs,
    notifications_tx: &mpsc::Sender<SubscriptionNotification>,
    completed: &mut HashSet<u64>,
) -> TextOutcome {
    // Try a data notification first — that's the overwhelmingly common frame,
    // so the hot path parses the payload once.
    if let Some(n) = parse_notification::<K>(text, subs) {
        if notifications_tx
            .send(K::into_notification(n))
            .await
            .is_err()
        {
            return TextOutcome::ChannelClosed;
        }
        return TextOutcome::Continue;
    }

    // Otherwise it may be the terminal end-of-stream marker.
    if let Some(sub_id) = parse_completion(text)
        && subs.contains(&sub_id)
    {
        completed.insert(sub_id);
        if subs.iter().all(|id| completed.contains(id)) {
            return TextOutcome::AllComplete;
        }
    }
    TextOutcome::Continue
}

/// Parse a terminal end-of-stream marker, returning the subscription id it
/// targets. Shape: `{"method":"subscriptionComplete","params":{"subscription":N}}`.
fn parse_completion(text: &str) -> Option<u64> {
    #[derive(Deserialize)]
    struct Msg {
        method: String,
        params: Params,
    }
    #[derive(Deserialize)]
    struct Params {
        subscription: u64,
    }

    let msg: Msg = serde_json::from_str(text).ok()?;
    (msg.method == "subscriptionComplete").then_some(msg.params.subscription)
}

fn parse_notification<K: SubKind>(text: &str, subs: &Subs) -> Option<K::Notification> {
    #[derive(Deserialize)]
    #[serde(bound = "T: DeserializeOwned")]
    struct Msg<T> {
        method: String,
        params: Params<T>,
    }
    #[derive(Deserialize)]
    #[serde(bound = "T: DeserializeOwned")]
    struct Params<T> {
        subscription: u64,
        result: T,
    }

    let msg: Msg<K::Notification> = serde_json::from_str(text).ok()?;
    if msg.method != K::NOTIFICATION_METHOD {
        return None;
    }
    if !subs.contains(&msg.params.subscription) {
        return None;
    }
    Some(msg.params.result)
}

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

    #[test]
    fn parse_completion_extracts_subscription_id() {
        let text =
            r#"{"jsonrpc":"2.0","method":"subscriptionComplete","params":{"subscription":7}}"#;
        assert_eq!(parse_completion(text), Some(7));
    }

    #[test]
    fn parse_completion_ignores_other_messages() {
        let notification = r#"{"jsonrpc":"2.0","method":"transactionNotification","params":{"subscription":7,"result":{}}}"#;
        assert_eq!(parse_completion(notification), None);
        assert_eq!(parse_completion("not json"), None);
    }
}