wasmeld/crates/wasmeld-console/src/persistence.rs

//! Toasty models and transactional libSQL snapshots for Console control data.
//!
//! Component bytes and WIT packages remain filesystem artifacts. This module
//! persists service manifests, deployments, metrics, the bounded event
//! history, and service-scoped Host KV entries.

use std::{
    collections::BTreeSet,
    path::Path,
    sync::Arc,
    time::{SystemTime, UNIX_EPOCH},
};

use toasty::Db;
use toasty_driver_turso::Turso;
use tokio::sync::Mutex;

#[derive(Clone, Debug, toasty::Model)]
pub(crate) struct StoredService {
    #[key]
    pub service_key: String,
    pub manifest_toml: String,
    pub updated_at_ms: u64,
    pub calls: u64,
    pub errors: u64,
    pub last_latency_ms: Option<u64>,
}

#[derive(Clone, Debug, toasty::Model)]
pub(crate) struct StoredEvent {
    #[key]
    pub id: u64,
    pub timestamp_ms: u64,
    pub kind: String,
    pub service_id: Option<String>,
    pub revision: Option<String>,
    pub message: String,
}

#[derive(Clone, Debug, toasty::Model)]
pub(crate) struct StoredDeployment {
    #[key]
    pub service_id: String,
    pub active_revision: String,
    pub updated_at_ms: u64,
}

#[derive(Clone, Debug, toasty::Model)]
#[key(service_id, entry_key)]
pub(crate) struct StoredKvEntry {
    pub service_id: String,
    pub entry_key: String,
    pub value: Vec<u8>,
    pub updated_at_ms: u64,
}

/// One invocation's durable metric delta and event.
pub(crate) struct InvocationUpdate {
    pub service_key: String,
    pub failed: bool,
    pub updated_at_ms: u64,
    pub latency_ms: u64,
    pub event: StoredEvent,
}

/// Serialized access to the Toasty database connection.
#[derive(Clone)]
pub(crate) struct Persistence {
    db: Arc<Mutex<Db>>,
}

