gpu_handle_types/

wait_thread.rs

// SPDX-License-Identifier: MIT OR Apache-2.0

//! Waiter-thread infrastructure behind `async` waits on primitives that
//! only offer a blocking wait.
//!
//! The fast-path yield-spin inside [`SyncWaiter::wait_async`] catches the
//! sub-millisecond signal-already-here case with 64 cooperative yields.
//! Anything past that bounded spin must hand off to a **permanently-
//! attached** thread that issues the **native blocking-with-timeout**
//! primitive on 10 ms slices — without `tokio::spawn_blocking`, without
//! per-runtime timers, without per-operation thread spawn.
//!
//! ## Components
//!
//! - [`WaiterThread`] — owns the request channel + the OS thread.
//!   Typically one instance per backend kind (lazy-initialised via
//!   [`OnceLock`] inside each backend's waiter module); per-instance
//!   ownership is supported when callers want a deterministic
//!   shutdown via `Drop`.
//! - [`SliceFn`] — boxed closure each per-backend `wait_async`
//!   override builds at handoff time. The waiter thread calls it
//!   repeatedly with up to 10 ms of slice budget; it returns
//!   [`SliceOutcome::Signaled`] / `TimedOut` / `Failed`.
//! - [`BackendWaitFuture`] — the future returned to the executor.
//!   Its [`Drop`] flips a cancellation flag the waiter thread observes
//!   at slice boundaries; cancellation latency is therefore ≤ one
//!   slice (10 ms) — sized to fit inside a 60 fps frame budget.
//! - [`run_hybrid_wait`] — composes the fast-path yield-spin + the
//!   waiter-thread fallback (thread handoff). Backends with a native blocking primitive call
//!   this from their `wait_async` override and pass the slice
//!   closure.
//!
//! ## Why a per-process (per-backend-kind) static thread
//!
//! The waiter thread is **stateless** — every per-instance datum (device
//! handle, semaphore, fence, …) lives inside the [`SliceFn`] closure — so
//! one thread per backend kind serves every waiter of that kind. A second
//! per-instance thread would spend its life parked on `recv()` doing no
//! work, and a thread per `SyncWaiter` would spawn thousands per second on
//! a steady stream of frames.
//!
//! The trade-off: with a static singleton the thread is not joined on
//! `Drop` of any individual backend; it lives until process exit and
//! is reaped by the OS. The `Drop`-time join contract
//! still holds when callers construct a [`WaiterThread`] explicitly
//! and own it from a parent registry struct (the type supports that
//! shape — `new` + `Drop` join cleanly), which is how this crate's
//! tests exercise the lifecycle.

use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
// mpsc channel + join handle back the native waiter thread; wasm has none.
#[cfg(not(target_family = "wasm"))]
use std::sync::mpsc;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll, Waker};
#[cfg(not(target_family = "wasm"))]
use std::thread::JoinHandle;
use std::time::{Duration, Instant};

use crate::Error;

/// Per-slice native blocking-with-timeout budget. Sized to:
///
/// - Stay below the 60 fps frame budget (16.6 ms) so a future
///   dropped at the frame boundary releases its waiter slot inside
///   the same frame.
/// - Be coarse enough that the syscall overhead per slice is
///   amortised even on long waits — at 5 s of accumulated wait, the
///   thread re-enters the kernel ~500 times, not 5_000.
/// - Bound cancellation + deadline-exceeded latency to a single
///   slice.
///
/// Tunable as a `Duration` rather than a feature so the value can be
/// re-derived in tests without rebuilding the crate.
pub const WAITER_SLICE: Duration = Duration::from_millis(10);

/// Result of one slice-bounded native wait. Returned by the
/// [`SliceFn`] closure each per-backend `wait_async` override
/// supplies.
#[derive(Debug)]
pub enum SliceOutcome {
    /// The primitive reached its target value within the slice.
    /// Waiter thread resolves the future with `Ok(())`.
    Signaled,
    /// The slice's native timeout elapsed without signalling. Waiter
    /// thread re-issues a fresh slice (or transitions to
    /// `Err(Timeout)` if the per-request `deadline` has now passed).
    TimedOut,
    /// Driver-level failure (`VK_ERROR_DEVICE_LOST`,
    /// `DXGI_ERROR_DEVICE_REMOVED`, `CL_INVALID_EVENT`, …). Waiter
    /// thread resolves the future with this error directly.
    Failed(Error),
}

