wowlab_centrifuge/client/

mod.rs

// #t(file: rust_alloc_in_loop) client reconnect loop builds channel names per subscription

use std::{
    sync::{
        Arc,
        atomic::{AtomicBool, AtomicU32, Ordering},
    },
    time::{Duration, Instant},
};

use prost::Message;
use tokio::sync::{RwLock, mpsc, oneshot};
use tokio_util::sync::CancellationToken;
use wowlab_types::{constants::MS_PER_SECOND_U64, sensitive::Sensitive, sim::FastMap};

use crate::{
    backoff::Backoff,
    error::Error,
    proto,
    subscription::{SubscriptionInner, SubscriptionState},
    transport::Transport,
    types::{ClientEvent, ConnectResult},
};

mod commands;
mod config;
mod push;
#[cfg(test)]
mod tests;

pub use config::ClientConfig;

const MAX_SERVER_PING_DELAY: Duration = Duration::from_secs(10);
const EVENT_CHANNEL_BUFFER: usize = 64;
const COMMAND_CHANNEL_BUFFER: usize = 32;
const BACKOFF_RESET_THRESHOLD: Duration = Duration::from_secs(30);
const TOKEN_REFRESH_TIMEOUT: Duration = Duration::from_secs(10);

fn check_reply(reply: proto::Reply) -> Result<proto::Reply, Error> {
    if let Some(err) = reply.error {
        return Err(Error::from_proto(err));
    }

    Ok(reply)
}

fn format_ttl(ttl: u32, expires: bool) -> String {
    if expires && ttl > 0 {
        format!("{ttl}s")
    } else {
        "none".to_string()
    }
}

fn ttl_to_instant(ttl: u32, expires: bool) -> Option<Instant> {
    (expires && ttl > 0).then(|| {
        let ttl_ms = (u64::from(ttl) * MS_PER_SECOND_U64).min(i32::MAX as u64);

        Instant::now() + Duration::from_millis(ttl_ms)
    })
}

fn log_connection_lost(error: &Error) {
    if error.is_benign_disconnect() {
        tracing::debug!(%error, "Connection lost");
    } else {
        tracing::warn!(%error, "Connection lost");
    }
}

fn log_disconnected() {
    tracing::info!("Disconnected");
}

fn log_reconnect(delay: Duration) {
    tracing::info!(?delay, "Scheduling reconnect");
}

fn log_no_ping() {
    tracing::warn!("No ping from server; disconnecting");
}

fn log_token_refresh_failed(error: &Error) {
    tracing::warn!(%error, "Token refresh failed");
}

fn log_command_send_failed(error: &Error) {
    tracing::error!(%error, "Failed to send command");
}

/// Lifecycle state of a Centrifugo connection.
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
#[non_exhaustive]
pub enum ClientState {
    #[default]
    Disconnected,
    Connecting,
    Connected,
}

type PendingRequest = oneshot::Sender<Result<proto::Reply, Error>>;

struct ClientInner {
    config: ClientConfig,
    state: ClientState,
    subscriptions: FastMap<String, SubscriptionInner>,
    pending: FastMap<u32, PendingRequest>,
    next_id: AtomicU32,
    event_tx: mpsc::Sender<ClientEvent>,
    cmd_tx: Option<mpsc::Sender<CommandRequest>>,
    server_ping: Option<Duration>,
    send_pong: bool,
    refresh_required: bool,
}

struct CommandRequest {
    id: u32,
    cmd: proto::Command,
}

struct PendingRefresh {
    id: u32,
    token: Sensitive<String>,
    deadline: tokio::time::Instant,
}

struct DisconnectAdvice {
    code: u32,
    reason: String,
    reconnect: bool,
}

enum ConnectionExit {
    Shutdown,
    ServerDisconnect(DisconnectAdvice),
}

struct RunGuard(Arc<AtomicBool>);

impl Drop for RunGuard {
    fn drop(&mut self) {
        self.0.store(false, Ordering::Release);
    }
}

/// Centrifugo realtime client with auto-reconnect; clones share state.
#[derive(Clone)]
pub struct Client {
    inner: Arc<RwLock<ClientInner>>,
    shutdown: CancellationToken,
    run_active: Arc<AtomicBool>,
}

