modelshelf 0.1.0

//! The registry: loading, saving, and safely mutating `registry.json`.

pub mod lock;
pub mod schema;

use std::path::PathBuf;
use std::time::Duration;

pub use schema::{
    AppRef, Ecosystem, ExtraFields, GgufMeta, Location, LocationRole, ModelEntry, ModelId,
    Registry, Source, SCHEMA_VERSION,
};

use crate::paths::ShelfPaths;
use crate::{Error, Result};

/// Handle for reading and mutating the on-disk registry with correct locking.
#[derive(Debug, Clone)]
pub struct RegistryStore {
    paths: ShelfPaths,
    lock_timeout: Duration,
}

impl RegistryStore {
    /// Create a store over the given shelf paths.
    pub fn new(paths: ShelfPaths, lock_timeout: Duration) -> Self {
        RegistryStore {
            paths,
            lock_timeout,
        }
    }

    /// The shelf paths this store operates on.
    pub fn paths(&self) -> &ShelfPaths {
        &self.paths
    }

    /// Read the registry under a shared lock.
    ///
    /// A missing `registry.json` yields an empty registry (the shelf may not
    /// have been written to yet); a present-but-unparsable one is quarantined
    /// and reported as [`Error::RegistryCorrupt`].
    pub fn read(&self) -> Result<Registry> {
        let _guard = lock::shared(&self.paths.lock_file(), self.lock_timeout)?;
        self.load_locked()
    }

    /// Mutate the registry under an exclusive lock.
    ///
    /// The registry is re-read *under the lock* (read-modify-write), passed to
    /// `mutate`, then written back atomically. If `mutate` returns an error,
    /// nothing is written.
    pub fn update<T>(&self, mutate: impl FnOnce(&mut Registry) -> Result<T>) -> Result<T> {
        let _guard = lock::exclusive(&self.paths.lock_file(), self.lock_timeout)?;
        let mut reg = self.load_locked()?;
        let out = mutate(&mut reg)?;
        reg.schema_version = reg.schema_version.max(SCHEMA_VERSION);
        reg.generated_by = format!("modelshelf/{}", env!("CARGO_PKG_VERSION"));
        reg.updated_at = crate::time_util::now_rfc3339();
        let mut bytes = serde_json::to_vec_pretty(&reg)
            .map_err(|e| Error::io(self.paths.registry_json(), std::io::Error::other(e)))?;
        bytes.push(b'\n');
        lock::write_atomic(&self.paths.registry_json(), &bytes)?;
        Ok(out)
    }

    /// Discard a corrupt `registry.json` (already quarantined by a previous
    /// failed read) and start over with an empty registry. Explicit opt-in —
    /// the library never rebuilds silently.
    pub fn rebuild(&self) -> Result<Registry> {
        let _guard = lock::exclusive(&self.paths.lock_file(), self.lock_timeout)?;
        let reg = Registry::empty();
        let bytes = serde_json::to_vec_pretty(&reg)
            .map_err(|e| Error::io(self.paths.registry_json(), std::io::Error::other(e)))?;
        lock::write_atomic(&self.paths.registry_json(), &bytes)?;
        Ok(reg)
    }

