ARCADIA-Studio/crates/oa-server/src/store.rs

//! Single-owner on-disk envelopes. All methods run on a blocking worker.
use crate::error::{ApiError, ApiResult, issues, source_error};
use axum::http::StatusCode;
use oa_core::persistence::{DecodeResult, MigrationReport};
use oa_core::source_project::{self, CompiledSourceProject, ProjectLimits, SourceBundle};
use oa_core::{Id, Project};
use serde::{Deserialize, Serialize};
use std::{
    collections::BTreeMap,
    fs::{self, File, OpenOptions},
    io::{Read, Write},
    path::{Path, PathBuf},
    sync::{Mutex, MutexGuard},
};

#[cfg(test)]
#[path = "store_preparation_gate.rs"]
pub(crate) mod preparation_gate;

pub const MAX_MODEL_BYTES: usize = 8 * 1024 * 1024;
pub const MAX_SOURCE_PROJECT_BYTES: usize = 8 * 1024 * 1024;
pub const MAX_OBJECTS: usize = 5_000;
pub const MAX_PROJECTS: usize = 128;
pub const MAX_STORE_BYTES: usize = 256 * 1024 * 1024;
pub(crate) const MODEL_LIMITS: oa_core::language::ParseLimits = oa_core::language::ParseLimits {
    max_source_bytes: MAX_MODEL_BYTES,
    max_graph_objects: MAX_OBJECTS,
    max_records: 20_000,
    max_step_references: 50_000,
};
pub(crate) const SOURCE_LIMITS: ProjectLimits = ProjectLimits {
    max_files: 256,
    max_include_depth: 32,
    max_include_edges: 1_024,
    model: MODEL_LIMITS,
};
const MAX_ENVELOPE_BYTES: usize = MAX_MODEL_BYTES + MAX_SOURCE_PROJECT_BYTES + 512;

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct StoredProject {
    pub resource_id: Id,
    pub revision: u64,
    pub model: Project,
    pub source_project: SourceBundle,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub migration: Option<MigrationReport>,
}

impl StoredProject {
    pub fn etag(&self) -> String {
        format!("\"p-{}-r{}\"", self.resource_id, self.revision)
    }
}

#[derive(Serialize)]
pub(crate) struct ProjectSummary {
    pub resource_id: Id,
    pub revision: u64,
    pub model_id: Id,
    pub schema_version: u32,
    pub name: String,
    pub elements: usize,
    pub relations: usize,
    pub migration: Option<MigrationReport>,
}
impl From<&StoredProject> for ProjectSummary {
    fn from(stored: &StoredProject) -> Self {
        Self {
            resource_id: stored.resource_id,
            revision: stored.revision,
            model_id: stored.model.id,
            schema_version: stored.model.schema_version,
            name: stored.model.name.clone(),
            elements: stored.model.elements.len(),
            relations: stored.model.relations.len(),
            migration: stored.migration.clone(),
        }
    }
}

#[derive(Debug, thiserror::Error)]
pub enum StoreError {
    #[error("Cannot open data directory: {0}")]
    Io(#[from] std::io::Error),
    #[error("Data directory is already owned or cannot be locked: {0}")]
    Locked(String),
    #[error("Invalid stored project {path}: {detail}")]
    Corrupt { path: PathBuf, detail: String },
    #[error("Store capacity exceeded: {0}")]
    Capacity(String),
}

pub struct Store {
    inner: Mutex<StoreInner>,
    #[cfg(test)]
    preparation_gate: Mutex<Option<preparation_gate::Armed>>,
    // Keep the OS lock alive for the store's entire lifetime, including requests.
    _directory_lock: File,
}

struct StoreInner {
    directory: PathBuf,
    pub projects: BTreeMap<Id, StoredProject>,
    byte_sizes: BTreeMap<Id, usize>,
    total_bytes: usize,
}

/// Constructed only after complete candidate validation, outside the catalog
/// mutex. Commit methods cannot accidentally accept an unchecked candidate.
pub(crate) struct PreparedProject {
    compiled: CompiledSourceProject,
    migration: Option<MigrationReport>,
}
impl PreparedProject {
    fn new(model: Project) -> ApiResult<Self> {
        validate_model(&model)?;
        let prepared = Self::from_bundle(source_project::format_bundle(&model))?;
        if prepared.model() != &model {
            return Err(ApiError::internal(
                "Generated source does not reproduce the model.",
            ));
        }
        Ok(prepared)
    }

