gitdataai/libs/transport/lib.rs

use std::collections::HashMap;
use config::AppConfig;
use db::cache::AppCache;
use db::database::AppDatabase;
use service::AppService;
use crate::seq::SeqAllocator;

#[derive(Clone)]
pub struct AppTransport {
    pub db: AppDatabase,
    pub service: AppService,
    pub config: AppConfig,
    pub cache: AppCache,
    pub nats: Option<async_nats::Client>,
    pub seq: SeqAllocator,
    pub ack: ack::AckTracker,
    pub reconnect: reconnect::ReconnectManager,
    pub dedup: dedup::DeduplicationManager,
    pub rate_limiter: security::RateLimiter,
    pub csrf: security::CsrfProtection,
    pub circuit_breaker: circuit_breaker::CircuitBreaker,
}

pub mod ack;
pub mod bus;
pub mod cdn;
pub mod circuit_breaker;
pub mod dedup;
pub mod e2e;
pub mod envelope;
pub mod error;
pub mod event;
pub mod handler;
pub mod metrics;
pub mod pagination;
pub mod reconnect;
pub mod richtext;
pub mod search;
pub mod security;
pub mod seq;
pub mod token;
pub mod unread;

impl AppTransport {
    pub async fn new(
        service: AppService,
        config: AppConfig,
    ) -> Result<Self, crate::error::AppTransportError> {
        let nats = if config.nats_is_enabled() {
            let url = match config.nats_url() {
                Some(u) => u,
                None => {
                    tracing::warn!("NATS_URL not configured, running without NATS transport");
                    return Ok(Self {
                        db: service.db.clone(),
                        cache: service.cache.clone(),
                        seq: SeqAllocator::new(service.cache.clone(), service.db.clone()),
                        ack: ack::AckTracker::new(service.cache.clone()),
                        reconnect: reconnect::ReconnectManager::new(service.cache.clone(), service.db.clone()),
                        dedup: dedup::DeduplicationManager::new(service.cache.clone()),
                        rate_limiter: security::RateLimiter::new(service.cache.clone()),
                        csrf: security::CsrfProtection::new(service.cache.clone()),
                        circuit_breaker: circuit_breaker::CircuitBreaker::new(),
                        service,
                        config,
                        nats: None,
                    });
                }
            };
            let token = match config.nats_token() {
                Some(t) => t,
                None => {
                    tracing::warn!("NATS_TOKEN not configured, running without NATS transport");
                    return Ok(Self {
                        db: service.db.clone(),
                        cache: service.cache.clone(),
                        seq: SeqAllocator::new(service.cache.clone(), service.db.clone()),
                        ack: ack::AckTracker::new(service.cache.clone()),
                        reconnect: reconnect::ReconnectManager::new(service.cache.clone(), service.db.clone()),
                        dedup: dedup::DeduplicationManager::new(service.cache.clone()),
                        rate_limiter: security::RateLimiter::new(service.cache.clone()),
                        csrf: security::CsrfProtection::new(service.cache.clone()),
                        circuit_breaker: circuit_breaker::CircuitBreaker::new(),
                        service,
                        config,
                        nats: None,
                    });
                }
            };

            let opts = async_nats::ConnectOptions::with_token(token)
                .retry_on_initial_connect()
                .connection_timeout(std::time::Duration::from_secs(10))
                .reconnect_delay_callback(|attempts| {
                    let base = std::time::Duration::from_secs(1);
                    let delay = base.saturating_mul(2u32.saturating_pow(attempts as u32));
                    std::cmp::min(delay, std::time::Duration::from_secs(30))
                })
                .event_callback(|event| async move {
                    match event {
                        async_nats::Event::Connected => {
                            tracing::info!("NATS connected");
                        }
                        async_nats::Event::Disconnected => {
                            tracing::warn!("NATS disconnected, reconnecting");
                        }
                        async_nats::Event::ServerError(e) => {
                            tracing::warn!(error = %e, "NATS server error");
                        }
                        _ => {}
                    }
                });

            match opts.connect(&url).await {
                Ok(client) => {
                    tracing::info!(url = %url, "NATS connected");
                    Some(client)
                }
                Err(e) => {
                    tracing::warn!(error = %e, url = %url, "NATS connect failed, running without NATS transport");
                    None
                }
            }
        } else {
            tracing::info!("NATS not enabled (NATS_URL or NATS_TOKEN not set)");
            None
        };

        Ok(Self {
            db: service.db.clone(),
            cache: service.cache.clone(),
            seq: SeqAllocator::new(
                service.cache.clone(),
                service.db.clone(),
            ),
            ack: ack::AckTracker::new(service.cache.clone()),
            reconnect: reconnect::ReconnectManager::new(service.cache.clone(), service.db.clone()),
            dedup: dedup::DeduplicationManager::new(service.cache.clone()),
            rate_limiter: security::RateLimiter::new(service.cache.clone()),
            csrf: security::CsrfProtection::new(service.cache.clone()),
            circuit_breaker: circuit_breaker::CircuitBreaker::new(),
            service,
            config,
            nats,
        })
    }

    pub fn nats_connected(&self) -> bool {
        self.nats.is_some()
    }

    pub async fn publish_nats(
        &self,
        subject: &str,
        payload: &[u8],
    ) -> Result<(), crate::error::AppTransportError> {
        match &self.nats {
            Some(client) => {
                let js = async_nats::jetstream::new(client.clone());
                js.publish(subject.to_string(), payload.to_vec().into())
                    .await
                    .map_err(|e| {
                        tracing::warn!(error = %e, subject = %subject, "NATS publish failed");
                        crate::error::AppTransportError::Internal
                    })?;
                Ok(())
            }
            None => {
                tracing::debug!(subject = %subject, "NATS not connected, skipping publish");
                Ok(())
            }
        }
    }

    pub async fn seq(&self, room: models::RoomId) -> Result<i64, crate::error::AppTransportError> {
        self.seq.seq(room).await
    }

    pub async fn bootstrap_seq(&self, room: models::RoomId) -> Result<i64, crate::error::AppTransportError> {
        self.seq.bootstrap(room).await
    }

    pub async fn bootstrap_all_seq(
        &self,
        rooms: Vec<models::RoomId>,
    ) -> Result<HashMap<models::RoomId, i64>, crate::error::AppTransportError> {
        self.seq.bootstrap_all(rooms).await
    }
}