foxhunt/trading_engine/src/persistence/postgres.rs

//! `PostgreSQL` connection pool and management for HFT trading operations
//!
//! This module provides high-performance `PostgreSQL` connectivity optimized for
//! sub-millisecond latency requirements in high-frequency trading.

use serde::{Deserialize, Serialize};
use sqlx::postgres::{PgConnectOptions, PgPool, PgPoolOptions};
use std::sync::Arc;
use std::time::{Duration, Instant};
use thiserror::Error;
use tokio::sync::RwLock;
use tracing::warn;

/// PostgreSQL-specific errors
#[derive(Debug, Error)]
pub enum PostgresError {
    #[error("Connection failed: {0}")]
    Connection(#[from] sqlx::Error),
    #[error("Query timeout: operation took {actual_ms}ms, max allowed {max_ms}ms")]
    QueryTimeout { actual_ms: u64, max_ms: u64 },
    #[error("Pool exhausted: no connections available")]
    PoolExhausted,
    #[error("Configuration error: {0}")]
    Configuration(String),
    #[error("Performance violation: {0}")]
    Performance(String),
}

/// `PostgreSQL` configuration optimized for HFT operations
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct PostgresConfig {
    /// Database connection URL
    pub url: String,
    /// Maximum number of connections in the pool
    pub max_connections: u32,
    /// Minimum number of connections to maintain
    pub min_connections: u32,
    /// Connection timeout in milliseconds (HFT optimized)
    pub connect_timeout_ms: u64,
    /// Query timeout in microseconds for HFT operations
    pub query_timeout_micros: u64,
    /// Connection acquire timeout in milliseconds
    pub acquire_timeout_ms: u64,
    /// Maximum connection lifetime in seconds
    pub max_lifetime_seconds: u64,
    /// Idle timeout in seconds
    pub idle_timeout_seconds: u64,
    /// Enable connection prewarming
    pub enable_prewarming: bool,
    /// Enable statement preparation
    pub enable_prepared_statements: bool,
    /// Enable query logging for slow queries
    pub enable_slow_query_logging: bool,
    /// Slow query threshold in microseconds
    pub slow_query_threshold_micros: u64,
}

impl Default for PostgresConfig {
    fn default() -> Self {
        Self {
            url: "postgresql://foxhunt:password@localhost:5432/foxhunt".to_owned(),
            max_connections: 50,
            min_connections: 10,
            connect_timeout_ms: 100,    // Fast connection establishment
            query_timeout_micros: 800,  // <1ms for HFT operations
            acquire_timeout_ms: 50,     // Fast pool acquisition
            max_lifetime_seconds: 3600, // 1 hour connection lifetime
            idle_timeout_seconds: 300,  // 5 minutes idle timeout
            enable_prewarming: true,
            enable_prepared_statements: true,
            enable_slow_query_logging: true,
            slow_query_threshold_micros: 1000, // Log queries >1ms
        }
    }
}

/// `PostgreSQL` connection pool with HFT optimizations
pub struct PostgresPool {
    pool: PgPool,
    config: PostgresConfig,
    metrics: Arc<RwLock<PostgresMetrics>>,
}

impl PostgresPool {
    /// Create a new `PostgreSQL` connection pool optimized for HFT
    pub async fn new(config: PostgresConfig) -> Result<Self, PostgresError> {
        // Parse and configure connection options for optimal performance
        let mut connect_options: PgConnectOptions = config
            .url
            .parse()
            .map_err(|e| PostgresError::Configuration(format!("Invalid URL: {}", e)))?;

        // Configure connection-level optimizations
        connect_options = connect_options
            .application_name("foxhunt-hft")
            .statement_cache_capacity(1000); // Cache prepared statements

        // Enable detailed logging for development/debugging
        if config.enable_slow_query_logging {
            // Note: SQLx logging configuration removed as it depends on the log crate
            // Consider using tracing-based alternatives if needed
        }

        // Create connection pool with HFT-optimized settings
        let pool = PgPoolOptions::new()
            .max_connections(config.max_connections)
            .min_connections(config.min_connections)
            .acquire_timeout(Duration::from_millis(config.acquire_timeout_ms))
            .max_lifetime(Duration::from_secs(config.max_lifetime_seconds))
            .idle_timeout(Duration::from_secs(config.idle_timeout_seconds))
            .test_before_acquire(true) // Ensure connections are healthy
            .after_connect(|conn, _meta| {
                Box::pin(async move {
                    // Optimize each connection for HFT performance
                    sqlx::query("SET synchronous_commit = OFF")
                        .execute(&mut *conn)
                        .await?;
                    sqlx::query("SET wal_writer_delay = '10ms'")
                        .execute(&mut *conn)
                        .await?;
                    sqlx::query("SET commit_delay = 0")
                        .execute(&mut *conn)
                        .await?;
                    sqlx::query("SET commit_siblings = 5")
                        .execute(&mut *conn)
                        .await?;
                    sqlx::query("SET tcp_keepalives_idle = 60")
                        .execute(&mut *conn)
                        .await?;
                    sqlx::query("SET tcp_keepalives_interval = 10")
                        .execute(&mut *conn)
                        .await?;
                    sqlx::query("SET tcp_keepalives_count = 3")
                        .execute(&mut *conn)
                        .await?;
                    Ok(())
                })
            })
            .connect_with(connect_options)
            .await
            .map_err(PostgresError::Connection)?;

        // Pre-warm connections if enabled
        if config.enable_prewarming {
            for _ in 0..config.min_connections {
                let _conn = pool.acquire().await.map_err(PostgresError::Connection)?;
                // Execute a simple query to warm up the connection
                sqlx::query("SELECT 1")
                    .fetch_one(&pool)
                    .await
                    .map_err(PostgresError::Connection)?;
            }
        }

        let metrics = Arc::new(RwLock::new(PostgresMetrics::new()));

        Ok(Self {
            pool,
            config,
            metrics,
        })
    }