    pub(crate) fn from_bundle(bundle: SourceBundle) -> ApiResult<Self> {
        check_source_size(&bundle)?;
        let compiled =
            source_project::compile_bundle(&bundle, SOURCE_LIMITS).map_err(source_error)?;
        Self::from_compiled(compiled)
    }

    fn from_compiled(compiled: CompiledSourceProject) -> ApiResult<Self> {
        check_model_size(&compiled.decoded.project)?;
        check_source_size(&compiled.bundle)?;
        let migration = compiled.decoded.migration.clone();
        Ok(Self {
            compiled,
            migration,
        })
    }

    pub(crate) fn from_command(
        current: &StoredProject,
        prepared: oa_core::PreparedCommand,
    ) -> ApiResult<Self> {
        check_model_size(&prepared.project)?;
        let compiled = source_project::compile_bundle(&current.source_project, SOURCE_LIMITS)
            .map_err(source_error)?;
        // Model equality alone is not a source no-op: delete/recreate and
        // identity swaps can change declaration ownership without changing the
        // final semantic snapshot.
        let compiled = source_project::update_bundle_with_identity_map(
            &compiled,
            &prepared.project,
            &prepared.identity_map,
            SOURCE_LIMITS,
        )
        .map_err(source_error)?;
        Self::from_compiled(compiled)
    }

    fn replacement(current: &StoredProject, decoded: DecodeResult) -> ApiResult<Self> {
        validate_model(&decoded.project)?;
        let compiled = source_project::compile_bundle(&current.source_project, SOURCE_LIMITS)
            .map_err(source_error)?;
        let compiled = source_project::update_bundle(&compiled, &decoded.project, SOURCE_LIMITS)
            .map_err(source_error)?;
        let mut prepared = Self::from_compiled(compiled)?;
        prepared.migration = decoded.migration;
        Ok(prepared)
    }

    pub(crate) fn model(&self) -> &Project {
        &self.compiled.decoded.project
    }

    pub(crate) fn bundle(&self) -> &SourceBundle {
        &self.compiled.bundle
    }

    pub(crate) fn migration(&self) -> Option<&MigrationReport> {
        self.migration.as_ref()
    }

    pub(crate) fn source_issues(&self) -> Vec<crate::error::IssueDto> {
        self.compiled
            .diagnostics
            .iter()
            .cloned()
            .map(crate::error::source_issue)
            .collect()
    }

    pub(crate) fn changed(&self, current: &StoredProject) -> bool {
        self.model() != &current.model || self.bundle() != &current.source_project
    }
}

impl Store {
    fn lock(&self) -> ApiResult<MutexGuard<'_, StoreInner>> {
        self.inner.lock().map_err(|_| {
            ApiError::new(
                StatusCode::SERVICE_UNAVAILABLE,
                "storage_unavailable",
                "Storage worker needs a restart.",
            )
        })
    }

    pub(crate) fn get(&self, id: Id) -> ApiResult<StoredProject> {
        Ok(self.lock()?.get(id)?.clone())
    }

    pub(crate) fn checked(&self, id: Id, expected: &str) -> ApiResult<StoredProject> {
        Ok(self.lock()?.checked(id, expected)?.clone())
    }

    pub(crate) fn catalog(&self) -> ApiResult<Vec<ProjectSummary>> {
        Ok(self
            .lock()?
            .projects
            .values()
            .map(ProjectSummary::from)
            .collect())
    }

    pub(crate) fn create(&self, model: Project) -> ApiResult<StoredProject> {
        let prepared = PreparedProject::new(model)?;
        self.lock()?.create(prepared)
    }

