1use std::any::Any;
9use std::collections::HashMap;
10use std::ffi::{c_void, CStr};
11use std::future::Future;
12use std::pin::Pin;
13use std::sync::atomic::{AtomicUsize, Ordering};
14use std::sync::{Arc, Condvar, Mutex, OnceLock, PoisonError};
15use std::task::{Context, Poll, Waker};
16use std::time::{Duration, Instant};
17
18use crate::utils::panic_safe::catch_user_panic;
19
20pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(30);
22
23pub const TIMEOUT_ENV_VAR: &str = "SCREENCAPTUREKIT_COMPLETION_TIMEOUT_SECS";
25
26pub const TIMEOUT_MESSAGE_PREFIX: &str = "screencapturekit: completion callback did not fire";
28
29#[must_use]
31pub fn is_timeout_error(message: &str) -> bool {
32 message.starts_with(TIMEOUT_MESSAGE_PREFIX)
33}
34
35#[must_use]
37pub fn default_timeout() -> Option<Duration> {
38 static TIMEOUT: OnceLock<Option<Duration>> = OnceLock::new();
39 *TIMEOUT.get_or_init(|| {
40 std::env::var(TIMEOUT_ENV_VAR).map_or(Some(DEFAULT_TIMEOUT), |raw| {
41 raw.trim()
42 .parse::<u64>()
43 .map_or(Some(DEFAULT_TIMEOUT), |seconds| {
44 (seconds != 0).then(|| Duration::from_secs(seconds))
45 })
46 })
47 })
48}
49
50#[must_use]
52pub fn timed_out_context_count() -> usize {
53 TIMED_OUT_CONTEXTS.load(Ordering::Relaxed)
54}
55
56static TIMED_OUT_CONTEXTS: AtomicUsize = AtomicUsize::new(0);
57static NEXT_CONTEXT_ID: AtomicUsize = AtomicUsize::new(1);
58static CONTEXTS: Mutex<Option<HashMap<usize, Box<dyn Any + Send>>>> = Mutex::new(None);
59static NEXT_TIMEOUT_ID: AtomicUsize = AtomicUsize::new(1);
60static TIMEOUT_SCHEDULER: OnceLock<Arc<TimeoutScheduler>> = OnceLock::new();
61
62pub type SyncCompletionPtr = *mut c_void;
64
65#[allow(clippy::significant_drop_tightening)]
66fn register_context<T>(context: T) -> (SyncCompletionPtr, usize)
67where
68 T: Any + Send,
69{
70 let mut contexts = CONTEXTS.lock().unwrap_or_else(PoisonError::into_inner);
71 let contexts = contexts.get_or_insert_with(HashMap::new);
72
73 loop {
74 let id = NEXT_CONTEXT_ID.fetch_add(1, Ordering::Relaxed);
75 if id != 0 && !contexts.contains_key(&id) {
76 contexts.insert(id, Box::new(context));
77 return (id as SyncCompletionPtr, id);
78 }
79 }
80}
81
82fn take_context<T>(context: SyncCompletionPtr) -> Option<T>
83where
84 T: Any + Send,
85{
86 let id = context as usize;
87 if id == 0 {
88 return None;
89 }
90
91 let entry = CONTEXTS
92 .lock()
93 .unwrap_or_else(PoisonError::into_inner)
94 .as_mut()?
95 .remove(&id)?;
96 entry.downcast::<T>().ok().map(|entry| *entry)
97}
98
99fn remove_context(id: usize) -> bool {
100 let removed = {
101 let mut contexts = CONTEXTS.lock().unwrap_or_else(PoisonError::into_inner);
102 contexts.as_mut().and_then(|contexts| contexts.remove(&id))
103 };
104 let existed = removed.is_some();
105 drop(removed);
106 existed
107}
108
109fn timeout_message(timeout: Duration) -> String {
110 format!("{TIMEOUT_MESSAGE_PREFIX} within {timeout:?}")
111}
112
113struct TimeoutEntry {
114 id: usize,
115 deadline: Instant,
116 action: Option<Box<dyn FnOnce() + Send>>,
117}
118
119struct TimeoutScheduler {
120 entries: Mutex<Vec<TimeoutEntry>>,
121 changed: Condvar,
122}
123
124impl TimeoutScheduler {
125 fn shared() -> &'static Arc<Self> {
126 TIMEOUT_SCHEDULER.get_or_init(|| {
127 let scheduler = Arc::new(Self {
128 entries: Mutex::new(Vec::new()),
129 changed: Condvar::new(),
130 });
131 let worker = Arc::clone(&scheduler);
132 std::thread::Builder::new()
133 .name("screencapturekit-completions".to_string())
134 .spawn(move || worker.run())
135 .expect("failed to start completion timeout thread");
136 scheduler
137 })
138 }
139
140 fn schedule(timeout: Duration, action: impl FnOnce() + Send + 'static) -> usize {
141 let Some(deadline) = Instant::now().checked_add(timeout) else {
142 return 0;
143 };
144 let scheduler = Self::shared();
145 let mut entries = scheduler
146 .entries
147 .lock()
148 .unwrap_or_else(PoisonError::into_inner);
149 let id = loop {
150 let id = NEXT_TIMEOUT_ID.fetch_add(1, Ordering::Relaxed);
151 if id != 0 && !entries.iter().any(|entry| entry.id == id) {
152 break id;
153 }
154 };
155 entries.push(TimeoutEntry {
156 id,
157 deadline,
158 action: Some(Box::new(action)),
159 });
160 drop(entries);
161 scheduler.changed.notify_one();
162 id
163 }
164
165 fn cancel(id: usize) {
166 if id == 0 {
167 return;
168 }
169 let Some(scheduler) = TIMEOUT_SCHEDULER.get() else {
170 return;
171 };
172 let removed = {
173 let mut entries = scheduler
174 .entries
175 .lock()
176 .unwrap_or_else(PoisonError::into_inner);
177 entries
178 .iter()
179 .position(|entry| entry.id == id)
180 .map(|index| entries.swap_remove(index))
181 };
182 drop(removed);
183 scheduler.changed.notify_one();
184 }
185
186 fn run(self: Arc<Self>) -> ! {
187 loop {
188 let mut entries = self.entries.lock().unwrap_or_else(PoisonError::into_inner);
189 while entries.is_empty() {
190 entries = self
191 .changed
192 .wait(entries)
193 .unwrap_or_else(PoisonError::into_inner);
194 }
195
196 let (index, deadline) = entries
197 .iter()
198 .enumerate()
199 .min_by_key(|(_, entry)| entry.deadline)
200 .map(|(index, entry)| (index, entry.deadline))
201 .expect("the timeout queue is non-empty");
202 let now = Instant::now();
203 if deadline > now {
204 let (guard, _) = self
205 .changed
206 .wait_timeout(entries, deadline - now)
207 .unwrap_or_else(PoisonError::into_inner);
208 drop(guard);
209 continue;
210 }
211
212 let mut entry = entries.swap_remove(index);
213 drop(entries);
214 if let Some(action) = entry.action.take() {
215 catch_user_panic("async completion timeout", action);
216 }
217 }
218 }
219}
220
221fn cancel_async_timeout<T>(inner: &AsyncCompletionInner<T>) {
222 let timeout_id = inner.timeout_id.swap(0, Ordering::AcqRel);
223 TimeoutScheduler::cancel(timeout_id);
224}
225
226struct SyncCompletionInner<T> {
227 result: Mutex<Option<Result<T, String>>>,
228 cvar: Condvar,
229}
230
231pub struct SyncCompletion<T: Send + 'static> {
233 inner: Arc<SyncCompletionInner<T>>,
234 context_id: usize,
235}
236
237impl<T: Send + 'static> std::fmt::Debug for SyncCompletion<T> {
238 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
239 let completed = self
240 .inner
241 .result
242 .lock()
243 .unwrap_or_else(PoisonError::into_inner)
244 .is_some();
245 f.debug_struct("SyncCompletion")
246 .field("completed", &completed)
247 .finish_non_exhaustive()
248 }
249}
250
251impl<T: Send + 'static> SyncCompletion<T> {
252 #[must_use]
254 pub fn new() -> (Self, SyncCompletionPtr) {
255 let inner = Arc::new(SyncCompletionInner {
256 result: Mutex::new(None),
257 cvar: Condvar::new(),
258 });
259 let (context, context_id) = register_context(Arc::clone(&inner));
260 (Self { inner, context_id }, context)
261 }
262
263 pub fn wait(self) -> Result<T, String> {
269 match default_timeout() {
270 Some(timeout) => self.wait_timeout(timeout),
271 None => self.wait_forever(),
272 }
273 }
274
275 #[allow(clippy::significant_drop_tightening)]
281 pub fn wait_timeout(self, timeout: Duration) -> Result<T, String> {
282 let guard = self
283 .inner
284 .result
285 .lock()
286 .unwrap_or_else(PoisonError::into_inner);
287 let (mut guard, wait_result) = self
288 .inner
289 .cvar
290 .wait_timeout_while(guard, timeout, |result| result.is_none())
291 .unwrap_or_else(PoisonError::into_inner);
292
293 if wait_result.timed_out() && guard.is_none() {
294 TIMED_OUT_CONTEXTS.fetch_add(1, Ordering::Relaxed);
295 return Err(timeout_message(timeout));
296 }
297
298 guard
299 .take()
300 .unwrap_or_else(|| Err("completion signalled without a result".to_string()))
301 }
302
303 #[allow(clippy::significant_drop_tightening)]
309 pub fn wait_forever(self) -> Result<T, String> {
310 let guard = self
311 .inner
312 .result
313 .lock()
314 .unwrap_or_else(PoisonError::into_inner);
315 let mut guard = self
316 .inner
317 .cvar
318 .wait_while(guard, |result| result.is_none())
319 .unwrap_or_else(PoisonError::into_inner);
320 guard
321 .take()
322 .unwrap_or_else(|| Err("completion signalled without a result".to_string()))
323 }
324
325 pub unsafe fn complete_ok(context: SyncCompletionPtr, value: T) {
332 unsafe { Self::complete_with_result(context, Ok(value)) };
333 }
334
335 pub unsafe fn complete_err(context: SyncCompletionPtr, error: String) {
342 unsafe { Self::complete_with_result(context, Err(error)) };
343 }
344
345 pub unsafe fn complete_with_result(context: SyncCompletionPtr, result: Result<T, String>) {
354 let Some(inner) = take_context::<Arc<SyncCompletionInner<T>>>(context) else {
355 return;
356 };
357
358 {
359 let mut slot = inner.result.lock().unwrap_or_else(PoisonError::into_inner);
360 *slot = Some(result);
361 }
362 inner.cvar.notify_all();
363 }
364}
365
366impl<T: Send + 'static> Default for SyncCompletion<T> {
367 fn default() -> Self {
368 Self::new().0
369 }
370}
371
372impl<T: Send + 'static> Drop for SyncCompletion<T> {
373 fn drop(&mut self) {
374 remove_context(self.context_id);
375 }
376}
377
378struct AsyncCompletionState<T> {
379 result: Option<Result<T, String>>,
380 waker: Option<Waker>,
381 completion_hook: Option<AsyncCompletionHook<T>>,
382}
383
384type AsyncCompletionHook<T> = Box<dyn FnOnce(&Result<T, String>) + Send>;
385
386struct AsyncCompletionInner<T> {
387 state: Mutex<AsyncCompletionState<T>>,
388 timeout_id: AtomicUsize,
389}
390
391pub struct AsyncCompletion<T: Send + 'static> {
393 _marker: std::marker::PhantomData<T>,
394}
395
396impl<T: Send + 'static> std::fmt::Debug for AsyncCompletion<T> {
397 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
398 f.debug_struct("AsyncCompletion").finish_non_exhaustive()
399 }
400}
401
402pub struct AsyncCompletionFuture<T: Send + 'static> {
404 inner: Arc<AsyncCompletionInner<T>>,
405 context_id: usize,
406 cancel_on_drop: bool,
407}
408
409impl<T: Send + 'static> std::fmt::Debug for AsyncCompletionFuture<T> {
410 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
411 f.debug_struct("AsyncCompletionFuture")
412 .finish_non_exhaustive()
413 }
414}
415
416impl<T: Send + 'static> AsyncCompletion<T> {
417 #[must_use]
422 pub fn create() -> (AsyncCompletionFuture<T>, SyncCompletionPtr) {
423 Self::create_inner(None, true, default_timeout())
424 }
425
426 #[cfg(feature = "async")]
427 pub(crate) fn create_with_hook(
428 hook: impl FnOnce(&Result<T, String>) + Send + 'static,
429 ) -> (AsyncCompletionFuture<T>, SyncCompletionPtr) {
430 Self::create_inner(Some(Box::new(hook)), false, default_timeout())
431 }
432
433 #[cfg(feature = "async")]
434 pub(crate) fn create_unbounded() -> (AsyncCompletionFuture<T>, SyncCompletionPtr) {
435 Self::create_inner(None, true, None)
436 }
437
438 #[cfg(feature = "async")]
439 pub(crate) fn create_unbounded_with_hook(
440 hook: impl FnOnce(&Result<T, String>) + Send + 'static,
441 ) -> (AsyncCompletionFuture<T>, SyncCompletionPtr) {
442 let (mut future, context) = Self::create_unbounded();
443 future
444 .inner
445 .state
446 .lock()
447 .unwrap_or_else(PoisonError::into_inner)
448 .completion_hook = Some(Box::new(hook));
449 future.cancel_on_drop = false;
450 (future, context)
451 }
452
453 fn create_inner(
454 completion_hook: Option<AsyncCompletionHook<T>>,
455 cancel_on_drop: bool,
456 timeout: Option<Duration>,
457 ) -> (AsyncCompletionFuture<T>, SyncCompletionPtr) {
458 let inner = Arc::new(AsyncCompletionInner {
459 state: Mutex::new(AsyncCompletionState {
460 result: None,
461 waker: None,
462 completion_hook,
463 }),
464 timeout_id: AtomicUsize::new(0),
465 });
466 let (context, context_id) = register_context(Arc::clone(&inner));
467 if let Some(timeout) = timeout {
468 let context_address = context as usize;
469 let timeout_id = TimeoutScheduler::schedule(timeout, move || unsafe {
470 Self::complete_err(
471 context_address as SyncCompletionPtr,
472 timeout_message(timeout),
473 );
474 });
475 inner.timeout_id.store(timeout_id, Ordering::Release);
476 }
477 (
478 AsyncCompletionFuture {
479 inner,
480 context_id,
481 cancel_on_drop,
482 },
483 context,
484 )
485 }
486
487 pub unsafe fn complete_ok(context: SyncCompletionPtr, value: T) {
494 unsafe { Self::complete_with_result(context, Ok(value)) };
495 }
496
497 pub unsafe fn complete_err(context: SyncCompletionPtr, error: String) {
504 unsafe { Self::complete_with_result(context, Err(error)) };
505 }
506
507 pub unsafe fn complete_with_result(context: SyncCompletionPtr, result: Result<T, String>) {
515 Self::complete_if_pending(context, result);
516 }
517
518 pub(crate) fn complete_if_pending(
519 context: SyncCompletionPtr,
520 result: Result<T, String>,
521 ) -> bool {
522 let Some(inner) = take_context::<Arc<AsyncCompletionInner<T>>>(context) else {
523 return false;
524 };
525 cancel_async_timeout(&inner);
526
527 let completion_hook = {
528 let mut state = inner.state.lock().unwrap_or_else(PoisonError::into_inner);
529 state.completion_hook.take()
530 };
531 if let Some(completion_hook) = completion_hook {
532 catch_user_panic("async completion hook", || completion_hook(&result));
533 }
534 let waker = {
535 let mut state = inner.state.lock().unwrap_or_else(PoisonError::into_inner);
536 state.result = Some(result);
537 state.waker.take()
538 };
539 if let Some(waker) = waker {
540 waker.wake();
541 }
542 true
543 }
544}
545
546impl<T: Send + 'static> Future for AsyncCompletionFuture<T> {
547 type Output = Result<T, String>;
548
549 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
550 let waker = cx.waker().clone();
551 let mut state = self
552 .inner
553 .state
554 .lock()
555 .unwrap_or_else(PoisonError::into_inner);
556
557 if let Some(result) = state.result.take() {
558 drop(state);
559 drop(waker);
560 return Poll::Ready(result);
561 }
562
563 let replaced = match state.waker.as_ref() {
564 Some(existing) if existing.will_wake(&waker) => Some(waker),
565 _ => state.waker.replace(waker),
566 };
567 drop(state);
568 drop(replaced);
569 Poll::Pending
570 }
571}
572
573impl<T: Send + 'static> Drop for AsyncCompletionFuture<T> {
574 fn drop(&mut self) {
575 if self.cancel_on_drop {
576 remove_context(self.context_id);
577 cancel_async_timeout(&self.inner);
578 return;
579 }
580
581 let waker = self
582 .inner
583 .state
584 .lock()
585 .unwrap_or_else(PoisonError::into_inner)
586 .waker
587 .take();
588 drop(waker);
589 }
590}
591
592#[must_use]
598pub unsafe fn error_from_cstr(msg: *const i8) -> String {
599 if msg.is_null() {
600 "Unknown error".to_string()
601 } else {
602 unsafe { CStr::from_ptr(msg) }
603 .to_str()
604 .map_or_else(|_| "Unknown error".to_string(), String::from)
605 }
606}
607
608pub type UnitCompletion = SyncCompletion<()>;
610
611impl UnitCompletion {
612 #[allow(clippy::not_unsafe_ptr_arg_deref)]
614 pub extern "C" fn callback(context: SyncCompletionPtr, success: bool, msg: *const i8) {
615 catch_user_panic("UnitCompletion::callback", || {
616 if success {
617 unsafe { Self::complete_ok(context, ()) };
618 } else {
619 let error = unsafe { error_from_cstr(msg) };
620 unsafe { Self::complete_err(context, error) };
621 }
622 });
623 }
624}
625
626#[cfg(all(test, feature = "async"))]
627mod tests {
628 use super::*;
629 use std::sync::atomic::AtomicBool;
630
631 #[test]
632 fn completion_hook_runs_after_future_is_dropped() {
633 let called = Arc::new(AtomicBool::new(false));
634 let observed = Arc::clone(&called);
635 let (future, context) = AsyncCompletion::<()>::create_with_hook(move |result| {
636 assert!(result.is_ok());
637 observed.store(true, Ordering::Release);
638 });
639
640 drop(future);
641 unsafe { AsyncCompletion::complete_ok(context, ()) };
642
643 assert!(called.load(Ordering::Acquire));
644 }
645
646 #[test]
647 fn async_completion_times_out_and_wakes() {
648 struct WakeSignal(std::sync::mpsc::Sender<()>);
649 impl std::task::Wake for WakeSignal {
650 fn wake(self: Arc<Self>) {
651 self.0
652 .send(())
653 .expect("timeout wake receiver should remain available");
654 }
655 }
656
657 let (future, _context) =
658 AsyncCompletion::<()>::create_inner(None, true, Some(Duration::from_millis(100)));
659 let mut future = Box::pin(future);
660 let (wake_sender, wake_receiver) = std::sync::mpsc::channel();
661 let waker = Waker::from(Arc::new(WakeSignal(wake_sender)));
662 let mut context = Context::from_waker(&waker);
663 assert!(future.as_mut().poll(&mut context).is_pending());
664
665 wake_receiver
666 .recv_timeout(Duration::from_secs(5))
667 .expect("completion did not wake after its deadline");
668 let Poll::Ready(Err(error)) = future.as_mut().poll(&mut context) else {
669 panic!("completion did not resolve after its deadline");
670 };
671 assert!(is_timeout_error(&error));
672 }
673
674 #[test]
675 fn unbounded_completion_does_not_register_a_deadline() {
676 let (future, context) = AsyncCompletion::<()>::create_unbounded();
677 assert_eq!(future.inner.timeout_id.load(Ordering::Acquire), 0);
678
679 drop(future);
680 unsafe { AsyncCompletion::complete_ok(context, ()) };
681 }
682
683 #[test]
684 fn unbounded_start_completion_does_not_register_a_deadline() {
685 let called = Arc::new(AtomicBool::new(false));
686 let observed = Arc::clone(&called);
687 let (future, context) = AsyncCompletion::<()>::create_unbounded_with_hook(move |result| {
688 assert!(result.is_ok());
689 observed.store(true, Ordering::Release);
690 });
691 assert_eq!(future.inner.timeout_id.load(Ordering::Acquire), 0);
692
693 drop(future);
694 assert!(AsyncCompletion::complete_if_pending(context, Ok(())));
695 assert!(called.load(Ordering::Acquire));
696 }
697
698 #[test]
699 fn late_start_completion_reports_a_timeout() {
700 let (sender, receiver) = std::sync::mpsc::channel();
701 let (future, context) = AsyncCompletion::<()>::create_inner(
702 Some(Box::new(move |result| {
703 sender.send(result.is_err()).expect("receiver is open");
704 })),
705 false,
706 Some(Duration::from_millis(10)),
707 );
708 assert!(receiver
709 .recv_timeout(Duration::from_secs(5))
710 .expect("the timeout hook did not run"));
711 assert!(!AsyncCompletion::complete_if_pending(context, Ok(())));
712 drop(future);
713 }
714
715 #[test]
716 fn unrepresentable_deadline_registers_no_timeout() {
717 let (future, context) =
718 AsyncCompletion::<()>::create_inner(None, true, Some(Duration::MAX));
719 assert_eq!(future.inner.timeout_id.load(Ordering::Acquire), 0);
720
721 drop(future);
722 unsafe { AsyncCompletion::complete_ok(context, ()) };
723 }
724}