obscura-server 0.3.11

A server for relaying secure messages using the Signal Protocol
Documentation
use crate::config::HealthConfig;
use crate::storage::DbPool;
use aws_sdk_s3::Client;
use opentelemetry::{KeyValue, global, metrics::Gauge};
use std::time::Duration;
use tokio::time::timeout;

#[derive(Clone)]
pub struct HealthMetrics {
    pub health_status: Gauge<i64>,
}

impl HealthMetrics {
    pub fn new() -> Self {
        let meter = global::meter("obscura-server");
        Self {
            health_status: meter
                .i64_gauge("health_status")
                .with_description("Status of health checks (1 for ok, 0 for error)")
                .build(),
        }
    }
}

impl Default for HealthMetrics {
    fn default() -> Self {
        Self::new()
    }
}

#[derive(Clone)]
pub struct HealthService {
    pool: DbPool,
    s3_client: Client,
    storage_bucket: String,
    config: HealthConfig,
    metrics: HealthMetrics,
}

impl HealthService {
    pub fn new(pool: DbPool, s3_client: Client, storage_bucket: String, config: HealthConfig) -> Self {
        Self {
            pool,
            s3_client,
            storage_bucket,
            config,
            metrics: HealthMetrics::new(),
        }
    }

    pub async fn check_db(&self) -> Result<(), String> {
        let db_timeout = Duration::from_millis(self.config.db_timeout_ms);
        
        match timeout(db_timeout, sqlx::query("SELECT 1").execute(&self.pool)).await {
            Ok(Ok(_)) => {
                self.metrics.health_status.record(1, &[KeyValue::new("component", "database")]);
                Ok(())
            }
            Ok(Err(e)) => {
                self.metrics.health_status.record(0, &[KeyValue::new("component", "database")]);
                Err(format!("Database connection failed: {:?}", e))
            }
            Err(_) => {
                self.metrics.health_status.record(0, &[KeyValue::new("component", "database")]);
                Err("Database connection timed out".to_string())
            }
        }
    }

    pub async fn check_storage(&self) -> Result<(), String> {
        let storage_timeout = Duration::from_millis(self.config.storage_timeout_ms);

        match timeout(storage_timeout, self.s3_client.head_bucket().bucket(&self.storage_bucket).send()).await {
            Ok(Ok(_)) => {
                self.metrics.health_status.record(1, &[KeyValue::new("component", "storage")]);
                Ok(())
            }
            Ok(Err(e)) => {
                self.metrics.health_status.record(0, &[KeyValue::new("component", "storage")]);
                Err(format!("Storage connection failed for bucket {}: {:?}", self.storage_bucket, e))
            }
            Err(_) => {
                self.metrics.health_status.record(0, &[KeyValue::new("component", "storage")]);
                Err("Storage connection timed out".to_string())
            }
        }
    }
}