    pub(crate) fn create_decoded(&self, decoded: DecodeResult) -> ApiResult<StoredProject> {
        let mut prepared = PreparedProject::new(decoded.project)?;
        prepared.migration = decoded.migration;
        self.lock()?.create(prepared)
    }
    pub(crate) fn create_source_project(&self, bundle: SourceBundle) -> ApiResult<StoredProject> {
        let prepared = PreparedProject::from_bundle(bundle)?;
        self.lock()?.create(prepared)
    }
    pub(crate) fn replace_decoded(
        &self,
        id: Id,
        expected: &str,
        decoded: DecodeResult,
    ) -> ApiResult<(StoredProject, bool)> {
        let current = self.checked(id, expected)?;
        let prepared = PreparedProject::replacement(&current, decoded)?;
        #[cfg(test)]
        self.wait_after_preparation(id, expected, &prepared)?;
        self.lock()?.replace(id, expected, prepared)
    }

    pub(crate) fn replace(
        &self,
        id: Id,
        expected: &str,
        prepared: PreparedProject,
    ) -> ApiResult<(StoredProject, bool)> {
        self.lock()?.replace(id, expected, prepared)
    }

    pub(crate) fn delete(&self, id: Id, expected: &str) -> ApiResult<()> {
        self.lock()?.delete(id, expected)
    }

    /// The caller must keep this synchronous disk work off async runtime threads.
    pub fn open(directory: impl AsRef<Path>) -> Result<Self, StoreError> {
        fs::create_dir_all(directory.as_ref())?;
        let directory = fs::canonicalize(directory.as_ref())?;
        let directory_lock = OpenOptions::new()
            .create(true)
            .truncate(false)
            .read(true)
            .write(true)
            .open(directory.join(".arcadia.lock"))?;
        directory_lock
            .try_lock()
            .map_err(|error| StoreError::Locked(error.to_string()))?;
        let mut projects = BTreeMap::new();
        let mut byte_sizes = BTreeMap::new();
        let mut total_bytes: usize = 0;
        for entry in fs::read_dir(&directory)? {
            let entry = entry?;
            let path = entry.path();
            // Aborted sibling temporary files are not committed envelopes.
            if path.extension().is_none_or(|ext| ext != "json") {
                continue;
            }
            let corrupt = |detail: String| StoreError::Corrupt {
                path: path.clone(),
                detail,
            };
            if !entry.file_type()?.is_file() {
                return Err(corrupt(
                    "Envelope must be a regular file, not a link/directory.".into(),
                ));
            }
            let filename_id: Id = path
                .file_stem()
                .and_then(|s| s.to_str())
                .ok_or_else(|| corrupt("Filename is not a UUID.".into()))?
                .parse()
                .map_err(|_| corrupt("Filename is not a UUID.".into()))?;
            if path.file_stem().and_then(|stem| stem.to_str())
                != Some(filename_id.to_string().as_str())
            {
                return Err(corrupt(
                    "Filename must use the canonical lowercase hyphenated resource UUID.".into(),
                ));
            }
            let mut bytes = Vec::new();
            File::open(&path)?
                .take((MAX_ENVELOPE_BYTES + 1) as u64)
                .read_to_end(&mut bytes)?;
            if bytes.len() > MAX_ENVELOPE_BYTES {
                return Err(corrupt("Envelope exceeds the size limit.".into()));
            }
            if projects.len() >= MAX_PROJECTS {
                return Err(StoreError::Capacity(
                    "At most 128 models and 256 MiB of serialized envelopes are supported.".into(),
                ));
            }
            #[derive(Deserialize)]
            #[serde(deny_unknown_fields)]
            struct Envelope {
                resource_id: Id,
                revision: u64,
                model: Box<serde_json::value::RawValue>,
                #[serde(default, deserialize_with = "present_source_project")]
                source_project: Option<SourceBundle>,
            }
            let envelope: Envelope =
                serde_json::from_slice(&bytes).map_err(|error| corrupt(error.to_string()))?;
            let decoded = oa_core::persistence::decode_with_report_limited(
                envelope.model.get().as_bytes(),
                MODEL_LIMITS,
            )
            .map_err(|error| corrupt(error.to_string()))?;
            let prepared = if let Some(bundle) = envelope.source_project {
                let prepared = PreparedProject::from_bundle(bundle.clone())
                    .map_err(|error| corrupt(error.detail))?;
                // Explicit source upgrades may change tokens and generate new
                // records. They must still agree with independently migrated
                // JSON. Current-format unstabilized/inconsistent pairs reject.
                if (prepared.bundle() != &bundle && prepared.migration.is_none())
                    || prepared.model() != &decoded.project
                {
                    return Err(corrupt(
                        "Stored source and model disagree, or source lacks stable identities."
                            .into(),
                    ));
                }
                prepared
            } else {
                // Old model-only envelopes gain deterministic source in memory.
                // Opening the store never saves or advances their revision.
                PreparedProject::new(decoded.project.clone())
                    .map_err(|error| corrupt(error.detail))?
            };
            let stored = StoredProject {
                resource_id: envelope.resource_id,
                revision: envelope.revision,
                model: decoded.project,
                source_project: prepared.compiled.bundle,
                migration: decoded.migration.or(prepared.migration),
            };
            if stored.resource_id != filename_id || stored.revision == 0 {
                return Err(corrupt(
                    "Resource ID/filename mismatch or zero revision.".into(),
                ));
            }
            let size = envelope_bytes(&stored)
                .map_err(|error| corrupt(error.detail))?
                .len();
            // Capacity charges the complete accepted pair, including an old
            // envelope's in-memory source upgrade before its next changed save.
            total_bytes = total_bytes.saturating_add(size);
            if total_bytes > MAX_STORE_BYTES {
                return Err(StoreError::Capacity(
                    "Complete accepted envelopes exceed 256 MiB.".into(),
                ));
            }
            byte_sizes.insert(stored.resource_id, size);
            if projects.insert(stored.resource_id, stored).is_some() {
                return Err(corrupt("Duplicate hosted resource ID.".into()));
            }
        }
        Ok(Self {
            #[cfg(test)]
            preparation_gate: Mutex::new(None),
            inner: Mutex::new(StoreInner {
                directory,
                projects,
                byte_sizes,
                total_bytes,
            }),
            _directory_lock: directory_lock,
        })
    }
}