/// Boxed closure each per-backend `wait_async` override builds at
/// handoff time. The waiter thread calls it on its own thread with
/// the current slice budget; the closure must:
///
/// - Issue the native blocking-with-timeout primitive
///   (`vkWaitSemaphores`, `WaitForSingleObject`,
///   `MTLSharedEvent::waitUntilSignaledValue:timeoutMS:`,
///   `clWaitForEvents`, `device.poll(Wait { … })`, …) with the slice
///   budget.
///
/// - Return one of the three [`SliceOutcome`] variants.
///
/// `Send` is required because the closure runs on the waiter thread,
/// not on the executor's poll thread. `'static` because the slice
/// closure is consumed asynchronously and outlives the `await`
/// position that built it.
// `Send` off wasm (the slice runs on the waiter thread). On wasm there is
// no waiter thread — `enqueue` reports `NotSupported` — and the
// slice closure may close over thread-affine `wgpu` handles, so the bound
// is dropped. An explicit `+ Send` on a `dyn FnMut` cannot be spelled with
// a non-auto marker trait, so the alias is cfg-split directly.
#[cfg(not(target_family = "wasm"))]
pub type SliceFn = Box<dyn FnMut(Duration) -> SliceOutcome + Send + 'static>;
#[cfg(target_family = "wasm")]
pub type SliceFn = Box<dyn FnMut(Duration) -> SliceOutcome + 'static>;

/// Shared state between [`BackendWaitFuture`] (on the executor) and
/// the waiter thread.
struct WaitCompletion {
    result: Option<Result<(), Error>>,
    waker: Option<Waker>,
}

// Native-only: the request that crosses the channel to the waiter thread.
// wasm has no waiter thread, so this is never constructed there.
#[cfg(not(target_family = "wasm"))]
struct WaitRequest {
    slice_fn: SliceFn,
    deadline: Option<Instant>,
    cancelled: Arc<AtomicBool>,
    completion: Arc<Mutex<WaitCompletion>>,
}

/// Permanently-attached waiter thread + request channel.
///
/// Cheap to construct (one `mpsc::channel` + one `thread::spawn`).
/// Joined cleanly on `Drop`: the channel sender is dropped, the
/// thread's `recv()` returns `Err`, the loop exits, and the join
/// handle is consumed.
///
/// Every request that is in flight or still queued when `Drop` runs is
/// resolved with [`Error::Cancelled`] before the thread exits.
/// [`BackendWaitFuture`] borrows nothing from this type, so a caller may
/// hold one across the drop; without that resolution its `await` would
/// pend forever.
///
/// Per-instance ownership is supported, but the usual shape is a
/// per-backend-kind static singleton in each backend's waiter module.
/// See module-level docs.
pub struct WaiterThread {
    // Native-only: the channel + join handle to the waiter thread. On wasm
    // these are absent so `WaiterThread` stays `Send + Sync` (a
    // `Sender<WaitRequest>` carrying a `!Send` `SliceFn` would otherwise be
    // `!Sync`, breaking the process-`static` singletons that hold it).
    #[cfg(not(target_family = "wasm"))]
    sender: Option<mpsc::Sender<WaitRequest>>,
    #[cfg(not(target_family = "wasm"))]
    join: Option<JoinHandle<()>>,
    /// Thread-wide shutdown flag observed at slice boundaries.
    /// Flipped in `Drop` so the thread exits even mid-request — closing
    /// the channel alone wouldn't suffice when a long-running request is
    /// in flight (the slice loop never returns to `recv()`).
    shutdown: Arc<AtomicBool>,
}

impl std::fmt::Debug for WaiterThread {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        #[cfg(not(target_family = "wasm"))]
        let alive = self.sender.is_some();
        // wasm has no waiter thread — never "alive".
        #[cfg(target_family = "wasm")]
        let alive = false;
        f.debug_struct("WaiterThread").field("alive", &alive).finish()
    }
}