impl Client {
    /// Creates a disconnected client with shared connection state.
    #[must_use]
    pub fn new(config: ClientConfig) -> Self {
        let (event_tx, _) = mpsc::channel(EVENT_CHANNEL_BUFFER);

        Self {
            inner: Arc::new(RwLock::new(ClientInner {
                config,
                state: ClientState::Disconnected,
                subscriptions: FastMap::default(),
                pending: FastMap::default(),
                next_id: AtomicU32::new(1),
                event_tx,
                cmd_tx: None,
                server_ping: None,
                send_pong: false,
                refresh_required: false,
            })),
            shutdown: CancellationToken::new(),
            run_active: Arc::new(AtomicBool::new(false)),
        }
    }

    /// Replaces the client event receiver and returns the new receiver; install before [`Self::connect`].
    pub async fn events(&self) -> mpsc::Receiver<ClientEvent> {
        let (tx, rx) = mpsc::channel(EVENT_CHANNEL_BUFFER);

        self.inner.write().await.event_tx = tx;

        rx
    }

    /// Returns the current connection lifecycle state.
    pub async fn state(&self) -> ClientState {
        self.inner.read().await.state
    }

    /// Returns whether the client completed the Centrifugo handshake.
    pub async fn is_connected(&self) -> bool {
        self.inner.read().await.state == ClientState::Connected
    }

    /// Spawns a background task that maintains the connection with auto-reconnect.
    pub fn connect(&self) {
        let client = self.clone();

        tokio::spawn(async move {
            client.run().await;
        });
    }

    /// Maintains the connection and reconnects until [`Self::disconnect`] is called.
    pub async fn run(&self) {
        if self
            .run_active
            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
            .is_err()
        {
            tracing::debug!("Connection loop is already running");

            return;
        }

        let _run_guard = RunGuard(Arc::clone(&self.run_active));
        let mut backoff = {
            let inner = self.inner.read().await;

            Backoff::new(
                inner.config.min_reconnect_delay,
                inner.config.max_reconnect_delay,
            )
        };

        loop {
            if self.shutdown.is_cancelled() {
                return;
            }

            self.inner.write().await.state = ClientState::Connecting;
            self.emit_client_event(ClientEvent::Connecting).await;

            let start = Instant::now();

            let disconnect = match self.connect_and_run().await {
                Ok(ConnectionExit::Shutdown) => {
                    log_disconnected();

                    DisconnectAdvice {
                        code: 0,
                        reason: "Disconnected".to_string(),
                        reconnect: false,
                    }
                }
                Ok(ConnectionExit::ServerDisconnect(disconnect)) => {
                    log_disconnected();

                    disconnect
                }
                Err(error) => {
                    log_connection_lost(&error);

                    if error.requires_token_refresh() {
                        self.inner.write().await.refresh_required = true;
                    }

                    self.emit_client_event(ClientEvent::Error(error.to_string()))
                        .await;

                    DisconnectAdvice {
                        code: 0,
                        reason: "Connection closed".to_string(),
                        reconnect: true,
                    }
                }
            };

            let disconnected_event = ClientEvent::Disconnected {
                code: disconnect.code,
                reason: disconnect.reason,
                reconnect: disconnect.reconnect,
            };

            let mut inner = self.inner.write().await;

            inner.state = ClientState::Disconnected;
            inner.cmd_tx = None;

            inner.pending.drain().for_each(|(_, reply_tx)| {
                let _ = reply_tx.send(Err(Error::connection_closed()));
            });

            inner
                .subscriptions
                .values_mut()
                .for_each(|sub| sub.state = SubscriptionState::Unsubscribed);
            drop(inner);
            self.emit_client_event(disconnected_event).await;

            if self.shutdown.is_cancelled() || !disconnect.reconnect {
                return;
            }

            if start.elapsed() > BACKOFF_RESET_THRESHOLD {
                backoff.reset();
            }

            let delay = backoff.next_delay();

            log_reconnect(delay);

            tokio::select! {
                () = tokio::time::sleep(delay) => {}
                () = self.shutdown.cancelled() => return,
            }

            tokio::task::yield_now().await;
        }
    }

    async fn emit_client_event(&self, event: ClientEvent) {
        let event_tx = self.inner.read().await.event_tx.clone();

        tokio::select! {
            biased;
            () = self.shutdown.cancelled() => {
                if let Err(error) = event_tx.try_send(event) {
                    tracing::debug!(%error, "Client event receiver unavailable during shutdown");
                }
            }
            result = event_tx.send(event.clone()) => {
                if let Err(error) = result {
                    tracing::debug!(%error, "Client event receiver closed");
                }
            }
        }
    }

