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, AsyncRead, 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                    crate::timing::SIMULATION_SETUP_SECONDS,
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        self.call_bounded(method, params, seconds, token, 512 * 1024 * 1024)
179            .await
180    }
181    /// JSON-only commands reject any announced binary body before allocating it.
182    pub async fn call_json(
183        &mut self,
184        method: &str,
185        params: Value,
186        seconds: f64,
187        token: &CancellationToken,
188    ) -> Result<Value> {
189        self.call_bounded(method, params, seconds, token, 0)
190            .await
191            .map(|(value, _)| value)
192    }
193    /// Bound an image reply by the caller's admitted ROI before allocating pixels.
194    /// A zero-length streaming poll remains valid; capture checks exact length.
195    pub async fn call_image(
196        &mut self,
197        method: &str,
198        params: Value,
199        seconds: f64,
200        token: &CancellationToken,
201        maximum_bytes: usize,
202    ) -> Result<(Value, Vec<u8>)> {
203        ensure!(
204            matches!(method, "download" | "stream-download" | "stream-poll")
205                && (1..=512 * 1024 * 1024).contains(&maximum_bytes),
206            Failure::Invalid("Invalid worker image admission".into())
207        );
208        self.call_bounded(method, params, seconds, token, maximum_bytes)
209            .await
210    }
211    async fn call_bounded(
212        &mut self,
213        method: &str,
214        params: Value,
215        seconds: f64,
216        token: &CancellationToken,
217        maximum_bytes: usize,
218    ) -> Result<(Value, Vec<u8>)> {
219        self.call_until(
220            method,
221            params,
222            tokio::time::Instant::now() + Duration::from_secs_f64(seconds),
223            token,
224            maximum_bytes,
225        )
226        .await
227    }
228    /// Persistent camera controls share absolute-deadline/framed-write tracking.
229    pub(crate) async fn control_call(
230        &mut self,
231        method: &str,
232        params: Value,
233        deadline: tokio::time::Instant,
234        token: &CancellationToken,
235    ) -> Result<(Value, Vec<u8>)> {
236        self.call_until(method, params, deadline, token, 0).await
237    }
238    async fn call_until(
239        &mut self,
240        method: &str,
241        params: Value,
242        deadline: tokio::time::Instant,
243        token: &CancellationToken,
244        maximum_bytes: usize,
245    ) -> Result<(Value, Vec<u8>)> {
246        let control =
247            method == "set" && matches!(params["control"].as_i64(), Some(0 | 5 | 16 | 17));
248        let mut dispatched = false;
249        let result = tokio::select! {
250            biased;
251            _=token.cancelled()=>Err(Failure::Cancelled.into()),
252            result=tokio::time::timeout_at(deadline,self.exchange(method,params,maximum_bytes,&mut dispatched,deadline))=>result.unwrap_or_else(|_|{
253                if control && !dispatched {Err(crate::cooling::CoolingError::Expired.into())}
254                else {Err(anyhow::anyhow!("Worker {method} acknowledgement timed out"))}
255            }),
256        };
257        let result = result.map_err(|error| {
258            if control
259                && dispatched
260                && !matches!(
261                    error.downcast_ref::<Failure>(),
262                    Some(Failure::Worker { .. } | Failure::UncertainControl { .. })
263                )
264            {
265                // The process may have consumed a partial/complete command. A
266                // timeout, cancellation or malformed/lost reply cannot prove
267                // this persistent hardware write was not applied.
268                Failure::UncertainControl {
269                    message: format!("{error:#}"),
270                    code: None,
271                }
272                .into()
273            } else {
274                error
275            }
276        });
277        if result.is_err()
278            && !(!dispatched
279                && result.as_ref().err().is_some_and(|error| {
280                    matches!(
281                        error.downcast_ref::<crate::cooling::CoolingError>(),
282                        Some(crate::cooling::CoolingError::Expired)
283                    )
284                }))
285            && !matches!(
286                result
287                    .as_ref()
288                    .err()
289                    .and_then(|e| e.downcast_ref::<Failure>()),
290                Some(Failure::Worker { .. })
291            )
292        {
293            self.kill().await;
294        }
295        result
296    }
297    async fn exchange(
298        &mut self,
299        method: &str,
300        params: Value,
301        maximum_bytes: usize,
302        dispatched: &mut bool,
303        deadline: tokio::time::Instant,
304    ) -> Result<(Value, Vec<u8>)> {
305        let id = self
306            .id
307            .checked_add(1)
308            .ok_or_else(|| invalid("Worker command counter exhausted"))?;
309        let bytes =
310            serde_json::to_vec(&json!({"version":1,"id":id,"method":method,"params":params}))?;
311        ensure!(
312            bytes.len() <= 65536,
313            Failure::Invalid("Command too large".into())
314        );
315        ensure!(
316            tokio::time::Instant::now() < deadline,
317            crate::cooling::CoolingError::Expired
318        );
319        self.id = id;
320        *dispatched = true;
321        self.input.write_u32_le(bytes.len() as u32).await?;
322        self.input.write_all(&bytes).await?;
323        self.input.flush().await?;
324        read_reply(&mut self.output, method, self.id, maximum_bytes).await
325    }
326}
327
328async fn read_reply<R: AsyncRead + Unpin>(
329    output: &mut R,
330    method: &str,
331    id: u64,
332    maximum_bytes: usize,
333) -> Result<(Value, Vec<u8>)> {
334    let length = output.read_u32_le().await? as usize;
335    ensure!(
336        (1..=65536).contains(&length),
337        Failure::Invalid("Invalid worker response length".into())
338    );
339    let mut header = vec![0; length];
340    output.read_exact(&mut header).await?;
341    let reply: Value =
342        serde_json::from_slice(&header).map_err(|_| invalid("Malformed worker JSON"))?;
343    ensure!(
344        reply["id"] == id && reply["version"] == 1,
345        Failure::Invalid("Stale worker response".into())
346    );
347    let count = reply["binaryLength"]
348        .as_u64()
349        .ok_or_else(|| invalid("Missing frame length"))?;
350    ensure!(
351        count <= maximum_bytes as u64
352            && count <= 512 * 1024 * 1024
353            && (matches!(method, "download" | "stream-download" | "stream-poll") || count == 0),
354        Failure::Invalid("Invalid worker image length".into())
355    );
356    if reply["ok"] == false {
357        ensure!(
358            count == 0,
359            Failure::Invalid("Error response contains pixels".into())
360        );
361        let message = reply["error"].as_str().unwrap_or("Worker error").into();
362        let code = reply["sdkCode"]
363            .as_i64()
364            .and_then(|v| i32::try_from(v).ok());
365        return Err(if reply["controlUncertain"] == true {
366            Failure::UncertainControl { message, code }
367        } else {
368            Failure::Worker { message, code }
369        }
370        .into());
371    }
372    ensure!(
373        reply["ok"] == true,
374        Failure::Invalid("Missing worker status".into())
375    );
376    let mut pixels = Vec::new();
377    pixels
378        .try_reserve_exact(count as usize)
379        .context("Worker image allocation failed")?;
380    pixels.resize(count as usize, 0);
381    output.read_exact(&mut pixels).await?;
382    // Move the result out, avoiding a second decoded metadata tree.
383    let mut reply = reply;
384    Ok((reply["result"].take(), pixels))
385}
386impl Drop for Worker {
387    fn drop(&mut self) {
388        self.diagnostic.abort();
389    }
390}
391
392#[cfg(test)]
393mod tests {
394    use super::*;
395    fn simulated_runtime(settings: Value) -> Runtime {
396        Runtime {
397            directory: std::env::var_os("REGAIN_TEST_WORKERS")
398                .map(PathBuf::from)
399                .unwrap_or_else(|| {
400                    PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../target/debug")
401                }),
402            sdk: "unused".into(),
403            simulate: true,
404            sdk_simulation: Some(settings),
405        }
406    }
407
408    fn response(count: u64, ok: bool) -> Vec<u8> {
409        let header = serde_json::to_vec(&json!({
410            "version":1,"id":7,"ok":ok,"binaryLength":count,
411            "result":{"width":3,"height":2,"readRecoveries":2},
412            "error":"private SDK error","sdkCode":11,
413        }))
414        .unwrap();
415        let mut bytes = (header.len() as u32).to_le_bytes().to_vec();
416        bytes.extend(header);
417        bytes
418    }
419
420    #[tokio::test]
421    async fn json_only_admission_rejects_announced_pixels_before_reading_a_body() {
422        let zero = response(0, true);
423        assert!(
424            read_reply(&mut zero.as_slice(), "list", 7, 0)
425                .await
426                .unwrap()
427                .1
428                .is_empty()
429        );
430        // No body follows this header. Reject its announced size rather than
431        // attempting an allocation or reporting a later partial-body failure.
432        let binary = response(1, true);
433        let error = read_reply(&mut binary.as_slice(), "download", 7, 0)
434            .await
435            .unwrap_err();
436        assert!(error.to_string().contains("Invalid worker image length"));
437    }
438
439    #[tokio::test]
440    async fn timestamped_control_observations_cross_production_worker_framing() {
441        for direct in [false, true] {
442            let runtime = simulated_runtime(json!({"instant":true}));
443            let token = CancellationToken::new();
444            let mut worker = runtime
445                .spawn(direct, std::sync::Arc::new(|_, _, _| {}))
446                .await
447                .unwrap();
448            worker
449                .call(
450                    "open",
451                    json!({"name":if direct {"ZWO ASI585MM Pro"} else {"ZWO Simulated"}}),
452                    15.,
453                    &token,
454                )
455                .await
456                .unwrap();
457            worker
458                .call("set", json!({"control":0,"value":123}), 15., &token)
459                .await
460                .unwrap();
461            for control in [0, 8, 15, 16, 17] {
462                let requested = tokio::time::Instant::now();
463                let (reply, pixels) = worker
464                    .call("get-observation", json!({"control":control}), 15., &token)
465                    .await
466                    .unwrap();
467                let observation: crate::ControlObservationReply =
468                    serde_json::from_value(reply).unwrap();
469                assert!(pixels.is_empty());
470                assert!(observation.age_seconds.is_finite() && observation.age_seconds >= 0.);
471                assert!(observation.observed_at(requested).unwrap() <= requested);
472                assert_eq!(
473                    worker
474                        .call("get", json!({"control":control}), 15., &token)
475                        .await
476                        .unwrap()
477                        .0,
478                    observation.value
479                );
480                if control == 0 {
481                    assert_eq!(observation.value, 123);
482                }
483            }
484            let pid = worker.pid();
485            for control in [json!(999), json!("8"), json!(4294967296u64)] {
486                assert!(
487                    worker
488                        .call("get-observation", json!({"control":control}), 15., &token)
489                        .await
490                        .is_err()
491                );
492            }
493            assert_eq!(worker.pid(), pid);
494            worker
495                .call("get-observation", json!({"control":0}), 15., &token)
496                .await
497                .unwrap();
498            worker
499                .call("close", Value::Null, 15., &token)
500                .await
501                .unwrap();
502            worker.kill().await;
503        }
504    }
505
506    #[tokio::test]
507    async fn persistent_control_transport_loss_retires_simulated_worker_without_retry() {
508        let runtime = simulated_runtime(json!({"instant":true,"fault":"hang"}));
509        for control in [0, 5, 16, 17] {
510            for cancel_after_dispatch in [false, true] {
511                let token = CancellationToken::new();
512                let mut worker = runtime
513                    .spawn(false, std::sync::Arc::new(|_, _, _| {}))
514                    .await
515                    .unwrap();
516                worker
517                    .call("open", json!({"name":"ZWO Simulated"}), 15., &token)
518                    .await
519                    .unwrap();
520                worker.call("start",json!({"width":64,"height":64,"bin":1,"x":0,"y":0,"microseconds":1000,"dark":false}),15.,&token).await.unwrap();
521                // Deliberately park the simulated owner in download. Its serial
522                // command queue then cannot acknowledge the following control write.
523                worker.id += 1;
524                let header = serde_json::to_vec(
525                    &json!({"version":1,"id":worker.id,"method":"download","params":null}),
526                )
527                .unwrap();
528                worker
529                    .input
530                    .write_u32_le(header.len() as u32)
531                    .await
532                    .unwrap();
533                worker.input.write_all(&header).await.unwrap();
534                worker.input.flush().await.unwrap();
535                let mut call =
536                    Box::pin(worker.call("set", json!({"control":control,"value":0}), 1., &token));
537                if cancel_after_dispatch {
538                    // Poll through write admission before cancellation, rather than
539                    // using a scheduling-sensitive sleep to guess dispatch timing.
540                    std::future::poll_fn(|cx| {
541                        assert!(std::future::Future::poll(call.as_mut(), cx).is_pending());
542                        std::task::Poll::Ready(())
543                    })
544                    .await;
545                    token.cancel();
546                }
547                let error = call.await.unwrap_err();
548                assert!(
549                    matches!(
550                        error.downcast_ref::<Failure>(),
551                        Some(Failure::UncertainControl { .. })
552                    ),
553                    "{error:#}"
554                );
555                assert!(!crate::retryable(&error));
556                assert!(worker.child.try_wait().unwrap().is_some());
557            }
558        }
559        // Before-dispatch cancellation is still distinguishable from a write
560        // whose acknowledgement was lost. No command ID is consumed.
561        let mut worker = runtime
562            .spawn(false, std::sync::Arc::new(|_, _, _| {}))
563            .await
564            .unwrap();
565        let id = worker.id;
566        let token = CancellationToken::new();
567        token.cancel();
568        let error = worker
569            .call("set", json!({"control":17,"value":0}), 1., &token)
570            .await
571            .unwrap_err();
572        assert!(matches!(
573            error.downcast_ref::<Failure>(),
574            Some(Failure::Cancelled)
575        ));
576        assert_eq!(worker.id, id);
577        assert!(worker.child.try_wait().unwrap().is_some());
578    }
579
580    #[tokio::test]
581    async fn expired_control_write_preserves_worker_and_framing_without_dispatch() {
582        let mut worker = simulated_runtime(json!({"instant":true}))
583            .spawn(false, std::sync::Arc::new(|_, _, _| {}))
584            .await
585            .unwrap();
586        let token = CancellationToken::new();
587        let id = worker.id;
588        for control in [0, 5, 16, 17] {
589            let error = worker
590                .control_call(
591                    "set",
592                    json!({"control":control,"value":0}),
593                    tokio::time::Instant::now(),
594                    &token,
595                )
596                .await
597                .unwrap_err();
598            assert!(matches!(
599                error.downcast_ref::<crate::cooling::CoolingError>(),
600                Some(crate::cooling::CoolingError::Expired)
601            ));
602            assert_eq!(worker.id, id);
603            assert!(worker.child.try_wait().unwrap().is_none());
604        }
605        worker
606            .call("open", json!({"name":"ZWO Simulated"}), 15., &token)
607            .await
608            .unwrap();
609        worker
610            .call("close", Value::Null, 15., &token)
611            .await
612            .unwrap();
613        worker.kill().await;
614    }
615
616    #[tokio::test]
617    async fn uncertain_control_reply_is_typed_and_never_retryable() {
618        for code in [Some(11), None] {
619            let header=serde_json::to_vec(&json!({"version":1,"id":7,"ok":false,
620                "binaryLength":0,"error":"private cooling failure","sdkCode":code,"controlUncertain":true})).unwrap();
621            let mut bytes = (header.len() as u32).to_le_bytes().to_vec();
622            bytes.extend(header);
623            let error = read_reply(&mut bytes.as_slice(), "set", 7, 0)
624                .await
625                .unwrap_err();
626            assert!(
627                matches!(error.downcast_ref::<Failure>(),Some(Failure::UncertainControl {code:actual,message}) if *actual==code && message=="private cooling failure")
628            );
629            assert!(!crate::retryable(&error));
630        }
631    }
632
633    #[tokio::test]
634    async fn oversized_admitted_reply_is_rejected_before_reading_any_body() {
635        let (mut writer, mut reader) = tokio::io::duplex(1024);
636        writer.write_all(&response(13, true)).await.unwrap();
637        // Writer stays open and sends no body: waiting for pixels would time out.
638        let error = tokio::time::timeout(
639            Duration::from_secs(1),
640            read_reply(&mut reader, "download", 7, 12),
641        )
642        .await
643        .unwrap()
644        .unwrap_err();
645        assert!(
646            matches!(error.downcast_ref::<Failure>(), Some(Failure::Invalid(message)) if message == "Invalid worker image length")
647        );
648        assert!(!crate::retryable(&error));
649    }
650
651    #[tokio::test]
652    async fn bounded_reply_preserves_pixels_metadata_and_empty_stream_poll() {
653        let mut bytes = response(12, true);
654        bytes.extend(0u8..12);
655        let (metadata, pixels) = read_reply(&mut bytes.as_slice(), "download", 7, 12)
656            .await
657            .unwrap();
658        assert_eq!(pixels, (0u8..12).collect::<Vec<_>>());
659        assert_eq!(metadata["readRecoveries"], 2);
660        let bytes = response(0, true);
661        assert!(
662            read_reply(&mut bytes.as_slice(), "stream-poll", 7, 12)
663                .await
664                .unwrap()
665                .1
666                .is_empty()
667        );
668    }
669
670    #[tokio::test]
671    async fn bounded_reply_keeps_error_codes_and_rejects_stale_or_non_image_pixels() {
672        let bytes = response(0, false);
673        let error = read_reply(&mut bytes.as_slice(), "download", 7, 12)
674            .await
675            .unwrap_err();
676        assert!(
677            matches!(error.downcast_ref::<Failure>(), Some(Failure::Worker {code:Some(11), message}) if message == "private SDK error")
678        );
679        for (count, ok, method, id) in [
680            (1, false, "download", 7),
681            (1, true, "status", 7),
682            (0, true, "download", 8),
683        ] {
684            let bytes = response(count, ok);
685            let error = read_reply(&mut bytes.as_slice(), method, id, 12)
686                .await
687                .unwrap_err();
688            assert!(matches!(
689                error.downcast_ref::<Failure>(),
690                Some(Failure::Invalid(_))
691            ));
692        }
693        let bytes = response(12, true);
694        let error = read_reply(&mut bytes.as_slice(), "download", 7, 12)
695            .await
696            .unwrap_err();
697        assert_eq!(
698            error.downcast_ref::<std::io::Error>().unwrap().kind(),
699            std::io::ErrorKind::UnexpectedEof
700        );
701    }
702}
regain_core/worker.rs

regain_core/
worker.rs
