haematite-cli 0.21.8

use std::collections::BTreeMap;
use std::error::Error;
use std::num::{NonZeroU32, NonZeroUsize};
use std::sync::{
    Barrier,
    atomic::{AtomicUsize, Ordering},
};
use std::thread;

use super::{Running, Slots, StartError};

type TestResult = Result<(), Box<dyn Error>>;

#[test]
fn full_admission_refuses_and_dropping_each_guard_returns_exactly_one_slot() -> TestResult {
    let slots = Slots::new(NonZeroUsize::MIN.saturating_add(1));
    let first = slots.try_acquire()?;
    let second = slots.try_acquire()?;
    let failure = slots.try_acquire().err().ok_or("third slot accepted")?;
    assert_eq!(failure.code, "query_busy");
    drop(first);
    let replacement = slots.try_acquire()?;
    assert!(slots.try_acquire().is_err());
    drop(second);
    drop(replacement);
    let first = slots.try_acquire()?;
    let second = slots.try_acquire()?;
    drop((first, second));
    Ok(())
}

#[test]
fn simultaneous_admission_never_exceeds_the_limit() -> TestResult {
    let slots = Slots::new(NonZeroUsize::MIN);
    let barrier = Barrier::new(2);
    let accepted = AtomicUsize::new(0);
    thread::scope(|scope| -> TestResult {
        let mut threads = Vec::new();
        for _ in 0..2 {
            threads.push(scope.spawn(|| {
                barrier.wait();
                let slot = slots.try_acquire();
                if slot.is_ok() {
                    accepted.fetch_add(1, Ordering::Relaxed);
                }
                barrier.wait();
                slot.map(drop)
            }));
        }
        for thread in threads {
            if let Err(failure) = thread.join().map_err(|_| "admission thread panicked")? {
                assert_eq!(failure.code, "query_busy");
            }
        }
        assert_eq!(accepted.load(Ordering::Relaxed), 1);
        Ok(())
    })?;
    drop(slots.try_acquire()?);
    Ok(())
}

#[test]
fn absent_settings_refuse_without_acquiring_a_slot_or_starting_a_worker() -> TestResult {
    let slots = Slots::new(NonZeroUsize::MIN);
    let failure = Running::start(
        &slots,
        None,
        NonZeroU32::MIN,
        ".".to_owned(),
        None,
        BTreeMap::new(),
    )
    .err()
    .ok_or("absent settings accepted")?;
    match failure {
        StartError::Admission(failure) => assert_eq!(failure.code, "query_engine_unavailable"),
        StartError::Worker(error) => return Err(error),
    }
    drop(slots.try_acquire()?);
    Ok(())
}

/// A stand-in worker: a script that answers Ready as this crate's worker would,
/// then for the first Input writes one Output and Done, and waits for its input
/// to close. It lives in the system temporary directory under this process's id.
struct StandIn {
    script: std::path::PathBuf,
    frames: std::path::PathBuf,
}

impl StandIn {
    fn new(name: &str) -> Result<Self, Box<dyn Error>> {
        use haematite_query::{Ready, Tag, frame};
        use std::os::unix::fs::PermissionsExt;
        let base =
            std::env::temp_dir().join(format!("haematite-standin-{}-{name}", std::process::id()));
        let (frames, script) = (base.with_extension("frames"), base.with_extension("sh"));
        let ready = Ready {
            protocol: haematite_query::PROTOCOL,
            worker: env!("CARGO_PKG_VERSION").to_owned(),
            core_limit: 0,
            descriptors_closed: 0,
            environment: 0,
            engine: "stand-in".to_owned(),
            jaq_core: "0".to_owned(),
            jaq_std: "0".to_owned(),
            jaq_json: "0".to_owned(),
            reading: Some(haematite_query::Reading::Candidate),
        };
        let mut bytes = Vec::new();
        frame::write(
            &mut bytes,
            Tag::Ready as u8,
            &serde_json::to_vec(&ready)?,
            NonZeroU32::MAX,
        )?;
        frame::write(&mut bytes, Tag::Output as u8, b"1", NonZeroU32::MAX)?;
        frame::write(
            &mut bytes,
            Tag::Done as u8,
            br#"{"outputs":1}"#,
            NonZeroU32::MAX,
        )?;
        std::fs::write(&frames, bytes)?;
        std::fs::write(
            &script,
            format!(
                "#!/bin/sh\n/bin/cat '{}'\nexec /bin/cat >/dev/null\n",
                frames.display()
            ),
        )?;
        std::fs::set_permissions(&script, std::fs::Permissions::from_mode(0o700))?;
        let stand_in = Self { script, frames };
        let first = std::process::Command::new(&stand_in.script)
            .env_clear()
            .current_dir("/")
            .stdin(std::process::Stdio::null())
            .stdout(std::process::Stdio::null())
            .stderr(std::process::Stdio::piped())
            .output()
            .map_err(|error| format!("query stand-in prelaunch failed: {error}"))?;
        if !first.status.success() {
            return Err(format!(
                "query stand-in prelaunch ended with {}; stderr: {}",
                first.status,
                String::from_utf8_lossy(&first.stderr)
            )
            .into());
        }
        Ok(stand_in)
    }

