trusty-memory 0.28.2

//! Tests for `spawn_startup_tasks` (#474).
//!
//! Why: `spawn_startup_tasks` populates `AppState::pin_project_map` in the
//! background so handlers can resolve a palace id to a project path without
//! a filesystem walk at request time (issue #470). Moved out of `main.rs`
//! into its own file (pure code motion, no logic change) to keep `main.rs`
//! under the 500 SLOC production cap — see issue #2522 review.
//!
//! What: exercises the full wiring — a real `AppState` with a real temp
//! search root, a pinned project, and the async background task. The test
//! owns its runtime and its environment: `spawn_startup_tasks` fans out the
//! whole daemon startup sequence, and two parts of it outlive the assertions
//! (#5937), while a third reads an override the calling shell may already
//! export (#5821).
//!
//! Test: this file.

use super::*;
use std::ffi::{OsStr, OsString};
use std::fs;
use std::time::Duration;
use trusty_memory::project_root::{write_project_pin, ProjectPin};

/// RAII guard that pins one environment variable for the duration of a test
/// and restores the prior value on drop, including on panic.
///
/// Why: this binary's only env-mutating test set `HOME` and
/// `TRUSTY_SKIP_PALACE_ENFORCEMENT` and restored `HOME` on the success path
/// alone, so a failing assertion left the other twelve tests in this binary
/// pointed at a deleted tempdir (#5821).
/// What: `set` installs a value and `clear` removes one; both capture the
/// prior value and `Drop` puts it back.
/// Test: `spawn_startup_tasks_populates_pin_map`.
struct EnvGuard {
    key: &'static str,
    prev: Option<OsString>,
}

impl EnvGuard {
    fn set<V: AsRef<OsStr>>(key: &'static str, value: V) -> Self {
        let prev = std::env::var_os(key);
        // SAFETY: the caller is `#[serial_test::serial]`, so no sibling test
        // reads or writes the process environment concurrently.
        unsafe { std::env::set_var(key, value) };
        Self { key, prev }
    }

    fn clear(key: &'static str) -> Self {
        let prev = std::env::var_os(key);
        // SAFETY: see `set`.
        unsafe { std::env::remove_var(key) };
        Self { key, prev }
    }
}

impl Drop for EnvGuard {
    fn drop(&mut self) {
        // SAFETY: see `set`.
        unsafe {
            match self.prev.take() {
                Some(v) => std::env::set_var(self.key, v),
                None => std::env::remove_var(self.key),
            }
        }
    }
}

/// Outer wall-clock bound on [`spawn_startup_tasks_populates_pin_map`].
///
/// #5937: a hang guard, not a timing assertion. The scan finishes in
/// milliseconds; this only turns a wedged runtime into a failure.
const PIN_MAP_HANG_GUARD: Duration = Duration::from_secs(60);

/// Why: the pin scan inside `spawn_startup_tasks` must populate
/// `AppState::pin_project_map` so handlers can resolve a palace id to a
/// project path without a filesystem walk at request time (issue #470).
/// This test verifies the full wiring: a real `AppState` with a real temp
/// search root, a pinned project, and the async background task.
/// What: creates a project with a pin file under a temp search root, runs
/// `spawn_startup_tasks` on a runtime owned by a dedicated thread, and waits
/// at most [`PIN_MAP_HANG_GUARD`] for the pin map to gain the entry. The bound
/// lives outside the runtime so it fires even when the runtime is wedged.
/// Test: itself (issue #474 regression guard, plus #5937 and #5821).
#[serial_test::serial]
#[test]
fn spawn_startup_tasks_populates_pin_map() {
    // #5937: the fan-out warms the shared embedder, and a cold init downloads
    // the ONNX model on a blocking-pool thread. Seeding the mock first makes
    // `shared_embedder()` resolve from an already-initialised OnceCell, so no
    // download starts. Must happen before the runtime exists.
    trusty_common::memory_core::retrieval::seed_shared_embedder_with_mock();

    // Build a temp search root with one pinned project.
    let tmp = tempfile::tempdir().expect("tempdir");
    let search_root = tmp.path().join("Projects");
    let project_dir = search_root.join("my-project");
    fs::create_dir_all(&project_dir).expect("create project dir");
    write_project_pin(&project_dir, &ProjectPin::new("my-palace".to_string())).expect("write pin");

    // `HOME` points `default_search_dirs()` at the temp root, and the
    // enforcement bypass keeps `AppState::new` off a real project root.
    let _home = EnvGuard::set("HOME", tmp.path());
    let _enforcement = EnvGuard::set("TRUSTY_SKIP_PALACE_ENFORCEMENT", "1");
    // #5821: the #880 isolation branch skips the pin scan outright when a
    // data-dir override is active, so a shell exporting one made this test
    // assert against a map the scan is designed to leave empty — it reported
    // "0 pin(s) discovered" and panicked on the 500 ms deadline.
    let _data_dir = EnvGuard::clear(trusty_common::DATA_DIR_OVERRIDE_ENV);
    // #5937: the update check resolves crates.io on the blocking pool, which
    // is the second task free to outlive this test.
    let _update = EnvGuard::set(trusty_common::update::NO_UPDATE_CHECK_ENV, "1");

    let state_root = tmp.path().join("data");
    fs::create_dir_all(&state_root).expect("create data dir");

    // #5937: the runtime lives on its own thread so the bound below does not
    // depend on that runtime being able to poll a timer.
    let (map_tx, map_rx) = std::sync::mpsc::channel();
    let (found_tx, found_rx) = std::sync::mpsc::channel();
    let runtime_thread = std::thread::Builder::new()
        .name("pin-map-runtime".to_string())
        .spawn(move || {
            let rt = tokio::runtime::Builder::new_current_thread()
                .enable_all()
                .build()
                .expect("build test runtime");
            let found = rt.block_on(async move {
                let state = AppState::new(state_root);
                let _ = map_tx.send(std::sync::Arc::clone(&state.pin_project_map));
                spawn_startup_tasks(&state);
                loop {
                    if let Some(entry) = state.pin_project_map.get("my-palace") {
                        return entry.clone();
                    }
                    tokio::time::sleep(Duration::from_millis(10)).await;
                }
            });
            // #5937: shut down without waiting. Dropping a runtime blocks until
            // every blocking-pool task it started has finished, and the startup
            // fan-out is free to leave one in flight.
            rt.shutdown_timeout(Duration::ZERO);
            let _ = found_tx.send(found);
        })
        .expect("spawn the test runtime thread");

    let found = match found_rx.recv_timeout(PIN_MAP_HANG_GUARD) {
        Ok(found) => found,
        // The runtime thread panicked before sending: re-raise its panic.
        Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => std::panic::resume_unwind(
            runtime_thread
                .join()
                .expect_err("the runtime thread dropped its sender without panicking"),
        ),
        Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {
            let keys: Vec<String> = map_rx
                .try_recv()
                .map(|map| map.iter().map(|e| e.key().clone()).collect())
                .unwrap_or_default();
            panic!(
                "#5937 hang guard: pin_project_map had no 'my-palace' {PIN_MAP_HANG_GUARD:?} \
                 after spawn_startup_tasks. Keys present: {keys:?}. HOME is now {:?}; \
                 the scan must read {:?}.",
                std::env::var_os("HOME"),
                tmp.path(),
            );
        }
    };
    // Canonicalize to handle macOS /private symlinks.
    let actual = fs::canonicalize(found).expect("canonicalize actual");
    let expected = fs::canonicalize(&project_dir).expect("canonicalize expected");
    assert_eq!(
        actual, expected,
        "pin_project_map entry must point to the project directory"
    );
}