    /// Load the registry file. Caller must hold the lock.
    fn load_locked(&self) -> Result<Registry> {
        let path = self.paths.registry_json();
        let bytes = match std::fs::read(&path) {
            Ok(b) => b,
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
                return Ok(Registry::empty());
            }
            Err(e) => return Err(Error::io(&path, e)),
        };
        let reg: Registry = match serde_json::from_slice(&bytes) {
            Ok(r) => r,
            Err(e) => {
                let quarantined_to = self.quarantine(&path)?;
                return Err(Error::RegistryCorrupt {
                    path,
                    quarantined_to,
                    reason: e.to_string(),
                });
            }
        };
        if reg.schema_version > SCHEMA_VERSION {
            return Err(Error::UnsupportedSchema {
                found: reg.schema_version,
                supported: SCHEMA_VERSION,
            });
        }
        Ok(reg)
    }

    /// Preserve an unreadable registry file for post-mortem, out of the way
    /// of future reads.
    fn quarantine(&self, path: &std::path::Path) -> Result<PathBuf> {
        let mut dest = path.with_file_name(format!(
            "registry.json.corrupt-{}",
            crate::time_util::now_compact()
        ));
        // Avoid clobbering an earlier quarantine from the same second.
        let mut n = 1;
        while dest.exists() {
            dest = path.with_file_name(format!(
                "registry.json.corrupt-{}-{n}",
                crate::time_util::now_compact()
            ));
            n += 1;
        }
        std::fs::rename(path, &dest).map_err(|e| Error::io(path, e))?;
        Ok(dest)
    }
}

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

    fn store(dir: &std::path::Path) -> RegistryStore {
        RegistryStore::new(ShelfPaths::at(dir), Duration::from_secs(10))
    }

    #[test]
    fn missing_file_reads_as_empty() {
        let tmp = tempfile::tempdir().unwrap();
        let reg = store(tmp.path()).read().unwrap();
        assert_eq!(reg.schema_version, SCHEMA_VERSION);
        assert!(reg.models.is_empty());
    }

    #[test]
    fn update_persists_and_stamps_writer() {
        let tmp = tempfile::tempdir().unwrap();
        let s = store(tmp.path());
        s.update(|reg| {
            reg.extra
                .insert("marker".into(), serde_json::json!("hello"));
            Ok(())
        })
        .unwrap();
        let reg = s.read().unwrap();
        assert_eq!(reg.extra["marker"], serde_json::json!("hello"));
        assert!(reg.generated_by.starts_with("modelshelf/"));
    }

    #[test]
    fn failed_mutation_writes_nothing() {
        let tmp = tempfile::tempdir().unwrap();
        let s = store(tmp.path());
        let err = s
            .update(|_| -> Result<()> { Err(Error::Refused("nope".into())) })
            .unwrap_err();
        assert!(matches!(err, Error::Refused(_)));
        assert!(!tmp.path().join("registry.json").exists());
    }

    #[test]
    fn corrupt_file_is_quarantined_not_rebuilt() {
        let tmp = tempfile::tempdir().unwrap();
        let s = store(tmp.path());
        std::fs::write(tmp.path().join("registry.json"), b"{not json!").unwrap();
        let err = s.read().unwrap_err();
        let Error::RegistryCorrupt { quarantined_to, .. } = &err else {
            panic!("expected RegistryCorrupt, got {err:?}");
        };
        assert!(quarantined_to.exists());
        assert!(!tmp.path().join("registry.json").exists());
        // Explicit rebuild recovers.
        s.rebuild().unwrap();
        assert!(s.read().unwrap().models.is_empty());
    }

    #[test]
    fn newer_schema_version_is_refused() {
        let tmp = tempfile::tempdir().unwrap();
        let s = store(tmp.path());
        std::fs::write(
            tmp.path().join("registry.json"),
            br#"{"schema_version": 999, "generated_by": "future/1.0", "updated_at": "2030-01-01T00:00:00Z", "models": []}"#,
        )
        .unwrap();
        let err = s.read().unwrap_err();
        assert!(matches!(
            err,
            Error::UnsupportedSchema {
                found: 999,
                supported: SCHEMA_VERSION
            }
        ));
    }

    #[test]
    fn concurrent_read_modify_write_loses_no_updates() {
        let tmp = tempfile::tempdir().unwrap();
        let dir = tmp.path().to_path_buf();
        const THREADS: u64 = 16;
        const ROUNDS: u64 = 50;

        let handles: Vec<_> = (0..THREADS)
            .map(|_| {
                let dir = dir.clone();
                std::thread::spawn(move || {
                    let s = RegistryStore::new(ShelfPaths::at(&dir), Duration::from_secs(60));
                    for _ in 0..ROUNDS {
                        s.update(|reg| {
                            let n = reg
                                .extra
                                .get("counter")
                                .and_then(|v| v.as_u64())
                                .unwrap_or(0);
                            reg.extra.insert("counter".into(), serde_json::json!(n + 1));
                            Ok(())
                        })
                        .unwrap();
                    }
                })
            })
            .collect();
        for h in handles {
            h.join().unwrap();
        }

        let s = store(&dir);
        let reg = s.read().unwrap();
        assert_eq!(
            reg.extra["counter"].as_u64().unwrap(),
            THREADS * ROUNDS,
            "read-modify-write under exclusive lock must not lose updates"
        );
    }
}