impl WaiterThread {
    /// Spawn a new waiter thread. `name` is appended to `"gpu-wait-"` for
    /// the thread name visible to debuggers and `ps -L` /
    /// `Process Explorer`. Pick a short, stable backend-identifying string
    /// (`"vulkan"`, `"wgpu"`, `"d3d12"`, `"metal"`, `"opencl"`, …) —
    /// Linux truncates thread names to 15 bytes.
    pub fn new(name: &str) -> Self {
        // wasm has no OS threads. Construct an inert `WaiterThread`
        // (`sender = None`): `enqueue` then takes its existing
        // "already shut down" branch and reports `NotSupported`. Blocking
        // waits are unavailable on wasm — callers use a non-blocking
        // path; the `run_hybrid_wait` fast-path spin still runs.
        #[cfg(target_family = "wasm")]
        {
            let _ = name;
            Self { shutdown: Arc::new(AtomicBool::new(false)) }
        }
        #[cfg(not(target_family = "wasm"))]
        {
            let (tx, rx) = mpsc::channel::<WaitRequest>();
            let shutdown = Arc::new(AtomicBool::new(false));
            let shutdown_thread = shutdown.clone();
            let join = std::thread::Builder::new()
                .name(format!("gpu-wait-{name}"))
                .spawn(move || waiter_loop(&rx, &shutdown_thread))
                .expect("gpu-handle-types: failed to spawn waiter thread");
            Self { sender: Some(tx), join: Some(join), shutdown }
        }
    }

    /// Enqueue a wait request. Returns a [`BackendWaitFuture`] the
    /// executor can `.await`; `Drop` of the future signals
    /// cancellation to the waiter thread.
    ///
    /// `deadline = None` means "wait forever" — the slice loop never
    /// transitions to `Err(Timeout)` and only exits on
    /// `Signaled` / `Failed` / `cancelled`.
    #[cfg_attr(target_family = "wasm", allow(unused_variables, clippy::needless_pass_by_value))]
    pub fn enqueue(&self, slice_fn: SliceFn, deadline: Option<Instant>) -> BackendWaitFuture {
        let cancelled = Arc::new(AtomicBool::new(false));
        let completion = Arc::new(Mutex::new(WaitCompletion { result: None, waker: None }));
        // wasm: no waiter thread — blocking waits are unavailable.
        // Resolve the future immediately to `NotSupported`; callers use a
        // non-blocking path. The `run_hybrid_wait` fast-path spin
        // already handles the already-signaled case before reaching here.
        #[cfg(target_family = "wasm")]
        finish_completion(
            &completion,
            Err(Error::NotSupported(
                "blocking waits are unavailable on wasm (no waiter thread); drive the async wait \
                 path instead"
                    .into(),
            )),
        );
        #[cfg(not(target_family = "wasm"))]
        {
            let req = WaitRequest { slice_fn, deadline, cancelled: cancelled.clone(), completion: completion.clone() };
            match self.sender.as_ref() {
                Some(s) => {
                    if s.send(req).is_err() {
                        finish_completion(
                            &completion,
                            Err(Error::NotSupported("waiter thread died before request was accepted".into())),
                        );
                    }
                }
                None => {
                    finish_completion(&completion, Err(Error::NotSupported("waiter thread already shut down".into())));
                }
            }
        }
        BackendWaitFuture { completion, cancelled }
    }
}

impl Drop for WaiterThread {
    fn drop(&mut self) {
        // Two-phase shutdown: (1) flip the thread-wide flag so the
        // slice loop exits at the next slice boundary even if the
        // request is still nominally in flight; (2) drop the sender
        // so the outer `recv()` returns `Err` after the current
        // request (if any) finishes. (wasm holds neither the sender nor
        // the join handle — nothing to tear down beyond the flag.)
        self.shutdown.store(true, Ordering::Release);
        #[cfg(not(target_family = "wasm"))]
        {
            self.sender.take();
            if let Some(h) = self.join.take() {
                let _ = h.join();
            }
        }
    }
}

