1use 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 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 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 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 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}