khadhroony-solana-project/crates/ksp-onchain-transport-lib/src/ws_subscription.rs

// file: crates/ksp-onchain-transport-lib/src/ws_subscription.rs
// version: 6

/// Typed handle for one logical WebSocket subscription.
///
/// The handle owns the bounded typed notification receiver while the physical session actor owns the remote subscription identity and socket. The remote
/// numeric subscription identifier is intentionally never exposed because it is transient and is remapped by the session actor after reconnect.
pub struct WsSubscription<T> {
    session_id: crate::WsSessionId,
    id: crate::WsSubscriptionId,
    kind: crate::WsSubscriptionKind,
    notification_rx: tokio::sync::mpsc::Receiver<ksp_core_lib::Result<T>>,
    state_rx: tokio::sync::watch::Receiver<crate::WsSubscriptionState>,
    terminal_error_rx: tokio::sync::watch::Receiver<std::option::Option<ksp_core_lib::ErrorCode>>,
    command_tx: tokio::sync::mpsc::Sender<crate::WsSessionCommand>,
    command_timeout: std::time::Duration,
}

impl<T> WsSubscription<T> {
    /// Creates a typed subscription handle from one actor registration.
    pub(crate) fn new(
        session_id: crate::WsSessionId,
        registration: WsSubscriptionRegistration,
        notification_rx: tokio::sync::mpsc::Receiver<ksp_core_lib::Result<T>>,
        command_tx: tokio::sync::mpsc::Sender<crate::WsSessionCommand>,
        command_timeout: std::time::Duration,
    ) -> Self {
        return Self {
            session_id,
            id: registration.id,
            kind: registration.kind,
            notification_rx,
            state_rx: registration.state_rx,
            terminal_error_rx: registration.terminal_error_rx,
            command_tx,
            command_timeout,
        };
    }

    /// Returns the stable local logical subscription identity.
    #[must_use]
    pub const fn id(&self) -> crate::WsSubscriptionId {
        return self.id;
    }

    /// Returns the logical WebSocket subscription family.
    #[must_use]
    pub const fn kind(&self) -> crate::WsSubscriptionKind {
        return self.kind;
    }

    /// Returns the latest observable lifecycle state for this logical subscription.
    #[must_use]
    pub fn state(&self) -> crate::WsSubscriptionState {
        return *self.state_rx.borrow();
    }

    /// Returns the safe terminal error code when this logical subscription failed.
    ///
    /// Successful local cancellation and normal completion use `None`. The value never contains remote payloads or endpoint credentials.
    #[must_use]
    pub fn terminal_error_code(&self) -> std::option::Option<ksp_core_lib::ErrorCode> {
        return *self.terminal_error_rx.borrow();
    }

    /// Receives the next typed notification or terminal typed-decoding error.
    ///
    /// The underlying queue is bounded by `WsSessionSettings::notification_queue_capacity`. `None` means the actor closed this logical subscription and no
    /// further notifications can arrive.
    pub async fn recv(&mut self) -> std::option::Option<ksp_core_lib::Result<T>> {
        return self.notification_rx.recv().await;
    }

    /// Cancels this logical subscription and sends the matching protocol unsubscribe request when a remote binding still exists.
    ///
    /// The returned boolean preserves the standard Solana unsubscribe result when the current remote binding is reachable. Local cancellation is terminal
    /// for this handle; during reconnect it wins before resubscribe selection, and a late remote acknowledgement is cleaned up best-effort without
    /// reactivation.
    pub async fn unsubscribe(&mut self) -> ksp_core_lib::Result<bool> {
        match self.state() {
            crate::WsSubscriptionState::Closed => return std::result::Result::Ok(false),
            crate::WsSubscriptionState::Failed => {
                return std::result::Result::Err(subscription_closed_error(self.session_id, self.id, "WebSocket subscription is already failed"));
            },
            _ => {},
        }
        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
        let command = crate::WsSessionCommand::Unsubscribe { subscription_id: self.id, response_tx };
        let send_wait = tokio::time::timeout(self.command_timeout, self.command_tx.send(command)).await;
        match send_wait {
            std::result::Result::Ok(std::result::Result::Ok(())) => {},
            std::result::Result::Ok(std::result::Result::Err(_)) => {
                return std::result::Result::Err(subscription_closed_error(self.session_id, self.id, "WebSocket session command channel is closed"));
            },
            std::result::Result::Err(_) => {
                return std::result::Result::Err(subscription_timeout_error(
                    self.session_id,
                    self.id,
                    "WebSocket unsubscribe command queue remained unavailable until timeout",
                ));
            },
        }
        return match response_rx.await {
            std::result::Result::Ok(result) => result,
            std::result::Result::Err(_) => {
                std::result::Result::Err(subscription_closed_error(self.session_id, self.id, "WebSocket session ended before unsubscribe completion"))
            },
        };
    }
}

impl<T> std::fmt::Debug for WsSubscription<T> {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        return formatter
            .debug_struct("WsSubscription")
            .field("session_id", &self.session_id)
            .field("id", &self.id)
            .field("kind", &self.kind)
            .field("state", &self.state())
            .finish();
    }
}

/// Internal result of dispatching one decoded wire notification into a bounded typed channel.
pub(crate) enum WsNotificationDispatchOutcome {
    Delivered,
    DeliveredTerminal,
    ReceiverClosed,
    QueueFull,
    DecodeFailed { code: ksp_core_lib::ErrorCode },
}

/// Type-erased actor-owned dispatcher for one heterogeneous typed notification channel.
pub(crate) type WsNotificationDispatcher = std::boxed::Box<dyn Fn(serde_json::Value) -> WsNotificationDispatchOutcome + std::marker::Send + std::marker::Sync>;