/// Why (#7106): startup hydration opened every palace on disk with nothing
/// bounding how many were open at once. Each open hydrates a drawer table, an
/// HNSW graph and a KG adjacency (~90 MB) plus three redb page caches, so on the
/// reporter's 94-palace install the peak was the whole estate. The bound is the
/// fix, and nothing else in the run reports it — a bounded hydration and an
/// unbounded one load the same palaces and log the same count.
/// What: seeds eight palaces on disk, hydrates through an `AppState` whose gate
/// is pinned to a limit of two, and asserts every palace loaded AND that the
/// gate's high-water mark never exceeded the limit. Hydration itself is serial,
/// so the mark it sets alone is 1 — the limit is the ceiling this job shares
/// with the two BM25 sweeps, not a concurrency it introduces. Hydrating four at
/// a time was measured at a 336 MB → 409 MB `phys_footprint_peak` regression
/// over a 94-palace copy of the reporter's store, which is the resource #7106
/// is about.
/// Test: this test.
// #5937: serial because it sets `HOME`. Run in parallel, it swapped `HOME`
// under `spawn_startup_tasks_populates_pin_map`, whose scan then read the
// wrong root and found 0 pins.
#[serial_test::serial]
#[test]
fn hydration_never_exceeds_the_startup_open_limit() {
    let tmp = tempfile::tempdir().expect("tempdir");
    let _home = EnvGuard::set("HOME", tmp.path());
    let _enforcement = EnvGuard::set("TRUSTY_SKIP_PALACE_ENFORCEMENT", "1");
    let _update = EnvGuard::set(trusty_common::update::NO_UPDATE_CHECK_ENV, "1");

    let state_root = tmp.path().join("data");
    fs::create_dir_all(&state_root).expect("create data dir");

    const PALACES: usize = 8;
    const LIMIT: usize = 2;
    for i in 0..PALACES {
        let id = format!("palace-{i}");
        let palace = trusty_common::memory_core::palace::Palace {
            id: trusty_common::memory_core::palace::PalaceId::new(id.clone()),
            name: id.clone(),
            description: None,
            created_at: chrono::Utc::now(),
            data_dir: state_root.join(&id),
        };
        trusty_common::memory_core::store::PalaceStore::save_palace(&palace)
            .expect("seed palace metadata");
    }

    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()
        .expect("build test runtime");

    let (loaded, peak, limit) = rt.block_on(async {
        let mut state = AppState::new(state_root);
        // Pin the budget rather than reading the ambient env, so the assertion
        // is about the bound and not about the developer's shell.
        state.startup_gate = trusty_memory::startup_budget::StartupOpenGate::with_limit(LIMIT);
        let gate = state.startup_gate.clone();
        let loaded = state
            .load_palaces_from_disk()
            .await
            .expect("hydration must not fail");
        (loaded, gate.peak_concurrent(), gate.limit())
    });
    rt.shutdown_timeout(Duration::ZERO);

    assert_eq!(limit, LIMIT, "the pinned limit must survive into the gate");
    assert_eq!(
        loaded, PALACES,
        "every seeded palace must still hydrate — the bound throttles, it does not skip"
    );
    assert!(
        peak <= LIMIT,
        "#7106: hydration held {peak} palaces open at once against a limit of {LIMIT}"
    );
    assert!(
        peak > 0,
        "every open must take a permit, or the ceiling governs nothing"
    );
}