impl StoreInner {
    pub fn get(&self, id: Id) -> ApiResult<&StoredProject> {
        self.projects.get(&id).ok_or_else(ApiError::missing)
    }

    pub fn checked(&self, id: Id, expected: &str) -> ApiResult<&StoredProject> {
        let current = self.get(id)?;
        if current.etag() != expected {
            let mut error = ApiError::new(
                StatusCode::PRECONDITION_FAILED,
                "revision_conflict",
                "Project changed. Reload and review the current model before retrying.",
            );
            error.current_revision = Some(current.revision);
            return Err(error);
        }
        Ok(current)
    }

    fn create(&mut self, prepared: PreparedProject) -> ApiResult<StoredProject> {
        if self.projects.len() >= MAX_PROJECTS {
            return Err(capacity_error());
        }
        let stored = StoredProject {
            resource_id: Id::new_v4(),
            revision: 1,
            model: prepared.compiled.decoded.project,
            source_project: prepared.compiled.bundle,
            migration: prepared.migration,
        };
        let size = self.write(&stored, true)?;
        self.adopt(stored.clone(), size);
        self.sync_directory()?;
        Ok(stored)
    }

    pub fn replace(
        &mut self,
        id: Id,
        expected: &str,
        prepared: PreparedProject,
    ) -> ApiResult<(StoredProject, bool)> {
        let current = self.checked(id, expected)?;
        if !prepared.changed(current) {
            let mut unchanged = current.clone();
            if prepared.migration.is_some() {
                unchanged.migration = prepared.migration;
            }
            return Ok((unchanged, false));
        }
        let revision = current.revision.checked_add(1).ok_or_else(|| {
            ApiError::new(
                StatusCode::CONFLICT,
                "revision_exhausted",
                "Project revision counter is exhausted.",
            )
        })?;
        let stored = StoredProject {
            resource_id: id,
            revision,
            model: prepared.compiled.decoded.project,
            source_project: prepared.compiled.bundle,
            migration: prepared.migration,
        };
        let size = self.write(&stored, false)?;
        self.adopt(stored.clone(), size);
        self.sync_directory()?;
        Ok((stored, true))
    }