impl Persistence {
    /// Opens the database and installs or advances its schema.
    pub async fn open(path: &Path) -> toasty::Result<Self> {
        let is_new_database = !path.exists();
        let db = Db::builder()
            .models(toasty::models!(
                StoredService,
                StoredEvent,
                StoredDeployment,
                StoredKvEntry
            ))
            .build(Turso::file(path))
            .await?;
        if is_new_database {
            db.push_schema().await?;
        } else {
            // This project predates Toasty's migration history. Keep the
            // additive bootstrap explicit so existing databases can be opened
            // without recreating the service and event tables.
            let mut db = db;
            toasty::sql::statement(
                r#"CREATE TABLE IF NOT EXISTS "stored_deployments" (
                    "service_id" TEXT NOT NULL,
                    "active_revision" TEXT NOT NULL,
                    "updated_at_ms" INTEGER NOT NULL,
                    PRIMARY KEY ("service_id")
                )"#,
            )
            .exec(&mut db)
            .await?;
            toasty::sql::statement(
                r#"CREATE TABLE IF NOT EXISTS "stored_kv_entries" (
                    "service_id" TEXT NOT NULL,
                    "entry_key" TEXT NOT NULL,
                    "value" BLOB NOT NULL,
                    "updated_at_ms" INTEGER NOT NULL,
                    PRIMARY KEY ("service_id", "entry_key")
                )"#,
            )
            .exec(&mut db)
            .await?;
            return Ok(Self {
                db: Arc::new(Mutex::new(db)),
            });
        }
        Ok(Self {
            db: Arc::new(Mutex::new(db)),
        })
    }

    /// Loads the complete service and deployment sets plus retained events.
    pub async fn load(
        &self,
    ) -> toasty::Result<(Vec<StoredService>, Vec<StoredEvent>, Vec<StoredDeployment>)> {
        let mut db = self.db.lock().await;
        let services = StoredService::all().exec(&mut *db).await?;
        let events = StoredEvent::all().exec(&mut *db).await?;
        let deployments = StoredDeployment::all().exec(&mut *db).await?;
        Ok((services, events, deployments))
    }

    /// Atomically upserts a Console snapshot and prunes expired events.
    pub async fn sync(
        &self,
        services: Vec<StoredService>,
        events: Vec<StoredEvent>,
        deployments: Vec<StoredDeployment>,
    ) -> toasty::Result<()> {
        let mut db = self.db.lock().await;
        let mut tx = db.transaction().await?;

        let stored_services = StoredService::all().exec(&mut tx).await?;
        let service_keys = services
            .iter()
            .map(|service| service.service_key.clone())
            .collect::<BTreeSet<_>>();
        let existing_service_keys = stored_services
            .iter()
            .map(|service| service.service_key.clone())
            .collect::<BTreeSet<_>>();
        for stored in stored_services {
            if !service_keys.contains(&stored.service_key) {
                StoredService::delete_by_service_key(&mut tx, &stored.service_key).await?;
            }
        }
        for service in services {
            if existing_service_keys.contains(&service.service_key) {
                // Invocation metrics are written as ordered deltas by
                // `record_invocations`; control snapshots must not overwrite
                // or double-count those concurrent updates.
                StoredService::filter_by_service_key(&service.service_key)
                    .update()
                    .manifest_toml(service.manifest_toml)
                    .updated_at_ms(service.updated_at_ms)
                    .exec(&mut tx)
                    .await?;
            } else {
                StoredService::create()
                    .service_key(service.service_key)
                    .manifest_toml(service.manifest_toml)
                    .updated_at_ms(service.updated_at_ms)
                    .calls(0)
                    .errors(0)
                    .last_latency_ms(None)
                    .exec(&mut tx)
                    .await?;
            }
        }

        if let Some(oldest_event) = events.last() {
            StoredEvent::filter(StoredEvent::fields().id().lt(oldest_event.id))
                .delete()
                .exec(&mut tx)
                .await?;
        }

        for event in events {
            StoredEvent::upsert_by_id(event.id)
                .timestamp_ms(event.timestamp_ms)
                .kind(event.kind)
                .service_id(event.service_id)
                .revision(event.revision)
                .message(event.message)
                .exec(&mut tx)
                .await?;
        }

        let deployment_ids = deployments
            .iter()
            .map(|deployment| deployment.service_id.clone())
            .collect::<BTreeSet<_>>();
        for stored in StoredDeployment::all().exec(&mut tx).await? {
            if !deployment_ids.contains(&stored.service_id) {
                StoredDeployment::delete_by_service_id(&mut tx, &stored.service_id).await?;
            }
        }
        for deployment in deployments {
            StoredDeployment::upsert_by_service_id(deployment.service_id)
                .active_revision(deployment.active_revision)
                .updated_at_ms(deployment.updated_at_ms)
                .exec(&mut tx)
                .await?;
        }

        tx.commit().await
    }

    /// Atomically records invocation deltas and their bounded event history.
    pub(crate) async fn record_invocations(
        &self,
        updates: Vec<InvocationUpdate>,
    ) -> toasty::Result<()> {
        let mut db = self.db.lock().await;
        let mut tx = db.transaction().await?;
        let mut newest_event_id = 0;

        for update in updates {
            if let Some(stored) = StoredService::filter_by_service_key(&update.service_key)
                .first()
                .exec(&mut tx)
                .await?
            {
                StoredService::filter_by_service_key(&update.service_key)
                    .update()
                    .calls(stored.calls.saturating_add(1))
                    .errors(stored.errors.saturating_add(u64::from(update.failed)))
                    .updated_at_ms(update.updated_at_ms)
                    .last_latency_ms(Some(update.latency_ms))
                    .exec(&mut tx)
                    .await?;
            }

            newest_event_id = newest_event_id.max(update.event.id);
            StoredEvent::upsert_by_id(update.event.id)
                .timestamp_ms(update.event.timestamp_ms)
                .kind(update.event.kind)
                .service_id(update.event.service_id)
                .revision(update.event.revision)
                .message(update.event.message)
                .exec(&mut tx)
                .await?;
        }

        let oldest_retained_event_id = newest_event_id.saturating_sub(255);
        if oldest_retained_event_id > 0 {
            StoredEvent::filter(StoredEvent::fields().id().lt(oldest_retained_event_id))
                .delete()
                .exec(&mut tx)
                .await?;
        }

        tx.commit().await
    }

    pub(crate) async fn kv_get(
        &self,
        service_id: &str,
        entry_key: &str,
    ) -> toasty::Result<Option<Vec<u8>>> {
        let mut db = self.db.lock().await;
        Ok(
            StoredKvEntry::filter_by_service_id_and_entry_key(service_id, entry_key)
                .first()
                .exec(&mut *db)
                .await?
                .map(|entry| entry.value),
        )
    }

    pub(crate) async fn kv_set(
        &self,
        service_id: String,
        entry_key: String,
        value: Vec<u8>,
    ) -> toasty::Result<()> {
        let mut db = self.db.lock().await;
        StoredKvEntry::upsert_by_service_id_and_entry_key(service_id, entry_key)
            .value(value)
            .updated_at_ms(unix_time_ms())
            .exec(&mut *db)
            .await?;
        Ok(())
    }

    pub(crate) async fn kv_delete(&self, service_id: &str, entry_key: &str) -> toasty::Result<()> {
        let mut db = self.db.lock().await;
        StoredKvEntry::delete_by_service_id_and_entry_key(&mut *db, service_id, entry_key).await
    }
}

fn unix_time_ms() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|duration| u64::try_from(duration.as_millis()).unwrap_or(u64::MAX))
        .unwrap_or_default()
}