    fn settings(&self) -> super::Settings {
        super::Settings {
            time_ms: std::num::NonZeroU64::MIN.saturating_add(59_999),
            memory_bytes: std::num::NonZeroU64::MIN.saturating_add(999_999),
            outputs_per_record: std::num::NonZeroU64::MIN.saturating_add(6),
            scan_records: std::num::NonZeroU64::MIN.saturating_add(99),
            max_running: NonZeroUsize::MIN,
            worker: self.script.clone(),
        }
    }
}

impl Drop for StandIn {
    fn drop(&mut self) {
        let _script: std::io::Result<()> = std::fs::remove_file(&self.script);
        let _frames: std::io::Result<()> = std::fs::remove_file(&self.frames);
    }
}

/// Records in order, or a named failure in place of the first.
struct Source {
    records: std::collections::VecDeque<haematite_query::page::Record>,
    fail: bool,
}

impl haematite_query::page::RecordSource for Source {
    fn next_record(
        &mut self,
    ) -> Result<Option<haematite_query::page::Record>, haematite_query::handle::Error> {
        if self.fail {
            return Err(haematite_query::handle::Error::new(
                haematite_query::Refusal::QueryRuntime,
                "the source failed",
            ));
        }
        Ok(self.records.pop_front())
    }
}

fn source(ids: &[&str], fail: bool) -> Result<Source, serde_json::Error> {
    let raw = |json: &str| serde_json::value::RawValue::from_string(json.to_owned());
    let records = ids
        .iter()
        .map(|id| {
            Ok(haematite_query::page::Record {
                value: raw("1")?,
                origin: haematite_query::page::Origin {
                    collection: "patients".to_owned(),
                    id: (*id).to_owned(),
                    version: raw(r#"{"generation":7,"sequence":11}"#)?,
                    score: None,
                },
            })
        })
        .collect::<Result<_, serde_json::Error>>()?;
    Ok(Source { records, fail })
}

fn alive(pid: u32) -> Result<bool, std::io::Error> {
    std::process::Command::new("/bin/kill")
        .args(["-0", &pid.to_string()])
        .status()
        .map(|status| status.success())
}

/// One row per page, so the second record cuts the page.
fn one_row() -> haematite_query::page::Allowance {
    haematite_query::page::Allowance {
        rows: std::num::NonZeroU64::MIN,
        bytes: std::num::NonZeroU64::MIN.saturating_add(9_999),
        scan_records: 10,
    }
}

/// Test (d), a page cut: the page is read through the guard, which keeps the
/// worker for the whole call. The cut stops and reaps the worker while the guard
/// still holds the slot; dropping the guard then frees it.
#[test]
fn a_page_cut_through_the_guard_reaps_the_worker_before_the_slot_is_free() -> TestResult {
    let stand_in = StandIn::new("cut")?;
    let settings = stand_in.settings();
    let slots = Slots::new(NonZeroUsize::MIN);
    let frame = NonZeroU32::MIN.saturating_add(8191);
    let mut running = Running::start(
        &slots,
        Some(&settings),
        frame,
        ".".to_owned(),
        None,
        BTreeMap::new(),
    )?;
    let worker = running.worker()?;
    let pid = worker.pid();
    let mut records = source(&["a", "b"], false)?;
    let page = haematite_query::page::read_page(worker, &mut records, &one_row(), None)?;
    assert!(page.next.is_some(), "control: the page was not cut");
    assert!(!alive(pid)?, "the cut left the worker alive");
    assert!(
        slots.try_acquire().is_err(),
        "the slot was freed while the guard lives"
    );
    drop(running);
    drop(slots.try_acquire()?);
    Ok(())
}

/// Test (d), an error: a source failure leaves the worker to the guard, which
/// still holds the slot; dropping the guard reaps the worker and then frees it.
#[test]
fn an_error_through_the_guard_reaps_the_worker_before_the_slot_is_free() -> TestResult {
    let stand_in = StandIn::new("error")?;
    let settings = stand_in.settings();
    let slots = Slots::new(NonZeroUsize::MIN);
    let frame = NonZeroU32::MIN.saturating_add(8191);
    let mut running = Running::start(
        &slots,
        Some(&settings),
        frame,
        ".".to_owned(),
        None,
        BTreeMap::new(),
    )?;
    let worker = running.worker()?;
    let pid = worker.pid();
    let mut records = source(&["a"], true)?;
    let error = haematite_query::page::read_page(worker, &mut records, &one_row(), None)
        .err()
        .ok_or("a failing source gave a page")?;
    assert_eq!(error.refusal.code, haematite_query::Refusal::QueryRuntime);
    assert!(
        slots.try_acquire().is_err(),
        "the slot was freed while the guard lives"
    );
    drop(running);
    assert!(!alive(pid)?, "dropping the guard left the worker alive");
    drop(slots.try_acquire()?);
    Ok(())
}