karl/searchi
1
0 Fork 0
Code Issues Pull Requests Packages Projects Releases 10 Wiki Activity

Harden worker IPC and recoverable runtime staging

Browse Source
This commit is contained in:
karl committed 2026-10-12 23:46:01 +02:00
1 parent 0bba1a0aa7
commit b3ebd7a894
12 files changed
+1152 -75
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

crates/searchi-assets/src/lib.rs
+47
View File
@@ -94,6 +94,27 @@ pub fn default_model_directory() -> Result<PathBuf> {
.join("embeddinggemma2"))
}
/// Cheap location selection for optional bundled assets. This checks regular
/// files and expected lengths only; `verify_cached` still verifies SHA-256
/// before setup considers the model usable. Partial bundles use managed cache.
pub fn has_model_files(directory: &Path) -> Result<bool> {
model_files_present(directory, &manifest()?.files)
}
fn model_files_present(directory: &Path, artifacts: &[Artifact]) -> Result<bool> {
for artifact in artifacts {
let metadata = match fs::metadata(directory.join(&artifact.path)) {
Ok(metadata) => metadata,
Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(false),
Err(error) => return Err(error).context("inspect bundled model artifact"),
};
if !metadata.is_file() || metadata.len() != artifact.bytes {
return Ok(false);
}
}
Ok(true)
}
/// Offline validation includes file length and full SHA-256, not mere existence.
pub fn verify_cached(directory: &Path) -> Result<bool> {
for artifact in manifest()?.files {
@@ -383,6 +404,32 @@ mod tests {
);
}
#[test]
fn optional_bundle_requires_both_regular_files_at_expected_sizes() {
let directory = tempfile::tempdir().unwrap();
let artifacts = [artifact("one.gguf", b"abc"), artifact("two.gguf", b"xyz")];
assert!(!model_files_present(directory.path(), &artifacts).unwrap());
fs::write(directory.path().join("one.gguf"), b"abc").unwrap();
assert!(!model_files_present(directory.path(), &artifacts).unwrap());
fs::create_dir(directory.path().join("two.gguf")).unwrap();
assert!(!model_files_present(directory.path(), &artifacts).unwrap());
fs::remove_dir(directory.path().join("two.gguf")).unwrap();
fs::write(directory.path().join("two.gguf"), b"xyz").unwrap();
assert!(model_files_present(directory.path(), &artifacts).unwrap());
fs::write(directory.path().join("one.gguf"), b"ab").unwrap();
assert!(!model_files_present(directory.path(), &artifacts).unwrap());
}
#[test]
fn bundled_file_presence_never_substitutes_for_sha_verification() {
let directory = tempfile::tempdir().unwrap();
let artifact = artifact("one.gguf", b"abc");
let path = directory.path().join(&artifact.path);
fs::write(&path, b"xyz").unwrap();
assert!(model_files_present(directory.path(), std::slice::from_ref(&artifact)).unwrap());
assert!(!verified(&path, &artifact).unwrap());
}
#[test]
fn same_size_corruption_is_detected_and_missing_cache_is_offline() {
let directory = tempfile::tempdir().unwrap();
crates/searchi-semantic/src/gemma.rs
+233 -37
View File
@@ -60,7 +60,10 @@ fn query_operation(info: &WorkerInfo, code: bool) -> Result<&'static str> {
/// Persistent isolated native worker. No HTTP or remote model fallback.
pub struct GemmaEmbedder {
child: Child,
stdin: ChildStdin,
requests: Option<mpsc::SyncSender<Vec<u8>>>,
writes: Receiver<std::result::Result<(), String>>,
writer_thread: Option<thread::JoinHandle<()>>,
reader_thread: Option<thread::JoinHandle<()>>,
responses: Receiver<std::result::Result<Value, String>>,
pub info: WorkerInfo,
healthy: bool,
@@ -73,24 +76,48 @@ pub struct GemmaEmbedder {
impl Drop for GemmaEmbedder {
fn drop(&mut self) {
let _ = self.child.kill();
let _ = self.child.wait();
self.stop_worker();
}
}
type WriterTransport = (
mpsc::SyncSender<Vec<u8>>,
Receiver<std::result::Result<(), String>>,
thread::JoinHandle<()>,
);
fn spawn_writer(mut stdin: ChildStdin) -> WriterTransport {
let (requests, pending) = mpsc::sync_channel::<Vec<u8>>(1);
let (completed, writes) = mpsc::sync_channel(1);
let writer = thread::spawn(move || {
while let Ok(encoded) = pending.recv() {
let result = stdin
.write_all(&encoded)
.and_then(|_| stdin.write_all(b"\n"))
.and_then(|_| stdin.flush())
.map_err(|error| format!("Cannot write local Gemma worker request: {error}"));
let failed = result.is_err();
if completed.try_send(result).is_err() || failed {
break;
}
}
});
(requests, writes, writer)
}
#[derive(Debug)]
struct StartupTransportFailure(String);
impl std::fmt::Display for StartupTransportFailure {
struct WorkerTransportFailure(String);
impl std::fmt::Display for WorkerTransportFailure {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str(&self.0)
}
}
impl std::error::Error for StartupTransportFailure {}
impl std::error::Error for WorkerTransportFailure {}
/// Retry only process/transport failure, never an incompatible protocol or model.
fn start_with_cpu_fallback<T>(cpu: bool, mut start: impl FnMut(bool) -> Result<T>) -> Result<T> {
match start(cpu) {
Err(error) if !cpu && error.downcast_ref::<StartupTransportFailure>().is_some() => {
Err(error) if !cpu && error.downcast_ref::<WorkerTransportFailure>().is_some() => {
start(true).with_context(|| {
format!("CPU initialization failed after automatic worker startup failure: {error}")
})
@@ -163,7 +190,7 @@ impl GemmaEmbedder {
command.creation_flags(0x08000000); // CREATE_NO_WINDOW
}
let mut child = command.spawn().map_err(|error| {
StartupTransportFailure(format!("Cannot start local Gemma worker: {error}"))
WorkerTransportFailure(format!("Cannot start local Gemma worker: {error}"))
})?;
let (Some(stdin), Some(stdout)) = (child.stdin.take(), child.stdout.take()) else {
let _ = child.kill();
@@ -171,12 +198,12 @@ impl GemmaEmbedder {
anyhow::bail!("Gemma worker protocol pipes missing");
};
let (sender, responses) = mpsc::sync_channel(1);
thread::spawn(move || {
let reader_thread = thread::spawn(move || {
let mut reader = BufReader::new(stdout);
loop {
let response = read_response(&mut reader);
let failed = response.is_err();
if sender.send(response).is_err() || failed {
if sender.try_send(response).is_err() || failed {
break;
}
}
@@ -186,13 +213,14 @@ impl GemmaEmbedder {
other => {
let _ = child.kill();
let _ = child.wait();
let _ = reader_thread.join();
// Malformed/oversized JSON is a protocol violation, not evidence
// that hardware acceleration failed. Do not mask it with retry.
if matches!(&other, Ok(Err(error)) if error.starts_with("Invalid Gemma worker JSON:") || error == "Gemma worker response exceeds protocol byte budget")
{
anyhow::bail!("Invalid Gemma worker startup protocol: {other:?}");
}
return Err(StartupTransportFailure(format!(
return Err(WorkerTransportFailure(format!(
"Gemma worker did not initialize: {other:?}"
))
.into());
@@ -203,6 +231,7 @@ impl GemmaEmbedder {
Err(error) => {
let _ = child.kill();
let _ = child.wait();
let _ = reader_thread.join();
return Err(error).context("Invalid Gemma worker handshake");
}
};
@@ -216,11 +245,16 @@ impl GemmaEmbedder {
{
let _ = child.kill();
let _ = child.wait();
let _ = reader_thread.join();
anyhow::bail!("Gemma worker handshake violates configured model/device contract");
}
let (requests, writes, writer_thread) = spawn_writer(stdin);
Ok(Self {
child,
stdin,
requests: Some(requests),
writes,
writer_thread: Some(writer_thread),
reader_thread: Some(reader_thread),
responses,
info,
healthy: true,
@@ -246,7 +280,10 @@ impl GemmaEmbedder {
);
let response = match self.exchange(&encoded) {
Ok(response) => response,
Err(error) if can_retry_cpu(&self.info.device, self.force_cpu) => {
Err(error)
if can_retry_cpu(&self.info.device, self.force_cpu)
&& error.downcast_ref::<WorkerTransportFailure>().is_some() =>
{
// GPU runtimes can abort rather than return an encoder error.
// Restart once as CPU and retry this same bounded request.
let _ = self.child.wait();
@@ -296,33 +333,76 @@ impl GemmaEmbedder {
Ok(response)
}
fn stop_worker(&mut self) {
self.requests.take();
// The shipped native worker does not spawn descendants. Killing/reaping
// it closes its pipe ends and releases any blocked writer/read thread.
let _ = self.child.kill();
let _ = self.child.wait();
if let Some(writer) = self.writer_thread.take() {
let _ = writer.join();
}
if let Some(reader) = self.reader_thread.take() {
let _ = reader.join();
}
}
fn exchange(&mut self, encoded: &[u8]) -> Result<Value> {
self.exchange_with_timeout(encoded, Duration::from_secs(180))
}
fn exchange_with_timeout(&mut self, encoded: &[u8], timeout: Duration) -> Result<Value> {
let deadline = Instant::now() + timeout;
let write_started = self.query_profile.as_ref().map(|_| Instant::now());
let sent = (|| {
self.stdin.write_all(encoded)?;
self.stdin.write_all(b"\n")?;
self.stdin.flush()
})()
.context("Cannot write local Gemma worker request");
self.record_query_duration("request_write_us", write_started);
if let Err(error) = sent {
let result = (|| {
self.requests
.as_ref()
.ok_or_else(|| WorkerTransportFailure("Gemma worker writer stopped".into()))?
.try_send(encoded.to_vec())
.map_err(|_| {
WorkerTransportFailure("Gemma worker request queue unavailable".into())
})?;
match self
.writes
.recv_timeout(deadline.saturating_duration_since(Instant::now()))
{
Ok(Ok(())) => (),
other => {
return Err(WorkerTransportFailure(format!(
"Gemma worker request write failed or timed out: {other:?}"
))
.into())
}
}
self.record_query_duration("request_write_us", write_started);
let wait_started = self.query_profile.as_ref().map(|_| Instant::now());
let response = match self
.responses
.recv_timeout(deadline.saturating_duration_since(Instant::now()))
{
Ok(Ok(value)) => value,
Ok(Err(error))
if error.starts_with("Invalid Gemma worker JSON:")
|| error == "Gemma worker response exceeds protocol byte budget" =>
{
anyhow::bail!("Invalid Gemma worker response protocol: {error}")
}
other => {
return Err(WorkerTransportFailure(format!(
"Gemma worker request failed or timed out: {other:?}"
))
.into())
}
};
self.record_query_duration("worker_wait_and_response_decode_us", wait_started);
Ok(response)
})();
if result.is_err() {
self.healthy = false;
self.info.ready = false;
let _ = self.child.kill();
return Err(error);
self.stop_worker();
}
let wait_started = self.query_profile.as_ref().map(|_| Instant::now());
let response = match self.responses.recv_timeout(Duration::from_secs(180)) {
Ok(Ok(value)) => value,
other => {
self.healthy = false;
self.info.ready = false;
let _ = self.child.kill();
anyhow::bail!("Gemma worker request failed: {other:?}");
}
};
self.record_query_duration("worker_wait_and_response_decode_us", wait_started);
Ok(response)
result
}
pub fn split_windows(&mut self, text: &str, title: &str) -> Result<Vec<TextWindow>> {
@@ -710,6 +790,122 @@ fn verify_artifact(path: &Path, bytes: u64, sha: &str) -> Result<()> {
#[cfg(test)]
mod tests {
use super::*;
// This helper is an inert pipe-protocol fixture, never an embedding provider.
fn stalled_protocol_helper(read_request: bool) -> GemmaEmbedder {
#[cfg(windows)]
let mut command = {
use std::os::windows::process::CommandExt;
let mut command = Command::new("powershell.exe");
command.args([
"-NoProfile",
"-NonInteractive",
"-Command",
if read_request {
"[Console]::In.ReadLine() | Out-Null; Start-Sleep -Seconds 20"
} else {
"Start-Sleep -Seconds 20"
},
]);
command.creation_flags(0x08000000);
command
};
#[cfg(not(windows))]
let mut command = {
let mut command = Command::new("sh");
command.args([
"-c",
if read_request {
"read line; exec sleep 20"
} else {
"exec sleep 20"
},
]);
command
};
let mut child = command
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::null())
.spawn()
.unwrap();
let (requests, writes, writer_thread) = spawn_writer(child.stdin.take().unwrap());
let stdout = child.stdout.take().unwrap();
let (sender, responses) = mpsc::sync_channel(1);
let reader_thread = thread::spawn(move || {
let response = read_response(&mut BufReader::new(stdout));
let _ = sender.try_send(response);
});
GemmaEmbedder {
child,
requests: Some(requests),
writes,
writer_thread: Some(writer_thread),
reader_thread: Some(reader_thread),
responses,
info: WorkerInfo {
ready: true,
model: "pipe protocol fixture".into(),
dimensions: 0,
device: "cpu".into(),
threads: 1,
runtime: "fixture".into(),
media_supported: false,
code_query_supported: false,
},
healthy: true,
force_cpu: true,
model_dir: PathBuf::new(),
worker_path: PathBuf::new(),
projector_verified: false,
query_profile: None,
}
}
#[test]
fn disconnected_writer_queue_is_transport_failure_and_reaps_helper() {
let mut fixture = stalled_protocol_helper(false);
let (sender, receiver) = mpsc::sync_channel(1);
drop(receiver);
fixture.requests = Some(sender);
let error = fixture
.exchange_with_timeout(b"{}", Duration::from_millis(300))
.unwrap_err();
assert!(error.downcast_ref::<WorkerTransportFailure>().is_some());
assert!(fixture.child.try_wait().unwrap().is_some());
assert!(!fixture.healthy && fixture.requests.is_none());
assert!(fixture.writer_thread.is_none() && fixture.reader_thread.is_none());
}
#[test]
fn entire_request_deadline_kills_and_joins_worker_that_never_reads_stdin() {
let mut fixture = stalled_protocol_helper(false);
let started = Instant::now();
let error = fixture
.exchange_with_timeout(&vec![b'x'; MAX_REQUEST - 1], Duration::from_millis(300))
.unwrap_err();
assert!(error.downcast_ref::<WorkerTransportFailure>().is_some());
assert!(error.to_string().contains("write failed or timed out"));
assert!(started.elapsed() < Duration::from_secs(5));
assert!(fixture.child.try_wait().unwrap().is_some());
assert!(!fixture.healthy && !fixture.info.ready);
assert!(fixture.writer_thread.is_none() && fixture.reader_thread.is_none());
}
#[test]
fn response_wait_uses_same_deadline_and_reaps_silent_reader() {
let mut fixture = stalled_protocol_helper(true);
let started = Instant::now();
let error = fixture
.exchange_with_timeout(b"{}", Duration::from_millis(300))
.unwrap_err();
assert!(error.downcast_ref::<WorkerTransportFailure>().is_some());
assert!(error.to_string().contains("request failed or timed out"));
assert!(started.elapsed() < Duration::from_secs(5));
assert!(fixture.child.try_wait().unwrap().is_some());
assert!(fixture.requests.is_none());
assert!(fixture.writer_thread.is_none() && fixture.reader_thread.is_none());
}
#[test]
fn automatic_startup_transport_failure_retries_cpu_once() {
let mut attempts = Vec::new();
@@ -718,7 +914,7 @@ mod tests {
if cpu {
Ok("cpu ready")
} else {
Err(StartupTransportFailure("worker closed its output before ready".into()).into())
Err(WorkerTransportFailure("worker closed its output before ready".into()).into())
}
})
.unwrap();
@@ -732,7 +928,7 @@ mod tests {
let mut attempts = Vec::new();
let result: Result<()> = start_with_cpu_fallback(cpu, |selected_cpu| {
attempts.push(selected_cpu);
Err(StartupTransportFailure("worker exited".into()).into())
Err(WorkerTransportFailure("worker exited".into()).into())
});
let error = format!("{:#}", result.unwrap_err());
assert_eq!(attempts, if cpu { vec![true] } else { vec![false, true] });