#[cfg(not(target_family = "wasm"))]
fn waiter_loop(rx: &mpsc::Receiver<WaitRequest>, shutdown: &AtomicBool) {
    while let Ok(mut req) = rx.recv() {
        if shutdown.load(Ordering::Acquire) {
            abandon_on_shutdown(&req);
            break;
        }
        // Set when the per-request loop exited because the thread is
        // shutting down, so the outer loop stops pulling new work
        // instead of blocking in `recv()` behind a sender that
        // `WaiterThread::drop` may not have released yet.
        let mut shutting_down = false;
        loop {
            // Thread-wide shutdown — `WaiterThread::drop` was called.
            // Exit the per-request loop and the outer recv loop.
            if shutdown.load(Ordering::Acquire) {
                abandon_on_shutdown(&req);
                shutting_down = true;
                break;
            }
            // Cancellation — future dropped. Abandon the request;
            // the native slice just issued (if any) has already
            // returned, and the closure may keep state across
            // returns that dropping it mid-slice would corrupt.
            // Slice granularity (10 ms) is the cancellation latency
            // bound by design.
            //
            // No `finish_completion` here, deliberately: the only way
            // this flag is set is [`BackendWaitFuture::drop`], so there
            // is no future left to resolve. Contrast the shutdown arm
            // above, where the future is typically still alive.
            if req.cancelled.load(Ordering::Acquire) {
                break;
            }
            // Deadline check.
            let remaining = match req.deadline {
                Some(d) => match d.checked_duration_since(Instant::now()) {
                    Some(r) => r,
                    None => {
                        finish_completion(&req.completion, Err(Error::Timeout));
                        break;
                    }
                },
                None => Duration::MAX,
            };
            let slice = std::cmp::min(WAITER_SLICE, remaining);
            match (req.slice_fn)(slice) {
                SliceOutcome::Signaled => {
                    finish_completion(&req.completion, Ok(()));
                    break;
                }
                SliceOutcome::TimedOut => continue,
                SliceOutcome::Failed(e) => {
                    finish_completion(&req.completion, Err(e));
                    break;
                }
            }
        }
        if shutting_down {
            break;
        }
    }
    // Requests still queued when the thread stops are resolved too, for
    // the same reason: their futures are alive and nothing else will
    // ever complete them. The sender has been dropped (or is about to
    // be) by `WaiterThread::drop`, so this drains and terminates.
    while let Ok(req) = rx.try_recv() {
        abandon_on_shutdown(&req);
    }
}

/// Resolve a request the waiter thread is giving up on because the
/// thread itself is being torn down.
///
/// This is **not** optional bookkeeping. [`BackendWaitFuture`] borrows
/// nothing from [`WaiterThread`] — it holds only the shared completion
/// cell — so `let f = t.enqueue(..); drop(t); f.await` is well-typed and
/// a caller may legitimately outlive the thread it enqueued on. Leaving
/// the completion unresolved makes that `await` pend forever with no
/// waker left in existence to rescue it.
///
/// [`Error::Cancelled`] rather than a failure kind: tearing the waiter
/// thread down is a caller-initiated stop, exactly the "not a failure,
/// the operation was asked to stop" case that
/// [`crate::ErrorKind::Cancelled`] exists to classify. Callers void it
/// instead of surfacing a fault.
#[cfg(not(target_family = "wasm"))]
fn abandon_on_shutdown(req: &WaitRequest) {
    finish_completion(&req.completion, Err(Error::Cancelled));
}

fn finish_completion(completion: &Arc<Mutex<WaitCompletion>>, result: Result<(), Error>) {
    let mut c = completion.lock().expect("WaitCompletion mutex poisoned");
    // Don't clobber an existing terminal outcome (the cancellation
    // path lets the future drop without resolving, but a follow-up
    // poll would still observe the result — keep the first one).
    if c.result.is_none() {
        c.result = Some(result);
    }
    let waker = c.waker.take();
    drop(c);
    if let Some(w) = waker {
        w.wake();
    }
}