    /// Get the underlying connection pool
    pub const fn pool(&self) -> &PgPool {
        &self.pool
    }

    /// Get current configuration
    pub const fn config(&self) -> &PostgresConfig {
        &self.config
    }

    /// Execute a query with performance monitoring
    pub async fn execute_monitored<'query, A>(
        &self,
        query: sqlx::query::Query<'query, sqlx::Postgres, A>,
    ) -> Result<sqlx::postgres::PgQueryResult, PostgresError>
    where
        A: 'query + sqlx::IntoArguments<'query, sqlx::Postgres>,
    {
        let start = Instant::now();

        // Execute query with timeout
        let result = tokio::time::timeout(
            Duration::from_micros(self.config.query_timeout_micros),
            query.execute(&self.pool),
        )
        .await;

        let elapsed = start.elapsed();

        // Update metrics
        self.update_metrics(elapsed, result.is_ok()).await;

        // Check for performance violations
        if elapsed.as_micros() > self.config.query_timeout_micros as u128 {
            return Err(PostgresError::QueryTimeout {
                actual_ms: elapsed.as_millis() as u64,
                max_ms: self.config.query_timeout_micros / 1000,
            });
        }

        match result {
            Ok(Ok(query_result)) => Ok(query_result),
            Ok(Err(e)) => Err(PostgresError::Connection(e)),
            Err(_) => Err(PostgresError::QueryTimeout {
                actual_ms: elapsed.as_millis() as u64,
                max_ms: self.config.query_timeout_micros / 1000,
            }),
        }
    }