    async fn emit_subscription_event(&self, channel: &str, event: crate::types::SubscriptionEvent) {
        let event_tx = {
            let inner = self.inner.read().await;

            inner
                .subscriptions
                .get(channel)
                .map(|sub| sub.event_tx.clone())
        };

        let Some(event_tx) = event_tx else {
            return;
        };

        tokio::select! {
            biased;
            () = self.shutdown.cancelled() => {}
            result = event_tx.send(event) => {
                if let Err(error) = result {
                    tracing::debug!(%channel, %error, "Subscription event receiver closed");
                }
            }
        }
    }

    async fn connection_config(&self) -> Result<Option<ClientConfig>, Error> {
        let mut config = self.inner.read().await.config.clone();
        let refresh_required = self.inner.read().await.refresh_required;

        if !refresh_required {
            return Ok(Some(config));
        }

        let Some(ref get_token) = config.get_token else {
            return Ok(Some(config));
        };

        let token_result = tokio::select! {
            biased;
            () = self.shutdown.cancelled() => return Ok(None),
            result = tokio::time::timeout(TOKEN_REFRESH_TIMEOUT, get_token()) => {
                result.map_err(|_elapsed| Error::protocol("Token callback timed out"))?
            }
        };

        match token_result {
            Ok(new_token) => {
                config.token = new_token.clone();
                let mut inner = self.inner.write().await;

                inner.config.token = new_token;
                inner.refresh_required = false;
            }
            Err(error) => {
                tracing::error!(%error, "Token refresh failed");

                return Err(error);
            }
        }

        Ok(Some(config))
    }

    async fn establish_connection(
        &self,
        config: &ClientConfig,
    ) -> Result<(Transport, proto::ConnectResult), Error> {
        let mut transport = Transport::connect(&config.url).await?;
        let subs = self.build_recovery_subs().await;
        let connect_req = proto::ConnectRequest {
            token: config.token.expose().clone(),
            name: config.name.clone(),
            version: config.version.clone(),
            data: config.data.clone().unwrap_or_default(),
            headers: config.headers.clone().into_iter().collect(),
            subs: subs.into_iter().collect(),
            ..Default::default()
        };
        let cmd = proto::Command {
            id: 1,
            connect: Some(connect_req),
            ..Default::default()
        };
        let reply = check_reply(transport.send_command(cmd).await?)?;
        let connect_result = reply
            .connect
            .ok_or_else(|| Error::protocol("No connect result"))?;

        Ok((transport, connect_result))
    }

