foxhunt/services/trading_service/src/async_audit_queue.rs

// Async Audit Queue Implementation
// Purpose: Reduce E2E latency by moving audit writes off critical path
// Target: 300μs synchronous write → 5μs async queue send (-98.3%)

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use sqlx::PgPool;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tokio::fs::OpenOptions;
use tokio::io::AsyncWriteExt;
use tokio::sync::mpsc;
use tokio::task::JoinHandle;
use tracing::{debug, error, info, warn};

/// Audit event to be logged
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AuditEvent {
    pub timestamp: DateTime<Utc>,
    pub user_id: String,
    pub action: String,
    pub details: serde_json::Value,
    pub ip_address: Option<String>,
    pub session_id: Option<String>,
}

/// Configuration for async audit queue
#[derive(Debug, Clone)]
pub struct AuditQueueConfig {
    /// Maximum number of events to buffer before backpressure
    pub buffer_size: usize,
    /// Number of events to batch before writing to database
    pub batch_size: usize,
    /// Maximum time to wait before flushing partial batch
    pub flush_interval: Duration,
    /// Path to fallback disk file if database unavailable
    pub fallback_path: String,
    /// Maximum retries for database writes
    pub max_retries: u32,
}

impl Default for AuditQueueConfig {
    fn default() -> Self {
        Self {
            buffer_size: 10_000,
            batch_size: 100,
            flush_interval: Duration::from_secs(1),
            fallback_path: "/var/lib/foxhunt/audit_fallback.jsonl".to_string(),
            max_retries: 3,
        }
    }
}

/// Metrics for monitoring audit queue performance
#[derive(Debug, Default)]
pub struct AuditQueueMetrics {
    pub events_sent: AtomicU64,
    pub events_written: AtomicU64,
    pub events_failed: AtomicU64,
    pub events_fallback: AtomicU64,
    pub batch_writes: AtomicU64,
    pub queue_depth: AtomicU64,
}

impl AuditQueueMetrics {
    pub fn record_send(&self) {
        self.events_sent.fetch_add(1, Ordering::Relaxed);
    }

    pub fn record_write(&self, count: u64) {
        self.events_written.fetch_add(count, Ordering::Relaxed);
        self.batch_writes.fetch_add(1, Ordering::Relaxed);
    }

    pub fn record_failure(&self) {
        self.events_failed.fetch_add(1, Ordering::Relaxed);
    }

    pub fn record_fallback(&self) {
        self.events_fallback.fetch_add(1, Ordering::Relaxed);
    }

    pub fn update_queue_depth(&self, depth: u64) {
        self.queue_depth.store(depth, Ordering::Relaxed);
    }
}

/// Async audit queue for non-blocking audit logging
pub struct AsyncAuditQueue {
    sender: mpsc::Sender<AuditEvent>,
    metrics: Arc<AuditQueueMetrics>,
    worker_handle: JoinHandle<()>,
}

impl AsyncAuditQueue {
    /// Create new async audit queue with background worker
    pub async fn new(pool: PgPool, config: AuditQueueConfig) -> Self {
        let (sender, receiver) = mpsc::channel(config.buffer_size);
        let metrics = Arc::new(AuditQueueMetrics::default());

        // Spawn background worker task
        let worker_handle = tokio::spawn(audit_worker(
            receiver,
            pool,
            config,
            Arc::clone(&metrics),
        ));

        info!(
            "Async audit queue started - buffer: {}, batch: {}, flush_interval: {:?}",
            config.buffer_size, config.batch_size, config.flush_interval
        );

        Self {
            sender,
            metrics,
            worker_handle,
        }
    }

    /// Log audit event (non-blocking, returns immediately)
    ///
    /// Target latency: <10μs
    pub async fn log_event(&self, event: AuditEvent) -> Result<(), &'static str> {
        self.metrics.record_send();