    pub fn delete(&mut self, id: Id, expected: &str) -> ApiResult<()> {
        self.checked(id, expected)?;
        fs::remove_file(self.path(id)).map_err(ApiError::storage)?;
        self.projects.remove(&id);
        self.total_bytes -= self.byte_sizes.remove(&id).unwrap_or(0);
        self.sync_directory()
    }

    fn path(&self, id: Id) -> PathBuf {
        self.directory.join(format!("{id}.json"))
    }

    fn adopt(&mut self, stored: StoredProject, size: usize) {
        let previous = self
            .byte_sizes
            .insert(stored.resource_id, size)
            .unwrap_or(0);
        self.total_bytes = self.total_bytes - previous + size;
        self.projects.insert(stored.resource_id, stored);
    }

    fn write(&self, stored: &StoredProject, create: bool) -> ApiResult<usize> {
        let bytes = envelope_bytes(stored)?;
        let previous = self
            .byte_sizes
            .get(&stored.resource_id)
            .copied()
            .unwrap_or(0);
        if self.total_bytes - previous + bytes.len() > MAX_STORE_BYTES {
            return Err(capacity_error());
        }
        let mut temporary = tempfile::Builder::new()
            .prefix(".arcadia-write-")
            .tempfile_in(&self.directory)
            .map_err(ApiError::storage)?;
        temporary.write_all(&bytes).map_err(ApiError::storage)?;
        temporary.as_file().sync_all().map_err(ApiError::storage)?;
        if create {
            temporary.persist_noclobber(self.path(stored.resource_id))
        } else {
            temporary.persist(self.path(stored.resource_id))
        }
        .map_err(|error| ApiError::storage(error.error))?;
        Ok(bytes.len())
    }

    fn sync_directory(&self) -> ApiResult<()> {
        // The installed file and in-memory head already agree if this sync fails.
        // Clients receive an uncertain write outcome and must reload.
        #[cfg(unix)]
        File::open(&self.directory)
            .and_then(|file| file.sync_all())
            .map_err(ApiError::storage)?;
        Ok(())
    }
}

fn capacity_error() -> ApiError {
    ApiError::new(
        StatusCode::INSUFFICIENT_STORAGE,
        "store_capacity",
        "Store supports at most 128 models and 256 MiB of serialized envelopes. Export/delete unneeded models before adding more data.",
    )
}

fn present_source_project<'de, D: serde::Deserializer<'de>>(
    deserializer: D,
) -> Result<Option<SourceBundle>, D::Error> {
    SourceBundle::deserialize(deserializer).map(Some)
}

fn envelope_bytes(stored: &StoredProject) -> ApiResult<Vec<u8>> {
    // Migration observations are response metadata, never persisted facts.
    #[derive(Serialize)]
    struct Envelope<'a> {
        resource_id: Id,
        revision: u64,
        model: &'a Project,
        source_project: &'a SourceBundle,
    }
    let bytes = serde_json::to_vec(&Envelope {
        resource_id: stored.resource_id,
        revision: stored.revision,
        model: &stored.model,
        source_project: &stored.source_project,
    })
    .map_err(|_| ApiError::internal("Could not serialize accepted project."))?;
    if bytes.len() > MAX_ENVELOPE_BYTES {
        return Err(ApiError::new(
            StatusCode::PAYLOAD_TOO_LARGE,
            "project_too_large",
            "Stored project exceeds the envelope limit.",
        ));
    }
    Ok(bytes)
}

pub(crate) fn check_source_size(bundle: &SourceBundle) -> ApiResult<()> {
    let bytes = serde_json::to_vec(bundle)
        .map_err(|_| ApiError::bad_request("Source project cannot be serialized."))?;
    if bytes.len() > MAX_SOURCE_PROJECT_BYTES {
        return Err(ApiError::new(
            StatusCode::PAYLOAD_TOO_LARGE,
            "source_project_too_large",
            "Serialized source project exceeds 8 MiB after identity stabilization.",
        ));
    }
    Ok(())
}

