rightkit_gpu/
exit.rs

1//! Leaving: no GPU work may be in flight when the process ends (issue #620). Adapted from Ember `crates/gpu`.
2//!
3//! The devices are statics that are never dropped, and the threads that render (preview workers,
4//! export, denoise, the device warm-up) are not the one that ends the process. Exiting runs the
5//! driver's own teardown while those threads keep running: a render that is inside the driver at
6//! that moment crashes there (seen as a SIGSEGV in the NVIDIA Vulkan driver), where neither the
7//! panic hook nor `catch_unwind` can help.
8//!
9//! So every entry into a device goes through [`enter`], and the host ends the process with
10//! [`begin_shutdown`] (nothing new enters) and [`wait_idle`] (what is in flight
11//! finishes, up to a deadline).
12
13use std::sync::{Condvar, Mutex, MutexGuard};
14use std::time::{Duration, Instant};
15
16struct State {
17    closed: bool,
18    /// [`Work`] in flight, on any thread.
19    busy: usize,
20}
21
22pub struct GpuGate {
23    state: Mutex<State>,
24    idle: Condvar,
25}
26
27/// GPU work in flight (until dropped).
28pub struct Work<'a>(&'a GpuGate);
29
30static GATE: GpuGate = GpuGate::new();
31
32/// Start GPU work, or `None` when the process is ending: the caller then does without the GPU.
33/// Work may nest (a render that creates the device).
34pub fn enter() -> Option<Work<'static>> {
35    GATE.enter()
36}
37
38pub fn begin_shutdown() {
39    GATE.close();
40}
41
42pub fn shutting_down() -> bool {
43    GATE.lock().closed
44}
45
46pub fn wait_idle(timeout: Duration) -> bool {
47    GATE.wait_idle(timeout)
48}
49
50impl Default for GpuGate {
51    fn default() -> Self {
52        Self::new()
53    }
54}
55
56impl GpuGate {
57    pub const fn new() -> GpuGate {
58        GpuGate {
59            state: Mutex::new(State {
60                closed: false,
61                busy: 0,
62            }),
63            idle: Condvar::new(),
64        }
65    }
66
67    fn lock(&self) -> MutexGuard<'_, State> {
68        self.state.lock().unwrap_or_else(|e| e.into_inner())
69    }
70
71    pub fn enter(&self) -> Option<Work<'_>> {
72        let mut s = self.lock();
73        if s.closed {
74            return None;
75        }
76        s.busy = s.busy.saturating_add(1);
77        Some(Work(self))
78    }
79
80    /// Nothing enters from now on.
81    pub fn close(&self) {
82        self.lock().closed = true;
83    }
84
85    /// Wait up to `timeout` until no work is in flight; `false` when some still is.
86    pub fn wait_idle(&self, timeout: Duration) -> bool {
87        let t0 = Instant::now();
88        let mut s = self.lock();
89        while s.busy > 0 {
90            let Some(left) = timeout.checked_sub(t0.elapsed()).filter(|d| !d.is_zero()) else {
91                return false;
92            };
93            s = self
94                .idle
95                .wait_timeout(s, left)
96                .unwrap_or_else(|e| e.into_inner())
97                .0;
98        }
99        true
100    }
101}
102
103impl Drop for Work<'_> {
104    fn drop(&mut self) {
105        let mut s = self.0.lock();
106        s.busy = s.busy.saturating_sub(1);
107        if s.busy == 0 {
108            drop(s);
109            self.0.idle.notify_all();
110        }
111    }
112}
rightkit_gpu/exit.rs

rightkit_gpu/
exit.rs
