297 lines
11 KiB
Rust
297 lines
11 KiB
Rust
|
//! Abrupt termination of an owned real CLI process tree, not a returned error.
|
||
|
|
//! Requires verified local Google assets/decoders. This is not power-loss proof.
|
||
|
|
use rusqlite::Connection;
|
||
|
|
use serde_json::Value;
|
||
|
|
use std::{
|
||
|
|
fs::{self, File},
|
||
|
|
io::Write,
|
||
|
|
path::{Path, PathBuf},
|
||
|
|
process::{Child, Command, Stdio},
|
||
|
|
thread,
|
||
|
|
time::{Duration, Instant},
|
||
|
|
};
|
||
|
|
|
||
|
|
struct OwnedCli {
|
||
|
|
child: Child,
|
||
|
|
}
|
||
|
|
impl OwnedCli {
|
||
|
|
fn spawn(mut command: Command, stdout: &Path, stderr: &Path) -> Self {
|
||
|
|
#[cfg(unix)]
|
||
|
|
{
|
||
|
|
use std::os::unix::process::CommandExt;
|
||
|
|
// Every worker/decoder inherits this exclusive process group.
|
||
|
|
command.process_group(0);
|
||
|
|
}
|
||
|
|
command
|
||
|
|
.env("SEARCHI_FORCE_CPU", "1")
|
||
|
|
.stdin(Stdio::null())
|
||
|
|
.stdout(File::create(stdout).unwrap())
|
||
|
|
.stderr(File::create(stderr).unwrap());
|
||
|
|
Self {
|
||
|
|
child: command.spawn().unwrap(),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
fn terminate(&mut self) -> std::io::Result<()> {
|
||
|
|
if self.child.try_wait()?.is_some() {
|
||
|
|
return Ok(());
|
||
|
|
}
|
||
|
|
#[cfg(windows)]
|
||
|
|
let status = Command::new("taskkill.exe")
|
||
|
|
.args(["/PID", &self.child.id().to_string(), "/T", "/F"])
|
||
|
|
.stdout(Stdio::null())
|
||
|
|
.stderr(Stdio::null())
|
||
|
|
.status()?;
|
||
|
|
#[cfg(unix)]
|
||
|
|
let status = Command::new("/bin/kill")
|
||
|
|
.args(["-KILL", "--", &format!("-{}", self.child.id())])
|
||
|
|
.status()?;
|
||
|
|
if !status.success() {
|
||
|
|
return Err(std::io::Error::other(
|
||
|
|
"Could not terminate owned CLI process tree",
|
||
|
|
));
|
||
|
|
}
|
||
|
|
self.child.wait()?;
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
fn finish(&mut self, timeout: Duration, stderr: &Path) {
|
||
|
|
let deadline = Instant::now() + timeout;
|
||
|
|
loop {
|
||
|
|
if let Some(status) = self.child.try_wait().unwrap() {
|
||
|
|
assert!(
|
||
|
|
status.success(),
|
||
|
|
"CLI failed: {}",
|
||
|
|
fs::read_to_string(stderr).unwrap()
|
||
|
|
);
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
assert!(Instant::now() < deadline, "Owned CLI exceeded deadline");
|
||
|
|
thread::sleep(Duration::from_millis(20));
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
impl Drop for OwnedCli {
|
||
|
|
fn drop(&mut self) {
|
||
|
|
let _ = self.terminate();
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn cli(db: &Path, command: &str) -> Command {
|
||
|
|
let mut cmd = Command::new(env!("CARGO_BIN_EXE_searchi"));
|
||
|
|
cmd.arg(command).arg("--db").arg(db);
|
||
|
|
cmd
|
||
|
|
}
|
||
|
|
fn run(command: Command, directory: &Path, label: &str) -> Value {
|
||
|
|
let stdout = directory.join(format!("{label}.stdout"));
|
||
|
|
let stderr = directory.join(format!("{label}.stderr"));
|
||
|
|
let mut child = OwnedCli::spawn(command, &stdout, &stderr);
|
||
|
|
child.finish(Duration::from_secs(900), &stderr);
|
||
|
|
serde_json::from_slice(&fs::read(stdout).unwrap()).unwrap()
|
||
|
|
}
|
||
|
|
fn embed(db: &Path, model: &Path, worker: &Path, ffmpeg: &Path, ffprobe: &Path) -> Command {
|
||
|
|
let mut command = cli(db, "embed");
|
||
|
|
command
|
||
|
|
.arg("--model-dir")
|
||
|
|
.arg(model)
|
||
|
|
.arg("--embedding-worker")
|
||
|
|
.arg(worker)
|
||
|
|
.arg("--cpu-threads")
|
||
|
|
.arg("1")
|
||
|
|
.arg("--ffmpeg")
|
||
|
|
.arg(ffmpeg)
|
||
|
|
.arg("--ffprobe")
|
||
|
|
.arg(ffprobe);
|
||
|
|
command
|
||
|
|
}
|
||
|
|
fn index(db: &Path, root: &Path, directory: &Path) {
|
||
|
|
let stdout = directory.join("index.stdout");
|
||
|
|
let stderr = directory.join("index.stderr");
|
||
|
|
let mut command = cli(db, "index");
|
||
|
|
command.arg("--root").arg(root);
|
||
|
|
let mut child = OwnedCli::spawn(command, &stdout, &stderr);
|
||
|
|
child.finish(Duration::from_secs(60), &stderr);
|
||
|
|
}
|
||
|
|
fn required(name: &str) -> PathBuf {
|
||
|
|
PathBuf::from(std::env::var_os(name).unwrap_or_else(|| panic!("{name} required")))
|
||
|
|
}
|
||
|
|
fn active(db: &Connection) -> i64 {
|
||
|
|
db.query_row(
|
||
|
|
"SELECT active_generation FROM semantic_meta WHERE singleton=1",
|
||
|
|
[],
|
||
|
|
|r| r.get(0),
|
||
|
|
)
|
||
|
|
.unwrap()
|
||
|
|
}
|
||
|
|
type CommittedChunk = (String, i64, String, Vec<u8>);
|
||
|
|
fn completed(db: &Connection, generation: i64) -> Vec<CommittedChunk> {
|
||
|
|
let mut statement = db.prepare("SELECT c.task_key,c.ordinal,c.path,c.vector FROM semantic_chunks c JOIN semantic_tasks t ON t.generation=c.generation AND t.task_key=c.task_key WHERE c.generation=?1 AND t.state='done' ORDER BY c.task_key,c.ordinal").unwrap();
|
||
|
|
statement
|
||
|
|
.query_map([generation], |r| {
|
||
|
|
Ok((r.get(0)?, r.get(1)?, r.get(2)?, r.get(3)?))
|
||
|
|
})
|
||
|
|
.unwrap()
|
||
|
|
.collect::<rusqlite::Result<_>>()
|
||
|
|
.unwrap()
|
||
|
|
}
|
||
|
|
fn write_audio(path: &Path) {
|
||
|
|
// Standard PCM WAV: three overlapping audio windows under normal limits.
|
||
|
|
let samples = 16_000_u32 * 65;
|
||
|
|
let mut wav = File::create(path).unwrap();
|
||
|
|
for bytes in [
|
||
|
|
b"RIFF".as_slice(),
|
||
|
|
&(36 + samples * 2).to_le_bytes(),
|
||
|
|
b"WAVEfmt ",
|
||
|
|
&16_u32.to_le_bytes(),
|
||
|
|
&1_u16.to_le_bytes(),
|
||
|
|
&1_u16.to_le_bytes(),
|
||
|
|
&16_000_u32.to_le_bytes(),
|
||
|
|
&32_000_u32.to_le_bytes(),
|
||
|
|
&2_u16.to_le_bytes(),
|
||
|
|
&16_u16.to_le_bytes(),
|
||
|
|
b"data",
|
||
|
|
&(samples * 2).to_le_bytes(),
|
||
|
|
] {
|
||
|
|
wav.write_all(bytes).unwrap();
|
||
|
|
}
|
||
|
|
let mut pcm = Vec::with_capacity(samples as usize * 2);
|
||
|
|
for n in 0..samples {
|
||
|
|
let value = ((n as f64 * 440.0 * std::f64::consts::TAU / 16_000.0).sin() * 4096.0) as i16;
|
||
|
|
pcm.extend(value.to_le_bytes());
|
||
|
|
}
|
||
|
|
wav.write_all(&pcm).unwrap();
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
#[ignore = "Requires verified Google assets/CPU auto worker and FFmpeg; kills only owned CLI tree"]
|
||
|
|
fn real_google_cli_kill_preserves_active_and_resumes_only_complete_groups() {
|
||
|
|
assert_eq!(std::env::var("SEARCHI_FORCE_CPU").unwrap(), "1");
|
||
|
|
let model = required("SEARCHI_GEMMA_MODEL_DIR");
|
||
|
|
let worker = required("SEARCHI_GEMMA_WORKER");
|
||
|
|
let ffmpeg = required("SEARCHI_TEST_FFMPEG");
|
||
|
|
let ffprobe = required("SEARCHI_TEST_FFPROBE");
|
||
|
|
let temp = tempfile::tempdir().unwrap();
|
||
|
|
let root = temp.path().join("sources");
|
||
|
|
fs::create_dir(&root).unwrap();
|
||
|
|
fs::write(
|
||
|
|
root.join("agreement.txt"),
|
||
|
|
"The supplier partnership can end with sixty days written notice.",
|
||
|
|
)
|
||
|
|
.unwrap();
|
||
|
|
let path = temp.path().join("index.db");
|
||
|
|
index(&path, &root, temp.path());
|
||
|
|
run(
|
||
|
|
embed(&path, &model, &worker, &ffmpeg, &ffprobe),
|
||
|
|
temp.path(),
|
||
|
|
"initial",
|
||
|
|
);
|
||
|
|
let db = Connection::open(&path).unwrap();
|
||
|
|
db.busy_timeout(Duration::from_secs(2)).unwrap();
|
||
|
|
let original = active(&db);
|
||
|
|
let old_vectors = completed(&db, original);
|
||
|
|
fs::write(
|
||
|
|
root.join("retry.rs"),
|
||
|
|
"pub fn exponential_backoff(attempt: u32) -> u64 { 100 * 2_u64.pow(attempt.min(8)) }\n",
|
||
|
|
)
|
||
|
|
.unwrap();
|
||
|
|
write_audio(&root.join("tone.wav"));
|
||
|
|
index(&path, &root, temp.path());
|
||
|
|
let stdout = temp.path().join("interrupted.stdout");
|
||
|
|
let stderr = temp.path().join("interrupted.stderr");
|
||
|
|
let mut child = OwnedCli::spawn(
|
||
|
|
embed(&path, &model, &worker, &ffmpeg, &ffprobe),
|
||
|
|
&stdout,
|
||
|
|
&stderr,
|
||
|
|
);
|
||
|
|
let deadline = Instant::now() + Duration::from_secs(300);
|
||
|
|
let generation = loop {
|
||
|
|
assert!(
|
||
|
|
child.child.try_wait().unwrap().is_none(),
|
||
|
|
"CLI finished before interruption boundary: {}",
|
||
|
|
fs::read_to_string(&stderr).unwrap()
|
||
|
|
);
|
||
|
|
let boundary: Option<i64> = db.query_row("SELECT g.id FROM semantic_generations g WHERE g.id!=?1 AND (SELECT count(*) FROM semantic_tasks t WHERE t.generation=g.id AND t.state='done')>=2 AND EXISTS(SELECT 1 FROM semantic_tasks t JOIN semantic_chunks c ON c.generation=t.generation AND c.path=t.path AND c.passage_id=t.passage_id WHERE t.generation=g.id AND t.state='running' AND t.source_type='media' AND c.task_key IS NULL AND c.source_type!='text') LIMIT 1", [original], |r| r.get(0)).ok();
|
||
|
|
if let Some(generation) = boundary {
|
||
|
|
break generation;
|
||
|
|
}
|
||
|
|
assert!(
|
||
|
|
Instant::now() < deadline,
|
||
|
|
"No committed partial media boundary within deadline"
|
||
|
|
);
|
||
|
|
thread::sleep(Duration::from_millis(20));
|
||
|
|
};
|
||
|
|
child.terminate().unwrap();
|
||
|
|
assert_eq!(
|
||
|
|
active(&db),
|
||
|
|
original,
|
||
|
|
"Abrupt kill must not publish staging"
|
||
|
|
);
|
||
|
|
assert_eq!(completed(&db, original), old_vectors);
|
||
|
|
let saved = completed(&db, generation);
|
||
|
|
assert_eq!(saved.len(), 2);
|
||
|
|
let partial: usize = db.query_row("SELECT count(*) FROM semantic_chunks c JOIN semantic_tasks t ON t.generation=c.generation AND t.path=c.path AND t.passage_id=c.passage_id WHERE c.generation=?1 AND t.state='running' AND t.source_type='media' AND c.task_key IS NULL AND c.source_type!='text'",[generation],|r|r.get(0)).unwrap();
|
||
|
|
assert!(
|
||
|
|
partial > 0,
|
||
|
|
"Need actual committed incomplete Google vectors"
|
||
|
|
);
|
||
|
|
let mut old_search = cli(&path, "search");
|
||
|
|
old_search
|
||
|
|
.arg("How can we stop working with the vendor?")
|
||
|
|
.arg("--json")
|
||
|
|
.arg("--model-dir")
|
||
|
|
.arg(&model)
|
||
|
|
.arg("--embedding-worker")
|
||
|
|
.arg(&worker)
|
||
|
|
.arg("--cpu-threads")
|
||
|
|
.arg("1");
|
||
|
|
let previous_results = run(old_search, temp.path(), "old-reader");
|
||
|
|
assert_eq!(previous_results[0]["filename"], "agreement.txt");
|
||
|
|
assert_eq!(
|
||
|
|
active(&db),
|
||
|
|
original,
|
||
|
|
"Searching must not activate unfinished work"
|
||
|
|
);
|
||
|
|
let report = run(
|
||
|
|
embed(&path, &model, &worker, &ffmpeg, &ffprobe),
|
||
|
|
temp.path(),
|
||
|
|
"resumed",
|
||
|
|
);
|
||
|
|
assert_eq!(report["resumed"], true);
|
||
|
|
assert_eq!(report["generation"], generation);
|
||
|
|
assert_eq!(report["document_inference_calls"], 0);
|
||
|
|
assert_eq!(report["reused_sources"], 2);
|
||
|
|
assert_eq!(report["encoded_sources"], 1);
|
||
|
|
assert_eq!(
|
||
|
|
report["audio_inference_calls"], 3,
|
||
|
|
"All incomplete media windows must be encoded afresh"
|
||
|
|
);
|
||
|
|
assert_eq!(report["media_failures"], 0);
|
||
|
|
assert_eq!(active(&db), generation);
|
||
|
|
let after = completed(&db, generation);
|
||
|
|
for chunk in saved {
|
||
|
|
assert!(
|
||
|
|
after.contains(&chunk),
|
||
|
|
"Committed real vector bytes/citations changed"
|
||
|
|
);
|
||
|
|
}
|
||
|
|
let (total,unique):(usize,usize)=db.query_row("SELECT count(*),count(DISTINCT task_key||':'||ordinal) FROM semantic_chunks WHERE generation=?1",[generation],|r|Ok((r.get(0)?,r.get(1)?))).unwrap();
|
||
|
|
assert_eq!(
|
||
|
|
(total, unique),
|
||
|
|
(5, 5),
|
||
|
|
"Partial media chunks must not leak or duplicate"
|
||
|
|
);
|
||
|
|
let mut search = cli(&path, "search");
|
||
|
|
search
|
||
|
|
.arg("retry failed requests with exponential delays")
|
||
|
|
.arg("--code-only")
|
||
|
|
.arg("--json")
|
||
|
|
.arg("--model-dir")
|
||
|
|
.arg(&model)
|
||
|
|
.arg("--embedding-worker")
|
||
|
|
.arg(&worker)
|
||
|
|
.arg("--cpu-threads")
|
||
|
|
.arg("1");
|
||
|
|
let results = run(search, temp.path(), "citation");
|
||
|
|
assert_eq!(results[0]["filename"], "retry.rs");
|
||
|
|
assert_eq!(results[0]["source"]["kind"], "line");
|
||
|
|
assert_eq!(results[0]["source"]["number"], 1);
|
||
|
|
eprintln!("Real CPU CLI crash recovery: generation={generation}, committed_sources=2, discarded_partial_chunks={partial}, resumed_audio_calls=3");
|
||
|
|
}
|