    async fn connect_and_run(&self) -> Result<ConnectionExit, Error> {
        let Some(config) = self.connection_config().await? else {
            return Ok(ConnectionExit::Shutdown);
        };

        let (mut transport, connect_result) = self.establish_connection(&config).await?;

        let server_ping =
            (connect_result.ping > 0).then(|| Duration::from_secs(u64::from(connect_result.ping)));
        let send_pong = connect_result.pong;

        let refresh_at = ttl_to_instant(connect_result.ttl, connect_result.expires);
        let ttl_display = format_ttl(connect_result.ttl, connect_result.expires);
        let ping_display =
            server_ping.map_or_else(|| "disabled".to_string(), |d| format!("{}s", d.as_secs()));

        tracing::info!(ping = %ping_display, token_ttl = %ttl_display, "Connected");

        let server_subs = connect_result.subs.clone();

        self.process_server_subs(server_subs).await;

        let (cmd_tx, mut cmd_rx) = mpsc::channel::<CommandRequest>(COMMAND_CHANNEL_BUFFER);

        let connected_event = ClientEvent::Connected(ConnectResult::from(connect_result));

        let mut inner = self.inner.write().await;

        inner.state = ClientState::Connected;
        inner.cmd_tx = Some(cmd_tx);
        inner.server_ping = server_ping;
        inner.send_pong = send_pong;
        drop(inner);
        self.emit_client_event(connected_event).await;

        self.resubscribe_all().await;

        let ping_timeout = server_ping.map(|p| p + MAX_SERVER_PING_DELAY);
        let mut last_data = Instant::now();
        let mut refresh_at = refresh_at;
        let mut pending_refresh: Option<PendingRefresh> = None;

        let exit = loop {
            let timeout_future = async {
                if let Some(timeout) = ping_timeout {
                    tokio::time::sleep_until((last_data + timeout).into()).await;
                } else {
                    std::future::pending::<()>().await;
                }
            };

            let refresh_future = async {
                if let Some(at) = refresh_at {
                    tokio::time::sleep_until(at.into()).await;
                } else {
                    std::future::pending::<()>().await;
                }
            };

            let refresh_timeout_future = async {
                if let Some(refresh) = &pending_refresh {
                    tokio::time::sleep_until(refresh.deadline).await;
                } else {
                    std::future::pending::<()>().await;
                }
            };

            tokio::select! {
                () = self.shutdown.cancelled() => {
                    break ConnectionExit::Shutdown;
                }

                () = timeout_future => {
                    log_no_ping();
                    return Err(Error::no_ping());
                }

                () = refresh_future => {
                    match self.begin_refresh(&mut transport).await {
                        Ok(Some(refresh)) => {
                            pending_refresh = Some(refresh);
                            refresh_at = None;
                        }
                        Ok(None) => refresh_at = None,
                        Err(error) => {
                            log_token_refresh_failed(&error);
                            self.inner.write().await.refresh_required = true;

                            return Err(error);
                        }
                    }
                }

                () = refresh_timeout_future => {
                    self.inner.write().await.refresh_required = true;

                    return Err(Error::protocol("Token refresh response timed out"));
                }

                Some(req) = cmd_rx.recv() => {
                    let mut cmd = req.cmd;
                    cmd.id = req.id;

                    let data = cmd.encode_length_delimited_to_vec();
                    if let Err(error) = transport.send_raw(data).await {
                        log_command_send_failed(&error);

                        return Err(error);
                    }
                }

                result = transport.read_message() => {
                    let reply = result?;
                    last_data = Instant::now();

                    if pending_refresh.as_ref().is_some_and(|refresh| refresh.id == reply.id) {
                        let refresh = pending_refresh.take().ok_or_else(|| {
                            Error::protocol("Refresh response arrived without a pending refresh")
                        })?;

                        match self.finish_refresh(refresh, reply).await {
                            Ok(new_refresh_at) => refresh_at = new_refresh_at,
                            Err(error) => {
                                log_token_refresh_failed(&error);
                                self.inner.write().await.refresh_required = true;

                                return Err(error);
                            }
                        }
                    } else if let Some(disconnect) = self.handle_reply(reply, &mut transport).await? {
                        break ConnectionExit::ServerDisconnect(disconnect);
                    }
                }
            }
        };

        transport.close().await;

        Ok(exit)
    }

    async fn begin_refresh(
        &self,
        transport: &mut Transport,
    ) -> Result<Option<PendingRefresh>, Error> {
        let get_token = {
            let inner = self.inner.read().await;

            inner.config.get_token.clone()
        };

        let Some(get_token) = get_token else {
            tracing::warn!("Token refresh requested but no get_token callback configured");

            return Ok(None);
        };

        let new_token = tokio::select! {
            biased;
            () = self.shutdown.cancelled() => return Err(Error::connection_closed()),
            result = tokio::time::timeout(TOKEN_REFRESH_TIMEOUT, get_token()) => {
                result
                    .map_err(|_elapsed| Error::protocol("Token callback timed out"))??
            }
        };

        let refresh_req = proto::RefreshRequest {
            token: new_token.expose().clone(),
        };

        let id = self
            .inner
            .read()
            .await
            .next_id
            .fetch_add(1, Ordering::Relaxed);
        let cmd = proto::Command {
            id,
            refresh: Some(refresh_req),
            ..Default::default()
        };

        let data = cmd.encode_length_delimited_to_vec();

        self.inner.write().await.refresh_required = true;
        transport.send_raw(data).await?;

        Ok(Some(PendingRefresh {
            id,
            token: new_token,
            deadline: tokio::time::Instant::now() + TOKEN_REFRESH_TIMEOUT,
        }))
    }

    async fn finish_refresh(
        &self,
        refresh: PendingRefresh,
        reply: proto::Reply,
    ) -> Result<Option<Instant>, Error> {
        let reply = check_reply(reply)?;

        let result = reply
            .refresh
            .ok_or_else(|| Error::protocol("No refresh result"))?;

        {
            let mut inner = self.inner.write().await;

            inner.config.token = refresh.token;
            inner.refresh_required = false;
        };

        let ttl_display = format_ttl(result.ttl, result.expires);

        tracing::info!(next_ttl = %ttl_display, "Token refreshed");

        Ok(ttl_to_instant(result.ttl, result.expires))
    }
}
impl std::fmt::Debug for Client {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Client")
            .field("inner", &"<ClientInner>")
            .field("shutdown", &"<CancellationToken>")
            .finish()
    }
}