loa_core/core/

database.rs

//! Database Actor - Manages tsink time-series database for local metrics storage
//!
//! This actor:
//! - Initializes tsink Storage with fixed 7-day retention
//! - Provides Arc<Storage> clones to metrics collectors
//! - Handles query operations
//! - Manages cleanup on shutdown

use ractor::{async_trait, Actor, ActorCell, ActorProcessingErr, ActorRef, SpawnErr};
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::oneshot;
use tsink::{DataPoint, Label, Storage, StorageBuilder, TimestampPrecision};

// ============================================================================
// STATE
// ============================================================================

pub struct DatabaseState {
    storage: Option<Arc<dyn Storage>>,
    metrics_path: PathBuf,
}

// ============================================================================
// MESSAGES
// ============================================================================

pub enum DatabaseMessage {
    /// Initialize the tsink database
    Initialize,
    /// Get a cloned Arc<Storage> for concurrent writes
    GetStorage {
        reply: oneshot::Sender<Arc<dyn Storage>>,
    },
    /// Query metrics data
    Query {
        metric: String,
        labels: Vec<Label>,
        start: i64,
        end: i64,
        reply: oneshot::Sender<Result<Vec<DataPoint>, String>>,
    },
    /// Close the database
    Close,
}

impl std::fmt::Debug for DatabaseMessage {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Initialize => write!(f, "DatabaseMessage::Initialize"),
            Self::GetStorage { .. } => write!(f, "DatabaseMessage::GetStorage {{ .. }}"),
            Self::Query { metric, labels, start, end, .. } => f
                .debug_struct("DatabaseMessage::Query")
                .field("metric", metric)
                .field("labels", labels)
                .field("start", start)
                .field("end", end)
                .finish(),
            Self::Close => write!(f, "DatabaseMessage::Close"),
        }
    }
}

// ============================================================================
// ACTOR IMPLEMENTATION
// ============================================================================

pub struct Database;

#[async_trait]
impl Actor for Database {
    type Msg = DatabaseMessage;
    type State = DatabaseState;
    type Arguments = PathBuf;

    async fn pre_start(
        &self,
        _myself: ActorRef<Self::Msg>,
        metrics_path: Self::Arguments,
    ) -> Result<Self::State, ActorProcessingErr> {
        tracing::debug!("Database actor starting");
        tracing::debug!("  Metrics path: {}", metrics_path.display());

        Ok(DatabaseState {
            storage: None,
            metrics_path,
        })
    }

    async fn post_start(
        &self,
        _myself: ActorRef<Self::Msg>,
        state: &mut Self::State,
    ) -> Result<(), ActorProcessingErr> {
        tracing::debug!(
            "Database actor initializing tsink storage at {}",
            state.metrics_path.display()
        );

        // Create directory if it doesn't exist
        if !state.metrics_path.exists() {
            std::fs::create_dir_all(&state.metrics_path)
                .map_err(|e| ActorProcessingErr::from(e.to_string()))?;
        }

        // Create tsink Storage with fixed configuration
        // Note: build() returns Arc<dyn Storage>, so no need to wrap in Arc::new()
        let storage = StorageBuilder::new()
            .with_data_path(&state.metrics_path)
            .with_retention(Duration::from_secs(7 * 24 * 3600)) // 7 days
            .with_partition_duration(Duration::from_secs(3600)) // 1 hour partitions
            .with_timestamp_precision(TimestampPrecision::Milliseconds)
            .with_write_timeout(Duration::from_secs(30))
            .with_wal_buffer_size(4096)
            .with_max_writers(2) // Allow 2 concurrent writers
            .build()
            .map_err(|e| ActorProcessingErr::from(e.to_string()))?;

        state.storage = Some(storage);
        tracing::info!("Database ready (7-day retention)");
        Ok(())
    }

    async fn handle(
        &self,
        _myself: ActorRef<Self::Msg>,
        message: Self::Msg,
        state: &mut Self::State,
    ) -> Result<(), ActorProcessingErr> {
        match message {
            DatabaseMessage::Initialize => {
                // No-op: Database now auto-initializes in post_start()
                tracing::debug!("DatabaseMessage::Initialize received (deprecated - auto-initializes in post_start)");
            }

            DatabaseMessage::GetStorage { reply } => {
                if let Some(storage) = &state.storage {
                    let _ = reply.send(storage.clone());
                    tracing::debug!("Provided Arc<Storage> clone to requester");
                } else {
                    tracing::warn!("GetStorage called before Initialize");
                }
            }

            DatabaseMessage::Query {
                metric,
                labels,
                start,
                end,
                reply,
            } => {
                if let Some(storage) = &state.storage {
                    match storage.select(&metric, &labels, start, end) {
                        Ok(datapoints) => {
                            tracing::debug!(
                                "Query successful: {} returned {} datapoints",
                                metric,
                                datapoints.len()
                            );
                            let _ = reply.send(Ok(datapoints));
                        }
                        Err(e) => {
                            tracing::error!("Query failed for {}: {}", metric, e);
                            let _ = reply.send(Err(e.to_string()));
                        }
                    }
                } else {
                    tracing::warn!("Query called before database initialized");
                    let _ = reply.send(Err("Database not initialized".to_string()));
                }
            }

            DatabaseMessage::Close => {
                if let Some(storage) = state.storage.take() {
                    // Close storage - Arc will handle final cleanup
                    if let Err(e) = storage.close() {
                        tracing::error!("Error closing database: {}", e);
                    } else {
                        tracing::info!("Database closed successfully");
                    }
                }
            }
        }

        Ok(())
    }

    async fn post_stop(
        &self,
        _myself: ActorRef<Self::Msg>,
        state: &mut Self::State,
    ) -> Result<(), ActorProcessingErr> {
        // Ensure storage is closed
        if let Some(storage) = state.storage.take() {
            let _ = storage.close();
        }
        tracing::info!("Database actor stopped");
        Ok(())
    }
}

// ============================================================================
// SPAWN HELPER
// ============================================================================

/// Spawn the Database actor
///
/// Returns an ActorCell for the supervisor to manage.
pub async fn spawn_database(
    supervisor_cell: ActorCell,
    child_id: String,
    storage_path: PathBuf,
) -> Result<ActorCell, SpawnErr> {
    tracing::debug!("Spawning Database actor as child '{}'", child_id);

    let metrics_path = storage_path.join("metrics");

    let (actor_ref, _handle) =
        Actor::spawn(Some("Database".to_string()), Database, metrics_path).await?;

    // Link to supervisor
    actor_ref.link(supervisor_cell.clone());

    tracing::debug!("✓ Database actor spawned successfully");

    Ok(actor_ref.get_cell())
}