screencapturekit/utils/
completion.rs

1//! Completion handles for Swift bridge callbacks.
2//!
3//! Synchronous waits are bounded by default. Callback contexts are opaque
4//! monotonic tokens backed by a process registry, not addresses. A callback
5//! that arrives after timeout/cancellation, or fires more than once, therefore
6//! finds no registry entry and returns without touching freed memory.
7
8use 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
20/// Default bound for synchronous waits and async completion futures.
21pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(30);
22
23/// Overrides [`DEFAULT_TIMEOUT`] in whole seconds; `0` disables the bound.
24pub const TIMEOUT_ENV_VAR: &str = "SCREENCAPTUREKIT_COMPLETION_TIMEOUT_SECS";
25
26/// Stable prefix for timeout errors.
27pub const TIMEOUT_MESSAGE_PREFIX: &str = "screencapturekit: completion callback did not fire";
28
29/// Whether an error came from a bounded wait expiring.
30#[must_use]
31pub fn is_timeout_error(message: &str) -> bool {
32    message.starts_with(TIMEOUT_MESSAGE_PREFIX)
33}
34
35/// Effective process-wide wait bound.
36#[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/// Number of synchronous operations that reached their wait deadline.
51#[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
62/// Opaque value passed through FFI callbacks.
63pub 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
231/// A blocking completion handler for asynchronous FFI callbacks.
232pub 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    /// Create a completion handle and its opaque callback context.
253    #[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    /// Wait until completion or the process-wide default deadline.
264    ///
265    /// # Errors
266    ///
267    /// Returns the callback error or a timeout error.
268    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    /// Wait with an explicit deadline.
276    ///
277    /// # Errors
278    ///
279    /// Returns the callback error or a timeout error.
280    #[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    /// Wait without a deadline.
304    ///
305    /// # Errors
306    ///
307    /// Returns the callback error.
308    #[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    /// Complete successfully.
326    ///
327    /// # Safety
328    ///
329    /// `context` must be the opaque token returned by [`Self::new`] for this
330    /// concrete `T`.
331    pub unsafe fn complete_ok(context: SyncCompletionPtr, value: T) {
332        unsafe { Self::complete_with_result(context, Ok(value)) };
333    }
334
335    /// Complete with an error.
336    ///
337    /// # Safety
338    ///
339    /// `context` must be the opaque token returned by [`Self::new`] for this
340    /// concrete `T`.
341    pub unsafe fn complete_err(context: SyncCompletionPtr, error: String) {
342        unsafe { Self::complete_with_result(context, Err(error)) };
343    }
344
345    /// Complete with a result.
346    ///
347    /// Duplicate, late, or cancelled callbacks are ignored safely.
348    ///
349    /// # Safety
350    ///
351    /// `context` must be the opaque token returned by [`Self::new`] for this
352    /// concrete `T`.
353    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
391/// Factory for future-based FFI completion handles.
392pub 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
402/// Future returned by [`AsyncCompletion::create`].
403pub 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    /// Create a future and its opaque callback context.
418    ///
419    /// The future resolves with a timeout error if the native callback does
420    /// not arrive within [`default_timeout`].
421    #[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    fn create_inner(
439        completion_hook: Option<AsyncCompletionHook<T>>,
440        cancel_on_drop: bool,
441        timeout: Option<Duration>,
442    ) -> (AsyncCompletionFuture<T>, SyncCompletionPtr) {
443        let inner = Arc::new(AsyncCompletionInner {
444            state: Mutex::new(AsyncCompletionState {
445                result: None,
446                waker: None,
447                completion_hook,
448            }),
449            timeout_id: AtomicUsize::new(0),
450        });
451        let (context, context_id) = register_context(Arc::clone(&inner));
452        if let Some(timeout) = timeout {
453            let context_address = context as usize;
454            let timeout_id = TimeoutScheduler::schedule(timeout, move || unsafe {
455                Self::complete_err(
456                    context_address as SyncCompletionPtr,
457                    timeout_message(timeout),
458                );
459            });
460            inner.timeout_id.store(timeout_id, Ordering::Release);
461        }
462        (
463            AsyncCompletionFuture {
464                inner,
465                context_id,
466                cancel_on_drop,
467            },
468            context,
469        )
470    }
471
472    /// Complete successfully.
473    ///
474    /// # Safety
475    ///
476    /// `context` must be the opaque token returned by [`Self::create`] for this
477    /// concrete `T`.
478    pub unsafe fn complete_ok(context: SyncCompletionPtr, value: T) {
479        unsafe { Self::complete_with_result(context, Ok(value)) };
480    }
481
482    /// Complete with an error.
483    ///
484    /// # Safety
485    ///
486    /// `context` must be the opaque token returned by [`Self::create`] for this
487    /// concrete `T`.
488    pub unsafe fn complete_err(context: SyncCompletionPtr, error: String) {
489        unsafe { Self::complete_with_result(context, Err(error)) };
490    }
491
492    /// Complete with a result. Duplicate, late, or cancelled callbacks are
493    /// ignored safely.
494    ///
495    /// # Safety
496    ///
497    /// `context` must be the opaque token returned by [`Self::create`] for this
498    /// concrete `T`.
499    pub unsafe fn complete_with_result(context: SyncCompletionPtr, result: Result<T, String>) {
500        let Some(inner) = take_context::<Arc<AsyncCompletionInner<T>>>(context) else {
501            return;
502        };
503        cancel_async_timeout(&inner);
504
505        let completion_hook = {
506            let mut state = inner.state.lock().unwrap_or_else(PoisonError::into_inner);
507            state.completion_hook.take()
508        };
509        if let Some(completion_hook) = completion_hook {
510            catch_user_panic("async completion hook", || completion_hook(&result));
511        }
512        let waker = {
513            let mut state = inner.state.lock().unwrap_or_else(PoisonError::into_inner);
514            state.result = Some(result);
515            state.waker.take()
516        };
517        if let Some(waker) = waker {
518            waker.wake();
519        }
520    }
521}
522
523impl<T: Send + 'static> Future for AsyncCompletionFuture<T> {
524    type Output = Result<T, String>;
525
526    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
527        let waker = cx.waker().clone();
528        let mut state = self
529            .inner
530            .state
531            .lock()
532            .unwrap_or_else(PoisonError::into_inner);
533
534        if let Some(result) = state.result.take() {
535            drop(state);
536            drop(waker);
537            return Poll::Ready(result);
538        }
539
540        let replaced = match state.waker.as_ref() {
541            Some(existing) if existing.will_wake(&waker) => Some(waker),
542            _ => state.waker.replace(waker),
543        };
544        drop(state);
545        drop(replaced);
546        Poll::Pending
547    }
548}
549
550impl<T: Send + 'static> Drop for AsyncCompletionFuture<T> {
551    fn drop(&mut self) {
552        if self.cancel_on_drop {
553            remove_context(self.context_id);
554            cancel_async_timeout(&self.inner);
555            return;
556        }
557
558        let waker = self
559            .inner
560            .state
561            .lock()
562            .unwrap_or_else(PoisonError::into_inner)
563            .waker
564            .take();
565        drop(waker);
566    }
567}
568
569/// Convert an optional NUL-terminated C string into an owned Rust string.
570///
571/// # Safety
572///
573/// `msg` must be null or point to a valid NUL-terminated string.
574#[must_use]
575pub unsafe fn error_from_cstr(msg: *const i8) -> String {
576    if msg.is_null() {
577        "Unknown error".to_string()
578    } else {
579        unsafe { CStr::from_ptr(msg) }
580            .to_str()
581            .map_or_else(|_| "Unknown error".to_string(), String::from)
582    }
583}
584
585/// Completion for operations that return only success or an error.
586pub type UnitCompletion = SyncCompletion<()>;
587
588impl UnitCompletion {
589    /// C callback for `(context, success, error_message)` operations.
590    #[allow(clippy::not_unsafe_ptr_arg_deref)]
591    pub extern "C" fn callback(context: SyncCompletionPtr, success: bool, msg: *const i8) {
592        catch_user_panic("UnitCompletion::callback", || {
593            if success {
594                unsafe { Self::complete_ok(context, ()) };
595            } else {
596                let error = unsafe { error_from_cstr(msg) };
597                unsafe { Self::complete_err(context, error) };
598            }
599        });
600    }
601}
602
603#[cfg(all(test, feature = "async"))]
604mod tests {
605    use super::*;
606    use std::sync::atomic::AtomicBool;
607
608    #[test]
609    fn completion_hook_runs_after_future_is_dropped() {
610        let called = Arc::new(AtomicBool::new(false));
611        let observed = Arc::clone(&called);
612        let (future, context) = AsyncCompletion::<()>::create_with_hook(move |result| {
613            assert!(result.is_ok());
614            observed.store(true, Ordering::Release);
615        });
616
617        drop(future);
618        unsafe { AsyncCompletion::complete_ok(context, ()) };
619
620        assert!(called.load(Ordering::Acquire));
621    }
622
623    #[test]
624    fn async_completion_times_out_and_wakes() {
625        struct WakeSignal(std::sync::mpsc::Sender<()>);
626        impl std::task::Wake for WakeSignal {
627            fn wake(self: Arc<Self>) {
628                self.0
629                    .send(())
630                    .expect("timeout wake receiver should remain available");
631            }
632        }
633
634        let (future, _context) =
635            AsyncCompletion::<()>::create_inner(None, true, Some(Duration::from_millis(100)));
636        let mut future = Box::pin(future);
637        let (wake_sender, wake_receiver) = std::sync::mpsc::channel();
638        let waker = Waker::from(Arc::new(WakeSignal(wake_sender)));
639        let mut context = Context::from_waker(&waker);
640        assert!(future.as_mut().poll(&mut context).is_pending());
641
642        wake_receiver
643            .recv_timeout(Duration::from_secs(5))
644            .expect("completion did not wake after its deadline");
645        let Poll::Ready(Err(error)) = future.as_mut().poll(&mut context) else {
646            panic!("completion did not resolve after its deadline");
647        };
648        assert!(is_timeout_error(&error));
649    }
650
651    #[test]
652    fn unbounded_completion_does_not_register_a_deadline() {
653        let (future, context) = AsyncCompletion::<()>::create_unbounded();
654        assert_eq!(future.inner.timeout_id.load(Ordering::Acquire), 0);
655
656        drop(future);
657        unsafe { AsyncCompletion::complete_ok(context, ()) };
658    }
659
660    #[test]
661    fn unrepresentable_deadline_registers_no_timeout() {
662        let (future, context) =
663            AsyncCompletion::<()>::create_inner(None, true, Some(Duration::MAX));
664        assert_eq!(future.inner.timeout_id.load(Ordering::Acquire), 0);
665
666        drop(future);
667        unsafe { AsyncCompletion::complete_ok(context, ()) };
668    }
669}
screencapturekit/utils/completion.rs

screencapturekit/utils/
completion.rs