        // Non-blocking send with timeout
        match tokio::time::timeout(Duration::from_micros(50), self.sender.send(event)).await {
            Ok(Ok(_)) => {
                let capacity = u64::try_from(self.sender.capacity()).unwrap_or(u64::MAX);
                self.metrics.update_queue_depth(capacity);
                Ok(())
            }
            Ok(Err(_)) => {
                self.metrics.record_failure();
                error!("Audit queue channel closed");
                Err("Audit queue channel closed")
            }
            Err(_) => {
                self.metrics.record_failure();
                warn!("Audit queue send timeout - queue may be full");
                Err("Audit queue timeout")
            }
        }
    }

    /// Get current metrics
    pub fn metrics(&self) -> &Arc<AuditQueueMetrics> {
        &self.metrics
    }

    /// Graceful shutdown - flush remaining events
    pub async fn shutdown(self) -> Result<(), &'static str> {
        info!("Shutting down audit queue - flushing remaining events");

        // Drop sender to signal worker to finish
        drop(self.sender);

        // Wait for worker to complete
        match tokio::time::timeout(Duration::from_secs(30), self.worker_handle).await {
            Ok(Ok(_)) => {
                info!("Audit queue shutdown complete");
                Ok(())
            }
            Ok(Err(e)) => {
                error!("Audit queue worker panicked: {:?}", e);
                Err("Worker panic during shutdown")
            }
            Err(_) => {
                error!("Audit queue shutdown timeout");
                Err("Shutdown timeout")
            }
        }
    }
}

/// Background worker task for batch writing audit events
async fn audit_worker(
    mut receiver: mpsc::Receiver<AuditEvent>,
    pool: PgPool,
    config: AuditQueueConfig,
    metrics: Arc<AuditQueueMetrics>,
) {
    let mut batch: Vec<AuditEvent> = Vec::with_capacity(config.batch_size);
    let mut flush_timer = tokio::time::interval(config.flush_interval);
    flush_timer.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);

    info!("Audit worker started");

    loop {
        tokio::select! {
            // Receive events from channel
            event = receiver.recv() => {
                match event {
                    Some(event) => {
                        batch.push(event);

                        // Write batch if full
                        if batch.len() >= config.batch_size {
                            write_batch(&pool, &mut batch, &config, &metrics).await;
                        }
                    }
                    None => {
                        // Channel closed - flush remaining events and exit
                        info!("Audit channel closed - flushing {} remaining events", batch.len());
                        if !batch.is_empty() {
                            write_batch(&pool, &mut batch, &config, &metrics).await;
                        }
                        break;
                    }
                }
            }

            // Periodic flush for partial batches
            _ = flush_timer.tick() => {
                if !batch.is_empty() {
                    debug!("Flushing partial batch: {} events", batch.len());
                    write_batch(&pool, &mut batch, &config, &metrics).await;
                }
            }
        }
    }

    info!("Audit worker stopped");
}

/// Write batch of audit events to database with retries and fallback
async fn write_batch(
    pool: &PgPool,
    batch: &mut Vec<AuditEvent>,
    config: &AuditQueueConfig,
    metrics: &Arc<AuditQueueMetrics>,
) {
    if batch.is_empty() {
        return;
    }

    let batch_size = batch.len();
    let mut retries = 0;

    loop {
        match write_batch_to_db(pool, batch).await {
            Ok(_) => {
                let batch_size_u64 = u64::try_from(batch_size).unwrap_or(u64::MAX);
                metrics.record_write(batch_size_u64);
                debug!("Batch write successful: {} events", batch_size);
                batch.clear();
                return;
            }
            Err(e) => {
                retries += 1;
                if retries > config.max_retries {
                    error!(
                        "Database write failed after {} retries: {:?} - writing to fallback",
                        config.max_retries, e
                    );
                    write_batch_to_fallback(batch, &config.fallback_path, metrics).await;
                    batch.clear();
                    return;
                }

                warn!(
                    "Database write failed (attempt {}/{}): {:?} - retrying",
                    retries, config.max_retries, e
                );
                let backoff_ms = u64::try_from(100u32.saturating_mul(retries)).unwrap_or(3000);
                tokio::time::sleep(Duration::from_millis(backoff_ms)).await;
            }
        }
    }
}

