regain_core/
worker.rs

1use crate::{Diagnostic, Failure, invalid, process::ProcessGuard};
2use anyhow::{Context, Result, ensure};
3use serde_json::{Value, json};
4use std::{path::PathBuf, process::Stdio, time::Duration};
5use tokio::{
6    io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader},
7    process::{Child, ChildStdin, ChildStdout, Command},
8};
9use tokio_util::sync::CancellationToken;
10
11#[derive(Clone)]
12pub struct Runtime {
13    pub directory: PathBuf,
14    pub sdk: PathBuf,
15    pub simulate: bool,
16    pub sdk_simulation: Option<Value>,
17}
18impl Runtime {
19    pub async fn usb_command(&self, args: &[&str], seconds: u64) -> Result<String> {
20        ensure!(
21            !self.simulate,
22            "Simulation must not launch USB recovery helpers"
23        );
24        let mut command = Command::new(
25            self.directory
26                .join(format!("regain-device{}", std::env::consts::EXE_SUFFIX)),
27        );
28        command
29            .args(args)
30            .stdin(Stdio::null())
31            .stdout(Stdio::piped())
32            .stderr(Stdio::piped())
33            .kill_on_drop(true);
34        #[cfg(windows)]
35        command.creation_flags(0x08000000);
36        let output = tokio::time::timeout(Duration::from_secs(seconds), command.output())
37            .await
38            .context("USB helper deadline exceeded")??;
39        ensure!(
40            output.status.success(),
41            "USB helper failed ({}; Windows may require administrator approval): {}",
42            output.status,
43            String::from_utf8_lossy(&output.stderr).trim()
44        );
45        Ok(String::from_utf8(output.stdout)?.trim().to_owned())
46    }
47    pub async fn spawn(&self, direct: bool, log: Diagnostic) -> Result<Worker> {
48        let path = self
49            .directory
50            .join(format!("regain-device{}", std::env::consts::EXE_SUFFIX));
51        let mut command = Command::new(&path);
52        command
53            .args([
54                "zwo",
55                if direct {
56                    "camera-direct"
57                } else {
58                    "camera-sdk"
59                },
60            ])
61            .stdin(Stdio::piped())
62            .stdout(Stdio::piped())
63            .stderr(Stdio::piped())
64            .kill_on_drop(true)
65            .current_dir(&self.directory);
66        #[cfg(windows)]
67        command.creation_flags(0x08000000);
68        if direct {
69            command.arg("--serve");
70        }
71        if self.simulate {
72            command.arg("--simulate");
73        } else if !direct {
74            command.arg("--sdk").arg(&self.sdk);
75        }
76        let mut child = command
77            .spawn()
78            .with_context(|| format!("Start {}", path.display()))?;
79        let pid = child.id().context("Worker has no process ID")?;
80        let guard = ProcessGuard::attach(pid)?;
81        let stderr = child.stderr.take().context("Worker stderr unavailable")?;
82        let diagnostic = tokio::spawn(async move {
83            let mut lines = BufReader::new(stderr).lines();
84            while let Ok(Some(line)) = lines.next_line().await {
85                if let Some(text) = line.strip_prefix("REGAIN_DIAGNOSTIC ")
86                    && let Ok(record) = serde_json::from_str::<Value>(text)
87                {
88                    log(
89                        record["level"].as_str().unwrap_or("info"),
90                        record["event"].as_str().unwrap_or("worker"),
91                        record["message"].as_str().unwrap_or(&line),
92                    );
93                } else {
94                    log("info", "worker", &line);
95                }
96            }
97        });
98        let mut worker = Worker {
99            input: child.stdin.take().context("Worker stdin unavailable")?,
100            output: child.stdout.take().context("Worker stdout unavailable")?,
101            child,
102            _guard: guard,
103            diagnostic,
104            id: 0,
105        };
106        if self.simulate
107            && !direct
108            && let Some(settings) = &self.sdk_simulation
109        {
110            worker
111                .call(
112                    "simulation",
113                    settings.clone(),
114                    5.,
115                    &CancellationToken::new(),
116                )
117                .await?;
118        }
119        Ok(worker)
120    }
121    /// SDK discovery is active hardware probing, even in a short-lived worker:
122    /// ASIGetCameraProperty can internally open cameras owned by other processes.
123    /// Cache results and schedule SDK discovery before imaging, not as live polling.
124    /// Worker isolation bounds user-space calls, not side effects on shared USB.
125    pub async fn list(
126        &self,
127        direct: bool,
128        log: Diagnostic,
129        token: &CancellationToken,
130    ) -> Result<Value> {
131        let mut worker = self.spawn(direct, log).await?;
132        let result = worker
133            .call("list", Value::Null, 20., token)
134            .await
135            .map(|r| r.0);
136        worker.kill().await;
137        result
138    }
139}
140pub struct Worker {
141    child: Child,
142    input: ChildStdin,
143    output: ChildStdout,
144    _guard: ProcessGuard,
145    diagnostic: tokio::task::JoinHandle<()>,
146    id: u64,
147}
148impl Worker {
149    /// Query (None) or configure shared software white balance while idle.
150    /// Inspect the open response's `whiteBalance.supported` capability first.
151    pub async fn white_balance(
152        &mut self,
153        settings: Option<crate::white_balance::Settings>,
154        token: &CancellationToken,
155    ) -> Result<Value> {
156        if let Some(settings) = settings {
157            settings.gains.validate()?;
158        }
159        Ok(self
160            .call("white-balance", serde_json::to_value(settings)?, 15., token)
161            .await?
162            .0)
163    }
164    pub fn pid(&self) -> Option<u32> {
165        self.child.id()
166    }
167    pub async fn kill(&mut self) {
168        let _ = self.child.kill().await;
169        let _ = self.child.wait().await;
170    }
171    pub async fn call(
172        &mut self,
173        method: &str,
174        params: Value,
175        seconds: f64,
176        token: &CancellationToken,
177    ) -> Result<(Value, Vec<u8>)> {
178        let result = tokio::select! {
179            biased;
180            _=token.cancelled()=>Err(Failure::Cancelled.into()),
181            result=tokio::time::timeout(Duration::from_secs_f64(seconds),self.exchange(method,params))=>result.unwrap_or_else(|_|Err(anyhow::anyhow!("Worker {method} timed out after {seconds} seconds"))),
182        };
183        if result.is_err()
184            && !matches!(
185                result
186                    .as_ref()
187                    .err()
188                    .and_then(|e| e.downcast_ref::<Failure>()),
189                Some(Failure::Worker { .. })
190            )
191        {
192            self.kill().await;
193        }
194        result
195    }
196    async fn exchange(&mut self, method: &str, params: Value) -> Result<(Value, Vec<u8>)> {
197        self.id += 1;
198        let bytes =
199            serde_json::to_vec(&json!({"version":1,"id":self.id,"method":method,"params":params}))?;
200        ensure!(
201            bytes.len() <= 65536,
202            Failure::Invalid("Command too large".into())
203        );
204        self.input.write_u32_le(bytes.len() as u32).await?;
205        self.input.write_all(&bytes).await?;
206        self.input.flush().await?;
207        let length = self.output.read_u32_le().await? as usize;
208        ensure!(
209            (1..=65536).contains(&length),
210            Failure::Invalid("Invalid worker response length".into())
211        );
212        let mut header = vec![0; length];
213        self.output.read_exact(&mut header).await?;
214        let reply: Value =
215            serde_json::from_slice(&header).map_err(|_| invalid("Malformed worker JSON"))?;
216        ensure!(
217            reply["id"] == self.id && reply["version"] == 1,
218            Failure::Invalid("Stale worker response".into())
219        );
220        let count = reply["binaryLength"]
221            .as_u64()
222            .ok_or_else(|| invalid("Missing frame length"))?;
223        ensure!(
224            count <= 512 * 1024 * 1024
225                && (matches!(method, "download" | "stream-download" | "stream-poll") || count == 0),
226            Failure::Invalid("Invalid worker image length".into())
227        );
228        if reply["ok"] == false {
229            ensure!(
230                count == 0,
231                Failure::Invalid("Error response contains pixels".into())
232            );
233            return Err(Failure::Worker {
234                message: reply["error"].as_str().unwrap_or("Worker error").into(),
235                code: reply["sdkCode"].as_i64().map(|v| v as i32),
236            }
237            .into());
238        }
239        ensure!(
240            reply["ok"] == true,
241            Failure::Invalid("Missing worker status".into())
242        );
243        let mut pixels = vec![0; count as usize];
244        self.output.read_exact(&mut pixels).await?;
245        Ok((reply["result"].clone(), pixels))
246    }
247}
248impl Drop for Worker {
249    fn drop(&mut self) {
250        self.diagnostic.abort();
251    }
252}
regain_core/worker.rs

regain_core/
worker.rs