/// Creates one bounded typed notification receiver whose dispatcher can mark a successfully delivered value as terminal.
pub(crate) fn typed_notification_channel_with_completion<T, F, C>(
    capacity: usize,
    decoder: F,
    is_terminal: C,
) -> (WsNotificationDispatcher, tokio::sync::mpsc::Receiver<ksp_core_lib::Result<T>>)
where
    T: std::marker::Send + 'static,
    F: Fn(serde_json::Value) -> ksp_core_lib::Result<T> + std::marker::Send + std::marker::Sync + 'static,
    C: Fn(&T) -> bool + std::marker::Send + std::marker::Sync + 'static,
{
    let (notification_tx, notification_rx) = tokio::sync::mpsc::channel(capacity);
    let dispatcher = move |value: serde_json::Value| -> WsNotificationDispatchOutcome {
        let permit = match notification_tx.try_reserve() {
            std::result::Result::Ok(permit) => permit,
            std::result::Result::Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {
                return WsNotificationDispatchOutcome::ReceiverClosed;
            },
            std::result::Result::Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => {
                return WsNotificationDispatchOutcome::QueueFull;
            },
        };
        return match decoder(value) {
            std::result::Result::Ok(notification) => {
                let terminal = is_terminal(&notification);
                permit.send(std::result::Result::Ok(notification));
                if terminal {
                    return WsNotificationDispatchOutcome::DeliveredTerminal;
                }
                WsNotificationDispatchOutcome::Delivered
            },
            std::result::Result::Err(error) => {
                let code = error.code();
                permit.send(std::result::Result::Err(error));
                WsNotificationDispatchOutcome::DecodeFailed { code }
            },
        };
    };
    return (std::boxed::Box::new(dispatcher), notification_rx);
}

/// Internal registration returned after a remote subscribe acknowledgement becomes atomically bound.
pub(crate) struct WsSubscriptionRegistration {
    id: crate::WsSubscriptionId,
    kind: crate::WsSubscriptionKind,
    state_rx: tokio::sync::watch::Receiver<crate::WsSubscriptionState>,
    terminal_error_rx: tokio::sync::watch::Receiver<std::option::Option<ksp_core_lib::ErrorCode>>,
}

impl WsSubscriptionRegistration {
    /// Creates one successful actor registration without exposing the transient remote identifier.
    pub(crate) fn new(
        id: crate::WsSubscriptionId,
        kind: crate::WsSubscriptionKind,
        state_rx: tokio::sync::watch::Receiver<crate::WsSubscriptionState>,
        terminal_error_rx: tokio::sync::watch::Receiver<std::option::Option<ksp_core_lib::ErrorCode>>,
    ) -> Self {
        return Self { id, kind, state_rx, terminal_error_rx };
    }
}

/// Actor-owned runtime entry for one local logical WebSocket subscription.
pub(crate) struct WsSubscriptionRuntime {
    /// Stable local identity.
    pub(crate) id: crate::WsSubscriptionId,
    /// WebSocket subscription family.
    pub(crate) kind: crate::WsSubscriptionKind,
    /// Current logical lifecycle state.
    pub(crate) state: crate::WsSubscriptionState,
    /// Original subscribe parameters retained internally for deterministic resubscribe.
    pub(crate) params: std::vec::Vec<serde_json::Value>,
    /// Current transient remote subscription identity when bound.
    pub(crate) remote_id: std::option::Option<u64>,
    /// Lifecycle publisher observed by the public typed handle.
    pub(crate) state_tx: tokio::sync::watch::Sender<crate::WsSubscriptionState>,
    /// Safe terminal failure code publisher observed by the public typed handle.
    pub(crate) terminal_error_tx: tokio::sync::watch::Sender<std::option::Option<ksp_core_lib::ErrorCode>>,
    /// Type-erased dispatcher into the bounded typed notification channel.
    pub(crate) dispatcher: WsNotificationDispatcher,
}

impl WsSubscriptionRuntime {
    /// Builds the safe session-snapshot projection for this runtime entry.
    pub(crate) fn snapshot(&self) -> crate::WsSubscriptionSnapshot {
        return crate::WsSubscriptionSnapshot::new(self.id, self.kind, self.state, self.remote_id.is_some(), *self.terminal_error_tx.borrow());
    }

    /// Updates the runtime state and publishes it to the typed handle.
    pub(crate) fn set_state(&mut self, state: crate::WsSubscriptionState) {
        self.state = state;
        self.state_tx.send_replace(state);
    }

    /// Publishes a terminal failure code before moving the logical subscription to `Failed`.
    pub(crate) fn fail_with_code(&mut self, code: ksp_core_lib::ErrorCode) {
        self.terminal_error_tx.send_replace(std::option::Option::Some(code));
        self.set_state(crate::WsSubscriptionState::Failed);
    }
}

fn subscription_closed_error(session_id: crate::WsSessionId, subscription_id: crate::WsSubscriptionId, message: &'static str) -> ksp_core_lib::Error {
    return ksp_core_lib::Error::new(crate::ERROR_CODE_WS_SESSION_CLOSED, message)
        .with_context("session_id", session_id.get().to_string())
        .with_context("subscription_id", subscription_id.get().to_string());
}

fn subscription_timeout_error(session_id: crate::WsSessionId, subscription_id: crate::WsSubscriptionId, message: &'static str) -> ksp_core_lib::Error {
    return ksp_core_lib::Error::new(crate::ERROR_CODE_TIMEOUT, message)
        .with_context("session_id", session_id.get().to_string())
        .with_context("subscription_id", subscription_id.get().to_string());
}