/// Write batch to database as single transaction
async fn write_batch_to_db(
    pool: &PgPool,
    batch: &[AuditEvent],
) -> Result<(), sqlx::Error> {
    let mut tx = pool.begin().await?;

    for event in batch {
        sqlx::query(
            r#"
            INSERT INTO audit_log (timestamp, user_id, action, details, ip_address, session_id)
            VALUES ($1, $2, $3, $4, $5, $6)
            "#,
        )
        .bind(event.timestamp)
        .bind(&event.user_id)
        .bind(&event.action)
        .bind(&event.details)
        .bind(&event.ip_address)
        .bind(&event.session_id)
        .execute(&mut *tx)
        .await?;
    }

    tx.commit().await?;
    Ok(())
}

/// Fallback: Write batch to disk if database unavailable
async fn write_batch_to_fallback(
    batch: &[AuditEvent],
    fallback_path: &str,
    metrics: &Arc<AuditQueueMetrics>,
) {
    match write_to_disk(batch, fallback_path).await {
        Ok(_) => {
            metrics.record_fallback();
            warn!(
                fallback_path = %fallback_path,
                event_count = batch.len(),
                "Primary audit queue failed, writing to fallback file. Investigate queue health immediately."
            );
        }
        Err(e) => {
            metrics.record_failure();
            error!(
                "Failed to write to fallback file {}: {:?} - EVENTS LOST: {}",
                fallback_path,
                e,
                batch.len()
            );
        }
    }
}