    /// Fetch one row with performance monitoring
    pub async fn fetch_one_monitored<'query, A>(
        &self,
        query: sqlx::query::Query<'query, sqlx::Postgres, A>,
    ) -> Result<sqlx::postgres::PgRow, PostgresError>
    where
        A: 'query + sqlx::IntoArguments<'query, sqlx::Postgres>,
    {
        let start = Instant::now();

        let result = tokio::time::timeout(
            Duration::from_micros(self.config.query_timeout_micros),
            query.fetch_one(&self.pool),
        )
        .await;

        let elapsed = start.elapsed();
        self.update_metrics(elapsed, result.is_ok()).await;

        if elapsed.as_micros() > self.config.query_timeout_micros as u128 {
            return Err(PostgresError::QueryTimeout {
                actual_ms: elapsed.as_millis() as u64,
                max_ms: self.config.query_timeout_micros / 1000,
            });
        }

        match result {
            Ok(Ok(row)) => Ok(row),
            Ok(Err(e)) => Err(PostgresError::Connection(e)),
            Err(_) => Err(PostgresError::QueryTimeout {
                actual_ms: elapsed.as_millis() as u64,
                max_ms: self.config.query_timeout_micros / 1000,
            }),
        }
    }

    /// Health check for the `PostgreSQL` connection
    pub async fn health_check(&self) -> Result<(), PostgresError> {
        let start = Instant::now();

        let result = tokio::time::timeout(
            Duration::from_millis(100), // 100ms health check timeout
            sqlx::query("SELECT 1").fetch_one(&self.pool),
        )
        .await;

        let elapsed = start.elapsed();

        match result {
            Ok(Ok(_)) => {
                if elapsed.as_millis() > 10 {
                    warn!("PostgreSQL health check slow: {}ms", elapsed.as_millis());
                }
                Ok(())
            }
            Ok(Err(e)) => Err(PostgresError::Connection(e)),
            Err(_) => Err(PostgresError::QueryTimeout {
                actual_ms: elapsed.as_millis() as u64,
                max_ms: 100,
            }),
        }
    }

    /// Get current performance metrics
    pub async fn get_metrics(&self) -> Result<PostgresMetrics, PostgresError> {
        Ok(self.metrics.read().await.clone())
    }

    /// Get connection pool statistics
    pub async fn pool_stats(&self) -> PoolStats {
        PoolStats {
            size: self.pool.size(),
            idle: self.pool.num_idle() as u32,
            active: self.pool.size() - self.pool.num_idle() as u32,
            max_size: self.config.max_connections,
        }
    }

    /// Update internal metrics
    async fn update_metrics(&self, duration: Duration, success: bool) {
        let mut metrics = self.metrics.write().await;
        metrics.total_queries += 1;
        metrics.total_duration_micros += duration.as_micros() as u64;

        if success {
            metrics.successful_queries += 1;
        } else {
            metrics.failed_queries += 1;
        }

        if duration.as_micros() > self.config.query_timeout_micros as u128 {
            metrics.slow_queries += 1;
        }

        // Update latency percentiles (simplified)
        if duration.as_micros() < 500 {
            metrics.sub_500_micros += 1;
        } else if duration.as_micros() < 1000 {
            metrics.sub_1ms += 1;
        } else {
            metrics.over_1ms += 1;
        }
    }
}

/// `PostgreSQL` performance metrics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PostgresMetrics {
    pub total_queries: u64,
    pub successful_queries: u64,
    pub failed_queries: u64,
    pub slow_queries: u64,
    pub total_duration_micros: u64,
    pub sub_500_micros: u64,
    pub sub_1ms: u64,
    pub over_1ms: u64,
}

impl PostgresMetrics {
    const fn new() -> Self {
        Self {
            total_queries: 0,
            successful_queries: 0,
            failed_queries: 0,
            slow_queries: 0,
            total_duration_micros: 0,
            sub_500_micros: 0,
            sub_1ms: 0,
            over_1ms: 0,
        }
    }

    /// Calculate average query latency in microseconds
    pub fn average_latency_micros(&self) -> f64 {
        if self.total_queries == 0 {
            0.0
        } else {
            self.total_duration_micros as f64 / self.total_queries as f64
        }
    }

    /// Calculate success rate as percentage
    pub fn success_rate(&self) -> f64 {
        if self.total_queries == 0 {
            0.0
        } else {
            (self.successful_queries as f64 / self.total_queries as f64) * 100.0
        }
    }

    /// Calculate percentage of queries under 1ms
    pub fn sub_1ms_percentage(&self) -> f64 {
        if self.total_queries == 0 {
            0.0
        } else {
            ((self.sub_500_micros + self.sub_1ms) as f64 / self.total_queries as f64) * 100.0
        }
    }
}

/// Connection pool statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PoolStats {
    pub size: u32,
    pub idle: u32,
    pub active: u32,
    pub max_size: u32,
}

impl PoolStats {
    /// Calculate pool utilization percentage
    pub fn utilization_percentage(&self) -> f64 {
        (self.active as f64 / self.max_size as f64) * 100.0
    }

    /// Check if pool is healthy (not over-utilized)
    pub fn is_healthy(&self) -> bool {
        self.utilization_percentage() < 80.0 // Alert if >80% utilized
    }
}