pub(crate) fn check_model_size(model: &Project) -> ApiResult<()> {
    MODEL_LIMITS
        .check_model(model)
        .map_err(|detail| ApiError::new(StatusCode::UNPROCESSABLE_ENTITY, "model_limit", detail))?;
    let bytes = serde_json::to_vec(model)
        .map_err(|_| ApiError::bad_request("Model cannot be serialized."))?;
    if bytes.len() > MAX_MODEL_BYTES {
        return Err(ApiError::new(
            StatusCode::PAYLOAD_TOO_LARGE,
            "model_too_large",
            "Serialized model exceeds 8 MiB.",
        ));
    }
    Ok(())
}

pub(crate) fn validate_model(model: &Project) -> ApiResult<()> {
    check_model_size(model)?;
    let issues = issues(model);
    if issues.iter().any(|issue| issue.severity == "error") {
        let mut error = ApiError::new(
            StatusCode::UNPROCESSABLE_ENTITY,
            "invalid_model",
            "Model violates structural constraints.",
        );
        error.issues = issues;
        return Err(error);
    }
    Ok(())
}

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

    #[test]
    fn source_bytes_are_charged_on_legacy_open_update_and_delete() {
        let directory = tempfile::tempdir().unwrap();
        let resource_id = Id::new_v4();
        let model = Project::empty();
        let path = directory.path().join(format!("{resource_id}.json"));
        let legacy_bytes = serde_json::to_vec(
            &serde_json::json!({"resource_id":resource_id,"revision":4,"model":model}),
        )
        .unwrap();
        fs::write(&path, &legacy_bytes).unwrap();
        let store = Store::open(directory.path()).unwrap();
        let current = store.get(resource_id).unwrap();
        let accepted_size = envelope_bytes(&current).unwrap().len();
        assert!(accepted_size > legacy_bytes.len());
        assert_eq!(store.lock().unwrap().total_bytes, accepted_size);
        assert_eq!(fs::read(&path).unwrap(), legacy_bytes);
        let mut bundle = current.source_project.clone();
        bundle
            .files
            .get_mut(&bundle.entry)
            .unwrap()
            .push_str(&format!("// {}\n", "rationale ".repeat(512)));
        let candidate = PreparedProject::from_bundle(bundle).unwrap();
        let changed = store
            .replace(resource_id, &current.etag(), candidate)
            .unwrap()
            .0;
        let changed_size = envelope_bytes(&changed).unwrap().len();
        assert!(changed_size > accepted_size + 5_000);
        assert_eq!(store.lock().unwrap().total_bytes, changed_size);
        assert_eq!(fs::read(&path).unwrap().len(), changed_size);
        drop(store);
        let store = Store::open(directory.path()).unwrap();
        assert_eq!(store.lock().unwrap().total_bytes, changed_size);
        store.delete(resource_id, &changed.etag()).unwrap();
        assert_eq!(store.lock().unwrap().total_bytes, 0);
        assert!(store.lock().unwrap().byte_sizes.is_empty());
    }

    #[test]
    fn prepared_candidates_recheck_revision_at_commit_and_failed_loser_is_not_installed() {
        let directory = tempfile::tempdir().unwrap();
        let store = Store::open(directory.path()).unwrap();
        let first = store.create(Project::empty()).unwrap();
        let expected = first.etag();
        let mut winner = store.checked(first.resource_id, &expected).unwrap().model;
        let mut loser = winner.clone();
        winner.name = "Winning review".into();
        loser.name = "Losing stale review".into();
        // Both requests completed preparation against the same original head.
        let winner = PreparedProject::new(winner).unwrap();
        let loser = PreparedProject::new(loser).unwrap();
        let committed = store
            .lock()
            .unwrap()
            .replace(first.resource_id, &expected, winner)
            .unwrap()
            .0;
        let error = store
            .lock()
            .unwrap()
            .replace(first.resource_id, &expected, loser)
            .unwrap_err();
        assert_eq!(error.status, StatusCode::PRECONDITION_FAILED.as_u16());
        assert_eq!(error.current_revision, Some(2));
        let installed: StoredProject = serde_json::from_slice(
            &fs::read(directory.path().join(format!("{}.json", first.resource_id))).unwrap(),
        )
        .unwrap();
        assert_eq!(installed.model, committed.model);
        assert_eq!(installed.revision, 2);
    }
}