/// Write events to disk as JSONL (`JSON` Lines)
async fn write_to_disk(batch: &[AuditEvent], path: &str) -> Result<(), std::io::Error> {
    let mut file = OpenOptions::new()
        .create(true)
        .append(true)
        .open(path)
        .await?;

    for event in batch {
        let json = serde_json::to_string(event)?;
        file.write_all(json.as_bytes()).await?;
        file.write_all(b"\n").await?;
    }

    file.sync_all().await?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    fn create_test_event(user_id: &str, action: &str) -> AuditEvent {
        AuditEvent {
            timestamp: Utc::now(),
            user_id: user_id.to_string(),
            action: action.to_string(),
            details: json!({"test": true}),
            ip_address: Some("127.0.0.1".to_string()),
            session_id: Some("test-session".to_string()),
        }
    }

    #[tokio::test]
    async fn test_queue_send_latency() {
        // Test: Queue send latency should be <10μs
        let pool = match PgPool::connect("postgresql://foxhunt:foxhunt_dev_password@localhost:5432/foxhunt")
            .await
        {
            Ok(p) => p,
            Err(e) => panic!("Failed to connect to database: {}", e),
        };

        let config = AuditQueueConfig {
            buffer_size: 1000,
            batch_size: 10,
            ..Default::default()
        };

        let queue = AsyncAuditQueue::new(pool, config).await;

        // Measure latency of 1000 sends
        let start = std::time::Instant::now();
        for i in 0..1000 {
            let event = create_test_event(&format!("user_{}", i), "test_action");
            queue.log_event(event).await.unwrap_or_else(|e| panic!("Send failed: {}", e));
        }
        let elapsed = start.elapsed();

        let avg_latency = elapsed.as_micros() / 1000;
        println!("Average queue send latency: {}μs", avg_latency);

        assert!(
            avg_latency < 10,
            "Queue send latency {}μs exceeds 10μs target",
            avg_latency
        );

        // Shutdown and verify metrics
        queue.shutdown().await.unwrap_or_else(|e| panic!("Shutdown failed: {}", e));
    }

    #[tokio::test]
    async fn test_batch_writing() {
        // Test: Batch writes should group events efficiently
        let pool = match PgPool::connect("postgresql://foxhunt:foxhunt_dev_password@localhost:5432/foxhunt")
            .await
        {
            Ok(p) => p,
            Err(e) => panic!("Failed to connect to database: {}", e),
        };

        let config = AuditQueueConfig {
            buffer_size: 1000,
            batch_size: 100,
            flush_interval: Duration::from_millis(100),
            ..Default::default()
        };

        let queue = AsyncAuditQueue::new(pool, config).await;

        // Send 250 events (should create 3 batches: 100, 100, 50)
        for i in 0..250 {
            let event = create_test_event(&format!("user_{}", i), "batch_test");
            queue.log_event(event).await.unwrap_or_else(|e| panic!("Send failed: {}", e));
        }

        // Wait for batches to be written
        tokio::time::sleep(Duration::from_millis(500)).await;

        let metrics = queue.metrics();
        let batch_writes = metrics.batch_writes.load(Ordering::Relaxed);
        let events_written = metrics.events_written.load(Ordering::Relaxed);

        println!("Batch writes: {}", batch_writes);
        println!("Events written: {}", events_written);

        assert!(batch_writes >= 2, "Expected at least 2 batch writes");
        assert_eq!(
            events_written, 250,
            "Expected 250 events written, got {}",
            events_written
        );

        queue.shutdown().await.unwrap_or_else(|e| panic!("Shutdown failed: {}", e));
    }

    #[tokio::test]
    async fn test_no_event_loss_under_load() {
        // Test: No events should be lost under high load
        let pool = match PgPool::connect("postgresql://foxhunt:foxhunt_dev_password@localhost:5432/foxhunt")
            .await
        {
            Ok(p) => p,
            Err(e) => panic!("Failed to connect to database: {}", e),
        };

        let config = AuditQueueConfig {
            buffer_size: 10_000,
            batch_size: 100,
            flush_interval: Duration::from_millis(50),
            ..Default::default()
        };

        let queue = AsyncAuditQueue::new(pool, config).await;

        // Send 10,000 events rapidly
        for i in 0..10_000 {
            let event = create_test_event(&format!("user_{}", i), "load_test");
            queue.log_event(event).await.unwrap_or_else(|e| panic!("Send failed: {}", e));
        }

        // Shutdown and wait for all events to be written
        queue.shutdown().await.unwrap_or_else(|e| panic!("Shutdown failed: {}", e));

        let metrics = queue.metrics();
        let events_sent = metrics.events_sent.load(Ordering::Relaxed);
        let events_written = metrics.events_written.load(Ordering::Relaxed);
        let events_failed = metrics.events_failed.load(Ordering::Relaxed);

        println!("Events sent: {}", events_sent);
        println!("Events written: {}", events_written);
        println!("Events failed: {}", events_failed);

        assert_eq!(events_sent, 10_000, "Expected 10,000 events sent");
        assert_eq!(
            events_written, 10_000,
            "Expected 10,000 events written, got {}",
            events_written
        );
        assert_eq!(events_failed, 0, "No events should fail");
    }

    #[tokio::test]
    async fn test_fallback_on_db_failure() {
        // Test: Events should fallback to disk if database unavailable
        let fallback_path = "/tmp/test_audit_fallback.jsonl";

        // Remove existing fallback file
        let _ = tokio::fs::remove_file(fallback_path).await;

        // Use invalid connection string to simulate database failure
        let pool = match PgPool::connect("postgresql://invalid:invalid@localhost:9999/invalid")
            .await
        {
            Ok(p) => p,
            Err(e) => panic!("Expected connection to fail but got error: {}", e),
        };

        let config = AuditQueueConfig {
            buffer_size: 100,
            batch_size: 10,
            flush_interval: Duration::from_millis(100),
            fallback_path: fallback_path.to_string(),
            max_retries: 1,
        };

        let queue = AsyncAuditQueue::new(pool, config).await;

        // Send 10 events
        for i in 0..10 {
            let event = create_test_event(&format!("user_{}", i), "fallback_test");
            queue.log_event(event).await.unwrap_or_else(|e| panic!("Send failed: {}", e));
        }

        // Wait for fallback write
        tokio::time::sleep(Duration::from_millis(500)).await;
        queue.shutdown().await.unwrap_or_else(|e| panic!("Shutdown failed: {}", e));

        // Verify fallback file exists
        let fallback_exists = tokio::fs::metadata(fallback_path).await.is_ok();
        assert!(
            fallback_exists,
            "Fallback file should exist at {}",
            fallback_path
        );

        // Cleanup
        let _ = tokio::fs::remove_file(fallback_path).await;
    }
}