Skip to main content

regain_core/
usb_peripherals.rs

1//! Explicit frontend confirmation after a camera USB reset. The mailbox never
2//! takes the camera engine lock, so a capture can await NINA while stdin remains
3//! available to acknowledge the matching recovery episode.
4use crate::{CancellationToken, Failure, SharedStatus, invalid};
5use anyhow::{Result, ensure};
6use serde::Deserialize;
7use std::{
8    sync::{
9        Arc, Mutex,
10        atomic::{AtomicU64, Ordering},
11    },
12    time::Duration,
13};
14use tokio::{sync::oneshot, time::Instant};
15
16pub const WAIT_SECONDS: u64 = 120;
17static NEXT_ID: AtomicU64 = AtomicU64::new(0);
18
19fn next_id() -> Result<u64> {
20    // Unlike a per-Session counter, this cannot reuse an episode on reconnect.
21    // UTC microseconds also separate supervisor restarts; lifecycle records
22    // carry the supervisor PID so frontends can reject old-process records.
23    let seed = chrono::Utc::now().timestamp_micros().max(1) as u64;
24    let mut previous = NEXT_ID.load(Ordering::Relaxed);
25    loop {
26        let next = previous
27            .max(seed)
28            .checked_add(1)
29            .ok_or_else(|| invalid("USB peripheral recovery ID space exhausted"))?;
30        match NEXT_ID.compare_exchange_weak(previous, next, Ordering::Relaxed, Ordering::Relaxed) {
31            Ok(_) => return Ok(next),
32            Err(current) => previous = current,
33        }
34    }
35}
36
37#[derive(Debug, Deserialize)]
38#[serde(deny_unknown_fields)]
39pub struct Acknowledgement {
40    pub id: u64,
41    pub success: bool,
42    pub error: Option<String>,
43}
44
45struct Pending {
46    id: u64,
47    deadline: Option<Instant>,
48    reply: Option<oneshot::Sender<Acknowledgement>>,
49}
50#[derive(Default)]
51struct State {
52    enabled: bool,
53    retired: bool,
54    pending: Option<Pending>,
55    blocked: Option<String>,
56}
57
58#[derive(Clone)]
59pub struct PeripheralRecoveryHandle {
60    state: Arc<Mutex<State>>,
61    status: SharedStatus,
62}
63impl PeripheralRecoveryHandle {
64    pub(crate) fn new(status: SharedStatus) -> Self {
65        Self {
66            state: Arc::new(Mutex::new(State::default())),
67            status,
68        }
69    }
70    pub(crate) fn enable(&self, enabled: bool) -> Result<()> {
71        let mut state = self.state.lock().unwrap();
72        ensure!(
73            !state.retired && state.pending.is_none() && state.blocked.is_none(),
74            invalid("Cannot change USB peripheral recovery during or after recovery")
75        );
76        state.enabled = enabled;
77        Ok(())
78    }
79    /// Reject captures after an unconfirmed reset, including a later Start
80    /// request. An explicit disconnect/reconnect clears this safety fence.
81    pub(crate) fn ensure_ready(&self) -> Result<()> {
82        let state = self.state.lock().unwrap();
83        ensure!(!state.retired, invalid("Camera disconnected"));
84        if let Some(message) = &state.blocked {
85            return Err(invalid(message.clone()));
86        }
87        ensure!(
88            state.pending.is_none(),
89            invalid("USB peripheral recovery is still pending")
90        );
91        Ok(())
92    }
93    /// The Hub retains a Session across explicit disconnect/connect. Retiring
94    /// it rejects stale handles, while a subsequent explicit connect may clear
95    /// the fence. Calling connect alone during a failed episode cannot do so.
96    pub(crate) fn reconnect_after_close(&self) -> Result<()> {
97        {
98            let mut state = self.state.lock().unwrap();
99            if state.retired {
100                state.retired = false;
101                state.blocked = None;
102            }
103        }
104        self.ensure_ready()
105    }
106    /// Reserve the episode before reset dispatch; publish pending status only
107    /// after camera reappearance when the frontend may reconnect peripherals.
108    pub(crate) fn begin(&self) -> Result<Option<Recovery>> {
109        self.ensure_ready()?;
110        let mut state = self.state.lock().unwrap();
111        if !state.enabled {
112            return Ok(None);
113        }
114        let id = next_id()?;
115        let (reply, received) = oneshot::channel();
116        state.pending = Some(Pending {
117            id,
118            deadline: None,
119            reply: Some(reply),
120        });
121        Ok(Some(Recovery {
122            owner: self.clone(),
123            id,
124            received: Some(received),
125            finished: false,
126        }))
127    }
128    pub fn acknowledge(&self, acknowledgement: Acknowledgement) -> Result<()> {
129        ensure!(
130            acknowledgement.id > 0,
131            invalid("Invalid USB peripheral recovery ID")
132        );
133        if let Some(error) = &acknowledgement.error {
134            ensure!(
135                error.len() <= 2048 && !error.chars().any(|c| c.is_control()),
136                invalid("Invalid USB peripheral recovery error")
137            );
138            ensure!(
139                !acknowledgement.success,
140                invalid("Successful USB peripheral recovery cannot include an error")
141            );
142        }
143        let mut state = self.state.lock().unwrap();
144        let pending = state
145            .pending
146            .as_mut()
147            .ok_or_else(|| invalid("No USB peripheral recovery acknowledgement is pending"))?;
148        ensure!(
149            pending.id == acknowledgement.id,
150            invalid("USB peripheral recovery acknowledgement does not match the pending episode")
151        );
152        ensure!(
153            pending
154                .deadline
155                .is_some_and(|deadline| Instant::now() < deadline),
156            invalid("USB peripheral recovery is not awaiting acknowledgement or has expired")
157        );
158        let reply = pending
159            .reply
160            .take()
161            .ok_or_else(|| invalid("USB peripheral recovery was already acknowledged"))?;
162        reply
163            .send(acknowledgement)
164            .map_err(|_| invalid("USB peripheral recovery is no longer waiting"))?;
165        Ok(())
166    }
167    pub(crate) fn retire(&self) {
168        let mut state = self.state.lock().unwrap();
169        state.retired = true;
170        state.pending = None;
171        self.status.lock().unwrap().usb_peripheral_recovery_id = None;
172    }
173}
174
175pub(crate) struct Recovery {
176    owner: PeripheralRecoveryHandle,
177    pub(crate) id: u64,
178    received: Option<oneshot::Receiver<Acknowledgement>>,
179    finished: bool,
180}
181impl Recovery {
182    pub(crate) async fn wait(mut self, token: &CancellationToken) -> Result<()> {
183        let deadline = Instant::now() + Duration::from_secs(WAIT_SECONDS);
184        {
185            let mut state = self.owner.state.lock().unwrap();
186            let pending = state
187                .pending
188                .as_mut()
189                .filter(|p| p.id == self.id)
190                .ok_or_else(|| invalid("USB peripheral recovery was retired"))?;
191            pending.deadline = Some(deadline);
192            self.owner.status.lock().unwrap().usb_peripheral_recovery_id = Some(self.id);
193        }
194        let result = tokio::select! { biased;
195            _ = token.cancelled() => Err(Failure::Cancelled.into()),
196            _ = tokio::time::sleep_until(deadline) => Err(invalid("NINA peripheral recovery confirmation timed out")),
197            reply = self.received.take().unwrap() => match reply {
198                Ok(reply) if reply.success => Ok(()),
199                Ok(reply) => Err(invalid(format!("NINA peripheral recovery failed: {}",
200                    reply.error.as_deref().unwrap_or("peripherals were not confirmed ready")))),
201                Err(_) => Err(invalid("USB peripheral recovery was retired")),
202            }
203        };
204        self.finish(result.as_ref().err().map(|error| format!("{error:#}")));
205        result
206    }
207    fn finish(&mut self, failure: Option<String>) {
208        let mut state = self.owner.state.lock().unwrap();
209        if state.pending.as_ref().is_some_and(|p| p.id == self.id) {
210            state.pending = None;
211            state.blocked = failure.map(|message| format!(
212                "USB peripheral recovery was not confirmed ({message}); disconnect and reconnect the camera before another exposure"));
213            self.owner.status.lock().unwrap().usb_peripheral_recovery_id = None;
214        }
215        self.finished = true;
216    }
217}
218impl Drop for Recovery {
219    fn drop(&mut self) {
220        if !self.finished {
221            self.finish(Some("recovery did not complete".into()));
222        }
223    }
224}
225
226#[cfg(test)]
227mod tests {
228    use super::*;
229    use crate::Status;
230    fn handle() -> PeripheralRecoveryHandle {
231        let handle = PeripheralRecoveryHandle::new(Arc::new(Mutex::new(Status::default())));
232        handle.enable(true).unwrap();
233        handle
234    }
235    fn ack(id: u64, success: bool) -> Acknowledgement {
236        Acknowledgement {
237            id,
238            success,
239            error: None,
240        }
241    }
242    async fn waiting(
243        handle: &PeripheralRecoveryHandle,
244        token: CancellationToken,
245    ) -> (u64, tokio::task::JoinHandle<Result<()>>) {
246        let recovery = handle.begin().unwrap().unwrap();
247        let id = recovery.id;
248        assert!(
249            handle
250                .status
251                .lock()
252                .unwrap()
253                .usb_peripheral_recovery_id
254                .is_none()
255        );
256        assert!(
257            handle.acknowledge(ack(id, true)).is_err(),
258            "no ACK before camera return"
259        );
260        let wait = tokio::spawn(async move { recovery.wait(&token).await });
261        tokio::task::yield_now().await;
262        assert_eq!(
263            handle.status.lock().unwrap().usb_peripheral_recovery_id,
264            Some(id)
265        );
266        (id, wait)
267    }
268    #[tokio::test(start_paused = true)]
269    async fn replacement_waits_for_matching_success_and_repeated_old_future_ack_is_rejected() {
270        let handle = handle();
271        let (id, wait) = waiting(&handle, CancellationToken::new()).await;
272        tokio::time::advance(Duration::from_secs(119)).await;
273        assert!(
274            !wait.is_finished(),
275            "replacement must not start before confirmation"
276        );
277        assert!(handle.ensure_ready().is_err());
278        for wrong in [0, id - 1, id + 1] {
279            assert!(handle.acknowledge(ack(wrong, true)).is_err());
280        }
281        handle.acknowledge(ack(id, true)).unwrap();
282        assert!(handle.acknowledge(ack(id, true)).is_err());
283        wait.await.unwrap().unwrap();
284        handle.ensure_ready().unwrap();
285        assert!(
286            handle
287                .status
288                .lock()
289                .unwrap()
290                .usb_peripheral_recovery_id
291                .is_none()
292        );
293        let (next, wait) = waiting(&handle, CancellationToken::new()).await;
294        assert!(next > id);
295        assert!(handle.acknowledge(ack(id, true)).is_err());
296        handle.acknowledge(ack(next, true)).unwrap();
297        wait.await.unwrap().unwrap();
298    }
299    #[tokio::test(start_paused = true)]
300    async fn failed_ack_is_accepted_but_blocks_this_and_future_exposures() {
301        let handle = handle();
302        let (id, wait) = waiting(&handle, CancellationToken::new()).await;
303        handle
304            .acknowledge(Acknowledgement {
305                id,
306                success: false,
307                error: Some("Filter wheel did not reconnect".into()),
308            })
309            .unwrap();
310        let error = wait.await.unwrap().unwrap_err();
311        assert!(!crate::retryable(&error));
312        assert!(error.to_string().contains("Filter wheel"));
313        assert!(handle.ensure_ready().is_err());
314        assert!(handle.begin().is_err());
315        assert!(handle.acknowledge(ack(id, true)).is_err());
316    }
317    #[tokio::test(start_paused = true)]
318    async fn expiry_and_cancel_fence_future_starts_and_clear_pending_status() {
319        for cancel in [false, true] {
320            let handle = handle();
321            let token = CancellationToken::new();
322            let (id, wait) = waiting(&handle, token.clone()).await;
323            if cancel {
324                token.cancel();
325            } else {
326                tokio::time::advance(Duration::from_secs(WAIT_SECONDS)).await;
327            }
328            assert!(wait.await.unwrap().is_err());
329            assert!(
330                handle
331                    .status
332                    .lock()
333                    .unwrap()
334                    .usb_peripheral_recovery_id
335                    .is_none()
336            );
337            assert!(handle.ensure_ready().is_err());
338            assert!(handle.acknowledge(ack(id, true)).is_err());
339        }
340    }
341    #[tokio::test(start_paused = true)]
342    async fn cancelled_confirmation_wins_over_queued_success_and_retirement_rejects_ack() {
343        let first = handle();
344        let token = CancellationToken::new();
345        let (id, wait) = waiting(&first, token.clone()).await;
346        first.acknowledge(ack(id, true)).unwrap();
347        token.cancel();
348        assert!(wait.await.unwrap().unwrap_err().is::<Failure>());
349        assert!(first.ensure_ready().is_err());
350        let fresh = handle();
351        let (next, wait) = waiting(&fresh, CancellationToken::new()).await;
352        assert!(next > id, "new Session must not reuse episode IDs");
353        fresh.retire();
354        assert!(wait.await.unwrap().is_err());
355        assert!(fresh.acknowledge(ack(next, true)).is_err());
356    }
357    #[test]
358    fn default_is_disabled_and_incomplete_reset_is_a_safety_fence() {
359        let status = Arc::new(Mutex::new(Status::default()));
360        let handle = PeripheralRecoveryHandle::new(status);
361        assert!(handle.begin().unwrap().is_none());
362        handle.enable(true).unwrap();
363        drop(handle.begin().unwrap().unwrap());
364        assert!(handle.ensure_ready().is_err());
365        assert!(handle.enable(false).is_err());
366        assert!(
367            handle.reconnect_after_close().is_err(),
368            "connect alone must not clear a failed episode"
369        );
370        handle.retire();
371        handle.reconnect_after_close().unwrap();
372        assert!(
373            handle.begin().unwrap().is_some(),
374            "explicit reconnect preserves opt-in"
375        );
376    }
377    #[tokio::test(start_paused = true)]
378    async fn malformed_or_conflicting_ack_cannot_consume_a_valid_pending_reply() {
379        for value in [
380            serde_json::json!({"id":1,"success":"true"}),
381            serde_json::json!({"id":-1,"success":true}),
382            serde_json::json!({"id":1,"success":true,"other":true}),
383        ] {
384            assert!(serde_json::from_value::<Acknowledgement>(value).is_err());
385        }
386        let handle = handle();
387        let (id, wait) = waiting(&handle, CancellationToken::new()).await;
388        for error in [
389            "line\nbreak".into(),
390            "x".repeat(2049),
391            "unexpected success error".into(),
392        ] {
393            assert!(
394                handle
395                    .acknowledge(Acknowledgement {
396                        id,
397                        success: true,
398                        error: Some(error)
399                    })
400                    .is_err()
401            );
402        }
403        handle.acknowledge(ack(id, true)).unwrap();
404        wait.await.unwrap().unwrap();
405    }
406}