/// Future returned by [`WaiterThread::enqueue`]. Resolves when the
/// waiter thread reports `Signaled` / `Failed` / `Err(Timeout)` on
/// the deadline.
///
/// Drop semantics — `Drop` flips the cancellation flag the waiter
/// thread observes at the next slice boundary. The native wait
/// already in flight completes its current slice (up to
/// [`WAITER_SLICE`]) before the slot is released; cancellation
/// latency is bounded.
pub struct BackendWaitFuture {
    completion: Arc<Mutex<WaitCompletion>>,
    cancelled: Arc<AtomicBool>,
}

impl std::fmt::Debug for BackendWaitFuture {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("BackendWaitFuture")
            .field("cancelled", &self.cancelled.load(Ordering::Relaxed))
            .field("resolved", &self.completion.lock().map(|c| c.result.is_some()).unwrap_or(false))
            .finish()
    }
}

impl Future for BackendWaitFuture {
    type Output = Result<(), Error>;
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let mut c = self.completion.lock().expect("WaitCompletion mutex poisoned");
        if let Some(r) = c.result.take() {
            return Poll::Ready(r);
        }
        // Refresh the waker on every poll — executors that pass new
        // wakers per re-poll (futures::executor) need this.
        c.waker = Some(cx.waker().clone());
        Poll::Pending
    }
}

impl Drop for BackendWaitFuture {
    fn drop(&mut self) {
        self.cancelled.store(true, Ordering::Release);
    }
}

/// Compose the fast-path yield-spin (bounded) + the waiter-thread
/// fallback (thread handoff) into the canonical hybrid `wait_async` body.
///
/// Per-backend `SyncWaiter::wait_async` overrides call this with:
/// - `is_signaled`: closure issuing the backend's non-blocking probe
///   (`vkGetSemaphoreCounterValue`, `GetCompletedValue`,
///   `signaledValue() >= value`, `clGetEventInfo`, `device.poll`
///   with `Duration::ZERO`, …). Runs on the executor thread during
///   the fast-path yield-spin.
/// - `waiter_thread`: per-backend-kind [`WaiterThread`] reference.
/// - `timeout`: absolute timeout from the caller.
/// - `make_slice_fn`: builder for the [`SliceFn`] consumed on the
///   waiter-thread fallback. Lazy so backends that resolve in the
///   fast-path yield-spin skip closure construction entirely.
///
/// Returns:
/// - `Ok(())` on signal,
/// - `Err(Error::Timeout)` on deadline,
/// - `Err(Error::DeviceLost { .. })` / `Err(Error::NotSupported(_.into()))`
///   on driver-level failure surfaced by the slice closure.
pub async fn run_hybrid_wait<F, M>(
    is_signaled: F,
    waiter_thread: &WaiterThread,
    timeout: Duration,
    make_slice_fn: M,
) -> Result<(), Error>
where
    // `MaybeSend` = `Send` off wasm. On wasm the closures
    // may capture thread-affine `wgpu` handles; the fast-path spin loop
    // still runs on the executor, and the waiter-thread fallback reports
    // `NotSupported`.
    F: Fn() -> Result<bool, Error> + crate::MaybeSend,
    M: FnOnce() -> SliceFn + crate::MaybeSend,
{
    // Fast-path yield-spin — bounded yield-poll. Iteration count, not wall-clock
    // (pollster's re-poll-on-wake collapses wall-clock caps to a
    // CPU burn here).
    const SPIN_ITERATIONS: usize = 64;
    // `None` iff `timeout == Duration::MAX` (the wait-forever sentinel, which
    // `enqueue` also reads as "no deadline"); every finite timeout gets a bounded,
    // always-representable deadline, so an `Instant + Duration` overflow cannot
    // masquerade as "wait forever".
    let deadline = crate::wait_deadline(timeout);
    for _ in 0..SPIN_ITERATIONS {
        match is_signaled() {
            Ok(true) => return Ok(()),
            Ok(false) => {}
            Err(e) => return Err(e),
        }
        if let Some(d) = deadline
            && Instant::now() >= d
        {
            return Err(Error::Timeout);
        }
        crate::yield_once().await;
    }
    // Waiter-thread fallback — hand off to the waiter thread.
    let slice_fn = make_slice_fn();
    waiter_thread.enqueue(slice_fn, deadline).await
}