simulator_client/

subscriptions.rs

use std::{future::Future, time::Duration};

use futures::StreamExt;
use solana_client::{
    nonblocking::pubsub_client::PubsubClient,
    rpc_response::{Response, RpcLogsResponse},
};
use solana_commitment_config::CommitmentConfig;
use solana_rpc_client_api::config::{RpcTransactionLogsConfig, RpcTransactionLogsFilter};
use thiserror::Error;
use tokio::{
    sync::{oneshot, watch},
    task::JoinHandle,
};

use crate::urls::{UrlError, http_to_ws_url};

/// Error establishing a PubSub log subscription.
#[derive(Debug, Error)]
pub enum SubscriptionError {
    #[error(transparent)]
    InvalidUrl(#[from] UrlError),

    #[error("pubsub connect to {url} failed: {source}")]
    Connect {
        url: String,
        #[source]
        source: Box<dyn std::error::Error + Send + Sync>,
    },

    #[error("logs_subscribe failed: {source}")]
    Subscribe {
        #[source]
        source: Box<dyn std::error::Error + Send + Sync>,
    },

    #[error("subscription task exited unexpectedly before signaling ready")]
    TaskDropped,

    #[error("session has no rpc_endpoint (was the session created?)")]
    NoRpcEndpoint,
}

/// Handle for a running log subscription background task.
pub struct LogSubscriptionHandle {
    /// Background task that drives the subscription and spawns per-notification callbacks.
    ///
    /// Resolves after `stop.send(true)` is called, remaining buffered
    /// notifications are drained, and all spawned callback tasks complete.
    pub join_handle: JoinHandle<()>,

    /// Send `true` to signal the background task to stop accepting new
    /// notifications, drain remaining buffered ones, and exit cleanly.
    pub stop: watch::Sender<bool>,
}

/// Subscribe to program log notifications and invoke a callback for each one.
///
/// Spawns a background task that:
/// 1. Connects to the PubSub endpoint derived from `rpc_endpoint`.
/// 2. Subscribes to logs mentioning `program_id`.
/// 3. For each notification, spawns `on_notification(notification)` as a Tokio task.
/// 4. When `handle.stop.send(true)` is called, drains remaining buffered
///    notifications (up to 1s), waits for all spawned tasks, then returns.
///
/// Returns after the subscription is established. If setup fails, an error is
/// returned before any background task is left running.
///
/// ## Example
///
/// ```no_run
/// use std::sync::{Arc, Mutex};
/// use simulator_client::subscribe_program_logs;
/// use solana_commitment_config::CommitmentConfig;
///
/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
/// let handle = subscribe_program_logs(
///     "https://api.mainnet-beta.solana.com",
///     "TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA",
///     CommitmentConfig::confirmed(),
///     |notification| async move {
///         println!("sig: {}", notification.value.signature);
///     },
/// )
/// .await?;
///
/// // ... do other work ...
///
/// handle.stop.send(true).ok();
/// handle.join_handle.await.ok();
/// # Ok(())
/// # }
/// ```
pub async fn subscribe_program_logs<F, Fut>(
    rpc_endpoint: &str,
    program_id: &str,
    commitment: CommitmentConfig,
    on_notification: F,
) -> Result<LogSubscriptionHandle, SubscriptionError>
where
    F: Fn(Response<RpcLogsResponse>) -> Fut + Send + Sync + 'static,
    Fut: Future<Output = ()> + Send + 'static,
{
    let ws_url = http_to_ws_url(rpc_endpoint)?;
    let program_id = program_id.to_string();

    let (ready_tx, ready_rx) = oneshot::channel::<Result<(), SubscriptionError>>();
    let (stop_tx, mut stop_rx) = watch::channel(false);

    // PubsubClient::logs_subscribe borrows &self, so both the client and the
    // stream must live inside the spawned task. We report setup success/failure
    // back through a oneshot channel before entering the notification loop.
    let join_handle = tokio::spawn(async move {
        let client = match PubsubClient::new(&ws_url).await {
            Ok(c) => c,
            Err(e) => {
                let _ = ready_tx.send(Err(SubscriptionError::Connect {
                    url: ws_url,
                    source: Box::new(e),
                }));
                return;
            }
        };

        let (mut stream, _unsubscribe) = match client
            .logs_subscribe(
                RpcTransactionLogsFilter::Mentions(vec![program_id]),
                RpcTransactionLogsConfig {
                    commitment: Some(commitment),
                },
            )
            .await
        {
            Ok(s) => s,
            Err(e) => {
                let _ = ready_tx.send(Err(SubscriptionError::Subscribe {
                    source: Box::new(e),
                }));
                return;
            }
        };

        let _ = ready_tx.send(Ok(()));

        let mut tasks: Vec<JoinHandle<()>> = Vec::new();

        loop {
            if *stop_rx.borrow() {
                // Drain notifications that arrived just before stop was signaled,
                // bounded by a short total deadline to avoid running indefinitely
                // for high-traffic programs.
                let drain_until = tokio::time::Instant::now() + Duration::from_secs(1);
                while let Ok(Some(notification)) =
                    tokio::time::timeout_at(drain_until, stream.next()).await
                {
                    tasks.push(tokio::spawn(on_notification(notification)));
                }
                break;
            }

            let notification = tokio::select! {
                n = stream.next() => n,
                _ = stop_rx.changed() => continue,
            };

            match notification {
                Some(n) => tasks.push(tokio::spawn(on_notification(n))),
                None => break,
            }
        }

        // Wait for all in-flight callback tasks to complete.
        for task in tasks {
            let _ = task.await;
        }
    });

    match ready_rx.await {
        Ok(Ok(())) => Ok(LogSubscriptionHandle {
            join_handle,
            stop: stop_tx,
        }),
        Ok(Err(e)) => {
            join_handle.abort();
            Err(e)
        }
        Err(_) => {
            join_handle.abort();
            Err(SubscriptionError::TaskDropped)
        }
    }
}