Harden worker IPC and recoverable runtime staging

crates/searchi-semantic/src/gemma.rs

@@ -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] });